feat(forecast): Rust forecast crate, bit-for-bit with the Python serving path
CI / Test suite (push) Successful in 1m2s
CI / Format & lint (push) Successful in 22s
Rust / fmt, clippy, test, build (push) Successful in 9m53s
Security / Dependency vulnerabilities (push) Successful in 1m52s
Security / Static analysis (push) Successful in 20s
Security / License report (push) Successful in 1m41s

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.
This commit is contained in:
grabowski committed 2026-10-18 14:21:56 +07:00
1 parent 6d4beae65a
commit 5c4f464383
31 files changed
+2889 -29

No files matched your search

+106
View File
@@ -0,0 +1,106 @@
"""Model/config hand-over to the Rust forecaster (src/ml/export.py)."""
import gzip
import json
import numpy as np
import pytest
from scipy.special import expit
from sklearn.ensemble import (
HistGradientBoostingClassifier,
HistGradientBoostingRegressor,
)
from src.ml import export
def test_rust_flood_config_is_current():
"""rust/forecast/flood_config.json is compiled into ping-collector.
If this fails, a threshold/stage/lead/notify constant changed: run
`uv run python -m src.ml.export --config` and rebuild the collector.
"""
assert export.CONFIG_PATH.read_text(encoding="utf-8") == export.flood_config_text()
def _data(n=600, seed=7):
rng = np.random.default_rng(seed)
X = rng.normal(size=(n, 4))
X[rng.random(size=X.shape) < 0.1] = np.nan # NaN routing is exported too
y = np.nan_to_num(X[:, 0]) * 2 + np.nan_to_num(X[:, 1]) + rng.normal(size=n) * 0.1
return X, y
def _walk(head, X):
out = np.full(len(X), head["baseline"])
for tree in head["trees"]:
thr = [np.inf if t is None else t for t in tree["threshold"]]
for i, row in enumerate(X):
node = 0
while not tree["leaf"][node]:
v = row[tree["feature"][node]]
left = tree["missing_left"][node] if np.isnan(v) else v <= thr[node]
node = tree["left"][node] if left else tree["right"][node]
out[i] += tree["value"][node]
return out
def test_exported_trees_reproduce_sklearn_exactly():
X, y = _data()
reg = HistGradientBoostingRegressor(max_iter=30, random_state=0).fit(X, y)
clf = HistGradientBoostingClassifier(max_iter=30, random_state=0).fit(X, y > 0.5)
# through JSON, as the Rust side reads it
reg_head = json.loads(json.dumps(export.export_head(reg)))
clf_head = json.loads(json.dumps(export.export_head(clf)))
assert reg_head["kind"] == "regressor" and clf_head["kind"] == "classifier"
assert np.array_equal(_walk(reg_head, X), reg.predict(X))
assert np.array_equal(expit(_walk(clf_head, X)), clf.predict_proba(X)[:, 1])
def test_bundle_json_roundtrip(tmp_path):
X, y = _data()
reg = HistGradientBoostingRegressor(max_iter=5, random_state=0).fit(X, y)
bundle = {
"station_code": "P.1",
"model_version": "hgb-v3+test",
"trained_at": "2026-10-01T03:30:00",
"regression_target": "rise",
"feature_names": ["a", "b", "c", "d"],
"horizons": [6, 12, 24],
"sigma": {6: 0.15, 12: 0.15, 24: 0.2},
"thresholds": {"warning": 3.7, "danger": 4.2},
"heads": {"max_24": reg, "warn_24": None},
"skipped_heads": {"warn_24": "too few positives"},
}
path = tmp_path / "flood_P.1.json.gz"
export.write_bundle_json(bundle, path)
with gzip.open(path, "rt", encoding="utf-8") as fh:
doc = json.load(fh)
assert doc["format_version"] == export.FORMAT_VERSION
assert doc["sigma"] == {"6": 0.15, "12": 0.15, "24": 0.2}
assert doc["heads"]["warn_24"] is None
assert doc["heads"]["max_24"]["kind"] == "regressor"
assert not (tmp_path / "flood_P.1.json.gz.tmp").exists()
first = path.read_bytes()
export.write_bundle_json(bundle, path)
assert path.read_bytes() == first # reproducible bytes (gzip mtime 0)
def test_export_models_dir_skips_up_to_date(tmp_path):
joblib = pytest.importorskip("joblib")
X, y = _data()
reg = HistGradientBoostingRegressor(max_iter=3, random_state=0).fit(X, y)
bundle = {
"station_code": "P.20",
"model_version": "hgb-v3+test",
"feature_names": ["a", "b", "c", "d"],
"horizons": [6],
"sigma": {6: 0.15},
"thresholds": {"warning": 2.35, "danger": 2.8},
"heads": {"max_6": reg},
}
joblib.dump(bundle, tmp_path / "flood_P.20.joblib")
assert export.export_models_dir(tmp_path) == 1
assert (tmp_path / "flood_P.20.json.gz").exists()
assert export.export_models_dir(tmp_path) == 0
assert export.export_models_dir(tmp_path, force=True) == 1