grabowski
764764e07e
train_all() now raises RainUnavailableError when use_rain=True and the
Open-Meteo history cannot be loaded, instead of logging a warning and
writing gauge-only (v2) bundles over the deployed v3 set -- which is what
the 2026-09-01 server retrain did unnoticed. --no-rain remains the explicit
way to get v2. CLI exits 2 with a one-line error. Three tests cover the
guard, the opt-out, and the v3 happy path.
scripts/retrain.sh trains into models/.staging, refuses to promote unless
metrics.json shows hgb-v3+ and >=14 trained stations, then renames bundles
into place (previous generation kept in models/.previous). No API restart:
predict.py reloads by mtime on the hourly precompute.
water-monitor-retrain.{service,timer}: 1st of each month 03:30, Persistent,
OMP_NUM_THREADS=4, Nice=15, same sandbox as the API unit. install.sh now
does `uv sync` into .venv (one env rule; removes a stale venv/) and enables
the timer. water-monitor.service in the repo matched neither the deployed
unit nor the uv env; it now does (run.py --web-api, .venv, EnvironmentFile).
316 lines
13 KiB
Python
316 lines
13 KiB
Python
"""Tests for the flood forecast ML package. Synthetic data only -- no DB/network."""
|
|
|
|
import datetime
|
|
from typing import Dict, List, Optional
|
|
|
|
import joblib
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
from src.ml import features, predict, train
|
|
|
|
|
|
def make_synth(
|
|
n_hours: int,
|
|
stations: List[str],
|
|
seed: int = 0,
|
|
start: str = "2020-01-01",
|
|
pulses: Optional[Dict[str, List[tuple]]] = None,
|
|
missing_patches: Optional[Dict[str, List[tuple]]] = None,
|
|
) -> pd.DataFrame:
|
|
"""Generate a synthetic long measurement frame with smooth levels, flood pulses,
|
|
and optional missing patches, for `stations` over `n_hours` hourly steps.
|
|
|
|
pulses: {station: [(start_hour, width_hours, peak_add), ...]}
|
|
missing_patches: {station: [(start_hour, length_hours), ...]}
|
|
"""
|
|
rng = np.random.default_rng(seed)
|
|
idx = pd.date_range(start, periods=n_hours, freq="h")
|
|
rows = []
|
|
for station in stations:
|
|
base = 1.5 + 0.1 * np.sin(np.linspace(0, 6 * np.pi, n_hours))
|
|
noise = rng.normal(0, 0.02, n_hours)
|
|
level = base + noise
|
|
for pulse_start, width, peak_add in (pulses or {}).get(station, []):
|
|
t = np.arange(n_hours)
|
|
bump = peak_add * np.exp(-0.5 * ((t - (pulse_start + width / 2)) / (width / 4)) ** 2)
|
|
level = level + bump
|
|
discharge = 20.0 * level + rng.normal(0, 1.0, n_hours)
|
|
|
|
missing = np.zeros(n_hours, dtype=bool)
|
|
for patch_start, length in (missing_patches or {}).get(station, []):
|
|
missing[patch_start : patch_start + length] = True
|
|
|
|
for i in range(n_hours):
|
|
if missing[i]:
|
|
continue
|
|
rows.append(
|
|
{
|
|
"timestamp": idx[i],
|
|
"station_code": station,
|
|
"water_level": round(float(level[i]), 3),
|
|
"discharge": round(float(discharge[i]), 2),
|
|
}
|
|
)
|
|
return pd.DataFrame(rows)
|
|
|
|
|
|
def test_no_future_leakage():
|
|
stations = ["P.1", "P.20"]
|
|
df_a = make_synth(200, stations, seed=1, pulses={"P.1": [(150, 10, 3.0)]})
|
|
grid_a = features.make_hourly_grid(df_a)
|
|
feat_a = features.build_features(grid_a, "P.1")
|
|
|
|
t0 = grid_a.observed.index[120]
|
|
|
|
df_b = df_a.copy()
|
|
future_mask = df_b["timestamp"] > t0
|
|
df_b.loc[future_mask, "water_level"] = df_b.loc[future_mask, "water_level"] + 50.0
|
|
df_b.loc[future_mask, "discharge"] = df_b.loc[future_mask, "discharge"] + 500.0
|
|
grid_b = features.make_hourly_grid(df_b)
|
|
feat_b = features.build_features(grid_b, "P.1")
|
|
|
|
past_a = feat_a.loc[feat_a.index <= t0]
|
|
past_b = feat_b.loc[feat_b.index <= t0]
|
|
pd.testing.assert_frame_equal(past_a, past_b)
|
|
|
|
|
|
def test_label_alignment():
|
|
idx = pd.date_range("2020-01-01", periods=12, freq="h")
|
|
warn_thr, _danger_thr = features.get_thresholds("P.1")
|
|
peak = warn_thr + 0.3
|
|
levels = [1.0, 1.0, 1.0, 1.0, 1.0, peak, peak, 1.0, 1.0, 1.0, 1.0, 1.0]
|
|
df = pd.DataFrame(
|
|
{
|
|
"timestamp": idx,
|
|
"station_code": "P.1",
|
|
"water_level": levels,
|
|
"discharge": [20.0 * lvl for lvl in levels],
|
|
}
|
|
)
|
|
grid = features.make_hourly_grid(df)
|
|
labels = features.build_labels(grid, "P.1", horizons=(6,))
|
|
|
|
# Level crosses the warning threshold at t=5. A 6h forward window (t, t+6]
|
|
# first includes t=5 for t=0 .. t=4 (inclusive), so exceed_warn_6 should be
|
|
# 1 for t=0..4 and not (necessarily) for later rows in this hand-built series.
|
|
for t in range(5):
|
|
assert labels["exceed_warn_6"].iloc[t] == 1.0, f"t={t} expected warn exceedance"
|
|
|
|
# max_level_6 at t=0 covers hours 1..6 -> includes the peak.
|
|
assert labels["max_level_6"].iloc[0] == pytest.approx(peak)
|
|
|
|
|
|
def test_label_coverage_gate():
|
|
n = 40
|
|
idx = pd.date_range("2020-01-01", periods=n, freq="h")
|
|
levels = [1.0] * n
|
|
df = pd.DataFrame(
|
|
{"timestamp": idx, "station_code": "P.1", "water_level": levels, "discharge": [20.0] * n}
|
|
)
|
|
# Drop 70% of a future window (hours 21..26) for the row at t=20, no exceedance in it.
|
|
df_missing = df[~df["timestamp"].isin(idx[21:26])].copy()
|
|
grid = features.make_hourly_grid(df_missing)
|
|
labels = features.build_labels(grid, "P.1", horizons=(6,))
|
|
t20 = idx[20]
|
|
assert pd.isna(labels.loc[t20, "exceed_warn_6"])
|
|
|
|
# Same sparse window, but WITH an observed exceedance inside it -> must be 1, not NaN.
|
|
df_with_peak = df_missing.copy()
|
|
peak_row = pd.DataFrame(
|
|
[{"timestamp": idx[22], "station_code": "P.1", "water_level": 5.0, "discharge": 100.0}]
|
|
)
|
|
df_with_peak = pd.concat([df_with_peak, peak_row], ignore_index=True)
|
|
grid2 = features.make_hourly_grid(df_with_peak)
|
|
labels2 = features.build_labels(grid2, "P.1", horizons=(6,))
|
|
assert labels2.loc[t20, "exceed_warn_6"] == 1.0
|
|
|
|
|
|
def test_ffill_and_staleness():
|
|
n = 20
|
|
idx = pd.date_range("2020-01-01", periods=n, freq="h")
|
|
df = pd.DataFrame(
|
|
{
|
|
"timestamp": idx,
|
|
"station_code": "P.1",
|
|
"water_level": [1.0 + 0.01 * i for i in range(n)],
|
|
"discharge": [20.0] * n,
|
|
}
|
|
)
|
|
# Small gap: drop hours 5,6 (2h gap).
|
|
df_small_gap = df[~df["timestamp"].isin(idx[5:7])].copy()
|
|
grid = features.make_hourly_grid(df_small_gap)
|
|
feat = features.build_features(grid, "P.1")
|
|
assert feat.loc[idx[5], "obs_age_h"] == pytest.approx(1.0)
|
|
assert feat.loc[idx[6], "obs_age_h"] == pytest.approx(2.0)
|
|
|
|
# Large gap: drop hours 5..9 (5h gap) -> rows with age>3 dropped (NaN).
|
|
df_big_gap = df[~df["timestamp"].isin(idx[5:10])].copy()
|
|
grid2 = features.make_hourly_grid(df_big_gap)
|
|
feat2 = features.build_features(grid2, "P.1")
|
|
assert feat2.loc[idx[8], "obs_age_h"] != feat2.loc[idx[8], "obs_age_h"] # NaN
|
|
assert feat2.loc[idx[9], "obs_age_h"] != feat2.loc[idx[9], "obs_age_h"] # NaN
|
|
assert feat2.loc[idx[7], "obs_age_h"] == pytest.approx(3.0)
|
|
|
|
|
|
_FORECAST_KEYS = {
|
|
"station_code",
|
|
"horizon_hours",
|
|
"p_warning",
|
|
"p_danger",
|
|
"predicted_max_level",
|
|
"current_level",
|
|
"as_of",
|
|
"model_version",
|
|
"trained_at",
|
|
"source",
|
|
"threshold_warning",
|
|
"threshold_danger",
|
|
}
|
|
|
|
|
|
def _assert_valid_forecast_row(row: dict) -> None:
|
|
# "stages" is optional: model rows for stations in features.FLOOD_STAGES carry
|
|
# per-inundation-stage exceedance probabilities (currently P.1 only).
|
|
assert _FORECAST_KEYS <= set(row.keys())
|
|
assert set(row.keys()) - _FORECAST_KEYS <= {"stages"}
|
|
assert 0.0 <= row["p_warning"] <= 1.0
|
|
assert 0.0 <= row["p_danger"] <= 1.0
|
|
assert row["p_danger"] <= row["p_warning"]
|
|
assert row["predicted_max_level"] >= row["current_level"]
|
|
assert row["source"] in ("model", "heuristic")
|
|
for stage in row.get("stages", []):
|
|
assert 0.0 <= stage["p_exceed"] <= 1.0
|
|
assert stage["level"] > 0
|
|
|
|
|
|
def test_train_smoke_and_roundtrip(tmp_path):
|
|
# Include every station P.1's feature set actually references (its UPSTREAM_LEADS)
|
|
# so no upstream column is entirely NaN -- HistGradientBoosting's binning step
|
|
# cannot fit a fully-degenerate column (see train._safe_fit).
|
|
upstream = [code for code, _lead in features.UPSTREAM_LEADS["P.1"]]
|
|
data_stations = ["P.1"] + upstream
|
|
target_stations = ["P.1", "P.20"]
|
|
n = 700
|
|
pulses = {station: [(start, 20, 2.0) for start in range(50, n - 50, 110)] for station in data_stations}
|
|
df = make_synth(n, data_stations, seed=7, pulses=pulses)
|
|
|
|
metrics = train.train_all(
|
|
df, target_stations, models_dir=tmp_path, skip_eval=True, hgb_overrides={"max_iter": 20}, use_rain=False, use_dam=False
|
|
)
|
|
assert metrics["stations"]["P.1"]["status"] == "trained"
|
|
assert metrics["stations"]["P.20"]["status"] == "trained"
|
|
assert (tmp_path / "flood_P.1.joblib").exists()
|
|
assert (tmp_path / "metrics.json").exists()
|
|
|
|
readings_by_station = {
|
|
code: group[["timestamp", "water_level", "discharge"]].to_dict("records")
|
|
for code, group in df.groupby("station_code")
|
|
if code in target_stations
|
|
}
|
|
now = df["timestamp"].max()
|
|
forecasts = predict.get_forecasts(readings_by_station, models_dir=tmp_path, now=now)
|
|
|
|
assert len(forecasts) > 0
|
|
for row in forecasts:
|
|
_assert_valid_forecast_row(row)
|
|
assert any(row["source"] == "model" for row in forecasts)
|
|
|
|
|
|
def test_heuristic_fallback(tmp_path):
|
|
df = make_synth(50, ["P.1"], seed=3)
|
|
readings_by_station = {"P.1": df[["timestamp", "water_level", "discharge"]].to_dict("records")}
|
|
now = df["timestamp"].max()
|
|
|
|
forecasts = predict.get_forecasts(readings_by_station, models_dir=tmp_path, now=now)
|
|
|
|
assert len(forecasts) == len(predict.DEFAULT_HORIZONS)
|
|
for row in forecasts:
|
|
_assert_valid_forecast_row(row)
|
|
assert row["source"] == "heuristic"
|
|
assert row["model_version"] == "heuristic-v1"
|
|
assert row["trained_at"] is None
|
|
|
|
|
|
def _p1_synth(n: int = 300, seed: int = 11) -> pd.DataFrame:
|
|
upstream = [code for code, _lead in features.UPSTREAM_LEADS["P.1"]]
|
|
return make_synth(n, ["P.1"] + upstream, seed=seed, pulses={"P.1": [(100, 20, 2.0)]})
|
|
|
|
|
|
def test_train_refuses_silent_rain_downgrade(tmp_path, monkeypatch):
|
|
"""use_rain=True with no rain series must abort, not write v2 bundles.
|
|
|
|
Regression for the 2026-09-01 server retrain that overwrote v3 with v2
|
|
because the Open-Meteo archive fetch failed on a cache-less checkout.
|
|
"""
|
|
from src.ml import rain as rain_mod
|
|
|
|
df = _p1_synth()
|
|
overrides = {"max_iter": 10}
|
|
|
|
# Case 1: the loader returns None (archive unreachable, no cache file)
|
|
monkeypatch.setattr(rain_mod, "load_history", lambda *a, **k: None)
|
|
with pytest.raises(train.RainUnavailableError, match="--no-rain"):
|
|
train.train_all(df, ["P.1"], models_dir=tmp_path, skip_eval=True, hgb_overrides=overrides, use_rain=True, use_dam=False)
|
|
assert not (tmp_path / "flood_P.1.joblib").exists()
|
|
assert not (tmp_path / "metrics.json").exists()
|
|
|
|
# Case 2: the loader raises (network / parse error)
|
|
def boom(*a, **k):
|
|
raise ConnectionError("simulated Open-Meteo outage")
|
|
|
|
monkeypatch.setattr(rain_mod, "load_history", boom)
|
|
with pytest.raises(train.RainUnavailableError, match="simulated Open-Meteo outage"):
|
|
train.train_all(df, ["P.1"], models_dir=tmp_path, skip_eval=True, hgb_overrides=overrides, use_rain=True, use_dam=False)
|
|
assert not (tmp_path / "flood_P.1.joblib").exists()
|
|
|
|
# Explicit opt-out still produces v2 bundles as before
|
|
metrics = train.train_all(df, ["P.1"], models_dir=tmp_path, skip_eval=True, hgb_overrides=overrides, use_rain=False, use_dam=False)
|
|
assert metrics["model_version"].startswith("hgb-v2+")
|
|
assert (tmp_path / "flood_P.1.joblib").exists()
|
|
|
|
|
|
def test_train_with_rain_series_yields_v3(tmp_path, monkeypatch):
|
|
from src.ml import rain as rain_mod
|
|
|
|
df = _p1_synth()
|
|
idx = pd.date_range(df["timestamp"].min(), df["timestamp"].max(), freq="h")
|
|
fake_rain = pd.DataFrame({"a": np.linspace(0, 1, len(idx)), "b": 0.5}, index=idx)
|
|
monkeypatch.setattr(rain_mod, "load_history", lambda *a, **k: fake_rain)
|
|
|
|
metrics = train.train_all(df, ["P.1"], models_dir=tmp_path, skip_eval=True, hgb_overrides={"max_iter": 10}, use_rain=True, use_dam=False)
|
|
assert metrics["model_version"].startswith("hgb-v3+")
|
|
bundle = joblib.load(tmp_path / "flood_P.1.joblib")
|
|
assert set(features.RAIN_FEATURES) <= set(bundle["feature_names"])
|
|
|
|
|
|
def test_cli_exit_code_on_rain_failure(tmp_path, monkeypatch, caplog):
|
|
"""The console entry turns the guard into a one-line error and exit 2."""
|
|
from src.ml import rain as rain_mod
|
|
|
|
df = _p1_synth()
|
|
monkeypatch.setattr(rain_mod, "load_history", lambda *a, **k: None)
|
|
monkeypatch.setattr(train, "load_measurements", lambda *a, **k: df)
|
|
monkeypatch.setattr(train, "resolve_db_url", lambda *a, **k: None)
|
|
monkeypatch.setattr(
|
|
"sys.argv",
|
|
["train", "--stations", "P.1", "--models-dir", str(tmp_path), "--skip-eval"],
|
|
)
|
|
assert train.cli() == 2
|
|
assert "refusing to silently downgrade" in caplog.text
|
|
assert not (tmp_path / "metrics.json").exists()
|
|
|
|
|
|
def test_feature_name_stability(tmp_path):
|
|
upstream = [code for code, _lead in features.UPSTREAM_LEADS["P.1"]]
|
|
data_stations = ["P.1"] + upstream
|
|
df = make_synth(300, data_stations, seed=11, pulses={"P.1": [(100, 20, 2.0)]})
|
|
train.train_all(df, ["P.1"], models_dir=tmp_path, skip_eval=True, hgb_overrides={"max_iter": 10}, use_rain=False, use_dam=False)
|
|
# Safe: loading the bundle this same test just wrote to tmp_path, not an external file.
|
|
bundle = joblib.load(tmp_path / "flood_P.1.joblib")
|
|
|
|
grid = features.make_hourly_grid(df)
|
|
fresh_columns = list(features.build_features(grid, "P.1").columns)
|
|
assert fresh_columns == bundle["feature_names"]
|