diff --git a/docs/RUST_REWRITE.md b/docs/RUST_REWRITE.md new file mode 100644 index 0000000..db67b7b --- /dev/null +++ b/docs/RUST_REWRITE.md @@ -0,0 +1,105 @@ +# A Rust rewrite: assessment + +Status: evaluation, nothing decided. Spike: [spikes/rust-forecast](../spikes/rust-forecast/README.md). + +## Short version + +- **Feasible.** The risky part, serving the scikit-learn forecast from Rust with + no train/serve skew, works: the spike reproduces the pandas feature matrix + and the HistGradientBoosting regression output bit for bit (classifier heads + to the last ulp). +- **Training stays Python.** There is no Rust HGB that trains like + scikit-learn, and the evaluation harness (`scripts/evaluate_variants.py`) is + where model work happens. A "Rust rewrite" means a Rust service (collectors, + API, forecast serving, ntfy) plus a Python offline job (monthly retrain and + evaluation) that exports models for Rust. +- **Performance is not a reason.** The forecast takes 158 ms per station in Python + and 0.25 ms in Rust, but it runs once an hour. Load tests showed the site + is bandwidth-bound once caches are hit. +- **Real reasons, if any:** one process instead of leader-elected uvicorn workers + (the root of per-worker state bugs like `/scraping/status`), a fraction of the + memory (324 MB today), typed parsing and SQL, one deployable binary, and the + same stack as proplan (axum + sqlx). +- **Recommendation:** no big-bang rewrite. If Rust is wanted, start with the + collector + notifier daemon (phase 1 below), which stands alone and can be + abandoned without loss. The cheap Python alternative is at the end. + +## What exists (Python, `src/`) + +| Area | LOC | Fate in Rust | +| --- | ---: | --- | +| Collectors: RID hourly, HII/ThaiWater, RID dams, Open-Meteo | 2,500 | port; typed serde payloads | +| Web API (20 routes), schemas, health, metrics | 2,300 | port with identical JSON (golden-file tests) | +| DB adapters: Postgres, SQLite, MySQL, InfluxDB, VictoriaMetrics | 1,300 | Postgres only via sqlx, existing 13 tables, no data migration | +| ntfy notifications; Matrix alerting (CLI path) | 1,050 | port ntfy; Matrix only if still used | +| ML serving: predict, features, rain/dam/HII fill | ~1,300 | port; parity-tested (spike) | +| ML training + evaluation | ~1,700 | **stays Python** | +| CLI (`main.py`) | 650 | small clap CLI (`fill-gaps`, test cycle) | +| `static/dashboard.html` (164 KB) | | unchanged, served as a static file | +| tests | 3,900 | Rust unit tests + parity fixtures + API golden files | + +## Target shape + +One binary, one process (tokio): + +- tasks: RID / HII / dams / rain collectors on their schedules; forecast + precompute when new readings land; ntfy notifier. Shared state is an + `Arc>`, so there's no leader election, no status hand-off file and + no executor sizing. +- axum + tower-http: same routes and JSON, compression (gzip/brotli also covers + the 257 KB river GeoJSON), static dashboard, CORS; `utoipa` if `/docs` stays. +- sqlx against the existing schema. Python and Rust can run against the same + database during migration. +- models: `scripts/retrain.sh` also exports each bundle to JSON (spike format, + ~4.9 MB per station vs 6.2 MB joblib). Rust loads them at start and reloads on + change. A bundle whose `feature_names` Rust cannot build is refused and the + station falls back to the heuristic, as Python does today. + +## The permanent cost: the Python/Rust model contract + +Every feature added in `features.py` must be ported to Rust in the same change, +or the new bundle is refused. That has to be enforced by CI: Python writes a parity +fixture per station (as `export_fixture.py` does), and a Rust test asserts +bit-identical features and regression output. The spike found two traps that +silently flip tree branches (HGB thresholds often equal real data values): +serde_json's default float parsing (1 ulp off; enable `float_roundtrip`) and +pandas' Kahan-compensated rolling sums/means (replicated in `rolling_kahan()`). +This slows model iteration, which is the most active part of the repo. + +## Benefits and costs + +| For | Against | +| --- | --- | +| Single process: no leader election, no per-worker state | Two languages forever; parity CI is mandatory | +| Memory ~324 MB -> est. 60-80 MB (mostly tree arrays; less with u32/u16 nodes) | Weeks of work with little user-visible change | +| Typed RID/HII parsing, compile-checked SQL | Re-introduces bugs into recently hardened code (tz, alerts, security) | +| One static binary; no venv on the server for the API | Server still needs Python + sklearn for the monthly retrain | +| Same stack as proplan | Feature work = Python + Rust change every time | + +## Plan, if pursued (each phase ships on its own) + +0. **Spike: done.** P.1 features + HGB inference, bit-identical. +1. **Collector + notifier daemon** (~1.5k LOC): RID (the fixed semantics: store + late readings, retry while behind, re-fetch yesterday's holes), HII, dams, + Open-Meteo, ntfy, `notification_state`. The Python API runs with collection + off and becomes a stateless reader. Shadow one week first: compare stored + hours against RID upstream. +2. **Forecast crate** (~2k LOC): all 15 stations, dam/HII/rain inputs, heuristic + fallback, probabilities and stages, `forecast_history`; parity CI for every + station on a flood window and a dry window. +3. **Read API** (~2k LOC): routes recorded from the Python API as golden files, + then a Caddy switch. Python keeps only training/evaluation. +4. **Cleanup:** drop the multi-backend adapters, demo databases and the + Python web stack. + +Open questions before phase 1: is Matrix alerting still used? Must `/docs` +(OpenAPI UI) stay? Is SQLite dev mode needed? Where does the monthly retrain +run if the server stops having a Python environment for the API? + +## Cheaper alternative (stay Python) + +Most of the operational gain comes from taking collection out of the web +workers: run the collector as its own systemd unit (one process) and the +uvicorn workers as read-only API. That's about a day of work. It removes leader +election and per-worker state, but doesn't bring the memory, typing or +single-binary benefits. diff --git a/spikes/rust-forecast/.gitignore b/spikes/rust-forecast/.gitignore new file mode 100644 index 0000000..5d1f033 --- /dev/null +++ b/spikes/rust-forecast/.gitignore @@ -0,0 +1,2 @@ +target/ +data/ diff --git a/spikes/rust-forecast/Cargo.lock b/spikes/rust-forecast/Cargo.lock new file mode 100644 index 0000000..cdec5bd --- /dev/null +++ b/spikes/rust-forecast/Cargo.lock @@ -0,0 +1,132 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "num-traits", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "ping-forecast-spike" +version = "0.1.0" +dependencies = [ + "chrono", + "serde", + "serde_json", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/spikes/rust-forecast/Cargo.toml b/spikes/rust-forecast/Cargo.toml new file mode 100644 index 0000000..8b14f93 --- /dev/null +++ b/spikes/rust-forecast/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "ping-forecast-spike" +version = "0.1.0" +edition = "2024" +publish = false + +[dependencies] +serde = { version = "1", features = ["derive"] } +# float_roundtrip: exact float parsing. Tree thresholds often equal real data +# values, so a 1-ulp parse error flips `x <= threshold` (default parser is not exact). +serde_json = { version = "1", features = ["float_roundtrip"] } +chrono = { version = "0.4", default-features = false, features = ["std"] } + +[profile.release] +debug = false diff --git a/spikes/rust-forecast/README.md b/spikes/rust-forecast/README.md new file mode 100644 index 0000000..8174670 --- /dev/null +++ b/spikes/rust-forecast/README.md @@ -0,0 +1,43 @@ +# Rust forecast spike + +Answers one question: can the flood forecast be served from Rust without +train/serve skew, while training stays in Python (scikit-learn)? +Assessment and migration plan: [docs/RUST_REWRITE.md](../../docs/RUST_REWRITE.md). + +- `export/export_fixture.py`: dumps `models/flood_P.1.joblib` (all 9 HGB heads) as + flat tree arrays, plus two fixtures from `models/cache` (Aug-Oct 2024 record + flood, Jun-Aug 2026): the observed hourly grid of P.1 and its 7 upstream + gauges, rain, the pandas feature matrix and every head's sklearn prediction. +- `src/features.rs`: port of `src/ml/features.py::build_features` (P.1 path: + lags, rises, rolling stats, upstream leads, season, rain). +- `src/hgb.rs`: HistGradientBoosting inference (~90 lines). +- `src/main.rs`: compares bit-for-bit and times serving-shaped work. + +```bash +cd spikes/rust-forecast/export +uv run --project ../../.. python export_fixture.py # writes ../data/*.json (gitignored) +cd .. && cargo run --release +uv run --project ../.. python export/time_python.py # Python timing for comparison +``` + +Result (model hgb-v3+cbb3bf7, 3,600 hourly rows): + +| | Result | +| --- | --- | +| Features, 68 columns x 3,600 rows | 244,800 / 244,800 cells bit-identical | +| Regression heads (max_6/12/24) | bit-identical on every row | +| Classifier heads (warn/danger) | equal to the last ulp (<= 2.2e-16, `exp` vs scipy `expit`) | +| One station: 336 h window -> features + 9 heads | Rust 0.25 ms, Python 158 ms | +| Model load | Rust 16 ms (JSON), Python 199 ms (joblib) | + +Two traps the spike had to get past, both of which flip tree branches because +HGB thresholds often equal real data values exactly: + +1. `serde_json` parses floats best-effort (can be 1 ulp off) unless the + `float_roundtrip` feature is on. +2. pandas `rolling().sum()/.mean()` are incremental with Kahan compensation, + a "run of identical values" shortcut and a sign clamp. A plain window sum + differs in the last bits; `rolling_kahan()` replicates pandas exactly. + +Not covered yet: dam features (Mae Ngat), HII gap fill, Open-Meteo catchment +mean, the heuristic fallback, stage probabilities, the other 14 stations. diff --git a/spikes/rust-forecast/export/export_fixture.py b/spikes/rust-forecast/export/export_fixture.py new file mode 100644 index 0000000..66f302a --- /dev/null +++ b/spikes/rust-forecast/export/export_fixture.py @@ -0,0 +1,132 @@ +"""Export a trained flood bundle + a feature/prediction fixture for the Rust spike. + +Run from the monitor repo's environment: + uv run --project ../../.. python export_fixture.py (from this directory) + +Writes ../data/model_.json (all heads as flat tree arrays) and +../data/fixture_.json (observed hourly grid + rain, the Python feature +matrix, and every head's Python prediction for every row). +""" + +import json +import math +import sys +from pathlib import Path + +import joblib +import numpy as np +import pandas as pd + +REPO = Path(__file__).resolve().parents[3] # repo root +sys.path.insert(0, str(REPO)) + +from src.ml import features # noqa: E402 + +OUT = Path(__file__).resolve().parent.parent / "data" +OUT.mkdir(exist_ok=True) +STATION = "P.1" +WINDOWS = { + "flood2024": ("2024-08-15", "2024-10-31"), # record flood, high-level paths + "recent2026": ("2026-06-01", "2026-08-13"), # latest cached data +} + + +def clean(values): + return [None if (v is None or (isinstance(v, float) and math.isnan(v))) else float(v) for v in values] + + +def export_head(est): + preds = est._predictors + assert all(len(it) == 1 for it in preds), "multiclass not expected" + trees = [] + for (tp,) in preds: + n = tp.nodes + assert not n["is_categorical"].any() + assert np.isfinite(n["value"]).all() and not np.isneginf(n["num_threshold"]).any() + assert not np.isnan(n["num_threshold"]).any() + trees.append( + { + "feature": n["feature_idx"].astype(int).tolist(), + # +inf thresholds (a "missing vs everything" split) are not + # valid JSON: written as null, read back as +inf in Rust + "threshold": [ + None if np.isposinf(t) else float(t) for t in n["num_threshold"] + ], + "missing_left": n["missing_go_to_left"].astype(bool).tolist(), + "left": n["left"].astype(int).tolist(), + "right": n["right"].astype(int).tolist(), + "leaf": n["is_leaf"].astype(bool).tolist(), + "value": n["value"].astype(float).tolist(), + } + ) + kind = "classifier" if hasattr(est, "predict_proba") else "regressor" + return { + "kind": kind, + "baseline": float(np.ravel(est._baseline_prediction)[0]), + "trees": trees, + } + + +bundle = joblib.load(REPO / "models" / f"flood_{STATION}.joblib") +model = { + "station_code": STATION, + "model_version": bundle["model_version"], + "regression_target": bundle.get("regression_target"), + "feature_names": bundle["feature_names"], + "horizons": bundle["horizons"], + "sigma": {str(k): float(v) for k, v in bundle["sigma"].items()}, + "thresholds": bundle["thresholds"], + "heads": {name: export_head(est) for name, est in bundle["heads"].items()}, +} +(OUT / f"model_{STATION}.json").write_text(json.dumps(model)) +n_nodes = sum(len(t["value"]) for h in model["heads"].values() for t in h["trees"]) +print(f"model {model['model_version']}: {len(model['heads'])} heads, {n_nodes} nodes") + +stations = [STATION] + [u for u, _ in features.UPSTREAM_LEADS[STATION]] +long = pd.concat( + pd.read_csv(REPO / "models" / "cache" / f"{s}.csv.gz", parse_dates=["timestamp"]) + for s in stations +) +rain_all = pd.read_csv(REPO / "models" / "cache" / "rain_openmeteo.csv.gz") +print("rain columns:", list(rain_all.columns)) +ts_col = rain_all.columns[0] +rain_all[ts_col] = pd.to_datetime(rain_all[ts_col]) +rain_all = rain_all.set_index(ts_col)[rain_all.columns[1]].astype(float) + +for name, (start, end) in WINDOWS.items(): + df = long[(long.timestamp >= start) & (long.timestamp < end)] + grid = features.make_hourly_grid(df) + idx = grid.observed.index + rain = rain_all.reindex(idx) + # punch deterministic holes so NaN handling of the rolling sums is exercised + rain.iloc[::97] = np.nan + feats = features.build_features(grid, STATION, rain=rain)[model["feature_names"]] + + preds = {} + for head_name, est in bundle["heads"].items(): + if hasattr(est, "predict_proba"): + preds[head_name] = est.predict_proba(feats)[:, 1].tolist() + else: + preds[head_name] = est.predict(feats).tolist() + + fixture = { + "station": STATION, + "timestamps": [t.strftime("%Y-%m-%dT%H:%M:%S") for t in idx], + "observed": { + s: { + "water_level": clean(grid.observed[(s, "water_level")].tolist()), + "discharge": clean(grid.observed[(s, "discharge")].tolist()), + } + for s in stations + if (s, "water_level") in grid.observed.columns + }, + "rain": clean(rain.tolist()), + "feature_names": model["feature_names"], + "features": [clean(row) for row in feats.to_numpy().tolist()], + "predictions": preds, + } + (OUT / f"fixture_{name}.json").write_text(json.dumps(fixture)) + print( + f"{name}: {len(idx)} hours, {len(fixture['observed'])} stations, " + f"P.1 max {np.nanmax(grid.observed[(STATION, 'water_level')]):.2f} m" + ) diff --git a/spikes/rust-forecast/export/time_python.py b/spikes/rust-forecast/export/time_python.py new file mode 100644 index 0000000..8dcde03 --- /dev/null +++ b/spikes/rust-forecast/export/time_python.py @@ -0,0 +1,38 @@ +"""Python timing for the same serving-shaped work as the Rust spike.""" + +import sys +import time +from pathlib import Path + +import joblib +import numpy as np +import pandas as pd + +REPO = Path(__file__).resolve().parents[3] # repo root +sys.path.insert(0, str(REPO)) +from src.ml import features # noqa: E402 + +bundle = joblib.load(REPO / "models" / "flood_P.1.joblib") +stations = ["P.1"] + [u for u, _ in features.UPSTREAM_LEADS["P.1"]] +long = pd.concat( + pd.read_csv(REPO / "models" / "cache" / f"{s}.csv.gz", parse_dates=["timestamp"]) + for s in stations +) +df = long[(long.timestamp >= "2026-07-30") & (long.timestamp < "2026-08-13")] +rain_all = pd.read_csv(REPO / "models" / "cache" / "rain_openmeteo.csv.gz", parse_dates=["timestamp"]) +rain_all = rain_all.set_index("timestamp")["chiang_dao"].astype(float) + +reps = 20 +t = time.perf_counter() +for _ in range(reps): + grid = features.make_hourly_grid(df) + rain = rain_all.reindex(grid.observed.index) + row = features.build_features(grid, "P.1", rain=rain)[bundle["feature_names"]].iloc[[-1]] + out = { + k: (h.predict_proba(row)[0, 1] if hasattr(h, "predict_proba") else h.predict(row)[0]) + for k, h in bundle["heads"].items() + } +print(f"python: grid + features + 9 heads for the last hour: {(time.perf_counter() - t) / reps * 1000:.1f} ms per station") +t = time.perf_counter() +joblib.load(REPO / "models" / "flood_P.1.joblib") +print(f"python: bundle load {(time.perf_counter() - t) * 1000:.0f} ms") diff --git a/spikes/rust-forecast/src/features.rs b/spikes/rust-forecast/src/features.rs new file mode 100644 index 0000000..47994d6 --- /dev/null +++ b/spikes/rust-forecast/src/features.rs @@ -0,0 +1,318 @@ +//! Port of `src/ml/features.py::build_features` (pandas) to plain Rust. +//! +//! Pandas semantics reproduced: `ffill(limit)`, `shift(k)` (NaN-padded), +//! `rolling(w, min_periods=1)` aggregations that skip NaN and yield NaN for an +//! all-NaN window, NaN propagation through subtraction. + +use chrono::{Datelike, NaiveDateTime}; +use std::collections::HashMap; +use std::f64::consts::PI; + +pub const FFILL_LIMIT_H: usize = 3; +pub const BASIN_ANCHOR: &str = "P.1"; +const MONSOON_MONTHS: [u32; 5] = [6, 7, 8, 9, 10]; + +pub fn upstream_leads(station: &str) -> &'static [(&'static str, i64)] { + match station { + "P.1" => &[ + ("P.103", 1), + ("P.67", 7), + ("P.21", 9), + ("P.75", 12), + ("P.4A", 12), + ("P.92", 15), + ("P.20", 17), + ], + "P.21" => &[ + ("P.67", 1), + ("P.75", 3), + ("P.4A", 3), + ("P.92", 6), + ("P.20", 8), + ], + "P.67" => &[("P.75", 5), ("P.4A", 5), ("P.92", 8), ("P.20", 10)], + "P.75" | "P.4A" => &[("P.92", 3), ("P.20", 5)], + "P.92" => &[("P.20", 2)], + "P.5" => &[("P.1", 12), ("P.103", 13)], + "P.81" => &[("P.1", 4), ("P.103", 5)], + _ => &[], // spike: P.103's table omitted + } +} + +pub type Series = Vec; + +pub struct Grid { + pub times: Vec, + /// station -> observed (unfilled) hourly series, NaN where missing + pub level: HashMap, + pub discharge: HashMap, +} + +fn nan_series(n: usize) -> Series { + vec![f64::NAN; n] +} + +pub fn ffill(v: &[f64], limit: usize) -> Series { + let mut out = v.to_vec(); + let mut last = f64::NAN; + let mut run = 0; + for x in out.iter_mut() { + if x.is_nan() { + run += 1; + if !last.is_nan() && run <= limit { + *x = last; + } + } else { + last = *x; + run = 0; + } + } + out +} + +/// pandas `shift(k)`: positive k looks back, negative k looks ahead. +pub fn shift(v: &[f64], k: i64) -> Series { + let n = v.len() as i64; + (0..n) + .map(|i| { + let j = i - k; + if (0..n).contains(&j) { + v[j as usize] + } else { + f64::NAN + } + }) + .collect() +} + +fn sub(a: &[f64], b: &[f64]) -> Series { + a.iter().zip(b).map(|(x, y)| x - y).collect() +} + +#[derive(Clone, Copy, PartialEq)] +enum Agg { + Mean, + Sum, + Max, + Min, +} + +/// pandas `rolling(w, min_periods=1).()` over the trailing window [i-w+1, i]. +/// +/// Sum/Mean replicate pandas' `roll_sum`/`roll_mean` (window/aggregations.pyx) +/// operation for operation: an incremental window with separate Kahan +/// compensations for adds and removes, the "run of identical values" shortcut +/// and the sign clamp of the mean. A plain per-window sum differs in the last +/// bits, and HGB thresholds often sit exactly on real data values, so last-bit +/// differences flip tree branches. Max/Min involve no arithmetic: exact as is. +fn rolling(v: &[f64], w: usize, agg: Agg) -> Series { + match agg { + Agg::Sum | Agg::Mean => rolling_kahan(v, w, agg == Agg::Mean), + Agg::Max | Agg::Min => (0..v.len()) + .map(|i| { + let lo = (i + 1).saturating_sub(w); + v[lo..=i] + .iter() + .copied() + .filter(|x| !x.is_nan()) + .fold(f64::NAN, |acc, x| { + if acc.is_nan() { + x + } else if agg == Agg::Max { + acc.max(x) + } else { + acc.min(x) + } + }) + }) + .collect(), + } +} + +fn rolling_kahan(v: &[f64], w: usize, mean: bool) -> Series { + let (mut sum, mut comp_add, mut comp_remove) = (0.0f64, 0.0f64, 0.0f64); + let (mut nobs, mut neg_ct, mut same_run) = (0i64, 0i64, 0i64); + let mut prev = if v.is_empty() { f64::NAN } else { v[0] }; + let mut out = Vec::with_capacity(v.len()); + for i in 0..v.len() { + // window [s, e) for row i; previous window was [i-w, i) clipped + if i >= w { + let val = v[i - w]; + if !val.is_nan() { + nobs -= 1; + let y = -val - comp_remove; + let t = sum + y; + comp_remove = t - sum - y; + sum = t; + if val.is_sign_negative() { + neg_ct -= 1; + } + } + } + let val = v[i]; + if !val.is_nan() { + nobs += 1; + let y = val - comp_add; + let t = sum + y; + comp_add = t - sum - y; + sum = t; + if val.is_sign_negative() { + neg_ct += 1; + } + same_run = if val == prev { same_run + 1 } else { 1 }; + prev = val; + } + out.push(if nobs < 1 { + f64::NAN + } else if mean { + if same_run >= nobs { + prev + } else { + let r = sum / nobs as f64; + if neg_ct == 0 && r < 0.0 || neg_ct == nobs && r > 0.0 { + 0.0 + } else { + r + } + } + } else if same_run >= nobs { + prev * nobs as f64 + } else { + sum + }); + } + out +} + +/// Hours since the last observed value (0 when observed; NaN before the first). +fn hours_since_observed(mask: &[bool]) -> Series { + let mut last: Option = None; + mask.iter() + .enumerate() + .map(|(i, &m)| { + if m { + last = Some(i); + } + last.map_or(f64::NAN, |l| (i - l) as f64) + }) + .collect() +} + +/// Feature columns by name, in pandas' construction order. +pub fn build_features(grid: &Grid, station: &str, rain: Option<&[f64]>) -> Vec<(String, Series)> { + let n = grid.times.len(); + let filled = |s: &str| -> Series { + grid.level + .get(s) + .map_or_else(|| nan_series(n), |v| ffill(v, FFILL_LIMIT_H)) + }; + let mut cols: Vec<(String, Series)> = Vec::new(); + let mut put = |name: String, s: Series| cols.push((name, s)); + + let level = filled(station); + let discharge = grid + .discharge + .get(station) + .cloned() + .unwrap_or_else(|| nan_series(n)); + let mask: Vec = grid.level.get(station).map_or_else( + || vec![false; n], + |v| v.iter().map(|x| !x.is_nan()).collect(), + ); + + put("level".into(), level.clone()); + for k in [1, 2, 3, 6, 12, 24, 48, 72] { + put(format!("level_lag_{k}"), shift(&level, k)); + } + for k in [1, 3, 6, 12, 24] { + put(format!("rise_{k}"), sub(&level, &shift(&level, k))); + } + put("roll_mean_6".into(), rolling(&level, 6, Agg::Mean)); + put("roll_mean_24".into(), rolling(&level, 24, Agg::Mean)); + put("roll_max_6".into(), rolling(&level, 6, Agg::Max)); + put("roll_max_24".into(), rolling(&level, 24, Agg::Max)); + put("roll_max_72".into(), rolling(&level, 72, Agg::Max)); + put("roll_min_24".into(), rolling(&level, 24, Agg::Min)); + + put("discharge".into(), discharge.clone()); + put("discharge_lag_6".into(), shift(&discharge, 6)); + put("discharge_lag_24".into(), shift(&discharge, 24)); + put( + "discharge_rise_6".into(), + sub(&discharge, &shift(&discharge, 6)), + ); + + let age = hours_since_observed(&mask); + put( + "obs_age_h".into(), + age.iter() + .map(|&a| { + if a <= FFILL_LIMIT_H as f64 { + a + } else { + f64::NAN + } + }) + .collect(), + ); + let mask_f: Series = mask.iter().map(|&m| if m { 1.0 } else { 0.0 }).collect(); + put("cov_24h".into(), rolling(&mask_f, 24, Agg::Mean)); + + for &(up, lead) in upstream_leads(station) { + let u = filled(up); + let u_rise_6 = sub(&u, &shift(&u, 6)); + let u_rollmax_24 = rolling(&u, 24, Agg::Max); + let near = (lead - 3).max(0); + put(format!("{up}_level_lag_{near}"), shift(&u, near)); + put(format!("{up}_level_lag_{lead}"), shift(&u, lead)); + put(format!("{up}_level_lag_{}", lead + 3), shift(&u, lead + 3)); + put(format!("{up}_rise_6_lag_{lead}"), shift(&u_rise_6, lead)); + put( + format!("{up}_rollmax_24_lag_{near}"), + shift(&u_rollmax_24, near), + ); + } + + if station != BASIN_ANCHOR { + let p1 = filled(BASIN_ANCHOR); + put("P1_level".into(), p1.clone()); + put("P1_rollmax_24".into(), rolling(&p1, 24, Agg::Max)); + put("P1_rise_24".into(), sub(&p1, &shift(&p1, 24))); + } + + let doy: Vec = grid.times.iter().map(|t| t.ordinal() as f64).collect(); + put( + "doy_sin".into(), + doy.iter().map(|d| (2.0 * PI * d / 365.25).sin()).collect(), + ); + put( + "doy_cos".into(), + doy.iter().map(|d| (2.0 * PI * d / 365.25).cos()).collect(), + ); + put( + "is_monsoon".into(), + grid.times + .iter() + .map(|t| { + if MONSOON_MONTHS.contains(&t.month()) { + 1.0 + } else { + 0.0 + } + }) + .collect(), + ); + + if let Some(r) = rain { + put("rain_6h".into(), rolling(r, 6, Agg::Sum)); + put("rain_24h".into(), rolling(r, 24, Agg::Sum)); + put("rain_72h".into(), rolling(r, 72, Agg::Sum)); + // forward sum over (t, t+24] + let mut ahead = shift(r, -1); + ahead.reverse(); + let mut fc = rolling(&ahead, 24, Agg::Sum); + fc.reverse(); + put("rain_fc24".into(), fc); + } + cols +} diff --git a/spikes/rust-forecast/src/hgb.rs b/spikes/rust-forecast/src/hgb.rs new file mode 100644 index 0000000..f3fcf25 --- /dev/null +++ b/spikes/rust-forecast/src/hgb.rs @@ -0,0 +1,85 @@ +//! Inference for scikit-learn HistGradientBoosting models exported as flat +//! tree arrays (see export/export_fixture.py). Mirrors sklearn's +//! `_predictor.pyx`: NaN follows `missing_go_to_left`, otherwise +//! `x <= threshold` goes left; raw = baseline + sum of leaf values, summed in +//! tree order (same order as sklearn, so results are bit-identical). + +use serde::Deserialize; +use std::collections::BTreeMap; + +fn thresholds<'de, D: serde::Deserializer<'de>>(d: D) -> Result, D::Error> { + let raw: Vec> = Deserialize::deserialize(d)?; + Ok(raw + .into_iter() + .map(|t| t.unwrap_or(f64::INFINITY)) + .collect()) +} + +#[derive(Deserialize)] +pub struct Tree { + feature: Vec, + #[serde(deserialize_with = "thresholds")] + threshold: Vec, + missing_left: Vec, + left: Vec, + right: Vec, + leaf: Vec, + value: Vec, +} + +impl Tree { + #[inline] + fn predict(&self, x: &[f64]) -> f64 { + let mut n = 0; + while !self.leaf[n] { + let v = x[self.feature[n]]; + let go_left = if v.is_nan() { + self.missing_left[n] + } else { + v <= self.threshold[n] + }; + n = if go_left { self.left[n] } else { self.right[n] }; + } + self.value[n] + } +} + +#[derive(Deserialize, PartialEq, Clone, Copy)] +#[serde(rename_all = "lowercase")] +pub enum Kind { + Regressor, + Classifier, +} + +#[derive(Deserialize)] +pub struct Head { + pub kind: Kind, + baseline: f64, + trees: Vec, +} + +impl Head { + /// Regressor: the prediction. Classifier: P(class 1), like predict_proba[:, 1]. + pub fn predict(&self, x: &[f64]) -> f64 { + let mut raw = self.baseline; + for t in &self.trees { + raw += t.predict(x); + } + match self.kind { + Kind::Regressor => raw, + Kind::Classifier => 1.0 / (1.0 + (-raw).exp()), + } + } +} + +#[derive(Deserialize)] +pub struct Bundle { + pub station_code: String, + pub model_version: String, + pub regression_target: Option, + pub feature_names: Vec, + pub horizons: Vec, + pub sigma: BTreeMap, + pub thresholds: BTreeMap, + pub heads: BTreeMap, +} diff --git a/spikes/rust-forecast/src/main.rs b/spikes/rust-forecast/src/main.rs new file mode 100644 index 0000000..7dfc578 --- /dev/null +++ b/spikes/rust-forecast/src/main.rs @@ -0,0 +1,211 @@ +//! Parity + timing check: Rust features and HGB inference vs the Python output +//! recorded in data/fixture_*.json. + +mod features; +mod hgb; + +use chrono::NaiveDateTime; +use serde::Deserialize; +use std::collections::HashMap; +use std::time::Instant; + +#[derive(Deserialize)] +struct StationObs { + water_level: Vec>, + discharge: Vec>, +} + +#[derive(Deserialize)] +struct Fixture { + station: String, + timestamps: Vec, + observed: HashMap, + rain: Vec>, + feature_names: Vec, + features: Vec>>, + predictions: HashMap>, +} + +fn nanify(v: &[Option]) -> Vec { + v.iter().map(|x| x.unwrap_or(f64::NAN)).collect() +} + +fn same(a: f64, b: f64, tol: f64) -> bool { + (a.is_nan() && b.is_nan()) || (a - b).abs() <= tol * (1.0 + b.abs()) +} + +fn matrix(cols: &[(String, Vec)], names: &[String]) -> Vec> { + let by_name: HashMap<&str, &Vec> = cols.iter().map(|(k, v)| (k.as_str(), v)).collect(); + let picked: Vec<&Vec> = names + .iter() + .map(|n| { + *by_name + .get(n.as_str()) + .unwrap_or_else(|| panic!("missing feature {n}")) + }) + .collect(); + let rows = picked[0].len(); + (0..rows) + .map(|i| picked.iter().map(|c| c[i]).collect()) + .collect() +} + +fn main() { + let data = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("data"); + let t = Instant::now(); + let bundle: hgb::Bundle = + serde_json::from_slice(&std::fs::read(data.join("model_P.1.json")).unwrap()).unwrap(); + println!( + "model {} ({}): {} heads, {} features, target={:?}, horizons {:?}, sigma {:?}, thresholds {:?}; loaded in {:?}", + bundle.station_code, + bundle.model_version, + bundle.heads.len(), + bundle.feature_names.len(), + bundle.regression_target, + bundle.horizons, + bundle.sigma, + bundle.thresholds, + t.elapsed() + ); + + for name in ["flood2024", "recent2026"] { + let fx: Fixture = serde_json::from_slice( + &std::fs::read(data.join(format!("fixture_{name}.json"))).unwrap(), + ) + .unwrap(); + assert_eq!(fx.feature_names, bundle.feature_names); + let grid = features::Grid { + times: fx + .timestamps + .iter() + .map(|s| NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S").unwrap()) + .collect(), + level: fx + .observed + .iter() + .map(|(k, v)| (k.clone(), nanify(&v.water_level))) + .collect(), + discharge: fx + .observed + .iter() + .map(|(k, v)| (k.clone(), nanify(&v.discharge))) + .collect(), + }; + let rain = nanify(&fx.rain); + + // 1. feature parity + let t = Instant::now(); + let cols = features::build_features(&grid, &fx.station, Some(&rain)); + let feat_time = t.elapsed(); + let rust_x = matrix(&cols, &bundle.feature_names); + let py_x: Vec> = fx.features.iter().map(|r| nanify(r)).collect(); + let (mut cells, mut exact, mut close, mut worst) = + (0usize, 0usize, 0usize, (0.0f64, String::new())); + for (r, p) in rust_x.iter().zip(&py_x) { + for (j, (a, b)) in r.iter().zip(p).enumerate() { + cells += 1; + if a.to_bits() == b.to_bits() || (a.is_nan() && b.is_nan()) { + exact += 1; + } else if same(*a, *b, 1e-9) { + close += 1; + let d = (a - b).abs(); + if d > worst.0 { + worst = (d, bundle.feature_names[j].clone()); + } + } else { + panic!( + "{name}: feature {} differs: rust {a} vs python {b}", + bundle.feature_names[j] + ); + } + } + } + println!( + "\n[{name}] {} rows x {} features: {exact} bit-identical, {close} within 1e-9 (max |diff| {:.1e} in {}), 0 mismatches; Rust feature build {:?}", + rust_x.len(), + bundle.feature_names.len(), + worst.0, + if worst.1.is_empty() { "-" } else { &worst.1 }, + feat_time + ); + let _ = cells; + + // 2. inference parity, on Python's features and on Rust's own features + let t = Instant::now(); + let mut report = Vec::new(); + for (head_name, head) in &bundle.heads { + let want = &fx.predictions[head_name]; + let (mut max_py, mut max_rs, mut ident) = (0.0f64, 0.0f64, 0usize); + for i in 0..rust_x.len() { + let on_py = head.predict(&py_x[i]); + let on_rs = head.predict(&rust_x[i]); + if on_py.to_bits() == want[i].to_bits() { + ident += 1; + } + max_py = max_py.max((on_py - want[i]).abs()); + max_rs = max_rs.max((on_rs - want[i]).abs()); + } + report.push(format!( + " {head_name:<10} bit-identical {ident}/{} max|diff| on py-features {max_py:.1e}, on rust-features {max_rs:.1e}", + rust_x.len() + )); + } + let infer_time = t.elapsed(); + println!("{}", report.join("\n")); + println!( + " {} heads x {} rows x 2 inference passes in {:?}", + bundle.heads.len(), + rust_x.len(), + infer_time + ); + + // 3. serving-shaped work: 14-day window, forecast for the last hour + let tail = 336.min(grid.times.len()); + let start = grid.times.len() - tail; + let window = features::Grid { + times: grid.times[start..].to_vec(), + level: grid + .level + .iter() + .map(|(k, v)| (k.clone(), v[start..].to_vec())) + .collect(), + discharge: grid + .discharge + .iter() + .map(|(k, v)| (k.clone(), v[start..].to_vec())) + .collect(), + }; + let w_rain = rain[start..].to_vec(); + let reps = 200; + let t = Instant::now(); + let mut last = Vec::new(); + for _ in 0..reps { + let cols = features::build_features(&window, &fx.station, Some(&w_rain)); + let x = matrix(&cols, &bundle.feature_names); + let row = x.last().unwrap(); + last = bundle + .heads + .iter() + .map(|(k, h)| (k.clone(), h.predict(row))) + .collect(); + } + let per = t.elapsed() / reps; + let level_now = window.level[&fx.station] + .iter() + .rev() + .find(|x| !x.is_nan()) + .copied() + .unwrap(); + let rise24 = last.iter().find(|(k, _)| k == "max_24").unwrap().1; + let warn24 = last.iter().find(|(k, _)| k == "warn_24").unwrap().1; + println!( + " serving: 336 h window -> features + {} heads for the last hour: {:?} per station; at {} level {:.2} m, predicted 24 h max {:.2} m, warn_24 head p={:.3}", + bundle.heads.len(), + per, + window.times.last().unwrap(), + level_now, + (rise24 + level_now).max(level_now), + warn24 + ); + } +}