feat: add _calculate_discharge to estimate from water level when DB discharge is NULL

src/postgres_history.py

@@ -2,11 +2,35 @@
 
 import datetime
 import os
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 from sqlalchemy import create_engine, text
 
 
+# Stage-discharge rating curves: Q = a * (H - b)^c
+# Key: station_code, Value: (a, b, c)
+# Use linear fallback Q = slope * H if a curve is not defined.
+_RATING_CURVES: Dict[str, Tuple[float, float, float]] = {}
+_DEFAULT_LINEAR_SLOPE = 20.0  # m^3/s per meter
+
+
+def _calculate_discharge(water_level: Optional[float], station_code: str = None) -> Optional[float]:
+    """Estimate discharge from water level using a rating curve or linear fallback."""
+    if water_level is None:
+        return None
+
+    curve = _RATING_CURVES.get(station_code)
+    if curve:
+        a, b, c = curve
+        h_excess = water_level - b
+        if h_excess <= 0:
+            return 0.0
+        return round(a * (h_excess ** c), 2)
+    
+    # Linear fallback: Q = slope * H
+    return round(_DEFAULT_LINEAR_SLOPE * water_level, 2)
+
+
 class PostgresHistory:
     def __init__(self, connection_string: Optional[str] = None, engine=None):
         connection_string = connection_string or os.getenv("POSTGRES_CONNECTION_STRING")
@@ -54,12 +78,18 @@ class PostgresHistory:
                 timestamp = row[0]
                 if isinstance(timestamp, str):
                     timestamp = datetime.datetime.fromisoformat(timestamp)
+                station_code = row[1]
+                water_level = float(row[2]) if row[2] is not None else None
+                discharge = float(row[3]) if row[3] is not None else None
+                # Estimate discharge from water level if DB value is missing
+                if discharge is None and water_level is not None:
+                    discharge = _calculate_discharge(water_level, station_code)
                 result.append(
                     {
                         "timestamp": timestamp,
-                        "station_code": row[1],
+                        "station_code": station_code,
-                        "water_level": float(row[2]) if row[2] is not None else None,
+                        "water_level": water_level,
-                        "discharge": float(row[3]) if row[3] is not None else None,
+                        "discharge": discharge,
                         "discharge_percent": float(row[4]) if row[4] is not None else None,
                     }
                 )