grabowski
5c4f464383
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.
107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
"""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
|