diff --git a/.gitignore b/.gitignore index a4e7db0..a34c395 100644 --- a/.gitignore +++ b/.gitignore @@ -146,6 +146,7 @@ grafana_data/ # Trained flood-forecast model artifacts (produced on the server, ~100 MB; see docs/FLOOD_FORECASTING.md) models/*.joblib +models/*.json.gz models/cache/ models/metrics.json # scripts/retrain.sh working dirs (staging + one rollback generation) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index e558f30..7ea58a7 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "aho-corasick" version = "1.1.5" @@ -277,6 +283,15 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" +[[package]] +name = "crc32fast" +version = "1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01a7799fd6b852db0e61728dde9a204c423b44d689dbd432522543614b490e78" +dependencies = [ + "cfg-if", +] + [[package]] name = "crossbeam-queue" version = "0.3.14" @@ -400,6 +415,17 @@ version = "0.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aedcfb3409746eddb02b9e19ebda1c3394f759a152e48ee875a0844d1b955484" +[[package]] +name = "flate2" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" +dependencies = [ + "crc32fast", + "miniz_oxide", + "zlib-rs", +] + [[package]] name = "flume" version = "0.11.1" @@ -969,6 +995,16 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.2.4" @@ -1173,6 +1209,18 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "ping-forecast" +version = "0.1.0" +dependencies = [ + "anyhow", + "chrono", + "flate2", + "serde", + "serde_json", + "tracing", +] + [[package]] name = "pkcs1" version = "0.7.5" @@ -1393,7 +1441,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -1592,6 +1640,12 @@ dependencies = [ "rand_core", ] +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + [[package]] name = "slab" version = "0.4.12" @@ -1911,7 +1965,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2632,6 +2686,12 @@ dependencies = [ "syn 3.0.9", ] +[[package]] +name = "zlib-rs" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b268e58e7c693d7c271f93ffc4ba3b380412554231c85bf61ca7af91042a4112" + [[package]] name = "zmij" version = "1.0.23" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 612c2e9..33f5529 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,6 +1,6 @@ [workspace] resolver = "3" -members = ["collector"] +members = ["collector", "forecast"] [workspace.package] # 1.88: edition 2024 let-chains, and dependencies (home, icu_*) need it. Distro @@ -15,6 +15,7 @@ repository = "https://git.b4l.co.th/B4L/Northern-Thailand-Ping-River-Monitor" anyhow = "1" chrono = { version = "0.4", default-features = false, features = ["std", "clock", "serde"] } clap = { version = "4", features = ["derive"] } +flate2 = "1" log = "0.4" # native-tls (OpenSSL on Linux, SChannel on Windows), not rustls: app.rid.go.th (dams) # offers only TLS 1.2 CBC suites (ECDHE-RSA-AES256-SHA384), which rustls refuses. diff --git a/rust/forecast/Cargo.toml b/rust/forecast/Cargo.toml new file mode 100644 index 0000000..cbe7ef8 --- /dev/null +++ b/rust/forecast/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "ping-forecast" +version = "0.1.0" +description = "Ping River flood forecast inference (port of src/ml: features, HistGradientBoosting trees, heuristic fallback)" +edition.workspace = true +rust-version.workspace = true +license.workspace = true +publish.workspace = true +repository.workspace = true + +[dependencies] +anyhow.workspace = true +chrono.workspace = true +flate2.workspace = true +serde.workspace = true +serde_json.workspace = true +tracing.workspace = true + +[dev-dependencies] +flate2.workspace = true diff --git a/rust/forecast/flood_config.json b/rust/forecast/flood_config.json new file mode 100644 index 0000000..39482d4 --- /dev/null +++ b/rust/forecast/flood_config.json @@ -0,0 +1,348 @@ +{ + "thresholds": { + "*": [ + 3.0, + 4.5 + ], + "P.1": [ + 3.7, + 4.2 + ], + "P.103": [ + 5.95, + 6.75 + ], + "P.20": [ + 2.35, + 2.8 + ], + "P.21": [ + 3.2, + 3.6 + ], + "P.4A": [ + 3.4, + 3.9 + ], + "P.5": [ + 4.55, + 4.95 + ], + "P.67": [ + 2.45, + 2.9 + ], + "P.75": [ + 3.2, + 3.65 + ], + "P.76": [ + 5.35, + 5.45 + ], + "P.77": [ + 4.3, + 4.9 + ], + "P.81": [ + 5.15, + 6.3 + ], + "P.82": [ + 3.4, + 3.75 + ], + "P.84": [ + 3.45, + 3.9 + ], + "P.85": [ + 2.9, + 3.35 + ], + "P.87": [ + 3.75, + 4.05 + ], + "P.92": [ + 2.95, + 3.6 + ] + }, + "flood_stages": { + "P.1": [ + { + "stage": 1, + "level": 3.7, + "discharge_cms": 405.0 + }, + { + "stage": 2, + "level": 3.9, + "discharge_cms": 438.0 + }, + { + "stage": 3, + "level": 4.0, + "discharge_cms": 458.0 + }, + { + "stage": 4, + "level": 4.1, + "discharge_cms": 478.0 + }, + { + "stage": 5, + "level": 4.2, + "discharge_cms": 493.0 + }, + { + "stage": 6, + "level": 4.3, + "discharge_cms": 508.0 + }, + { + "stage": 7, + "level": 4.6, + "discharge_cms": 558.0 + } + ] + }, + "upstream_leads": { + "P.1": [ + [ + "P.103", + 1 + ], + [ + "P.67", + 7 + ], + [ + "P.21", + 9 + ], + [ + "P.75", + 12 + ], + [ + "P.4A", + 12 + ], + [ + "P.92", + 15 + ], + [ + "P.20", + 17 + ] + ], + "P.103": [ + [ + "P.67", + 6 + ], + [ + "P.21", + 8 + ], + [ + "P.75", + 11 + ], + [ + "P.4A", + 11 + ], + [ + "P.92", + 14 + ], + [ + "P.20", + 16 + ] + ], + "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.92", + 3 + ], + [ + "P.20", + 5 + ] + ], + "P.4A": [ + [ + "P.92", + 3 + ], + [ + "P.20", + 5 + ] + ], + "P.92": [ + [ + "P.20", + 2 + ] + ], + "P.20": [], + "P.5": [ + [ + "P.1", + 12 + ], + [ + "P.103", + 13 + ] + ], + "P.81": [ + [ + "P.1", + 4 + ], + [ + "P.103", + 5 + ] + ], + "P.82": [], + "P.84": [], + "P.87": [], + "P.77": [], + "P.85": [], + "P.76": [] + }, + "dam_stations": [ + "P.1", + "P.103", + "P.21", + "P.5", + "P.67", + "P.81" + ], + "monsoon_months": [ + 6, + 7, + 8, + 9, + 10 + ], + "ffill_limit_h": 3, + "basin_anchor": "P.1", + "forward_rain_h": 24, + "horizons": [ + 6, + 12, + 24 + ], + "stale_after_h": 6.0, + "heuristic_sigma": 0.3, + "heuristic_version": "heuristic-v1", + "latest_window_h": 336, + "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": [ + "P.1", + "P.103", + "P.20", + "P.4A", + "P.67", + "P.75", + "P.82", + "P.84", + "P.92" + ], + "hii_min_overlap_hours": 168, + "dam": { + "dam_id": "200103", + "report_hour": 7, + "ffill_limit_h": 48, + "serving_days": 21 + }, + "notify": { + "clear_margin_m": 0.1, + "capacity_guard_min_pct": 60.0, + "capacity_guard_exempt": [ + "P.1" + ], + "outlook_on": 0.5, + "outlook_off": 0.25, + "outlook_horizon": 24, + "feed_stale_after_h": 3.0, + "station_names": { + "P.1": "Nawarat Bridge, Chiang Mai city", + "P.103": "Ring Road Bridge 3, Chiang Mai", + "P.67": "Ban Tae (Mae Taeng)", + "P.21": "Ban Rim Tai (Mae Rim)", + "P.75": "Ban Chai Lat", + "P.92": "Ban Muang Aut", + "P.20": "Ban Chiang Dao", + "P.4A": "Ban Mae Taeng", + "P.5": "Tha Nang Bridge (downstream)", + "P.81": "Ban Pong (downstream)", + "P.82": "Ban Sob Win", + "P.84": "Ban Panton", + "P.87": "Ban Pa Sang", + "P.77": "Ban Sop Mae Sapuat", + "P.85": "Ban Lai Kaew", + "P.76": "Ban Mae I Hai" + } + } +} diff --git a/rust/forecast/src/config.rs b/rust/forecast/src/config.rs new file mode 100644 index 0000000..853ef46 --- /dev/null +++ b/rust/forecast/src/config.rs @@ -0,0 +1,106 @@ +//! Constants shared with Python, generated from it: `flood_config.json` is +//! written by `uv run python -m src.ml.export --config`, and +//! tests/test_model_export.py fails while it differs from the Python source. + +use serde::Deserialize; +use std::collections::BTreeMap; +use std::sync::OnceLock; + +#[derive(Debug, Deserialize)] +pub struct Stage { + pub stage: u32, + pub level: f64, + pub discharge_cms: f64, +} + +#[derive(Debug, Deserialize)] +pub struct DamConfig { + pub dam_id: String, + pub report_hour: u32, + pub ffill_limit_h: usize, + pub serving_days: i64, +} + +#[derive(Debug, Deserialize)] +pub struct NotifyConfig { + pub clear_margin_m: f64, + pub capacity_guard_min_pct: f64, + pub capacity_guard_exempt: Vec, + pub outlook_on: f64, + pub outlook_off: f64, + pub outlook_horizon: u32, + pub feed_stale_after_h: f64, + pub station_names: BTreeMap, +} + +#[derive(Debug, Deserialize)] +pub struct FloodConfig { + thresholds: BTreeMap, + pub flood_stages: BTreeMap>, + upstream_leads: BTreeMap>, + pub dam_stations: Vec, + pub monsoon_months: Vec, + pub ffill_limit_h: usize, + pub basin_anchor: String, + pub forward_rain_h: usize, + pub horizons: Vec, + pub stale_after_h: f64, + pub heuristic_sigma: f64, + pub heuristic_version: String, + pub latest_window_h: i64, + pub hii_fill_stations: Vec, + pub hii_exact_mirrors: Vec, + pub hii_min_overlap_hours: usize, + pub dam: DamConfig, + pub notify: NotifyConfig, +} + +impl FloodConfig { + /// (warning, danger) for a station, or the "*" default. + pub fn thresholds(&self, station: &str) -> (f64, f64) { + self.thresholds + .get(station) + .or_else(|| self.thresholds.get("*")) + .copied() + .unwrap_or((f64::INFINITY, f64::INFINITY)) + } + + /// (upstream station, lead hours) pairs, in Python's order. + pub fn upstream_leads(&self, station: &str) -> &[(String, i64)] { + self.upstream_leads.get(station).map_or(&[], Vec::as_slice) + } + + /// The stations the models are built for (Python's `UPSTREAM_LEADS` keys). + pub fn forecast_stations(&self) -> impl Iterator { + self.upstream_leads.keys().map(String::as_str) + } + + pub fn is_dam_station(&self, station: &str) -> bool { + self.dam_stations.iter().any(|s| s == station) + } +} + +pub fn config() -> &'static FloodConfig { + static CONFIG: OnceLock = OnceLock::new(); + CONFIG.get_or_init(|| { + serde_json::from_str(include_str!("../flood_config.json")) + .expect("flood_config.json is generated by src/ml/export.py and must parse") + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn generated_config_parses_and_has_the_basin() { + let c = config(); + assert_eq!(c.thresholds("P.1"), (3.7, 4.2)); + assert_eq!(c.thresholds("P.999"), (3.0, 4.5)); + assert_eq!(c.upstream_leads("P.1")[0], ("P.103".to_string(), 1)); + assert!(c.upstream_leads("P.82").is_empty()); + assert_eq!(c.forecast_stations().count(), 16); + assert_eq!(c.horizons, vec![6, 12, 24]); + assert!(c.is_dam_station("P.1") && !c.is_dam_station("P.20")); + } +} diff --git a/rust/forecast/src/dam.rs b/rust/forecast/src/dam.rs new file mode 100644 index 0000000..2744d92 --- /dev/null +++ b/rust/forecast/src/dam.rs @@ -0,0 +1,98 @@ +//! Mae Ngat reservoir state for the dam features (`src/ml/dam.py`). + +use crate::config::config; +use crate::rolling::ffill; +use chrono::{Duration, NaiveDate, NaiveDateTime}; +use std::collections::BTreeMap; + +/// One `rid_reservoir_daily` row; NaN = missing. +#[derive(Clone, Debug, PartialEq)] +pub struct DamDay { + pub date: NaiveDate, + pub storage_pct: f64, + pub inflow_mcm: f64, + pub outflow_mcm: f64, +} + +/// Hourly dam state: each daily report valid from 07:00 of its date, +/// forward-filled at most 48 h (`dam.hourly_frame`). +#[derive(Clone, Debug, Default)] +pub struct DamFrame { + pub hours: BTreeMap, +} + +impl DamFrame { + /// Columns storage_pct, inflow_mcm, outflow_mcm on the given hours + /// (`DataFrame.reindex`); an empty frame gives NaN columns. + pub fn reindex(&self, times: &[NaiveDateTime]) -> [Vec; 3] { + let col = |k: usize| { + times + .iter() + .map(|t| self.hours.get(t).map_or(f64::NAN, |r| r[k])) + .collect() + }; + [col(0), col(1), col(2)] + } +} + +/// `dam.hourly_frame`; `None` for no rows (Python then serves an empty frame). +pub fn hourly_frame(days: &[DamDay]) -> Option { + let cfg = &config().dam; + let mut by_time: BTreeMap = BTreeMap::new(); + for d in days { + let t = d.date.and_hms_opt(cfg.report_hour, 0, 0)?; + by_time.insert(t, [d.storage_pct, d.inflow_mcm, d.outflow_mcm]); // keep last + } + let first = *by_time.keys().next()?; + let last = *by_time.keys().next_back()? + Duration::hours(cfg.ffill_limit_h as i64); + let n = ((last - first).num_hours() + 1) as usize; + let times: Vec = (0..n).map(|i| first + Duration::hours(i as i64)).collect(); + let cols: Vec> = (0..3) + .map(|k| { + let raw: Vec = times + .iter() + .map(|t| by_time.get(t).map_or(f64::NAN, |r| r[k])) + .collect(); + ffill(&raw, cfg.ffill_limit_h) + }) + .collect(); + Some(DamFrame { + hours: times + .iter() + .enumerate() + .map(|(i, t)| (*t, [cols[0][i], cols[1][i], cols[2][i]])) + .collect(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn day(d: u32, storage: f64) -> DamDay { + DamDay { + date: NaiveDate::from_ymd_opt(2026, 10, d).unwrap(), + storage_pct: storage, + inflow_mcm: 1.0, + outflow_mcm: 2.0, + } + } + + #[test] + fn report_valid_from_seven_and_filled_two_days() { + let f = hourly_frame(&[day(1, 50.0), day(5, 60.0)]).unwrap(); + let at = |d: u32, h: u32| { + NaiveDate::from_ymd_opt(2026, 10, d) + .unwrap() + .and_hms_opt(h, 0, 0) + .unwrap() + }; + assert!(!f.hours.contains_key(&at(1, 6))); + assert_eq!(f.hours[&at(1, 7)][0], 50.0); + assert_eq!(f.hours[&at(3, 7)][0], 50.0); // +48 h + assert!(f.hours[&at(3, 8)][0].is_nan()); + assert_eq!(f.hours[&at(7, 7)][0], 60.0); + assert_eq!(f.hours.keys().next_back(), Some(&at(7, 7))); + assert!(hourly_frame(&[]).is_none()); + } +} diff --git a/rust/forecast/src/features.rs b/rust/forecast/src/features.rs new file mode 100644 index 0000000..105fb78 --- /dev/null +++ b/rust/forecast/src/features.rs @@ -0,0 +1,166 @@ +//! `features.build_features`: the feature columns for one target station, +//! in Python's construction order. Features are causal except `rain_fc24`, +//! the forward sum of the weather forecast. + +use crate::config::config; +use crate::dam::DamFrame; +use crate::grid::Grid; +use crate::rain::{RainSeries, reindex}; +use crate::rolling::*; +use chrono::{Datelike, Duration}; +use std::f64::consts::PI; + +pub type Columns = Vec<(String, Series)>; + +/// `rain`: `None` omits the rain columns (Python's `rain=None`), an empty +/// series gives NaN columns. Same for `dam`, which only DAM_STATIONS receive. +pub fn build_features( + grid: &Grid, + station: &str, + rain: Option<&RainSeries>, + dam: Option<&DamFrame>, +) -> Columns { + let cfg = config(); + let n = grid.len(); + let ffill_h = cfg.ffill_limit_h; + let filled = |s: &str| -> Series { + grid.level + .get(s) + .map_or_else(|| nan_series(n), |v| ffill(v, ffill_h)) + }; + let mut cols: Columns = Vec::with_capacity(80); + 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_mean(&level, 6)); + put("roll_mean_24".into(), rolling_mean(&level, 24)); + put("roll_max_6".into(), rolling_max(&level, 6)); + put("roll_max_24".into(), rolling_max(&level, 24)); + put("roll_max_72".into(), rolling_max(&level, 72)); + put("roll_min_24".into(), rolling_min(&level, 24)); + + 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)), + ); + + // hours since the last observed level, NaN beyond the ffill limit + let mut last: Option = None; + let age: Series = mask + .iter() + .enumerate() + .map(|(i, &m)| { + if m { + last = Some(i); + } + match last { + Some(l) if i - l <= ffill_h => (i - l) as f64, + _ => f64::NAN, + } + }) + .collect(); + put("obs_age_h".into(), age); + let mask_f: Series = mask.iter().map(|&m| if m { 1.0 } else { 0.0 }).collect(); + put("cov_24h".into(), rolling_mean(&mask_f, 24)); + + for (up, lead) in cfg.upstream_leads(station) { + let lead = *lead; + let u = filled(up); + let u_rise_6 = sub(&u, &shift(&u, 6)); + let u_rollmax_24 = rolling_max(&u, 24); + 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 != cfg.basin_anchor { + let p1 = filled(&cfg.basin_anchor); + let rollmax = rolling_max(&p1, 24); + let rise = sub(&p1, &shift(&p1, 24)); + put("P1_level".into(), p1); + put("P1_rollmax_24".into(), rollmax); + put("P1_rise_24".into(), rise); + } + + let mut put = |name: &str, s: Series| cols.push((name.to_string(), s)); + let doy: Vec = grid.times.iter().map(|t| t.ordinal() as f64).collect(); + put( + "doy_sin", + doy.iter().map(|d| (2.0 * PI * d / 365.25).sin()).collect(), + ); + put( + "doy_cos", + doy.iter().map(|d| (2.0 * PI * d / 365.25).cos()).collect(), + ); + put( + "is_monsoon", + grid.times + .iter() + .map(|t| { + if cfg.monsoon_months.contains(&t.month()) { + 1.0 + } else { + 0.0 + } + }) + .collect(), + ); + + if let Some(rain) = rain { + let r = reindex(rain, &grid.times); + put("rain_6h", rolling_sum(&r, 6)); + put("rain_24h", rolling_sum(&r, 24)); + put("rain_72h", rolling_sum(&r, 72)); + // forward sum over (t, t+24], on the grid extended past its last row + // so the newest row sees the weather forecast + let mut ext_times = grid.times.clone(); + if let Some(&last) = grid.times.last() { + ext_times.extend((1..=cfg.forward_rain_h as i64).map(|h| last + Duration::hours(h))); + } + let mut fc = forward_sum(&reindex(rain, &ext_times), cfg.forward_rain_h); + fc.truncate(grid.len()); + put("rain_fc24", fc); + } + + if let Some(dam) = dam + && cfg.is_dam_station(station) + { + let [storage, inflow, outflow] = dam.reindex(&grid.times); + put("dam_storage_pct", storage.clone()); + put("dam_storage_pct_d3", sub(&storage, &shift(&storage, 72))); + put("dam_inflow", inflow); + put("dam_outflow", outflow); + } + cols +} + +/// The row at index `i` as name -> value. +pub fn row_at(cols: &Columns, i: usize) -> std::collections::HashMap<&str, f64> { + cols.iter().map(|(k, v)| (k.as_str(), v[i])).collect() +} diff --git a/rust/forecast/src/grid.rs b/rust/forecast/src/grid.rs new file mode 100644 index 0000000..251f539 --- /dev/null +++ b/rust/forecast/src/grid.rs @@ -0,0 +1,86 @@ +//! Long readings -> complete hourly grid (`features.make_hourly_grid`). + +use crate::rolling::{Series, nan_series}; +use chrono::{Duration, NaiveDateTime, Timelike}; +use std::collections::BTreeMap; + +/// One reading in long format; NaN = missing. +#[derive(Clone, Debug, PartialEq)] +pub struct Row { + pub timestamp: NaiveDateTime, + pub station_code: String, + pub water_level: f64, + pub discharge: f64, +} + +pub fn floor_hour(t: NaiveDateTime) -> NaiveDateTime { + t.with_nanosecond(0) + .and_then(|t| t.with_second(0)) + .and_then(|t| t.with_minute(0)) + .unwrap_or(t) +} + +/// Every hour from the first to the last reading; per station the observed +/// (unfilled) level and discharge, NaN where nothing was recorded. +pub struct Grid { + pub times: Vec, + pub level: BTreeMap, + pub discharge: BTreeMap, +} + +impl Grid { + /// Timestamps are floored to the hour; for a duplicate (station, hour) + /// the later row wins (pandas `drop_duplicates(keep="last")`). + pub fn from_rows(rows: &[Row]) -> Grid { + let mut cells: BTreeMap<(String, NaiveDateTime), (f64, f64)> = BTreeMap::new(); + for r in rows { + cells.insert( + (r.station_code.clone(), floor_hour(r.timestamp)), + (r.water_level, r.discharge), + ); + } + let (Some(first), Some(last)) = ( + cells.keys().map(|k| k.1).min(), + cells.keys().map(|k| k.1).max(), + ) else { + return Grid { + times: vec![], + level: BTreeMap::new(), + discharge: BTreeMap::new(), + }; + }; + let n = ((last - first).num_hours() + 1) as usize; + let times: Vec = (0..n).map(|i| first + Duration::hours(i as i64)).collect(); + let mut level: BTreeMap = BTreeMap::new(); + let mut discharge: BTreeMap = BTreeMap::new(); + for ((code, t), (l, q)) in cells { + let i = (t - first).num_hours() as usize; + level.entry(code.clone()).or_insert_with(|| nan_series(n))[i] = l; + discharge.entry(code).or_insert_with(|| nan_series(n))[i] = q; + } + Grid { + times, + level, + discharge, + } + } + + pub fn len(&self) -> usize { + self.times.len() + } + + pub fn is_empty(&self) -> bool { + self.times.is_empty() + } + + pub fn index_of(&self, t: NaiveDateTime) -> Option { + let first = *self.times.first()?; + let i = (t - first).num_hours(); + (i >= 0 && (i as usize) < self.len() && self.times[i as usize] == t).then_some(i as usize) + } + + /// The station's last hour with an observed level. + pub fn last_observed(&self, station: &str) -> Option { + self.level.get(station)?.iter().rposition(|x| !x.is_nan()) + } +} diff --git a/rust/forecast/src/hgb.rs b/rust/forecast/src/hgb.rs new file mode 100644 index 0000000..04374b3 --- /dev/null +++ b/rust/forecast/src/hgb.rs @@ -0,0 +1,104 @@ +//! Inference for scikit-learn HistGradientBoosting heads exported as flat +//! tree arrays (src/ml/export.py). Mirrors sklearn's `_predictor.pyx`: NaN +//! follows `missing_go_to_left`, otherwise `x <= threshold` goes left; raw = +//! baseline + leaf values summed in tree order (sklearn's order, so the +//! regression output is bit-identical). + +use anyhow::{Result, bail}; +use serde::Deserialize; + +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(Debug, 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] + } + + fn validate(&self, n_features: usize) -> Result<()> { + let n = self.value.len(); + if n == 0 + || [ + self.feature.len(), + self.threshold.len(), + self.missing_left.len(), + self.left.len(), + self.right.len(), + self.leaf.len(), + ] + .iter() + .any(|&l| l != n) + { + bail!("tree arrays have different lengths"); + } + for i in 0..n { + if !self.leaf[i] + && (self.feature[i] >= n_features || self.left[i] >= n || self.right[i] >= n) + { + bail!("tree node {i} points outside the tree or the feature vector"); + } + } + Ok(()) + } +} + +#[derive(Debug, Deserialize, PartialEq, Eq, Clone, Copy)] +#[serde(rename_all = "lowercase")] +pub enum Kind { + Regressor, + Classifier, +} + +#[derive(Debug, Deserialize)] +pub struct Head { + pub kind: Kind, + baseline: f64, + trees: Vec, +} + +impl Head { + /// Regressor: the prediction. Classifier: P(class 1), as + /// `predict_proba(X)[:, 1]` (scipy's `expit`). + 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()), + } + } + + pub fn validate(&self, n_features: usize) -> Result<()> { + self.trees.iter().try_for_each(|t| t.validate(n_features)) + } +} diff --git a/rust/forecast/src/hii.rs b/rust/forecast/src/hii.rs new file mode 100644 index 0000000..22abb09 --- /dev/null +++ b/rust/forecast/src/hii.rs @@ -0,0 +1,224 @@ +//! Gap fill from the HII/ThaiWater mirror telemetry (`data.fill_from_hii`). +//! +//! For stations whose HII twin is the same sensor, hours missing from the RID +//! series (or carrying no level) are filled from HII, shifted by the median +//! offset between the two series over their overlap. This is also what keeps +//! the forecast current at night while RID's report lags. + +use crate::config::config; +use crate::grid::{Row, floor_hour}; +use chrono::NaiveDateTime; +use std::collections::{BTreeMap, BTreeSet, HashSet}; + +/// One `hii_waterlevel` row joined to its RID station code; NaN = missing. +#[derive(Clone, Debug, PartialEq)] +pub struct HiiRow { + pub timestamp: NaiveDateTime, + pub station_code: String, + pub wl_msl: f64, + pub discharge: f64, +} + +/// `data._prepare_hii_levels`: rows without a level dropped, one row per +/// station-hour, the latest reading of the hour winning (stable on ties). +pub fn prepare(rows: &[HiiRow]) -> Vec { + let mut kept: Vec<&HiiRow> = rows.iter().filter(|r| !r.wl_msl.is_nan()).collect(); + kept.sort_by_key(|r| r.timestamp); // stable + let mut by_hour: BTreeMap<(String, NaiveDateTime), HiiRow> = BTreeMap::new(); + for r in kept { + let t = floor_hour(r.timestamp); + by_hour.insert( + (r.station_code.clone(), t), + HiiRow { + timestamp: t, + ..r.clone() + }, + ); + } + by_hour.into_values().collect() +} + +/// The HII station codes worth fetching for these RID rows. +pub fn fill_candidates(rid: &[Row]) -> Vec { + let present: HashSet<&str> = rid.iter().map(|r| r.station_code.as_str()).collect(); + config() + .hii_fill_stations + .iter() + .filter(|c| present.contains(c.as_str())) + .cloned() + .collect() +} + +/// Per-station result of the fill, for logs and the notifier. +#[derive(Clone, Debug, PartialEq)] +pub struct FillInfo { + pub station_code: String, + pub offset: f64, + pub filled_hours: usize, +} + +/// pandas `Series.median` (numpy: mean of the two middle values when even). +fn median(mut v: Vec) -> f64 { + v.sort_by(f64::total_cmp); + let n = v.len(); + if n % 2 == 1 { + v[n / 2] + } else { + (v[n / 2 - 1] + v[n / 2]) / 2.0 + } +} + +/// `data.fill_from_hii`: `rid` as loaded (already one row per station-hour), +/// `hii` from [`prepare`]. Returns the merged rows (fill rows replace RID rows +/// that have no level) and what was filled. +pub fn fill(rid: &[Row], hii: &[HiiRow]) -> (Vec, Vec) { + let cfg = config(); + let candidates: BTreeSet = fill_candidates(rid).into_iter().collect(); + let mut mirrors: BTreeMap<&str, Vec<&HiiRow>> = BTreeMap::new(); + for h in hii { + if candidates.contains(&h.station_code) { + mirrors.entry(h.station_code.as_str()).or_default().push(h); + } + } + let mut fills: Vec = Vec::new(); + let mut info = Vec::new(); + for (code, mirror) in mirrors { + let base: BTreeMap = rid + .iter() + .filter(|r| r.station_code == code) + .map(|r| (r.timestamp, r.water_level)) + .collect(); + let diffs: Vec = mirror + .iter() + .filter_map(|m| match base.get(&m.timestamp) { + Some(level) if !level.is_nan() => Some(m.wl_msl - level), + _ => None, + }) + .collect(); + if diffs.len() < cfg.hii_min_overlap_hours { + continue; + } + let offset = median(diffs); + let exact = cfg.hii_exact_mirrors.iter().any(|c| c == code); + let before = fills.len(); + for m in mirror { + if base.get(&m.timestamp).is_some_and(|l| !l.is_nan()) { + continue; + } + fills.push(Row { + timestamp: m.timestamp, + station_code: code.to_string(), + water_level: m.wl_msl - offset, + discharge: if exact { m.discharge } else { f64::NAN }, + }); + } + if fills.len() > before { + info.push(FillInfo { + station_code: code.to_string(), + offset, + filled_hours: fills.len() - before, + }); + } + } + if fills.is_empty() { + return (rid.to_vec(), info); + } + // concat + drop_duplicates(keep="last"): a fill replaces a level-less RID row + let mut merged: BTreeMap<(String, NaiveDateTime), Row> = BTreeMap::new(); + for r in rid.iter().chain(fills.iter()) { + merged.insert((r.station_code.clone(), r.timestamp), r.clone()); + } + (merged.into_values().collect(), info) +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::{Duration, NaiveDate}; + + fn t0() -> NaiveDateTime { + NaiveDate::from_ymd_opt(2026, 10, 1) + .unwrap() + .and_hms_opt(0, 0, 0) + .unwrap() + } + + fn rid(code: &str, hours: std::ops::Range) -> Vec { + hours + .map(|h| Row { + timestamp: t0() + Duration::hours(h), + station_code: code.into(), + water_level: 1.0 + h as f64 * 0.01, + discharge: 50.0, + }) + .collect() + } + + fn mirror(code: &str, hours: std::ops::Range, offset: f64) -> Vec { + hours + .map(|h| HiiRow { + timestamp: t0() + Duration::hours(h), + station_code: code.into(), + wl_msl: 1.0 + h as f64 * 0.01 + offset, + discharge: 51.0, + }) + .collect() + } + + #[test] + fn trailing_hours_filled_with_the_overlap_offset() { + let r = rid("P.1", 0..200); + let h = prepare(&mirror("P.1", 0..203, 300.0)); + let (merged, info) = fill(&r, &h); + assert_eq!(merged.len(), 203); + assert_eq!(info[0].filled_hours, 3); + let last = merged.last().unwrap(); + assert_eq!(last.timestamp, t0() + Duration::hours(202)); + assert!((last.water_level - 3.02).abs() < 1e-9); + assert_eq!(last.discharge, 51.0); // exact mirror keeps discharge + } + + #[test] + fn short_overlap_or_other_sensor_is_not_merged() { + let r = rid("P.1", 0..100); + assert_eq!( + fill(&r, &prepare(&mirror("P.1", 0..103, 300.0))).0.len(), + 100 + ); + let r = rid("P.77", 0..200); + assert_eq!( + fill(&r, &prepare(&mirror("P.77", 0..203, 300.0))).0.len(), + 200 + ); + } + + #[test] + fn p81_is_bias_corrected_without_discharge() { + let (merged, _) = fill(&rid("P.81", 0..200), &prepare(&mirror("P.81", 0..201, 5.0))); + assert!(merged.last().unwrap().discharge.is_nan()); + } + + #[test] + fn latest_reading_of_an_hour_wins() { + let mut rows = mirror("P.1", 0..1, 300.0); + rows.push(HiiRow { + timestamp: t0() + Duration::minutes(40), + wl_msl: 302.0, + ..rows[0].clone() + }); + rows.push(HiiRow { + timestamp: t0() + Duration::minutes(10), + wl_msl: 301.5, + ..rows[0].clone() + }); + rows.push(HiiRow { + wl_msl: f64::NAN, + timestamp: t0() + Duration::minutes(50), + ..rows[0].clone() + }); + let p = prepare(&rows); + assert_eq!(p.len(), 1); + assert_eq!(p[0].wl_msl, 302.0); + assert_eq!(p[0].timestamp, t0()); + } +} diff --git a/rust/forecast/src/lib.rs b/rust/forecast/src/lib.rs new file mode 100644 index 0000000..a155bcf --- /dev/null +++ b/rust/forecast/src/lib.rs @@ -0,0 +1,24 @@ +//! Flood forecast inference: a port of the serving path of `src/ml` +//! (`features.py`, `predict.py`, the HII gap fill of `data.py`, `dam.py`). +//! +//! Training stays in Python (scikit-learn). `src/ml/export.py` writes each +//! trained bundle as `models/flood_.json.gz` and the shared constants +//! as `flood_config.json` (compiled in). Given the same inputs, this crate +//! reproduces Python's forecasts: features and levels bit for bit, +//! probabilities to the last bits of `exp` (see `tests/parity.rs`). + +pub mod config; +pub mod dam; +pub mod features; +pub mod grid; +pub mod hgb; +pub mod hii; +pub mod model; +pub mod predict; +pub mod rain; +pub mod rolling; + +pub use config::config; +pub use grid::Row; +pub use model::{Bundle, ModelSet, ModelStore}; +pub use predict::{Forecast, StageProbability, forecasts}; diff --git a/rust/forecast/src/model.rs b/rust/forecast/src/model.rs new file mode 100644 index 0000000..7a80244 --- /dev/null +++ b/rust/forecast/src/model.rs @@ -0,0 +1,165 @@ +//! Trained bundles exported by src/ml/export.py (`flood_.json.gz`), +//! reloaded when a retrain replaces a file. + +use crate::hgb::Head; +use anyhow::{Context, Result, bail}; +use serde::Deserialize; +use std::collections::BTreeMap; +use std::io::Read; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::SystemTime; + +pub const FORMAT_VERSION: u32 = 1; + +#[derive(Debug, Deserialize)] +pub struct BundleThresholds { + pub warning: f64, + pub danger: f64, +} + +#[derive(Debug, Deserialize)] +pub struct Bundle { + pub format_version: u32, + pub station_code: String, + pub model_version: String, + pub trained_at: Option, + pub regression_target: Option, + pub feature_names: Vec, + pub horizons: Vec, + /// Validation-residual sigma per horizon (keys are the hours as text). + pub sigma: BTreeMap, + pub thresholds: BundleThresholds, + /// `max_` regressors and `warn_` / `danger_` classifiers; a head + /// skipped at training time is null. + pub heads: BTreeMap>, +} + +impl Bundle { + pub fn head(&self, name: &str) -> Option<&Head> { + self.heads.get(name).and_then(Option::as_ref) + } + + pub fn sigma(&self, horizon: u32, default: f64) -> f64 { + self.sigma + .get(&horizon.to_string()) + .copied() + .unwrap_or(default) + } + + pub fn from_json(text: &str) -> Result { + let bundle: Bundle = serde_json::from_str(text)?; + if bundle.format_version != FORMAT_VERSION { + bail!( + "bundle format {} (this build reads {FORMAT_VERSION}): rebuild ping-collector", + bundle.format_version + ); + } + for (name, head) in &bundle.heads { + if let Some(h) = head { + h.validate(bundle.feature_names.len()) + .with_context(|| format!("head {name}"))?; + } + } + Ok(bundle) + } + + pub fn load(path: &Path) -> Result { + let raw = std::fs::read(path).with_context(|| format!("reading {}", path.display()))?; + let text = if path.extension().is_some_and(|e| e == "gz") { + let mut s = String::new(); + flate2::read::GzDecoder::new(&raw[..]) + .read_to_string(&mut s) + .with_context(|| format!("decompressing {}", path.display()))?; + s + } else { + String::from_utf8(raw)? + }; + Self::from_json(&text).with_context(|| format!("parsing {}", path.display())) + } +} + +/// The bundles available for one forecast run, by station. +pub type ModelSet = BTreeMap>; + +/// `models/flood_.json.gz`, cached by modification time. +pub struct ModelStore { + dir: PathBuf, + loaded: BTreeMap)>, +} + +fn station_of(path: &Path) -> Option { + let name = path.file_name()?.to_str()?; + let code = name.strip_prefix("flood_")?.strip_suffix(".json.gz")?; + (!code.is_empty()).then(|| code.to_string()) +} + +impl ModelStore { + pub fn new(dir: impl Into) -> Self { + Self { + dir: dir.into(), + loaded: BTreeMap::new(), + } + } + + pub fn dir(&self) -> &Path { + &self.dir + } + + /// Rescan the directory: load new or replaced bundles, drop removed ones. + /// A bundle that fails to load is skipped (that station falls back to the + /// heuristic, as in Python) and the previous copy, if any, is kept. + pub fn refresh(&mut self) -> Result { + let entries = std::fs::read_dir(&self.dir) + .with_context(|| format!("reading models directory {}", self.dir.display()))?; + let mut seen = BTreeMap::new(); + for entry in entries.flatten() { + let path = entry.path(); + let Some(code) = station_of(&path) else { + continue; + }; + let mtime = entry.metadata().and_then(|m| m.modified()).ok(); + seen.insert(code, (path, mtime)); + } + self.loaded.retain(|code, _| seen.contains_key(code)); + for (code, (path, mtime)) in seen { + let current = self.loaded.get(&code).map(|(t, _)| *t); + if mtime.is_some() && current == mtime { + continue; + } + self.warn_if_export_is_stale(&code, &path); + match Bundle::load(&path) { + Ok(b) if b.station_code == code => { + tracing::info!("model {code}: {} loaded", b.model_version); + self.loaded + .insert(code, (mtime.unwrap_or(SystemTime::UNIX_EPOCH), Arc::new(b))); + } + Ok(b) => tracing::error!( + "{}: bundle is for {}, not {code}; ignored", + path.display(), + b.station_code + ), + Err(e) => tracing::error!("model {code}: {e:#}"), + } + } + Ok(self + .loaded + .iter() + .map(|(k, (_, b))| (k.clone(), b.clone())) + .collect()) + } + + fn warn_if_export_is_stale(&self, code: &str, json: &Path) { + let joblib = self.dir.join(format!("flood_{code}.joblib")); + let modified = |p: &Path| std::fs::metadata(p).and_then(|m| m.modified()).ok(); + if let (Some(j), Some(e)) = (modified(&joblib), modified(json)) + && j > e + { + tracing::warn!( + "{} is newer than its export {}: run `python -m src.ml.export`", + joblib.display(), + json.display() + ); + } + } +} diff --git a/rust/forecast/src/predict.rs b/rust/forecast/src/predict.rs new file mode 100644 index 0000000..9058a57 --- /dev/null +++ b/rust/forecast/src/predict.rs @@ -0,0 +1,245 @@ +//! `predict.get_forecasts`: one forecast row per station x horizon, from the +//! trained bundle when it is usable, else the persistence heuristic. + +use crate::config::config; +use crate::dam::DamFrame; +use crate::features::{build_features, row_at}; +use crate::grid::{Grid, Row}; +use crate::model::{Bundle, ModelSet}; +use crate::rain::RainSeries; +use chrono::{Duration, NaiveDateTime}; +use serde::{Serialize, Serializer}; + +fn iso(t: &NaiveDateTime, s: S) -> Result { + s.serialize_str(&t.format("%Y-%m-%dT%H:%M:%S").to_string()) +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct StageProbability { + pub stage: u32, + pub level: f64, + pub p_exceed: f64, +} + +/// One `/forecast` row; field order and names as Python's dicts. +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct Forecast { + pub station_code: String, + pub horizon_hours: u32, + pub p_warning: f64, + pub p_danger: f64, + pub predicted_max_level: f64, + pub current_level: f64, + #[serde(serialize_with = "iso")] + pub as_of: NaiveDateTime, + pub model_version: String, + pub trained_at: Option, + pub source: String, + pub threshold_warning: f64, + pub threshold_danger: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub stages: Option>, +} + +fn sigmoid(predicted_max: f64, threshold: f64, sigma: f64) -> f64 { + 1.0 / (1.0 + (-(predicted_max - threshold) / sigma).exp()) +} + +fn clip(p: f64) -> f64 { + p.clamp(0.0, 1.0) +} + +struct Current { + station: String, + as_of: NaiveDateTime, + level: f64, + level_t_minus_3: Option, +} + +fn heuristic(c: &Current, horizons: &[u32]) -> Vec { + let cfg = config(); + let (warn, danger) = cfg.thresholds(&c.station); + let rate = c + .level_t_minus_3 + .map_or(0.0, |l3| ((c.level - l3) / 3.0).max(0.0)); + horizons + .iter() + .map(|&h| { + let predicted_max = (c.level + rate * h as f64 * 0.7).max(c.level); + let p_warning = clip(sigmoid(predicted_max, warn, cfg.heuristic_sigma)); + let p_danger = clip(sigmoid(predicted_max, danger, cfg.heuristic_sigma)).min(p_warning); + Forecast { + station_code: c.station.clone(), + horizon_hours: h, + p_warning, + p_danger, + predicted_max_level: predicted_max, + current_level: c.level, + as_of: c.as_of, + model_version: cfg.heuristic_version.clone(), + trained_at: None, + source: "heuristic".into(), + threshold_warning: warn, + threshold_danger: danger, + stages: None, + } + }) + .collect() +} + +/// `predict._model_forecast`: `None` when the bundle's features cannot be +/// built (the station then falls back to the heuristic); per horizon `None` +/// when that horizon's regressor was skipped at training time. +fn model_forecast( + c: &Current, + grid: &Grid, + bundle: &Bundle, + rain: Option<&RainSeries>, + dam: Option<&DamFrame>, +) -> Option>> { + let cfg = config(); + let (mut warn, mut danger) = (bundle.thresholds.warning, bundle.thresholds.danger); + // Thresholds changed since training: the classifier heads answer the old + // question, so use regression + sigma against the configured levels. + let configured = cfg.thresholds(&c.station); + let thresholds_stale = configured != (warn, danger); + if thresholds_stale { + tracing::warn!( + "{}: bundle thresholds ({warn}, {danger}) differ from configured {configured:?}; \ + using regression-derived probabilities until the model is retrained", + c.station + ); + (warn, danger) = configured; + } + let cols = build_features(grid, &c.station, rain, dam); + let i = grid.index_of(c.as_of)?; + let row = row_at(&cols, i); + let mut x = Vec::with_capacity(bundle.feature_names.len()); + for name in &bundle.feature_names { + match row.get(name.as_str()) { + Some(v) => x.push(*v), + None => { + tracing::error!( + "feature mismatch for {} (missing {name}); falling back to heuristic", + c.station + ); + return None; + } + } + } + let stages = cfg.flood_stages.get(&c.station); + let rows = bundle + .horizons + .iter() + .map(|&h| { + let reg = bundle.head(&format!("max_{h}"))?; + let mut raw = reg.predict(&x); + if bundle.regression_target.as_deref() == Some("rise") { + raw += c.level; // v2+ bundles predict the rise over the current level + } + let predicted_max = raw.max(c.level); + let sigma = bundle.sigma(h, cfg.heuristic_sigma); + // classifier OR regression-sigmoid, whichever is more alarmed + let head_p = |kind: &str| { + (!thresholds_stale) + .then(|| bundle.head(&format!("{kind}_{h}"))) + .flatten() + .map(|head| head.predict(&x)) + }; + let mut p_warning = sigmoid(predicted_max, warn, sigma); + if let Some(p) = head_p("warn") { + p_warning = p_warning.max(p); + } + let mut p_danger = sigmoid(predicted_max, danger, sigma); + if let Some(p) = head_p("danger") { + p_danger = p_danger.max(p); + } + let p_warning = clip(p_warning); + let p_danger = clip(p_danger).min(p_warning); + Some(Forecast { + station_code: c.station.clone(), + horizon_hours: h, + p_warning, + p_danger, + predicted_max_level: predicted_max, + current_level: c.level, + as_of: c.as_of, + model_version: bundle.model_version.clone(), + trained_at: bundle.trained_at.clone(), + source: "model".into(), + threshold_warning: warn, + threshold_danger: danger, + stages: stages.map(|stages| { + stages + .iter() + .map(|s| StageProbability { + stage: s.stage, + level: s.level, + p_exceed: clip(sigmoid(predicted_max, s.level, sigma)), + }) + .collect() + }), + }) + }) + .collect(); + Some(rows) +} + +fn forecast_station( + station: &str, + grid: &Grid, + models: &ModelSet, + now: NaiveDateTime, + rain: Option<&RainSeries>, + dam: Option<&DamFrame>, +) -> Vec { + let cfg = config(); + let Some(i) = grid.last_observed(station) else { + tracing::warn!("no observed readings for station {station}; omitting"); + return vec![]; + }; + let levels = &grid.level[station]; + let as_of = grid.times[i]; + let t3 = as_of - Duration::hours(3); + let c = Current { + station: station.to_string(), + as_of, + level: levels[i], + level_t_minus_3: grid.index_of(t3).map(|j| levels[j]).filter(|v| !v.is_nan()), + }; + let staleness_h = (now - as_of) + .num_nanoseconds() + .map_or(f64::INFINITY, |ns| ns as f64 / 1e9 / 3600.0); + let bundle = match models.get(station) { + Some(b) if staleness_h <= cfg.stale_after_h => b, + _ => return heuristic(&c, &cfg.horizons), + }; + let Some(rows) = model_forecast(&c, grid, bundle, rain, dam) else { + return heuristic(&c, &cfg.horizons); + }; + bundle + .horizons + .iter() + .zip(rows) + .flat_map(|(&h, row)| match row { + Some(r) => vec![r], + None => heuristic(&c, &[h]), + }) + .collect() +} + +/// Forecasts for every station in `rows` (the merged RID + HII-fill window), +/// stations in code order. `rain` and `dam` as `build_features`. +pub fn forecasts( + rows: &[Row], + models: &ModelSet, + now: NaiveDateTime, + rain: Option<&RainSeries>, + dam: Option<&DamFrame>, +) -> Vec { + let grid = Grid::from_rows(rows); + grid.level + .keys() + .flat_map(|station| forecast_station(station, &grid, models, now, rain, dam)) + .collect() +} diff --git a/rust/forecast/src/rain.rs b/rust/forecast/src/rain.rs new file mode 100644 index 0000000..2fea204 --- /dev/null +++ b/rust/forecast/src/rain.rs @@ -0,0 +1,27 @@ +//! Catchment rain series (`src/ml/rain.py`: `catchment_mean`). + +use chrono::NaiveDateTime; +use std::collections::BTreeMap; + +/// Hourly catchment-average precipitation (mm); NaN = no value. +pub type RainSeries = BTreeMap; + +/// pandas `mean(axis=1)` over the catchment points: the points that have a +/// value, summed in column order, divided by their count. +pub fn catchment_mean(points: &[Option]) -> Option { + let mut sum = 0.0; + let mut n = 0usize; + for v in points.iter().flatten() { + sum += *v; + n += 1; + } + (n > 0).then(|| sum / n as f64) +} + +/// The series on the given hours (`Series.reindex`), NaN where absent. +pub fn reindex(rain: &RainSeries, times: &[NaiveDateTime]) -> Vec { + times + .iter() + .map(|t| rain.get(t).copied().unwrap_or(f64::NAN)) + .collect() +} diff --git a/rust/forecast/src/rolling.rs b/rust/forecast/src/rolling.rs new file mode 100644 index 0000000..31b1c95 --- /dev/null +++ b/rust/forecast/src/rolling.rs @@ -0,0 +1,206 @@ +//! The pandas Series operations `features.py` uses, with pandas' +//! floating-point behaviour (pandas 2.0, `_libs/window/aggregations.pyx`). +//! NaN is "missing" throughout, as in pandas. + +pub type Series = Vec; + +pub fn nan_series(n: usize) -> Series { + vec![f64::NAN; n] +} + +/// `ffill(limit=limit)`: at most `limit` consecutive NaNs after a value are filled. +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 +} + +/// `shift(k)`: positive k looks back, negative k looks ahead; NaN-padded. +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() +} + +pub fn sub(a: &[f64], b: &[f64]) -> Series { + a.iter().zip(b).map(|(x, y)| x - y).collect() +} + +/// `rolling(w, min_periods=1).sum()`. +pub fn rolling_sum(v: &[f64], w: usize) -> Series { + kahan_window(v, w, false) +} + +/// `rolling(w, min_periods=1).mean()`. +pub fn rolling_mean(v: &[f64], w: usize) -> Series { + kahan_window(v, w, true) +} + +/// `rolling(w, min_periods=1).max()`. +pub fn rolling_max(v: &[f64], w: usize) -> Series { + extreme(v, w, f64::max) +} + +/// `rolling(w, min_periods=1).min()`. +pub fn rolling_min(v: &[f64], w: usize) -> Series { + extreme(v, w, f64::min) +} + +/// Sum over the next `w` rows, excluding the current one: +/// `s.shift(-1).iloc[::-1].rolling(w, min_periods=1).sum().iloc[::-1]`. +pub fn forward_sum(v: &[f64], w: usize) -> Series { + let mut ahead = shift(v, -1); + ahead.reverse(); + let mut out = rolling_sum(&ahead, w); + out.reverse(); + out +} + +fn extreme(v: &[f64], w: usize, pick: fn(f64, f64) -> f64) -> Series { + (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 { pick(acc, x) }, + ) + }) + .collect() +} + +/// `roll_sum` / `roll_mean` for a fixed window, operation for operation: the +/// window moves incrementally with separate Kahan compensations for adds and +/// removes, the "run of identical values" shortcut (GH#42064) and, for the +/// mean, the sign clamp. A plain per-window sum differs in the last bits, and +/// HGB thresholds can sit exactly on such values. +fn kahan_window(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 = v.first().copied().unwrap_or(f64::NAN); + let mut out = Vec::with_capacity(v.len()); + for i in 0..v.len() { + if w > 0 && 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 { + let r = sum / nobs as f64; + if same_run >= nobs { + prev + } else 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 +} + +#[cfg(test)] +mod tests { + use super::*; + + fn same(a: &[f64], b: &[f64]) -> bool { + a.len() == b.len() + && a.iter() + .zip(b) + .all(|(x, y)| x.to_bits() == y.to_bits() || (x.is_nan() && y.is_nan())) + } + + #[test] + fn ffill_respects_the_limit() { + let n = f64::NAN; + assert!(same( + &ffill(&[n, 1.0, n, n, n, n, 2.0], 3), + &[n, 1.0, 1.0, 1.0, 1.0, n, 2.0] + )); + } + + #[test] + fn shift_both_ways() { + let n = f64::NAN; + assert!(same(&shift(&[1.0, 2.0, 3.0], 1), &[n, 1.0, 2.0])); + assert!(same(&shift(&[1.0, 2.0, 3.0], -1), &[2.0, 3.0, n])); + } + + #[test] + fn rolling_skips_nan_and_needs_one_value() { + let n = f64::NAN; + let v = [n, 1.0, n, 3.0, n, n, n]; + assert!(same(&rolling_sum(&v, 2), &[n, 1.0, 1.0, 3.0, 3.0, n, n])); + assert!(same(&rolling_mean(&v, 3), &[n, 1.0, 1.0, 2.0, 3.0, 3.0, n])); + assert!(same(&rolling_max(&v, 3), &[n, 1.0, 1.0, 3.0, 3.0, 3.0, n])); + assert!(same(&rolling_min(&v, 3), &[n, 1.0, 1.0, 1.0, 3.0, 3.0, n])); + } + + #[test] + fn forward_sum_excludes_the_current_row() { + let n = f64::NAN; + assert!(same( + &forward_sum(&[1.0, 2.0, 3.0, 4.0], 2), + &[5.0, 7.0, 4.0, n] + )); + } + + #[test] + fn identical_run_shortcut_removes_float_residue() { + // a window of identical values yields value * count (sum) or the + // value itself (mean), whatever the Kahan accumulator holds + let v = [0.7, 0.1, 0.1, 0.1]; + assert_eq!(rolling_sum(&v, 3)[3], 0.1 * 3.0); + assert_eq!(rolling_mean(&v, 3)[3], 0.1); + } +} diff --git a/rust/forecast/testdata/fixture_flood2024.json.gz b/rust/forecast/testdata/fixture_flood2024.json.gz new file mode 100644 index 0000000..f15e536 Binary files /dev/null and b/rust/forecast/testdata/fixture_flood2024.json.gz differ diff --git a/rust/forecast/testdata/fixture_peak2024.json.gz b/rust/forecast/testdata/fixture_peak2024.json.gz new file mode 100644 index 0000000..3fc109e Binary files /dev/null and b/rust/forecast/testdata/fixture_peak2024.json.gz differ diff --git a/rust/forecast/testdata/fixture_recent2026.json.gz b/rust/forecast/testdata/fixture_recent2026.json.gz new file mode 100644 index 0000000..852e584 Binary files /dev/null and b/rust/forecast/testdata/fixture_recent2026.json.gz differ diff --git a/rust/forecast/testdata/fixture_stale2026.json.gz b/rust/forecast/testdata/fixture_stale2026.json.gz new file mode 100644 index 0000000..afd8297 Binary files /dev/null and b/rust/forecast/testdata/fixture_stale2026.json.gz differ diff --git a/rust/forecast/testdata/gen_forecast_fixtures.py b/rust/forecast/testdata/gen_forecast_fixtures.py new file mode 100644 index 0000000..2bbfcb1 --- /dev/null +++ b/rust/forecast/testdata/gen_forecast_fixtures.py @@ -0,0 +1,310 @@ +"""Parity fixtures for rust/forecast: inputs plus what Python computes from them. + +Run from the repo root (needs the trained bundles in models/ and the +models/cache/ snapshot, i.e. a machine that has trained or pulled models): + + uv run python rust/forecast/testdata/gen_forecast_fixtures.py + +Writes, next to this file: + models/flood_P.1.json.gz the P.1 bundle as src/ml/export.py writes it + fixture_.json.gz inputs (RID rows, HII mirror rows, rain points, + dam days, now) and Python's outputs: merged + rows, the feature row at each station's as_of, + forecasts with only P.1's model (checked in CI) + and with every local model (checked when + ../../models has the exports). + +The Python side is the production code path: data._normalize_long and +data.fill_from_hii (with the HII query replaced by the fixture rows), +rain.catchment_mean, dam.hourly_frame and predict.forecasts_from_long. +""" + +import datetime +import gzip +import json +import math +import shutil +import sys +import tempfile +from pathlib import Path +from unittest import mock + +import joblib +import numpy as np +import pandas as pd + +HERE = Path(__file__).resolve().parent +REPO = HERE.parents[2] +sys.path.insert(0, str(REPO)) + +from src.ml import dam as dam_mod # noqa: E402 +from src.ml import data, export, features, predict # noqa: E402 +from src.ml import rain as rain_mod # noqa: E402 + +MODELS = REPO / "models" +CACHE = MODELS / "cache" +FMT = "%Y-%m-%dT%H:%M:%S" + +# Synthetic MSL offsets of the HII twins (P.81 also gets a noisy bias). +HII_OFFSETS = { + "P.1": 300.0, + "P.103": 296.473, + "P.20": 410.31, + "P.4A": 352.09, + "P.67": 318.552, + "P.75": 389.11, + "P.82": 520.6, + "P.84": 288.125, + "P.92": 371.0, + "P.81": 281.402, + "P.77": 402.5, # a different sensor: must never be merged +} + + +def num(v): + if v is None: + return None + v = float(v) + return None if math.isnan(v) else v + + +def scenario_inputs(long, rain_points, dam_daily, now, rid_lag, null_levels): + end = now + start = now - datetime.timedelta(hours=336) + window = long[(long.timestamp >= start) & (long.timestamp <= end)].copy() + + # HII twins see every hour the gauges saw, including those RID "lags" on + rng = np.random.default_rng(42) + hii = [] + for code, offset in HII_OFFSETS.items(): + rows = window[(window.station_code == code) & window.water_level.notna()] + for ts, level, q in zip(rows.timestamp, rows.water_level, rows.discharge): + wl = level + offset + if code == "P.81": + wl += rng.normal(0, 0.05) + 0.19 + hii.append( + [ts, code, round(wl, 3), num(None if pd.isna(q) else round(q + 1.5, 2))] + ) + if code == "P.1" and ts.hour % 5 == 0: + # a second, later reading in the same hour wins; an earlier + # one and a level-less one must not + hii.append( + [ts + pd.Timedelta(minutes=40), code, round(wl + 0.004, 3), None] + ) + hii.append( + [ts + pd.Timedelta(minutes=10), code, round(wl - 0.25, 3), None] + ) + hii.append([ts + pd.Timedelta(minutes=50), code, None, None]) + hii = pd.DataFrame( + hii, columns=["timestamp", "station_code", "wl_msl", "discharge"] + ) + hii = hii[(hii.timestamp >= start) & (hii.timestamp <= end)] + # the query's ORDER BY m.timestamp, m.station_id (code stands in for the id) + hii = hii.sort_values(["timestamp", "station_code"], kind="stable").reset_index( + drop=True + ) + + # RID behind on some stations (the night case), NULL levels on others + newest = window.timestamp.max() + for code, hours in rid_lag.items(): + cut = newest - pd.Timedelta(hours=hours) + window = window[~((window.station_code == code) & (window.timestamp > cut))] + for code, every in null_levels.items(): + sel = (window.station_code == code) & (window.timestamp.dt.hour % every == 0) + window.loc[sel, "water_level"] = np.nan + window = window.sort_values(["timestamp", "station_code"], kind="stable") + + day0 = pd.Timestamp(now.date()) + rain_rows = rain_points[ + (rain_points.index >= day0 - pd.Timedelta(days=5)) + & (rain_points.index < day0 + pd.Timedelta(days=2)) + ] + dam_days = dam_daily[ + dam_daily.index >= pd.Timestamp(now.date() - datetime.timedelta(days=21)) + ] + return start, end, window, hii, rain_rows, dam_days + + +def python_side(window, hii, rain_rows, dam_days, start, end, now, models_dirs): + df = data._normalize_long(window) + + def fake_fetch(db_url, codes, s, e): + assert (s, e) == (start, end) + return data._prepare_hii_levels(hii[hii.station_code.isin(codes)]) + + with mock.patch.object(data, "_fetch_hii_levels", fake_fetch): + merged = data.fill_from_hii(df, "fixture://", start=start, end=end) + + rain = rain_mod.catchment_mean(rain_rows) + if rain is None: + rain = pd.Series(dtype=float) + dam_frame = dam_mod.hourly_frame(dam_days) + if dam_frame is None: + dam_frame = pd.DataFrame() + + grid = features.make_hourly_grid(merged) + feature_rows = {} + for code in sorted(merged.station_code.unique()): + observed = grid.observed[(code, "water_level")].dropna() + if observed.empty: + continue + as_of = observed.index.max() + X = features.build_features(grid, code, rain=rain, dam=dam_frame) + feature_rows[code] = { + "as_of": as_of.strftime(FMT), + "features": {k: num(v) for k, v in X.loc[as_of].items()}, + } + + forecasts = { + name: json.loads( + json.dumps( + predict.forecasts_from_long( + merged, models_dir=d, now=now, rain=rain, dam=dam_frame + ), + default=float, + ) + ) + for name, d in models_dirs.items() + } + merged_rows = [ + [t.strftime(FMT), c, num(l), num(q)] + for t, c, l, q in zip( + merged.timestamp, merged.station_code, merged.water_level, merged.discharge + ) + ] + merged_rows.sort(key=lambda r: (r[1], r[0])) + return merged_rows, feature_rows, forecasts + + +def main(): + long = data._read_cache(CACHE, None) + rain_points = pd.read_csv( + CACHE / "rain_openmeteo.csv.gz", index_col=0, parse_dates=True + ) + rain_points = rain_points[[p[0] for p in rain_mod.CATCHMENT_POINTS]] + dam_daily = pd.read_csv(CACHE / "dam_200103.csv.gz", index_col=0, parse_dates=True) + + out_models = HERE / "models" + out_models.mkdir(exist_ok=True) + p1 = joblib.load(MODELS / "flood_P.1.joblib") + export.write_bundle_json(p1, out_models / "flood_P.1.json.gz") + + tmp = Path(tempfile.mkdtemp()) + try: + p1_only = tmp / "p1" + p1_only.mkdir() + shutil.copy(MODELS / "flood_P.1.joblib", p1_only) + # the configured thresholds moved since training -> regression path + p1_moved = tmp / "p1_moved" + p1_moved.mkdir() + moved = dict(p1) + moved["thresholds"] = {"warning": 3.5, "danger": 4.0} + joblib.dump(moved, p1_moved / "flood_P.1.joblib") + + scenarios = { + # rising limb of the October 2024 record flood, RID 3 h behind + # on three HII-mirrored stations + "flood2024": dict( + now=datetime.datetime(2024, 10, 3, 1, 17, 31, 500000), + rid_lag={"P.1": 3, "P.103": 3, "P.67": 2}, + null_levels={"P.21": 7}, + ), + "peak2024": dict( + now=datetime.datetime(2024, 10, 5, 3, 5, 0), + rid_lag={"P.1": 1}, + null_levels={}, + ), + # quiet season, the newest cached data, with dam state + "recent2026": dict( + now=datetime.datetime(2026, 8, 13, 19, 40, 12), + rid_lag={}, + null_levels={"P.85": 5}, + ), + # every station > 6 h old: heuristic only + "stale2026": dict( + now=datetime.datetime(2026, 8, 14, 3, 45, 0), + rid_lag={"P.20": 30}, + null_levels={}, + ), + } + for name, sc in scenarios.items(): + start, end, window, hii, rain_rows, dam_days = scenario_inputs( + long, + rain_points, + dam_daily, + sc["now"], + sc["rid_lag"], + sc["null_levels"], + ) + dirs = {"p1_only": p1_only, "all_models": MODELS} + if name == "recent2026": + dirs["p1_moved_thresholds"] = p1_moved + merged, feats, fcs = python_side( + window, hii, rain_rows, dam_days, start, end, sc["now"], dirs + ) + fixture = { + "now": sc["now"].strftime(FMT + ".%f"), + "window": [start.strftime(FMT + ".%f"), end.strftime(FMT + ".%f")], + "rid_rows": [ + [t.strftime(FMT), c, num(l), num(q)] + for t, c, l, q in zip( + window.timestamp, + window.station_code, + window.water_level, + window.discharge, + ) + ], + "hii_rows": [ + [t.strftime(FMT), c, num(w), num(q)] + for t, c, w, q in zip( + hii.timestamp, hii.station_code, hii.wl_msl, hii.discharge + ) + ], + "rain_points": [ + [t.strftime(FMT), [num(v) for v in row]] + for t, row in zip(rain_rows.index, rain_rows.to_numpy().tolist()) + ], + "dam_days": [ + [d.strftime("%Y-%m-%d"), num(s), num(i), num(o)] + for d, s, i, o in zip( + dam_days.index, + dam_days.storage_pct, + dam_days.inflow_mcm, + dam_days.outflow_mcm, + ) + ], + "expected": { + "merged_rows": merged, + "features": feats, + "forecasts": fcs, + "all_models_versions": sorted( + { + joblib.load(p)["model_version"] + for p in MODELS.glob("flood_*.joblib") + } + ), + }, + } + path = HERE / f"fixture_{name}.json.gz" + with ( + open(path, "wb") as raw, + gzip.GzipFile(fileobj=raw, mode="wb", mtime=0) as gz, + ): + gz.write(json.dumps(fixture, separators=(",", ":")).encode()) + n_model = sum(r["source"] == "model" for r in fcs["p1_only"]) + p1_24 = [ + r + for r in fcs["all_models"] + if r["station_code"] == "P.1" and r["horizon_hours"] == 24 + ] + print( + f"{name}: {len(fixture['rid_rows'])} RID rows, {len(fixture['hii_rows'])} HII rows, " + f"{len(merged)} merged, {len(feats)} stations, model rows (P.1 only) {n_model}, " + f"P.1 24h {[(round(r['p_warning'], 3), round(r['predicted_max_level'], 2)) for r in p1_24]}" + ) + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +if __name__ == "__main__": + main() diff --git a/rust/forecast/testdata/models/flood_P.1.json.gz b/rust/forecast/testdata/models/flood_P.1.json.gz new file mode 100644 index 0000000..5666cc7 Binary files /dev/null and b/rust/forecast/testdata/models/flood_P.1.json.gz differ diff --git a/rust/forecast/tests/parity.rs b/rust/forecast/tests/parity.rs new file mode 100644 index 0000000..c14ce59 --- /dev/null +++ b/rust/forecast/tests/parity.rs @@ -0,0 +1,322 @@ +//! Replays the fixtures written by testdata/gen_forecast_fixtures.py (the +//! Python production path on the same inputs) through this crate. +//! +//! Exact equality for the merged readings, every feature and every level; +//! probabilities may differ in the last bits of `exp`, and `doy_sin`/`doy_cos` +//! in the last bits of the platform's `sin`/`cos` (numpy vs libm). + +use chrono::{NaiveDate, NaiveDateTime}; +use ping_forecast::dam::{DamDay, DamFrame, hourly_frame}; +use ping_forecast::features::{build_features, row_at}; +use ping_forecast::grid::{Grid, Row}; +use ping_forecast::hii::{self, HiiRow}; +use ping_forecast::model::{Bundle, ModelSet, ModelStore}; +use ping_forecast::rain::{RainSeries, catchment_mean}; +use ping_forecast::{Forecast, forecasts}; +use serde_json::Value; +use std::io::Read; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +const SCENARIOS: [&str; 4] = ["flood2024", "peak2024", "recent2026", "stale2026"]; +/// Max distance in units in the last place for transcendental results. +const ULPS: u64 = 4; + +fn testdata() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join("testdata") +} + +fn load(name: &str) -> Value { + let raw = std::fs::read(testdata().join(format!("fixture_{name}.json.gz"))).unwrap(); + let mut s = String::new(); + flate2::read::GzDecoder::new(&raw[..]) + .read_to_string(&mut s) + .unwrap(); + serde_json::from_str(&s).unwrap() +} + +fn ts(v: &Value) -> NaiveDateTime { + let s = v.as_str().unwrap(); + NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f").unwrap() +} + +fn f(v: &Value) -> f64 { + v.as_f64().unwrap_or(f64::NAN) +} + +fn rows(v: &Value) -> Vec { + v.as_array() + .unwrap() + .iter() + .map(|r| Row { + timestamp: ts(&r[0]), + station_code: r[1].as_str().unwrap().into(), + water_level: f(&r[2]), + discharge: f(&r[3]), + }) + .collect() +} + +struct Inputs { + now: NaiveDateTime, + rid: Vec, + hii: Vec, + rain: RainSeries, + dam: DamFrame, +} + +fn inputs(fx: &Value) -> Inputs { + let hii = fx["hii_rows"] + .as_array() + .unwrap() + .iter() + .map(|r| HiiRow { + timestamp: ts(&r[0]), + station_code: r[1].as_str().unwrap().into(), + wl_msl: f(&r[2]), + discharge: f(&r[3]), + }) + .collect(); + let rain = fx["rain_points"] + .as_array() + .unwrap() + .iter() + .map(|r| { + let points: Vec> = + r[1].as_array().unwrap().iter().map(Value::as_f64).collect(); + (ts(&r[0]), catchment_mean(&points).unwrap_or(f64::NAN)) + }) + .collect(); + let days: Vec = fx["dam_days"] + .as_array() + .unwrap() + .iter() + .map(|r| DamDay { + date: NaiveDate::parse_from_str(r[0].as_str().unwrap(), "%Y-%m-%d").unwrap(), + storage_pct: f(&r[1]), + inflow_mcm: f(&r[2]), + outflow_mcm: f(&r[3]), + }) + .collect(); + Inputs { + now: ts(&fx["now"]), + rid: rows(&fx["rid_rows"]), + hii, + rain, + dam: hourly_frame(&days).unwrap_or_default(), + } +} + +fn same_bits(a: f64, b: f64) -> bool { + a.to_bits() == b.to_bits() || (a.is_nan() && b.is_nan()) +} + +fn ulps(a: f64, b: f64) -> u64 { + if same_bits(a, b) || a == b { + return 0; + } + if a.is_sign_negative() != b.is_sign_negative() || a.is_nan() || b.is_nan() { + return u64::MAX; + } + a.to_bits().abs_diff(b.to_bits()) +} + +fn merged(inp: &Inputs) -> Vec { + let (mut m, _) = hii::fill(&inp.rid, &hii::prepare(&inp.hii)); + m.sort_by(|a, b| (&a.station_code, a.timestamp).cmp(&(&b.station_code, b.timestamp))); + m +} + +#[test] +fn hii_fill_matches_python() { + for name in SCENARIOS { + let fx = load(name); + let got = merged(&inputs(&fx)); + let want = rows(&fx["expected"]["merged_rows"]); + assert_eq!(got.len(), want.len(), "{name}: merged row count"); + for (g, w) in got.iter().zip(&want) { + assert!( + g.timestamp == w.timestamp + && g.station_code == w.station_code + && same_bits(g.water_level, w.water_level) + && same_bits(g.discharge, w.discharge), + "{name}: {g:?} != {w:?}" + ); + } + } +} + +#[test] +fn features_match_python_bit_for_bit() { + let mut checked = 0; + for name in SCENARIOS { + let fx = load(name); + let inp = inputs(&fx); + let grid = Grid::from_rows(&merged(&inp)); + for (code, exp) in fx["expected"]["features"].as_object().unwrap() { + let cols = build_features(&grid, code, Some(&inp.rain), Some(&inp.dam)); + let i = grid.index_of(ts(&exp["as_of"])).unwrap(); + assert_eq!(grid.last_observed(code), Some(i), "{name} {code}: as_of"); + let row = row_at(&cols, i); + let want = exp["features"].as_object().unwrap(); + let mut names: Vec<&str> = row.keys().copied().collect(); + names.sort_unstable(); + let mut want_names: Vec<&str> = want.keys().map(String::as_str).collect(); + want_names.sort_unstable(); + assert_eq!(names, want_names, "{name} {code}: feature set"); + for (k, v) in want { + let (g, w) = (row[k.as_str()], f(v)); + let ok = if k.starts_with("doy_") { + ulps(g, w) <= ULPS + } else { + same_bits(g, w) + }; + assert!(ok, "{name} {code} {k}: rust {g:?} python {w:?}"); + checked += 1; + } + } + } + assert!(checked > 3000, "only {checked} feature values compared"); +} + +fn compare(name: &str, got: &[Forecast], want: &Value) { + let want = want.as_array().unwrap(); + assert_eq!(got.len(), want.len(), "{name}: row count"); + for (g, w) in got.iter().zip(want) { + let ctx = format!("{name} {} {}h", g.station_code, g.horizon_hours); + assert_eq!(g.station_code, w["station_code"].as_str().unwrap(), "{ctx}"); + assert_eq!( + g.horizon_hours as u64, + w["horizon_hours"].as_u64().unwrap(), + "{ctx}" + ); + assert_eq!(g.as_of, ts(&w["as_of"]), "{ctx}"); + assert_eq!(g.source, w["source"].as_str().unwrap(), "{ctx}"); + assert_eq!( + g.model_version, + w["model_version"].as_str().unwrap(), + "{ctx}" + ); + assert_eq!(g.trained_at.as_deref(), w["trained_at"].as_str(), "{ctx}"); + for (field, gv) in [ + ("current_level", g.current_level), + ("predicted_max_level", g.predicted_max_level), + ("threshold_warning", g.threshold_warning), + ("threshold_danger", g.threshold_danger), + ] { + assert!( + same_bits(gv, f(&w[field])), + "{ctx} {field}: {gv} vs {}", + w[field] + ); + } + for (field, gv) in [("p_warning", g.p_warning), ("p_danger", g.p_danger)] { + let d = ulps(gv, f(&w[field])); + assert!(d <= ULPS, "{ctx} {field}: {gv} vs {} ({d} ulp)", w[field]); + } + match (&g.stages, w.get("stages")) { + (None, None) => {} + (Some(gs), Some(ws)) => { + let ws = ws.as_array().unwrap(); + assert_eq!(gs.len(), ws.len(), "{ctx}: stages"); + for (s, w) in gs.iter().zip(ws) { + assert_eq!(s.stage as u64, w["stage"].as_u64().unwrap()); + assert!(same_bits(s.level, f(&w["level"]))); + assert!( + ulps(s.p_exceed, f(&w["p_exceed"])) <= ULPS, + "{ctx} stage {}", + s.stage + ); + } + } + _ => panic!("{ctx}: stages present on one side only"), + } + } +} + +fn p1_models() -> ModelSet { + ModelStore::new(testdata().join("models")) + .refresh() + .unwrap() +} + +#[test] +fn forecasts_match_python_with_the_p1_model() { + let models = p1_models(); + assert_eq!(models.keys().collect::>(), ["P.1"]); + let mut model_rows = 0; + for name in SCENARIOS { + let fx = load(name); + let inp = inputs(&fx); + let got = forecasts( + &merged(&inp), + &models, + inp.now, + Some(&inp.rain), + Some(&inp.dam), + ); + model_rows += got.iter().filter(|r| r.source == "model").count(); + compare(name, &got, &fx["expected"]["forecasts"]["p1_only"]); + } + assert_eq!(model_rows, 9, "P.1 model path in the three fresh scenarios"); +} + +#[test] +fn moved_thresholds_use_the_regression_path_like_python() { + let path = testdata().join("models/flood_P.1.json.gz"); + let mut bundle = Bundle::load(&path).unwrap(); + bundle.thresholds.warning = 3.5; + bundle.thresholds.danger = 4.0; + let models: ModelSet = [("P.1".to_string(), Arc::new(bundle))].into(); + let fx = load("recent2026"); + let inp = inputs(&fx); + let got = forecasts( + &merged(&inp), + &models, + inp.now, + Some(&inp.rain), + Some(&inp.dam), + ); + compare( + "recent2026 moved", + &got, + &fx["expected"]["forecasts"]["p1_moved_thresholds"], + ); +} + +/// Every local model (needs `python -m src.ml.export` run in ../../models on +/// a machine with the same trained bundles as the fixture generator). +#[test] +fn forecasts_match_python_with_all_local_models() { + let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../models"); + let Ok(models) = ModelStore::new(&dir).refresh() else { + eprintln!("skipped: no {}", dir.display()); + return; + }; + let fx = load("recent2026"); + let want: Vec<&str> = fx["expected"]["all_models_versions"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap()) + .collect(); + let mut have: Vec<&str> = models.values().map(|b| b.model_version.as_str()).collect(); + have.sort_unstable(); + have.dedup(); + if models.len() < 15 || have != want { + eprintln!("skipped: local models {have:?} are not the fixture's {want:?}"); + return; + } + for name in SCENARIOS { + let fx = load(name); + let inp = inputs(&fx); + let got = forecasts( + &merged(&inp), + &models, + inp.now, + Some(&inp.rain), + Some(&inp.dam), + ); + compare(name, &got, &fx["expected"]["forecasts"]["all_models"]); + } +} diff --git a/scripts/retrain.sh b/scripts/retrain.sh index 072387e..7cd5ff1 100644 --- a/scripts/retrain.sh +++ b/scripts/retrain.sh @@ -11,9 +11,9 @@ # * A run that produced gauge-only (v2) bundles, or trained too few stations, # must NOT replace the deployed models. train.py already aborts on a # missing rain series; this script re-checks the written metrics anyway. -# * No API restart is needed: predict.py reloads changed bundles on the next -# precompute (every scrape cycle, hourly), so the new models are live -# within an hour. Restart manually if you want them live immediately. +# * No restart is needed: predict.py (Python) and ping-collector (Rust, +# FORECAST_IN_COLLECTOR=true) reload changed bundles on their next +# forecast run, so the new models are live within an hour. # # Exit codes: 0 ok, 2 training refused (see log), 3 verification failed. set -euo pipefail @@ -79,7 +79,8 @@ fi # Promote: per-file rename is atomic; readers see either the old or the new # bundle, never a partial one. Keep one previous generation for rollback. mkdir -p "${MODELS_DIR}/.previous" -for f in "${STAGE_DIR}"/flood_*.joblib "${STAGE_DIR}/metrics.json"; do +# flood_*.json.gz is the Rust forecaster's export of the same bundle. +for f in "${STAGE_DIR}"/flood_*.joblib "${STAGE_DIR}"/flood_*.json.gz "${STAGE_DIR}/metrics.json"; do name="$(basename "${f}")" if [ -f "${MODELS_DIR}/${name}" ]; then mv -f "${MODELS_DIR}/${name}" "${MODELS_DIR}/.previous/${name}" diff --git a/src/ml/dam.py b/src/ml/dam.py index ae6057e..01c826f 100644 --- a/src/ml/dam.py +++ b/src/ml/dam.py @@ -34,6 +34,7 @@ logger = logging.getLogger(__name__) REPORT_HOUR = 7 # daily value valid from 07:00 local on its own date FFILL_LIMIT_H = 48 # two missed daily reports -> NaN, not stale state +SERVING_DAYS = 21 # 336 h feature window + 72 h storage-delta lag, rounded up DAM_COLUMNS = ("storage_pct", "inflow_mcm", "outflow_mcm") CACHE_FILE = f"dam_{MAE_NGAT_DAM_ID}.csv.gz" @@ -101,9 +102,11 @@ def load_history(db_url: Optional[str] = None) -> Optional[pd.DataFrame]: def serving_frame( - db_url: Optional[str] = None, days: int = 21 + db_url: Optional[str] = None, + days: int = SERVING_DAYS, + today: Optional[datetime.date] = None, ) -> Optional[pd.DataFrame]: """Recent hourly dam state for inference (covers the 336 h feature window plus the 72 h storage-delta lag).""" - start = datetime.date.today() - datetime.timedelta(days=days) + start = (today or datetime.date.today()) - datetime.timedelta(days=days) return hourly_frame(load_daily(db_url, start=start)) diff --git a/src/ml/data.py b/src/ml/data.py index 1ce3b77..2661e28 100644 --- a/src/ml/data.py +++ b/src/ml/data.py @@ -143,6 +143,7 @@ def _fetch_hii_levels( query += f" AND s.rid_code IN ({placeholders})" for i, code in enumerate(stations): params[f"station_{i}"] = code + query += " ORDER BY m.timestamp, m.station_id" with engine.connect() as connection: df = pd.read_sql(text(query), connection, params=params) @@ -360,9 +361,10 @@ def load_latest( db_url: Optional[str] = None, hours: int = 336, stations: Optional[List[str]] = None, + end: Optional[datetime.datetime] = None, ) -> pd.DataFrame: """Load the last `hours` of history for all (or given) stations. Never cached to disk.""" - end = datetime.datetime.now() + end = end or datetime.datetime.now() start = end - datetime.timedelta(hours=hours) return load_measurements( db_url=db_url, diff --git a/src/ml/export.py b/src/ml/export.py new file mode 100644 index 0000000..a7bfa73 --- /dev/null +++ b/src/ml/export.py @@ -0,0 +1,207 @@ +"""Hand-over of trained models and forecast settings to the Rust forecaster. + +The Rust service (rust/forecast, run by ping-collector) cannot read +joblib/pickle. It reads: + +* ``models/flood_.json.gz`` -- the bundle's metadata plus every + HistGradientBoosting head as flat tree arrays (sklearn's predictor node + table). ``train_all()`` writes one next to each ``.joblib``; + ``python -m src.ml.export`` (re)exports the bundles already in ``models/``. +* ``rust/forecast/flood_config.json`` -- thresholds, flood stages, upstream + leads and the other constants features.py / predict.py / notify.py use, + compiled into the Rust binary. ``python -m src.ml.export --config`` + regenerates it; tests/test_model_export.py fails while it is out of date, + so a threshold change cannot reach Python without reaching Rust. + +Floats go through ``json`` (shortest round-trip repr) and the Rust side parses +with serde_json's ``float_roundtrip``, so every threshold and leaf value +arrives bit-identical. +""" + +import argparse +import gzip +import json +import os +import sys +from pathlib import Path +from typing import Optional + +import numpy as np + +from . import data as data_mod +from . import features, predict + +FORMAT_VERSION = 1 +REPO_ROOT = Path(__file__).resolve().parents[2] +CONFIG_PATH = REPO_ROOT / "rust" / "forecast" / "flood_config.json" + + +def export_head(est) -> dict: + """One fitted HistGradientBoosting{Regressor,Classifier} as flat arrays.""" + trees = [] + for per_iteration in est._predictors: + if len(per_iteration) != 1: + raise ValueError("multiclass heads are not supported") + nodes = per_iteration[0].nodes + if nodes["is_categorical"].any(): + raise ValueError("categorical splits are not supported") + if not np.isfinite(nodes["value"]).all(): + raise ValueError("non-finite leaf value") + thr = nodes["num_threshold"] + if np.isnan(thr).any() or np.isneginf(thr).any(): + raise ValueError("NaN or -inf split threshold") + trees.append( + { + "feature": nodes["feature_idx"].astype(int).tolist(), + # +inf ("missing vs everything" split) is not valid JSON: + # written as null, read back as +inf. + "threshold": [None if np.isposinf(t) else float(t) for t in thr], + "missing_left": nodes["missing_go_to_left"].astype(bool).tolist(), + "left": nodes["left"].astype(int).tolist(), + "right": nodes["right"].astype(int).tolist(), + "leaf": nodes["is_leaf"].astype(bool).tolist(), + "value": nodes["value"].astype(float).tolist(), + } + ) + return { + "kind": "classifier" if hasattr(est, "predict_proba") else "regressor", + "baseline": float(np.ravel(est._baseline_prediction)[0]), + "trees": trees, + } + + +def bundle_to_dict(bundle: dict) -> dict: + """The serving-relevant part of a train.py bundle, JSON-ready.""" + return { + "format_version": FORMAT_VERSION, + "station_code": bundle["station_code"], + "model_version": bundle["model_version"], + "trained_at": bundle.get("trained_at"), + "regression_target": bundle.get("regression_target"), + "feature_names": list(bundle["feature_names"]), + "horizons": [int(h) for h in bundle["horizons"]], + "sigma": {str(int(h)): float(v) for h, v in bundle["sigma"].items()}, + "thresholds": {k: float(v) for k, v in bundle["thresholds"].items()}, + # a head skipped at training time (too few positives) stays null + "heads": { + name: (None if est is None else export_head(est)) + for name, est in bundle["heads"].items() + }, + } + + +def write_bundle_json(bundle: dict, path: Path) -> None: + """Atomic write: readers see the old file or the complete new one.""" + path = Path(path) + tmp = path.with_name(path.name + ".tmp") + payload = json.dumps(bundle_to_dict(bundle), separators=(",", ":")) + with open(tmp, "wb") as raw: + # mtime=0: same model, same bytes + with gzip.GzipFile(fileobj=raw, mode="wb", mtime=0) as gz: + gz.write(payload.encode("utf-8")) + os.replace(tmp, path) + + +def json_path_for(joblib_path: Path) -> Path: + return joblib_path.with_name(joblib_path.stem + ".json.gz") + + +def export_models_dir(models_dir: Path, force: bool = False) -> int: + """Export every flood_*.joblib whose .json.gz is missing or older.""" + import joblib + + count = 0 + for path in sorted(Path(models_dir).glob("flood_*.joblib")): + out = json_path_for(path) + if not force and out.exists() and out.stat().st_mtime >= path.stat().st_mtime: + continue + # joblib.load runs pickle code: only ever our own train.py output + write_bundle_json(joblib.load(path), out) + count += 1 + print(f"exported {path.name} -> {out.name}") + return count + + +def flood_config() -> dict: + """Every constant the Rust forecaster and notifier share with Python.""" + from .. import notify + from . import dam as dam_mod + + return { + "thresholds": { + k: [float(w), float(d)] for k, (w, d) in features.THRESHOLDS.items() + }, + "flood_stages": { + code: [ + { + "stage": int(s["stage"]), + "level": float(s["level"]), + "discharge_cms": float(s["discharge_cms"]), + } + for s in stages + ] + for code, stages in features.FLOOD_STAGES.items() + }, + "upstream_leads": { + code: [[up, int(lead)] for up, lead in leads] + for code, leads in features.UPSTREAM_LEADS.items() + }, + "dam_stations": sorted(features.DAM_STATIONS), + "monsoon_months": sorted(features.MONSOON_MONTHS), + "ffill_limit_h": features.FFILL_LIMIT_H, + "basin_anchor": features.BASIN_ANCHOR, + "forward_rain_h": features.FORWARD_RAIN_H, + "horizons": list(predict.DEFAULT_HORIZONS), + "stale_after_h": float(predict.STALE_AFTER_H), + "heuristic_sigma": float(predict.HEURISTIC_SIGMA), + "heuristic_version": predict.HEURISTIC_VERSION, + "latest_window_h": predict.LATEST_WINDOW_H, + "hii_fill_stations": list(data_mod.HII_FILL_STATIONS), + "hii_exact_mirrors": sorted(data_mod._HII_EXACT_MIRRORS), + "hii_min_overlap_hours": data_mod._HII_MIN_OVERLAP_HOURS, + "dam": { + "dam_id": dam_mod.MAE_NGAT_DAM_ID, + "report_hour": dam_mod.REPORT_HOUR, + "ffill_limit_h": dam_mod.FFILL_LIMIT_H, + "serving_days": dam_mod.SERVING_DAYS, + }, + "notify": { + "clear_margin_m": float(notify.CLEAR_MARGIN_M), + "capacity_guard_min_pct": float(notify.CAPACITY_GUARD_MIN_PCT), + "capacity_guard_exempt": sorted(notify.CAPACITY_GUARD_EXEMPT), + "outlook_on": float(notify.OUTLOOK_ON), + "outlook_off": float(notify.OUTLOOK_OFF), + "outlook_horizon": int(notify.OUTLOOK_HORIZON), + "feed_stale_after_h": float(notify.FEED_STALE_AFTER_H), + "station_names": dict(notify.STATION_NAMES), + }, + } + + +def flood_config_text() -> str: + return json.dumps(flood_config(), indent=1, ensure_ascii=False) + "\n" + + +def main(argv: Optional[list] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument( + "models_dir", nargs="?", default=str(predict.DEFAULT_MODELS_DIR) + ) + parser.add_argument("--force", action="store_true", help="re-export everything") + parser.add_argument( + "--config", + action="store_true", + help=f"write {CONFIG_PATH.relative_to(REPO_ROOT)} instead of exporting models", + ) + args = parser.parse_args(argv) + if args.config: + CONFIG_PATH.write_text(flood_config_text(), encoding="utf-8", newline="\n") + print(f"wrote {CONFIG_PATH}") + return 0 + n = export_models_dir(Path(args.models_dir), force=args.force) + print(f"{n} bundle(s) exported") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/ml/predict.py b/src/ml/predict.py index 4cfe4b7..c35aa28 100644 --- a/src/ml/predict.py +++ b/src/ml/predict.py @@ -27,6 +27,7 @@ DEFAULT_HORIZONS: Tuple[int, ...] = (6, 12, 24) STALE_AFTER_H = 6.0 HEURISTIC_SIGMA = 0.3 HEURISTIC_VERSION = "heuristic-v1" +LATEST_WINDOW_H = 336 # history window a live forecast is computed from # Keyed by (path, mtime) so a retrained model (new mtime) invalidates the old entry. _MODEL_CACHE: Dict[Tuple[str, float], dict] = {} @@ -352,22 +353,14 @@ def get_forecasts( return results -def get_latest_forecasts( - db_url: Optional[str] = None, +def forecasts_from_long( + df_long: pd.DataFrame, models_dir: Union[str, Path] = DEFAULT_MODELS_DIR, - hours: int = 336, + now: Optional[Union[datetime.datetime, str]] = None, + rain: Optional[pd.Series] = None, + dam: Optional[pd.DataFrame] = None, ) -> List[dict]: - """Convenience wrapper for web_api: load the latest window from the DB/API and forecast. - - Raises FileNotFoundError when no trained model bundle exists at all, so the - API can 503 instead of serving purely heuristic output as if it were a forecast. - """ - from .data import load_latest - - if not sorted(Path(models_dir).glob("flood_*.joblib")): - raise FileNotFoundError(f"no trained model bundles in {models_dir}") - - df_long = load_latest(db_url=db_url, hours=hours) + """Forecasts from a loaded long window (load_latest's output).""" readings_by_station: Dict[str, List[dict]] = {} if not df_long.empty: for station_code, group in df_long.groupby("station_code"): @@ -380,13 +373,40 @@ def get_latest_forecasts( logger.warning( f"No recent data for station {missing_station}; omitting from forecasts" ) + return get_forecasts( + readings_by_station, models_dir=models_dir, now=now, rain=rain, dam=dam + ) + + +def get_latest_forecasts( + db_url: Optional[str] = None, + models_dir: Union[str, Path] = DEFAULT_MODELS_DIR, + hours: int = LATEST_WINDOW_H, + now: Optional[datetime.datetime] = None, + rain: Optional[pd.Series] = None, +) -> List[dict]: + """Convenience wrapper for web_api: load the latest window from the DB/API and forecast. + + Raises FileNotFoundError when no trained model bundle exists at all, so the + API can 503 instead of serving purely heuristic output as if it were a forecast. + ``now`` and ``rain`` are for reproducible comparisons + (scripts/forecast_parity.py); by default the clock and a live fetch. + """ + from .data import load_latest + + if not sorted(Path(models_dir).glob("flood_*.joblib")): + raise FileNotFoundError(f"no trained model bundles in {models_dir}") + + now = now or datetime.datetime.now() + df_long = load_latest(db_url=db_url, hours=hours, end=now) # Live rain: trailing days + next-48h forecast. On fetch failure pass an # EMPTY series (not None) so rain-trained bundles still find their columns # (as NaN) and serve model output instead of tripping the feature guard. - from .rain import serving_series + if rain is None: + from .rain import serving_series - rain = serving_series() + rain = serving_series() if rain is None: logger.warning("live rain unavailable; rain features will be NaN") rain = pd.Series(dtype=float) @@ -396,7 +416,7 @@ def get_latest_forecasts( from . import dam as dam_mod try: - dam = dam_mod.serving_frame(db_url=db_url) + dam = dam_mod.serving_frame(db_url=db_url, today=now.date()) except Exception as error: logger.warning(f"dam serving frame failed: {error}") dam = None @@ -404,4 +424,6 @@ def get_latest_forecasts( logger.warning("dam state unavailable; dam features will be NaN") dam = pd.DataFrame() - return get_forecasts(readings_by_station, models_dir=models_dir, rain=rain, dam=dam) + return forecasts_from_long( + df_long, models_dir=models_dir, now=now, rain=rain, dam=dam + ) diff --git a/src/ml/train.py b/src/ml/train.py index 0db4ea6..6aae442 100644 --- a/src/ml/train.py +++ b/src/ml/train.py @@ -554,6 +554,10 @@ def train_all( station_results[station] = station_metrics continue joblib.dump(bundle, models_dir / f"flood_{station}.joblib") + # the Rust forecaster's copy (src/ml/export.py) + from .export import write_bundle_json + + write_bundle_json(bundle, models_dir / f"flood_{station}.json.gz") logger.info( f"{station}: trained, {bundle['n_train_rows']} rows, " f"{len(bundle['skipped_heads'])} heads skipped" diff --git a/src/notify.py b/src/notify.py index 4144d33..ffdb1de 100644 --- a/src/notify.py +++ b/src/notify.py @@ -54,6 +54,8 @@ OUTLOOK_ON = 0.50 OUTLOOK_OFF = 0.25 # Below this the outlook is not announced at all (avoid "5 % chance" noise). OUTLOOK_HORIZON = 24 +# No reading this many hours old or newer: the feed is stale. +FEED_STALE_AFTER_H = 3.0 STATION_NAMES: Dict[str, str] = { "P.1": "Nawarat Bridge, Chiang Mai city", @@ -222,7 +224,7 @@ def evaluate( forecasts: Iterable[dict], state: NotificationState, publisher: NtfyPublisher, - stale_after_h: float = 3.0, + stale_after_h: float = FEED_STALE_AFTER_H, now: Optional[datetime.datetime] = None, ) -> List[Notification]: """Compare current readings/forecasts with last-sent state; publish transitions. diff --git a/tests/test_model_export.py b/tests/test_model_export.py new file mode 100644 index 0000000..ba1266f --- /dev/null +++ b/tests/test_model_export.py @@ -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