feat(forecast): Rust forecast crate, bit-for-bit with the Python serving path

Phase 2a of docs/RUST_REWRITE.md. rust/forecast (ping-forecast) ports the
serving half of src/ml: the HII mirror gap fill, the hourly grid, every
feature (pandas' Kahan rolling sums included), the Mae Ngat dam frame,
HistGradientBoosting tree inference, the heuristic fallback, stale-threshold
handling and the P.1 flood-stage probabilities.

Hand-over from Python (src/ml/export.py):
- models/flood_<station>.json.gz next to each joblib bundle; train_all() writes
  it, retrain.sh promotes it, `python -m src.ml.export` exports existing ones.
- rust/forecast/flood_config.json: thresholds, stages, leads and notify
  constants, compiled in. tests/test_model_export.py fails while it differs
  from the Python source.

Parity (rust/forecast/tests/parity.rs, fixtures from
testdata/gen_forecast_fixtures.py on the production code path): four
scenarios incl. the 2024 flood with RID lagging behind HII. Merged readings,
all features and all levels identical; with all 15 local models 377/384
probabilities bit-identical, the rest within 2 ulp (exp).

Python refactor without behaviour change: predict.forecasts_from_long, and
get_latest_forecasts/load_latest/serving_frame take now/rain for
reproducible comparisons; the HII query orders its rows.

src/ml/dam.py

@@ -34,6 +34,7 @@ logger = logging.getLogger(__name__)
 
 REPORT_HOUR = 7  # daily value valid from 07:00 local on its own date
 FFILL_LIMIT_H = 48  # two missed daily reports -> NaN, not stale state
+SERVING_DAYS = 21  # 336 h feature window + 72 h storage-delta lag, rounded up
 DAM_COLUMNS = ("storage_pct", "inflow_mcm", "outflow_mcm")
 CACHE_FILE = f"dam_{MAE_NGAT_DAM_ID}.csv.gz"
 
@@ -101,9 +102,11 @@ def load_history(db_url: Optional[str] = None) -> Optional[pd.DataFrame]:
 
 
 def serving_frame(
-    db_url: Optional[str] = None, days: int = 21
+    db_url: Optional[str] = None,
+    days: int = SERVING_DAYS,
+    today: Optional[datetime.date] = None,
 ) -> Optional[pd.DataFrame]:
     """Recent hourly dam state for inference (covers the 336 h feature window
     plus the 72 h storage-delta lag)."""
-    start = datetime.date.today() - datetime.timedelta(days=days)
+    start = (today or datetime.date.today()) - datetime.timedelta(days=days)
     return hourly_frame(load_daily(db_url, start=start))

src/ml/data.py

@@ -143,6 +143,7 @@ def _fetch_hii_levels(
     query += f" AND s.rid_code IN ({placeholders})"
     for i, code in enumerate(stations):
         params[f"station_{i}"] = code
+    query += " ORDER BY m.timestamp, m.station_id"
 
     with engine.connect() as connection:
         df = pd.read_sql(text(query), connection, params=params)
@@ -360,9 +361,10 @@ def load_latest(
     db_url: Optional[str] = None,
     hours: int = 336,
     stations: Optional[List[str]] = None,
+    end: Optional[datetime.datetime] = None,
 ) -> pd.DataFrame:
     """Load the last `hours` of history for all (or given) stations. Never cached to disk."""
-    end = datetime.datetime.now()
+    end = end or datetime.datetime.now()
     start = end - datetime.timedelta(hours=hours)
     return load_measurements(
         db_url=db_url,

src/ml/export.py

@@ -0,0 +1,207 @@
+"""Hand-over of trained models and forecast settings to the Rust forecaster.
+
+The Rust service (rust/forecast, run by ping-collector) cannot read
+joblib/pickle. It reads:
+
+* ``models/flood_<station>.json.gz`` -- the bundle's metadata plus every
+  HistGradientBoosting head as flat tree arrays (sklearn's predictor node
+  table). ``train_all()`` writes one next to each ``.joblib``;
+  ``python -m src.ml.export`` (re)exports the bundles already in ``models/``.
+* ``rust/forecast/flood_config.json`` -- thresholds, flood stages, upstream
+  leads and the other constants features.py / predict.py / notify.py use,
+  compiled into the Rust binary. ``python -m src.ml.export --config``
+  regenerates it; tests/test_model_export.py fails while it is out of date,
+  so a threshold change cannot reach Python without reaching Rust.
+
+Floats go through ``json`` (shortest round-trip repr) and the Rust side parses
+with serde_json's ``float_roundtrip``, so every threshold and leaf value
+arrives bit-identical.
+"""
+
+import argparse
+import gzip
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+
+from . import data as data_mod
+from . import features, predict
+
+FORMAT_VERSION = 1
+REPO_ROOT = Path(__file__).resolve().parents[2]
+CONFIG_PATH = REPO_ROOT / "rust" / "forecast" / "flood_config.json"
+
+
+def export_head(est) -> dict:
+    """One fitted HistGradientBoosting{Regressor,Classifier} as flat arrays."""
+    trees = []
+    for per_iteration in est._predictors:
+        if len(per_iteration) != 1:
+            raise ValueError("multiclass heads are not supported")
+        nodes = per_iteration[0].nodes
+        if nodes["is_categorical"].any():
+            raise ValueError("categorical splits are not supported")
+        if not np.isfinite(nodes["value"]).all():
+            raise ValueError("non-finite leaf value")
+        thr = nodes["num_threshold"]
+        if np.isnan(thr).any() or np.isneginf(thr).any():
+            raise ValueError("NaN or -inf split threshold")
+        trees.append(
+            {
+                "feature": nodes["feature_idx"].astype(int).tolist(),
+                # +inf ("missing vs everything" split) is not valid JSON:
+                # written as null, read back as +inf.
+                "threshold": [None if np.isposinf(t) else float(t) for t in thr],
+                "missing_left": nodes["missing_go_to_left"].astype(bool).tolist(),
+                "left": nodes["left"].astype(int).tolist(),
+                "right": nodes["right"].astype(int).tolist(),
+                "leaf": nodes["is_leaf"].astype(bool).tolist(),
+                "value": nodes["value"].astype(float).tolist(),
+            }
+        )
+    return {
+        "kind": "classifier" if hasattr(est, "predict_proba") else "regressor",
+        "baseline": float(np.ravel(est._baseline_prediction)[0]),
+        "trees": trees,
+    }
+
+
+def bundle_to_dict(bundle: dict) -> dict:
+    """The serving-relevant part of a train.py bundle, JSON-ready."""
+    return {
+        "format_version": FORMAT_VERSION,
+        "station_code": bundle["station_code"],
+        "model_version": bundle["model_version"],
+        "trained_at": bundle.get("trained_at"),
+        "regression_target": bundle.get("regression_target"),
+        "feature_names": list(bundle["feature_names"]),
+        "horizons": [int(h) for h in bundle["horizons"]],
+        "sigma": {str(int(h)): float(v) for h, v in bundle["sigma"].items()},
+        "thresholds": {k: float(v) for k, v in bundle["thresholds"].items()},
+        # a head skipped at training time (too few positives) stays null
+        "heads": {
+            name: (None if est is None else export_head(est))
+            for name, est in bundle["heads"].items()
+        },
+    }
+
+
+def write_bundle_json(bundle: dict, path: Path) -> None:
+    """Atomic write: readers see the old file or the complete new one."""
+    path = Path(path)
+    tmp = path.with_name(path.name + ".tmp")
+    payload = json.dumps(bundle_to_dict(bundle), separators=(",", ":"))
+    with open(tmp, "wb") as raw:
+        # mtime=0: same model, same bytes
+        with gzip.GzipFile(fileobj=raw, mode="wb", mtime=0) as gz:
+            gz.write(payload.encode("utf-8"))
+    os.replace(tmp, path)
+
+
+def json_path_for(joblib_path: Path) -> Path:
+    return joblib_path.with_name(joblib_path.stem + ".json.gz")
+
+
+def export_models_dir(models_dir: Path, force: bool = False) -> int:
+    """Export every flood_*.joblib whose .json.gz is missing or older."""
+    import joblib
+
+    count = 0
+    for path in sorted(Path(models_dir).glob("flood_*.joblib")):
+        out = json_path_for(path)
+        if not force and out.exists() and out.stat().st_mtime >= path.stat().st_mtime:
+            continue
+        # joblib.load runs pickle code: only ever our own train.py output
+        write_bundle_json(joblib.load(path), out)
+        count += 1
+        print(f"exported {path.name} -> {out.name}")
+    return count
+
+
+def flood_config() -> dict:
+    """Every constant the Rust forecaster and notifier share with Python."""
+    from .. import notify
+    from . import dam as dam_mod
+
+    return {
+        "thresholds": {
+            k: [float(w), float(d)] for k, (w, d) in features.THRESHOLDS.items()
+        },
+        "flood_stages": {
+            code: [
+                {
+                    "stage": int(s["stage"]),
+                    "level": float(s["level"]),
+                    "discharge_cms": float(s["discharge_cms"]),
+                }
+                for s in stages
+            ]
+            for code, stages in features.FLOOD_STAGES.items()
+        },
+        "upstream_leads": {
+            code: [[up, int(lead)] for up, lead in leads]
+            for code, leads in features.UPSTREAM_LEADS.items()
+        },
+        "dam_stations": sorted(features.DAM_STATIONS),
+        "monsoon_months": sorted(features.MONSOON_MONTHS),
+        "ffill_limit_h": features.FFILL_LIMIT_H,
+        "basin_anchor": features.BASIN_ANCHOR,
+        "forward_rain_h": features.FORWARD_RAIN_H,
+        "horizons": list(predict.DEFAULT_HORIZONS),
+        "stale_after_h": float(predict.STALE_AFTER_H),
+        "heuristic_sigma": float(predict.HEURISTIC_SIGMA),
+        "heuristic_version": predict.HEURISTIC_VERSION,
+        "latest_window_h": predict.LATEST_WINDOW_H,
+        "hii_fill_stations": list(data_mod.HII_FILL_STATIONS),
+        "hii_exact_mirrors": sorted(data_mod._HII_EXACT_MIRRORS),
+        "hii_min_overlap_hours": data_mod._HII_MIN_OVERLAP_HOURS,
+        "dam": {
+            "dam_id": dam_mod.MAE_NGAT_DAM_ID,
+            "report_hour": dam_mod.REPORT_HOUR,
+            "ffill_limit_h": dam_mod.FFILL_LIMIT_H,
+            "serving_days": dam_mod.SERVING_DAYS,
+        },
+        "notify": {
+            "clear_margin_m": float(notify.CLEAR_MARGIN_M),
+            "capacity_guard_min_pct": float(notify.CAPACITY_GUARD_MIN_PCT),
+            "capacity_guard_exempt": sorted(notify.CAPACITY_GUARD_EXEMPT),
+            "outlook_on": float(notify.OUTLOOK_ON),
+            "outlook_off": float(notify.OUTLOOK_OFF),
+            "outlook_horizon": int(notify.OUTLOOK_HORIZON),
+            "feed_stale_after_h": float(notify.FEED_STALE_AFTER_H),
+            "station_names": dict(notify.STATION_NAMES),
+        },
+    }
+
+
+def flood_config_text() -> str:
+    return json.dumps(flood_config(), indent=1, ensure_ascii=False) + "\n"
+
+
+def main(argv: Optional[list] = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    parser.add_argument(
+        "models_dir", nargs="?", default=str(predict.DEFAULT_MODELS_DIR)
+    )
+    parser.add_argument("--force", action="store_true", help="re-export everything")
+    parser.add_argument(
+        "--config",
+        action="store_true",
+        help=f"write {CONFIG_PATH.relative_to(REPO_ROOT)} instead of exporting models",
+    )
+    args = parser.parse_args(argv)
+    if args.config:
+        CONFIG_PATH.write_text(flood_config_text(), encoding="utf-8", newline="\n")
+        print(f"wrote {CONFIG_PATH}")
+        return 0
+    n = export_models_dir(Path(args.models_dir), force=args.force)
+    print(f"{n} bundle(s) exported")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/ml/predict.py

@@ -27,6 +27,7 @@ DEFAULT_HORIZONS: Tuple[int, ...] = (6, 12, 24)
 STALE_AFTER_H = 6.0
 HEURISTIC_SIGMA = 0.3
 HEURISTIC_VERSION = "heuristic-v1"
+LATEST_WINDOW_H = 336  # history window a live forecast is computed from
 
 # Keyed by (path, mtime) so a retrained model (new mtime) invalidates the old entry.
 _MODEL_CACHE: Dict[Tuple[str, float], dict] = {}
@@ -352,22 +353,14 @@ def get_forecasts(
     return results
 
 
-def get_latest_forecasts(
-    db_url: Optional[str] = None,
+def forecasts_from_long(
+    df_long: pd.DataFrame,
     models_dir: Union[str, Path] = DEFAULT_MODELS_DIR,
-    hours: int = 336,
+    now: Optional[Union[datetime.datetime, str]] = None,
+    rain: Optional[pd.Series] = None,
+    dam: Optional[pd.DataFrame] = None,
 ) -> 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)
+    """Forecasts from a loaded long window (load_latest's output)."""
     readings_by_station: Dict[str, List[dict]] = {}
     if not df_long.empty:
         for station_code, group in df_long.groupby("station_code"):
@@ -380,13 +373,40 @@ def get_latest_forecasts(
         logger.warning(
             f"No recent data for station {missing_station}; omitting from forecasts"
         )
+    return get_forecasts(
+        readings_by_station, models_dir=models_dir, now=now, rain=rain, dam=dam
+    )
+
+
+def get_latest_forecasts(
+    db_url: Optional[str] = None,
+    models_dir: Union[str, Path] = DEFAULT_MODELS_DIR,
+    hours: int = LATEST_WINDOW_H,
+    now: Optional[datetime.datetime] = None,
+    rain: Optional[pd.Series] = None,
+) -> 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.
+    ``now`` and ``rain`` are for reproducible comparisons
+    (scripts/forecast_parity.py); by default the clock and a live fetch.
+    """
+    from .data import load_latest
+
+    if not sorted(Path(models_dir).glob("flood_*.joblib")):
+        raise FileNotFoundError(f"no trained model bundles in {models_dir}")
+
+    now = now or datetime.datetime.now()
+    df_long = load_latest(db_url=db_url, hours=hours, end=now)
 
     # Live rain: trailing days + next-48h forecast. On fetch failure pass an
     # EMPTY series (not None) so rain-trained bundles still find their columns
     # (as NaN) and serve model output instead of tripping the feature guard.
-    from .rain import serving_series
+    if rain is None:
+        from .rain import serving_series
 
-    rain = serving_series()
+        rain = serving_series()
     if rain is None:
         logger.warning("live rain unavailable; rain features will be NaN")
         rain = pd.Series(dtype=float)
@@ -396,7 +416,7 @@ def get_latest_forecasts(
     from . import dam as dam_mod
 
     try:
-        dam = dam_mod.serving_frame(db_url=db_url)
+        dam = dam_mod.serving_frame(db_url=db_url, today=now.date())
     except Exception as error:
         logger.warning(f"dam serving frame failed: {error}")
         dam = None
@@ -404,4 +424,6 @@ def get_latest_forecasts(
         logger.warning("dam state unavailable; dam features will be NaN")
         dam = pd.DataFrame()
 
-    return get_forecasts(readings_by_station, models_dir=models_dir, rain=rain, dam=dam)
+    return forecasts_from_long(
+        df_long, models_dir=models_dir, now=now, rain=rain, dam=dam
+    )

src/ml/train.py

@@ -554,6 +554,10 @@ def train_all(
                 station_results[station] = station_metrics
                 continue
             joblib.dump(bundle, models_dir / f"flood_{station}.joblib")
+            # the Rust forecaster's copy (src/ml/export.py)
+            from .export import write_bundle_json
+
+            write_bundle_json(bundle, models_dir / f"flood_{station}.json.gz")
             logger.info(
                 f"{station}: trained, {bundle['n_train_rows']} rows, "
                 f"{len(bundle['skipped_heads'])} heads skipped"

src/notify.py

@@ -54,6 +54,8 @@ OUTLOOK_ON = 0.50
 OUTLOOK_OFF = 0.25
 # Below this the outlook is not announced at all (avoid "5 % chance" noise).
 OUTLOOK_HORIZON = 24
+# No reading this many hours old or newer: the feed is stale.
+FEED_STALE_AFTER_H = 3.0
 
 STATION_NAMES: Dict[str, str] = {
     "P.1": "Nawarat Bridge, Chiang Mai city",
@@ -222,7 +224,7 @@ def evaluate(
     forecasts: Iterable[dict],
     state: NotificationState,
     publisher: NtfyPublisher,
-    stale_after_h: float = 3.0,
+    stale_after_h: float = FEED_STALE_AFTER_H,
     now: Optional[datetime.datetime] = None,
 ) -> List[Notification]:
     """Compare current readings/forecasts with last-sent state; publish transitions.