"""Training CLI for the Ping River flood forecast models. Per station: build the feature/label matrix once, evaluate with a strict temporal holdout (Split B), then refit each head on the full record for the deployed artifact. Hyperparameters are fixed (chosen via an earlier Split A sweep, not repeated here) -- no random search, no shuffling, no sklearn early_stopping (its internal validation split is random and would leak across time). """ import argparse import datetime import json import logging import subprocess from pathlib import Path from typing import Dict, List, Optional, Tuple import joblib import numpy as np import pandas as pd import sklearn from sklearn.ensemble import ( HistGradientBoostingClassifier, HistGradientBoostingRegressor, ) from sklearn.metrics import ( average_precision_score, brier_score_loss, mean_absolute_error, mean_squared_error, ) from . import features from .data import DEFAULT_API_URL, load_measurements, resolve_db_url logger = logging.getLogger(__name__) HORIZONS: Tuple[int, ...] = (6, 12, 24) SPLIT_B_TRAIN_END = "2024-12-31" SPLIT_B_TEST_START = "2025-01-01" SPLIT_B_TEST_END = "2026-08-10" MIN_POSITIVES_FOR_CLASSIFIER = 30 MIN_SIGMA = 0.15 MIN_ROWS_TO_TRAIN = 200 MIN_ROWS_FOR_HEAD = 50 class RainUnavailableError(RuntimeError): """Raised when a rain-enabled training run cannot obtain the rain series. Training would otherwise fall through to gauge-only (v2) bundles and overwrite the deployed v3 artifacts without anyone noticing. """ HGB_PARAMS = { "max_iter": 300, "learning_rate": 0.06, "max_leaf_nodes": 31, "min_samples_leaf": 50, "l2_regularization": 1.0, "early_stopping": False, "random_state": 42, } def _git_short_sha() -> str: try: result = subprocess.run( ["git", "rev-parse", "--short", "HEAD"], capture_output=True, text=True, timeout=5, check=True, ) sha = result.stdout.strip() return sha or "nogit" except Exception: return "nogit" def _make_regressor(overrides: Optional[dict] = None) -> HistGradientBoostingRegressor: params = {**HGB_PARAMS, **(overrides or {})} return HistGradientBoostingRegressor(loss="squared_error", **params) def _make_classifier( overrides: Optional[dict] = None, ) -> HistGradientBoostingClassifier: params = {**HGB_PARAMS, **(overrides or {})} return HistGradientBoostingClassifier(**params) def _safe_fit( estimator, X: pd.DataFrame, y: pd.Series, head_key: str, skipped_heads: Dict[str, str], ): """Fit an estimator, converting any failure (e.g. HistGradientBoosting's binning step rejecting an all-NaN/constant feature column) into a recorded skip rather than a station-killing exception.""" try: estimator.fit(X, y) return estimator except Exception as error: skipped_heads[head_key] = f"fit failed: {error}" logger.warning(f"{head_key}: fit failed, skipping ({error})") return None def _recall_at_far( y_true: np.ndarray, y_score: np.ndarray, target_far: float ) -> Optional[float]: """Recall at the score threshold whose false-positive rate over true negatives is <= target_far.""" y_true = np.asarray(y_true) y_score = np.asarray(y_score) neg_scores = np.sort(y_score[y_true == 0])[::-1] n_pos = int((y_true == 1).sum()) n_neg = len(neg_scores) if n_pos == 0 or n_neg == 0: return None k = int(np.floor(target_far * n_neg)) threshold = neg_scores[k - 1] if k > 0 else neg_scores[0] + 1e-9 predicted_positive = y_score >= threshold tp = int(np.sum(predicted_positive & (y_true == 1))) return tp / n_pos def _p_warning_series( head, reg, X: pd.DataFrame, threshold: float, sigma: float ) -> pd.Series: """Model score if a classifier head exists, else the sigmoid-derived fallback probability.""" if head is not None: return pd.Series(head.predict_proba(X)[:, 1], index=X.index) # reg predicts the RISE over current level; add the level back predicted_max = pd.Series(reg.predict(X), index=X.index) + X["level"] return 1.0 / (1.0 + np.exp(-(predicted_max - threshold) / sigma)) def _find_events(observed_level: pd.Series, warn_thr: float) -> List[dict]: """Group contiguous observed hours >= warn_thr into flood events.""" above = observed_level >= warn_thr events: List[dict] = [] start = None prev_t = None for t, is_above in above.items(): if is_above and start is None: start = t elif not is_above and start is not None: window = observed_level.loc[start:prev_t] events.append( { "crossed_warn_at": start, "peak_time": window.idxmax(), "peak_level": float(window.max()), } ) start = None prev_t = t if start is not None: window = observed_level.loc[start:] events.append( { "crossed_warn_at": start, "peak_time": window.idxmax(), "peak_level": float(window.max()), } ) return events def _first_alert_at(p_series: pd.Series, crossed_at, lookback_h: int = 48): """Earliest time p_warning was sustained (>=0.5 for 2 consecutive hours) within the prior lookback_h.""" window = p_series.loc[crossed_at - pd.Timedelta(hours=lookback_h) : crossed_at] sustained = (window >= 0.5) & (window.shift(1) >= 0.5) hits = sustained[sustained].index if len(hits) == 0: return None return hits.min() - pd.Timedelta(hours=1) def _events_with_lead_time( observed_level_test: pd.Series, warn_thr: float, p_warning_test: pd.Series ) -> List[dict]: events = _find_events(observed_level_test, warn_thr) for event in events: first_alert_at = _first_alert_at(p_warning_test, event["crossed_warn_at"]) event["first_alert_at"] = ( first_alert_at.isoformat() if first_alert_at is not None else None ) if first_alert_at is not None: lead_hours = ( event["crossed_warn_at"] - first_alert_at ).total_seconds() / 3600.0 else: lead_hours = None event["lead_hours"] = lead_hours event["crossed_warn_at"] = event["crossed_warn_at"].isoformat() event["peak_time"] = event["peak_time"].isoformat() return events def train_station( df_long: pd.DataFrame, station: str, horizons: Tuple[int, ...] = HORIZONS, skip_eval: bool = False, hgb_overrides: Optional[dict] = None, split_train_end: str = SPLIT_B_TRAIN_END, split_test_start: str = SPLIT_B_TEST_START, split_test_end: str = SPLIT_B_TEST_END, rain: Optional[pd.Series] = None, dam: Optional[pd.DataFrame] = None, ) -> Tuple[Optional[dict], dict]: """Train every head for one station. Returns (bundle_or_None, station_metrics).""" X, Y, meta = features.build_matrix(df_long, station, horizons, rain=rain, dam=dam) if meta["n_rows"] < MIN_ROWS_TO_TRAIN: return None, { "status": "failed", "reason": f"only {meta['n_rows']} usable rows (< {MIN_ROWS_TO_TRAIN})", } warn_thr, danger_thr = features.get_thresholds(station) feature_names = list(X.columns) if skip_eval: train_mask = pd.Series(True, index=X.index) test_mask = pd.Series(False, index=X.index) else: train_mask = X.index <= pd.Timestamp(split_train_end) test_mask = (X.index >= pd.Timestamp(split_test_start)) & ( X.index <= pd.Timestamp(split_test_end) ) X_train, Y_train = X.loc[train_mask], Y.loc[train_mask] X_test, Y_test = X.loc[test_mask], Y.loc[test_mask] eval_X, eval_Y = (X, Y) if skip_eval else (X_train, Y_train) heads: Dict[str, object] = {} sigma: Dict[int, float] = {} skipped_heads: Dict[str, str] = {} per_horizon: Dict[int, dict] = {} observed_grid = features.make_hourly_grid(df_long).observed for h in horizons: max_col, warn_col, danger_col = ( f"max_level_{h}", f"exceed_warn_{h}", f"exceed_danger_{h}", ) horizon_metrics: dict = {} # --- regression head (rise to future max) --- # Target = future max MINUS current level ("rise"). Rises are far more # stationary than absolute stages, which softens the cannot-exceed- # training-max ceiling: on the rolling-origin harness (2026-08-12) the # rise target moved P.1 first-alert leads from +0h to +6/+46h and cut # the 2024 record-peak underprediction. Prediction = rise + level. reg_labeled = eval_Y[max_col].notna() reg = None if reg_labeled.sum() >= MIN_ROWS_FOR_HEAD: rise_target = ( eval_Y.loc[reg_labeled, max_col] - eval_X.loc[reg_labeled, "level"] ) reg = _safe_fit( _make_regressor(hgb_overrides), eval_X.loc[reg_labeled], rise_target, f"max_{h}", skipped_heads, ) else: skipped_heads[f"max_{h}"] = f"only {int(reg_labeled.sum())} labeled rows" sigma_h = MIN_SIGMA if reg is not None and not skip_eval: test_labeled = Y_test[max_col].notna() if test_labeled.sum() > 0: y_true = Y_test.loc[test_labeled, max_col] y_pred = ( reg.predict(X_test.loc[test_labeled]) + X_test.loc[test_labeled, "level"].to_numpy() ) residuals = y_true.to_numpy() - y_pred sigma_h = max(float(np.std(residuals)), MIN_SIGMA) horizon_metrics["n_test"] = int(test_labeled.sum()) horizon_metrics["mae"] = float(mean_absolute_error(y_true, y_pred)) horizon_metrics["rmse"] = float( np.sqrt(mean_squared_error(y_true, y_pred)) ) above_2m = y_true >= 2.0 horizon_metrics["mae_above_2m"] = ( float(mean_absolute_error(y_true[above_2m], y_pred[above_2m])) if above_2m.any() else None ) sigma[h] = sigma_h horizon_metrics["sigma"] = sigma_h # --- classification heads (warn / danger) --- p_warning_test = None for label_name, col, thr in ( ("warn", warn_col, warn_thr), ("danger", danger_col, danger_thr), ): train_labeled = eval_Y[col].notna() n_pos = ( int(eval_Y.loc[train_labeled, col].sum()) if train_labeled.any() else 0 ) head_key = f"{label_name}_{h}" clf = None if n_pos >= MIN_POSITIVES_FOR_CLASSIFIER: clf = _safe_fit( _make_classifier(hgb_overrides), eval_X.loc[train_labeled], eval_Y.loc[train_labeled, col], head_key, skipped_heads, ) else: skipped_heads[head_key] = ( f"only {n_pos} positives in train span (< {MIN_POSITIVES_FOR_CLASSIFIER})" ) heads[head_key] = clf if not skip_eval: test_labeled = Y_test[col].notna() horizon_metrics[f"base_rate_{label_name}"] = ( float(Y_test.loc[test_labeled, col].mean()) if test_labeled.any() else None ) if ( clf is not None and test_labeled.sum() > 0 and Y_test.loc[test_labeled, col].nunique() > 1 ): y_true = Y_test.loc[test_labeled, col] y_score = clf.predict_proba(X_test.loc[test_labeled])[:, 1] horizon_metrics[f"pr_auc_{label_name}"] = float( average_precision_score(y_true, y_score) ) horizon_metrics[f"brier_{label_name}"] = float( brier_score_loss(y_true, y_score) ) horizon_metrics[f"recall_{label_name}_at_far1pct"] = _recall_at_far( y_true, y_score, 0.01 ) horizon_metrics[f"recall_{label_name}_at_far5pct"] = _recall_at_far( y_true, y_score, 0.05 ) else: horizon_metrics[f"pr_auc_{label_name}"] = None horizon_metrics[f"brier_{label_name}"] = None horizon_metrics[f"recall_{label_name}_at_far1pct"] = None horizon_metrics[f"recall_{label_name}_at_far5pct"] = None if label_name == "warn" and not skip_eval and reg is not None: p_warning_test = _p_warning_series(clf, reg, X_test, thr, sigma_h) per_horizon[h] = horizon_metrics heads[f"max_{h}"] = reg if not skip_eval and reg is not None and p_warning_test is not None: observed_test_level = observed_grid.get((station, "water_level")) if observed_test_level is not None: observed_test_level = observed_test_level.loc[ observed_test_level.index.isin(X_test.index) ] per_horizon[h]["events"] = _events_with_lead_time( observed_test_level, warn_thr, p_warning_test ) # --- full refit on the ENTIRE record for the deployed artifact --- # This may include/exclude different heads than the eval-phase gate above (the # full record has more labeled rows), so skip reasons are re-derived here -- # skipped_heads must reflect what actually ends up in the saved bundle. final_heads: Dict[str, object] = {} for h in horizons: max_col, warn_col, danger_col = ( f"max_level_{h}", f"exceed_warn_{h}", f"exceed_danger_{h}", ) head_key = f"max_{h}" labeled = Y[max_col].notna() if labeled.sum() >= MIN_ROWS_FOR_HEAD: reg = _safe_fit( _make_regressor(hgb_overrides), X.loc[labeled], Y.loc[labeled, max_col] - X.loc[labeled, "level"], # rise target head_key, skipped_heads, ) final_heads[head_key] = reg if reg is not None: skipped_heads.pop(head_key, None) else: skipped_heads[head_key] = f"only {int(labeled.sum())} labeled rows" final_heads[head_key] = None for label_name, col in (("warn", warn_col), ("danger", danger_col)): head_key = f"{label_name}_{h}" train_labeled = Y[col].notna() n_pos = int(Y.loc[train_labeled, col].sum()) if train_labeled.any() else 0 if n_pos >= MIN_POSITIVES_FOR_CLASSIFIER: clf = _safe_fit( _make_classifier(hgb_overrides), X.loc[train_labeled], Y.loc[train_labeled, col], head_key, skipped_heads, ) final_heads[head_key] = clf if clf is not None: skipped_heads.pop(head_key, None) else: skipped_heads[head_key] = ( f"only {n_pos} positives in train span (< {MIN_POSITIVES_FOR_CLASSIFIER})" ) final_heads[head_key] = None # v4 = + Mae Ngat dam features; v3 = rise + rain; v2 = rise target only if "dam_storage_pct" in feature_names: version_prefix = "hgb-v4" elif "rain_24h" in feature_names: version_prefix = "hgb-v3" else: version_prefix = "hgb-v2" bundle = { "station_code": station, "model_version": f"{version_prefix}+{_git_short_sha()}", # v2+: regression heads predict the RISE over the current level; the # serving side must add the level back. Old v1 bundles lack this key. "regression_target": "rise", "trained_at": datetime.datetime.now().isoformat(), "sklearn_version": sklearn.__version__, "feature_names": feature_names, "horizons": list(horizons), "thresholds": {"warning": warn_thr, "danger": danger_thr}, "heads": final_heads, "sigma": sigma, "skipped_heads": skipped_heads, "train_span": meta["span"], "n_train_rows": meta["n_rows"], } station_metrics = {"status": "trained", "per_horizon": per_horizon} return bundle, station_metrics def train_all( df_long: pd.DataFrame, stations: List[str], horizons: Tuple[int, ...] = HORIZONS, models_dir: Path = Path("models"), skip_eval: bool = False, hgb_overrides: Optional[dict] = None, use_rain: bool = True, use_dam: bool = False, db_url: Optional[str] = None, ) -> dict: """Train and save every requested station's models. Returns the metrics.json payload.""" models_dir = Path(models_dir) models_dir.mkdir(parents=True, exist_ok=True) # Catchment rain (Open-Meteo archive, 2021+). A rain-less run produces v2 # bundles that serve fine but have measurably less flood lead (the 2024 # record flood: 13 h early with rain vs 18 h late without). The 2026-09-01 # server retrain hit exactly that -- the archive fetch failed on a checkout # with no models/cache/ and the run quietly wrote v2 over v3. So the # downgrade is now an error unless the caller opts out with use_rain=False # (the --no-rain flag), which is the only way to get v2 deliberately. rain_series = None if use_rain: try: from . import rain as rain_mod rain_series = rain_mod.catchment_mean(rain_mod.load_history()) except Exception as error: raise RainUnavailableError( f"rain history unavailable ({error}); refusing to silently " "downgrade to v2 bundles -- fix Open-Meteo access or restore " "models/cache/rain_openmeteo.csv.gz, or pass --no-rain to " "train gauge-only bundles on purpose" ) from error if rain_series is None: raise RainUnavailableError( "rain history unavailable (Open-Meteo archive unreachable and " "no models/cache/rain_openmeteo.csv.gz); refusing to silently " "downgrade to v2 bundles -- fix access, restore the cache file, " "or pass --no-rain to train gauge-only bundles on purpose" ) if rain_series is not None: logger.info( f"rain series: {rain_series.index.min()} .. {rain_series.index.max()}" ) # Mae Ngat reservoir state (rid_reservoir_daily, 2018+). OFF by default: # the 2026-08-13 backtest ablation showed every dam-feature subset COSTS # 1-3 h of first-alert lead on the 2024 record flood (the daily report # lags up to 31 h, so during fast onset the columns describe yesterday's # benign reservoir and damp the alarm). Kept as an opt-in for post-monsoon # re-evaluation once the 2026 season adds dam-era flood events. dam_frame = None if use_dam: try: from . import dam as dam_mod dam_frame = dam_mod.load_history(db_url=db_url) except Exception as error: logger.warning(f"dam history unavailable, training without it: {error}") if dam_frame is None: # load_history returns None (no raise) when both DB and cache # miss — an explicitly requested experiment must say so loudly. logger.warning( "--dam requested but no dam history available; " "training v3-style bundles WITHOUT dam features" ) if dam_frame is not None: logger.info(f"dam series: {dam_frame.index.min()} .. {dam_frame.index.max()}") # Run-level version: v4 only if some requested station actually receives # dam columns (they are gated to DAM_STATIONS; per-bundle versions are # derived from each station's own feature_names and remain authoritative). if dam_frame is not None and any(s in features.DAM_STATIONS for s in stations): version_prefix = "hgb-v4" elif rain_series is not None: version_prefix = "hgb-v3" else: version_prefix = "hgb-v2" model_version = f"{version_prefix}+{_git_short_sha()}" station_results: Dict[str, dict] = {} for station in stations: if station in features.NOT_TRAINABLE: reason = features.NOT_TRAINABLE[station] logger.info(f"{station}: heuristic ({reason})") station_results[station] = {"status": "heuristic", "reason": reason} continue try: bundle, station_metrics = train_station( df_long, station, horizons, skip_eval=skip_eval, hgb_overrides=hgb_overrides, rain=rain_series, dam=dam_frame, ) if bundle is None: logger.warning(f"{station}: failed ({station_metrics.get('reason')})") station_results[station] = station_metrics continue joblib.dump(bundle, models_dir / f"flood_{station}.joblib") logger.info( f"{station}: trained, {bundle['n_train_rows']} rows, " f"{len(bundle['skipped_heads'])} heads skipped" ) station_results[station] = station_metrics except Exception as error: logger.error(f"{station}: failed with exception: {error}") station_results[station] = {"status": "failed", "reason": str(error)} metrics_payload = { "generated_at": datetime.datetime.now().isoformat(), "model_version": model_version, "split": { "train_end": SPLIT_B_TRAIN_END, "test_start": SPLIT_B_TEST_START, "test_end": SPLIT_B_TEST_END, }, "stations": station_results, } with open(models_dir / "metrics.json", "w", encoding="utf-8") as handle: json.dump(metrics_payload, handle, indent=2, default=str) return metrics_payload def main(argv: Optional[List[str]] = None) -> None: logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s" ) parser = argparse.ArgumentParser( description="Train Ping River flood forecast models" ) parser.add_argument( "--stations", default="all", help="'all' or a comma-separated list of station codes", ) parser.add_argument("--models-dir", default="models") parser.add_argument("--db-url", default=None) parser.add_argument("--api-url", default=DEFAULT_API_URL) parser.add_argument( "--skip-eval", action="store_true", help="Refit-only fast path; skip Split B evaluation", ) parser.add_argument( "--start", default=None, help="ISO date; earliest measurement to load" ) parser.add_argument( "--end", default=None, help="ISO date; latest measurement to load" ) parser.add_argument( "--no-rain", action="store_true", help="DELIBERATELY train without the Open-Meteo rain features " "(v2-style bundles). Without this flag a missing rain series aborts " "the run instead of quietly downgrading the deployed model", ) parser.add_argument( "--dam", action="store_true", help="EXPERIMENTAL: include Mae Ngat reservoir features (v4 bundles); " "the 2026-08 ablation showed they cost 1-3 h of alert lead", ) args = parser.parse_args(argv) if args.stations == "all": stations = list(features.UPSTREAM_LEADS.keys()) else: stations = [s.strip() for s in args.stations.split(",") if s.strip()] start = datetime.datetime.fromisoformat(args.start) if args.start else None end = datetime.datetime.fromisoformat(args.end) if args.end else None logger.info(f"Loading measurements for {len(stations)} stations...") df_long = load_measurements( db_url=resolve_db_url(args.db_url), stations=None, start=start, end=end, api_url=args.api_url, ) logger.info( f"Loaded {len(df_long)} rows spanning {df_long['timestamp'].min()} .. {df_long['timestamp'].max()}" ) metrics_payload = train_all( df_long, stations, models_dir=Path(args.models_dir), skip_eval=args.skip_eval, use_rain=not args.no_rain, use_dam=args.dam, db_url=resolve_db_url(args.db_url), ) trained = sum( 1 for s in metrics_payload["stations"].values() if s["status"] == "trained" ) logger.info( f"Done: {trained}/{len(stations)} stations trained " f"({metrics_payload['model_version']}). " f"metrics.json written to {args.models_dir}" ) def cli() -> int: """Console entry: RainUnavailableError becomes a one-line error, exit 2.""" try: main() except RainUnavailableError as error: logger.error(str(error)) return 2 return 0 if __name__ == "__main__": raise SystemExit(cli())