diff --git a/src/ml/data.py b/src/ml/data.py index eb0834f..3699a12 100644 --- a/src/ml/data.py +++ b/src/ml/data.py @@ -97,6 +97,128 @@ def _fetch_from_db( return _normalize_long(df) +# Stations whose HII mirror is the SAME telemetry (corr ≈ 1.000, median diff +# == station offset exactly — validated 2026-08-11) plus P.81, where the HII +# twin reads the same river with a bias (corr 0.906, MAE 19 cm) that the +# dynamic overlap offset corrects. P.76/P.77/P.85/P.87 HII twins are DIFFERENT +# physical sensors (corr 0.25-0.62) and must never be merged into RID series. +HII_FILL_STATIONS = ( + "P.1", + "P.103", + "P.20", + "P.4A", + "P.67", + "P.75", + "P.82", + "P.84", + "P.92", + "P.81", +) +_HII_EXACT_MIRRORS = frozenset(HII_FILL_STATIONS) - {"P.81"} +_HII_MIN_OVERLAP_HOURS = 168 + + +def _fetch_hii_levels( + db_url: str, + stations: List[str], + start: Optional[datetime.datetime], + end: Optional[datetime.datetime], +) -> pd.DataFrame: + engine = create_engine(db_url, pool_pre_ping=True) + query = ( + "SELECT m.timestamp, s.rid_code AS station_code, m.wl_msl, m.discharge " + "FROM hii_waterlevel m JOIN hii_wl_stations s ON s.id = m.station_id " + "WHERE s.rid_code IS NOT NULL" + ) + params: Dict = {} + if start is not None: + query += " AND m.timestamp >= :start_time" + params["start_time"] = start + if end is not None: + query += " AND m.timestamp <= :end_time" + params["end_time"] = end + placeholders = ", ".join(f":station_{i}" for i in range(len(stations))) + query += f" AND s.rid_code IN ({placeholders})" + for i, code in enumerate(stations): + params[f"station_{i}"] = code + + with engine.connect() as connection: + df = pd.read_sql(text(query), connection, params=params) + df = df.dropna(subset=["wl_msl"]) + if df.empty: + return df + df["timestamp"] = pd.to_datetime(df["timestamp"]).dt.floor("h") + df["wl_msl"] = pd.to_numeric(df["wl_msl"], errors="coerce") + df["discharge"] = pd.to_numeric(df["discharge"], errors="coerce") + df = df.sort_values("timestamp").drop_duplicates( + subset=["station_code", "timestamp"], keep="last" + ) + return df + + +def fill_from_hii( + df: pd.DataFrame, + db_url: str, + start: Optional[datetime.datetime] = None, + end: Optional[datetime.datetime] = None, + stations: Optional[List[str]] = None, + min_overlap_hours: int = _HII_MIN_OVERLAP_HOURS, +) -> pd.DataFrame: + """Fill missing (station, hour) rows from the HII mirror telemetry. + + In-memory only — water_measurements is never written. Each station's + MSL→gauge offset is derived from the overlap between the two series + (median of wl_msl − water_level over ≥ `min_overlap_hours` shared hours), + which reproduces the published offset for exact mirrors and bias-corrects + P.81. Discharge is copied only for exact mirrors; P.81 fills get NaN + discharge (its discharge bias was never validated). Failures degrade to + returning `df` unchanged, so DBs without hii_* tables keep working. + """ + codes = [c for c in (stations or HII_FILL_STATIONS) if c in set(df["station_code"])] + if not codes: + return df + try: + hii = _fetch_hii_levels(db_url, codes, start, end) + except Exception as error: + logger.warning(f"HII gap-fill skipped (fetch failed): {error}") + return df + if hii.empty: + return df + + fills = [] + for code, mirror in hii.groupby("station_code"): + base = df[df["station_code"] == code] + overlap = base.merge( + mirror[["timestamp", "wl_msl"]], on="timestamp", how="inner" + ).dropna(subset=["water_level", "wl_msl"]) + if len(overlap) < min_overlap_hours: + continue + offset = (overlap["wl_msl"] - overlap["water_level"]).median() + # Hours the RID series lacks entirely OR carries only a NaN level; + # _normalize_long keeps the later (fill) row on collision. + present = base.loc[base["water_level"].notna(), "timestamp"] + missing = mirror[~mirror["timestamp"].isin(present)] + if missing.empty: + continue + fill = pd.DataFrame( + { + "timestamp": missing["timestamp"], + "station_code": code, + "water_level": missing["wl_msl"] - offset, + "discharge": missing["discharge"] + if code in _HII_EXACT_MIRRORS + else float("nan"), + } + ) + fills.append(fill) + logger.info( + f"HII gap-fill {code}: +{len(fill)} hours (offset {offset:.3f} m)" + ) + if not fills: + return df + return _normalize_long(pd.concat([df] + fills, ignore_index=True)) + + def _fetch_station_from_api( api_url: str, station_code: str, hours: int, limit: int = 100000 ) -> pd.DataFrame: @@ -177,18 +299,23 @@ def load_measurements( use_cache: bool = True, cache_dir: Path = CACHE_DIR, api_url: str = DEFAULT_API_URL, + hii_fill: bool = True, ) -> pd.DataFrame: """Load the long-format [timestamp, station_code, water_level, discharge] history. Tries PostgreSQL first, then the HTTP API, then the on-disk cache as a last resort. A successful DB/API fetch refreshes the cache; the cache itself is - never treated as a source of fresh data. + never treated as a source of fresh data. With `hii_fill` (DB path only), + gaps are patched in memory from the HII mirror telemetry — training and + serving both flow through here, so the two sides see identical series. """ resolved_db_url = resolve_db_url(db_url) if resolved_db_url: try: df = _fetch_from_db(resolved_db_url, stations, start, end) + if hii_fill: + df = fill_from_hii(df, resolved_db_url, start=start, end=end) if use_cache: _write_cache( df, cache_dir, source="postgres", discharge_maybe_synthetic=False diff --git a/tests/test_hii_fill.py b/tests/test_hii_fill.py new file mode 100644 index 0000000..f86cbab --- /dev/null +++ b/tests/test_hii_fill.py @@ -0,0 +1,124 @@ +"""Tests for the in-memory HII gap-fill in the ML data loader.""" + +import datetime + +import numpy as np +import pandas as pd +import pytest + +from src.hii_collector import HiiStore +from src.ml.data import fill_from_hii + +START = datetime.datetime(2024, 9, 1, 0, 0) + + +def _hours(n, offset=0): + return [START + datetime.timedelta(hours=offset + i) for i in range(n)] + + +def _base_frame(code="P.20", n=200, level=2.0): + return pd.DataFrame( + { + "timestamp": _hours(n), + "station_code": code, + "water_level": level, + "discharge": 100.0, + } + ) + + +@pytest.fixture +def hii_db(tmp_path): + """SQLite DB with hii_* tables; returns (db_url, store).""" + db_url = f"sqlite:///{tmp_path}/hii_fill.db" + store = HiiStore(db_url, "sqlite") + assert store.connect() + return db_url, store + +def _seed_mirror(store, station_id, rid_code, hours, wl_msl, discharge=250.0): + from sqlalchemy import text + + with store.engine.begin() as conn: + conn.execute( + text("INSERT INTO hii_wl_stations (id, rid_code) VALUES (:i, :c)"), + {"i": station_id, "c": rid_code}, + ) + conn.execute( + text( + "INSERT INTO hii_waterlevel (station_id, timestamp, wl_msl, discharge) " + "VALUES (:i, :t, :w, :d)" + ), + [ + {"i": station_id, "t": t, "w": w, "d": discharge} + for t, w in zip(hours, wl_msl) + ], + ) + + +class TestFillFromHii: + def test_exact_mirror_fills_missing_hours(self, hii_db): + db_url, store = hii_db + base = _base_frame("P.20", n=200, level=2.0) + # Mirror covers the base window plus 48 extra hours, at MSL offset +300 + _seed_mirror(store, 1, "P.20", _hours(248), [302.0] * 248) + + filled = fill_from_hii(base, db_url, min_overlap_hours=168) + p20 = filled[filled["station_code"] == "P.20"] + assert len(p20) == 248 + new_rows = p20[p20["timestamp"] >= START + datetime.timedelta(hours=200)] + assert len(new_rows) == 48 + # MSL converted back to gauge datum via the overlap-derived offset + assert new_rows["water_level"].round(6).eq(2.0).all() + # Exact mirrors copy discharge + assert new_rows["discharge"].eq(250.0).all() + + def test_p81_bias_corrected_without_discharge(self, hii_db): + db_url, store = hii_db + base = _base_frame("P.81", n=200, level=1.5) + # Biased mirror: offset 310.2 (not a published offset — derived only) + _seed_mirror(store, 2, "P.81", _hours(230), [311.7] * 230) + + filled = fill_from_hii(base, db_url, min_overlap_hours=168) + p81 = filled[filled["station_code"] == "P.81"] + new_rows = p81[p81["timestamp"] >= START + datetime.timedelta(hours=200)] + assert len(new_rows) == 30 + assert new_rows["water_level"].round(6).eq(1.5).all() + assert new_rows["discharge"].isna().all() + + def test_insufficient_overlap_skips_station(self, hii_db): + db_url, store = hii_db + base = _base_frame("P.20", n=50) # only 50 shared hours + _seed_mirror(store, 1, "P.20", _hours(100), [302.0] * 100) + + filled = fill_from_hii(base, db_url, min_overlap_hours=168) + assert len(filled) == len(base) + + def test_nan_level_rows_are_replaced(self, hii_db): + db_url, store = hii_db + base = _base_frame("P.20", n=200, level=2.0) + base.loc[10, "water_level"] = np.nan + _seed_mirror(store, 1, "P.20", _hours(200), [302.5] * 200) + + filled = fill_from_hii(base, db_url, min_overlap_hours=168) + p20 = filled[filled["station_code"] == "P.20"] + assert len(p20) == 200 # no duplicate hour + replaced = p20[p20["timestamp"] == START + datetime.timedelta(hours=10)] + assert replaced["water_level"].round(6).eq(2.0).all() + + def test_never_merged_station_untouched(self, hii_db): + db_url, store = hii_db + base = _base_frame("P.76", n=200, level=1.0) + _seed_mirror(store, 3, "P.76", _hours(300), [301.0] * 300) + + filled = fill_from_hii(base, db_url, min_overlap_hours=1) + assert len(filled) == len(base) # P.76 not in HII_FILL_STATIONS + + def test_missing_tables_degrade_gracefully(self, tmp_path): + base = _base_frame("P.20") + filled = fill_from_hii(base, f"sqlite:///{tmp_path}/empty.db") + assert filled.equals(base) + + def test_no_station_overlap_returns_input(self, hii_db): + db_url, _ = hii_db + base = _base_frame("P.5") # not fillable + assert fill_from_hii(base, db_url).equals(base)