"""Flood forecast inference. Integration contract (see get_forecasts / get_latest_forecasts): callers pass raw station readings, get back one forecast dict per station x horizon. A station with a stale, missing, or version-mismatched model transparently falls back to a simple persistence heuristic instead of raising -- this module must never crash the caller (e.g. the web API). """ import datetime import logging from pathlib import Path from typing import Dict, List, Optional, Tuple, Union import joblib import numpy as np import pandas as pd from . import features logger = logging.getLogger(__name__) DEFAULT_HORIZONS: Tuple[int, ...] = (6, 12, 24) STALE_AFTER_H = 6.0 HEURISTIC_SIGMA = 0.3 HEURISTIC_VERSION = "heuristic-v1" # Keyed by (path, mtime) so a retrained model (new mtime) invalidates the old entry. _MODEL_CACHE: Dict[Tuple[str, float], dict] = {} def _load_bundle(path: Path) -> dict: # joblib.load runs arbitrary pickle code; safe here because `path` is always # models/flood_{station}.joblib, an artifact this pipeline's own train.py wrote -- # never a user- or network-supplied file. key = (str(path), path.stat().st_mtime) cached = _MODEL_CACHE.get(key) if cached is not None: return cached bundle = joblib.load(path) for stale_key in [k for k in _MODEL_CACHE if k[0] == str(path)]: del _MODEL_CACHE[stale_key] _MODEL_CACHE[key] = bundle return bundle def _readings_to_long_df(readings_by_station: Dict[str, List[dict]]) -> pd.DataFrame: rows = [] for station_code, readings in readings_by_station.items(): for reading in readings: timestamp = reading.get("timestamp") if isinstance(timestamp, str): timestamp = pd.to_datetime(timestamp) rows.append( { "timestamp": timestamp, "station_code": station_code, "water_level": reading.get("water_level"), "discharge": reading.get("discharge"), } ) if not rows: return pd.DataFrame(columns=["timestamp", "station_code", "water_level", "discharge"]) df = pd.DataFrame(rows) return df.dropna(subset=["timestamp"]) def _clip_probability(value: float) -> float: return float(min(max(value, 0.0), 1.0)) def _sigmoid_probability(predicted_max: float, threshold: float, sigma: float) -> float: return 1.0 / (1.0 + np.exp(-(predicted_max - threshold) / sigma)) def _heuristic_forecast( station_code: str, as_of: pd.Timestamp, current_level: float, level_t_minus_3: Optional[float], warn_thr: float, danger_thr: float, horizons: Tuple[int, ...], ) -> List[dict]: if level_t_minus_3 is None: rate = 0.0 else: rate = max(0.0, (current_level - level_t_minus_3) / 3.0) results = [] for horizon_h in horizons: predicted_max = max(current_level + rate * horizon_h * 0.7, current_level) p_warning = _clip_probability(_sigmoid_probability(predicted_max, warn_thr, HEURISTIC_SIGMA)) p_danger = _clip_probability(_sigmoid_probability(predicted_max, danger_thr, HEURISTIC_SIGMA)) p_danger = min(p_danger, p_warning) results.append( { "station_code": station_code, "horizon_hours": horizon_h, "p_warning": p_warning, "p_danger": p_danger, "predicted_max_level": predicted_max, "current_level": current_level, "as_of": as_of.isoformat(), "model_version": HEURISTIC_VERSION, "trained_at": None, "source": "heuristic", "threshold_warning": warn_thr, "threshold_danger": danger_thr, } ) return results def _model_forecast( station_code: str, grid: features.HourlyGrid, bundle: dict, as_of: pd.Timestamp, current_level: float, ) -> List[dict]: warn_thr = bundle["thresholds"]["warning"] danger_thr = bundle["thresholds"]["danger"] feature_row = features.build_features(grid, station_code).loc[[as_of]] expected_columns = bundle["feature_names"] missing = [c for c in expected_columns if c not in feature_row.columns] if missing: logger.error(f"Feature mismatch for {station_code} (missing {missing}); falling back to heuristic") return None feature_row = feature_row[expected_columns] results = [] for horizon_h in bundle["horizons"]: reg = bundle["heads"].get(f"max_{horizon_h}") if reg is None: results.append(None) continue predicted_max = max(float(reg.predict(feature_row)[0]), current_level) sigma_h = bundle["sigma"].get(horizon_h, HEURISTIC_SIGMA) warn_head = bundle["heads"].get(f"warn_{horizon_h}") if warn_head is not None: p_warning = float(warn_head.predict_proba(feature_row)[0][1]) else: p_warning = _sigmoid_probability(predicted_max, warn_thr, sigma_h) danger_head = bundle["heads"].get(f"danger_{horizon_h}") if danger_head is not None: p_danger = float(danger_head.predict_proba(feature_row)[0][1]) else: p_danger = _sigmoid_probability(predicted_max, danger_thr, sigma_h) p_warning = _clip_probability(p_warning) p_danger = min(_clip_probability(p_danger), p_warning) row = { "station_code": station_code, "horizon_hours": horizon_h, "p_warning": p_warning, "p_danger": p_danger, "predicted_max_level": predicted_max, "current_level": current_level, "as_of": as_of.isoformat(), "model_version": bundle["model_version"], "trained_at": bundle["trained_at"], "source": "model", "threshold_warning": warn_thr, "threshold_danger": danger_thr, } stages = features.FLOOD_STAGES.get(station_code) if stages: # Exceedance probability per official inundation stage, from the # regression head and its validation-residual sigma. These are # threshold-agnostic, so no retraining is needed to serve them. row["stages"] = [ { "stage": s["stage"], "level": s["level"], "p_exceed": _clip_probability( _sigmoid_probability(predicted_max, s["level"], sigma_h) ), } for s in stages ] results.append(row) return results def _forecast_station( station_code: str, grid: features.HourlyGrid, models_dir: Path, now: pd.Timestamp, horizons: Tuple[int, ...], ) -> List[dict]: level_col = (station_code, "water_level") if level_col not in grid.observed.columns: logger.warning(f"No data for station {station_code}; omitting") return [] observed_level = grid.observed[level_col].dropna() if observed_level.empty: logger.warning(f"No observed readings for station {station_code}; omitting") return [] as_of = observed_level.index.max() current_level = float(observed_level.loc[as_of]) staleness_h = (pd.Timestamp(now) - as_of).total_seconds() / 3600.0 warn_thr, danger_thr = features.get_thresholds(station_code) t_minus_3 = as_of - pd.Timedelta(hours=3) level_t_minus_3 = float(observed_level.loc[t_minus_3]) if t_minus_3 in observed_level.index else None bundle_path = models_dir / f"flood_{station_code}.joblib" if not bundle_path.exists() or staleness_h > STALE_AFTER_H: return _heuristic_forecast( station_code, as_of, current_level, level_t_minus_3, warn_thr, danger_thr, horizons ) bundle = _load_bundle(bundle_path) model_results = _model_forecast(station_code, grid, bundle, as_of, current_level) if model_results is None: return _heuristic_forecast( station_code, as_of, current_level, level_t_minus_3, warn_thr, danger_thr, horizons ) # Per-horizon heads that were skipped at train time (e.g. too few positives) still # need a forecast row -- fall back to the single-horizon heuristic for just that row. filled = [] for horizon_h, row in zip(bundle["horizons"], model_results): if row is not None: filled.append(row) else: filled.extend( _heuristic_forecast( station_code, as_of, current_level, level_t_minus_3, warn_thr, danger_thr, (horizon_h,) ) ) return filled def get_forecasts( readings_by_station: Dict[str, List[dict]], models_dir: Union[str, Path] = "models", now: Optional[Union[datetime.datetime, str]] = None, ) -> List[dict]: """Produce flood forecasts for every station present in `readings_by_station`. Each reading dict needs at least {timestamp, water_level, discharge}; extra keys are ignored so raw API/DB rows can be passed straight through. At least 96 hours of span is required to populate every feature; 336 hours (14 days) is recommended. """ models_dir = Path(models_dir) if now is None: now = datetime.datetime.now() now = pd.Timestamp(now) df_long = _readings_to_long_df(readings_by_station) if df_long.empty: return [] grid = features.make_hourly_grid(df_long) results: List[dict] = [] for station_code in readings_by_station.keys(): try: results.extend(_forecast_station(station_code, grid, models_dir, now, DEFAULT_HORIZONS)) except Exception as error: logger.error(f"Forecast failed for station {station_code}: {error}") return results def get_latest_forecasts( db_url: Optional[str] = None, models_dir: Union[str, Path] = "models", hours: int = 336, ) -> List[dict]: """Convenience wrapper for web_api: load the latest window from the DB/API and forecast. Raises FileNotFoundError when no trained model bundle exists at all, so the API can 503 instead of serving purely heuristic output as if it were a forecast. """ from .data import load_latest if not sorted(Path(models_dir).glob("flood_*.joblib")): raise FileNotFoundError(f"no trained model bundles in {models_dir}") df_long = load_latest(db_url=db_url, hours=hours) readings_by_station: Dict[str, List[dict]] = {} if not df_long.empty: for station_code, group in df_long.groupby("station_code"): readings_by_station[station_code] = group[["timestamp", "water_level", "discharge"]].to_dict("records") expected_stations = set(features.UPSTREAM_LEADS.keys()) for missing_station in expected_stations - set(readings_by_station.keys()): logger.warning(f"No recent data for station {missing_station}; omitting from forecasts") return get_forecasts(readings_by_station, models_dir=models_dir)