diff --git a/src/postgres_history.py b/src/postgres_history.py index 1bae301..4e9b863 100644 --- a/src/postgres_history.py +++ b/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], - "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, + "station_code": station_code, + "water_level": water_level, + "discharge": discharge, "discharge_percent": float(row[4]) if row[4] is not None else None, } )