B4L/Northern-Thailand-Ping-River-Monitor
8
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 15 Wiki Activity

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

Browse Source 
Phase 2a of docs/RUST_REWRITE.md. rust/forecast (ping-forecast) ports the
serving half of src/ml: the HII mirror gap fill, the hourly grid, every
feature (pandas' Kahan rolling sums included), the Mae Ngat dam frame,
HistGradientBoosting tree inference, the heuristic fallback, stale-threshold
handling and the P.1 flood-stage probabilities.

Hand-over from Python (src/ml/export.py):
- models/flood_<station>.json.gz next to each joblib bundle; train_all() writes
  it, retrain.sh promotes it, `python -m src.ml.export` exports existing ones.
- rust/forecast/flood_config.json: thresholds, stages, leads and notify
  constants, compiled in. tests/test_model_export.py fails while it differs
  from the Python source.

Parity (rust/forecast/tests/parity.rs, fixtures from
testdata/gen_forecast_fixtures.py on the production code path): four
scenarios incl. the 2024 flood with RID lagging behind HII. Merged readings,
all features and all levels identical; with all 15 local models 377/384
probabilities bit-identical, the rest within 2 ulp (exp).

Python refactor without behaviour change: predict.forecasts_from_long, and
get_latest_forecasts/load_latest/serving_frame take now/rain for
reproducible comparisons; the HII query orders its rows.
This commit is contained in:
grabowski committed 2026-10-18 14:21:56 +07:00
1 parent 6d4beae65a
commit 5c4f464383
31 files changed
+2889 -29
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

rust/Cargo.lock
Generated
+62 -2
View File
@@ -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"
rust/Cargo.toml
+2 -1
View File
@@ -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.
rust/forecast/Cargo.toml
+20
View File
@@ -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
rust/forecast/flood_config.json
+348
View File
@@ -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"
}
}
}
rust/forecast/src/config.rs
+106
View File
@@ -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<String>,
pub outlook_on: f64,
pub outlook_off: f64,
pub outlook_horizon: u32,
pub feed_stale_after_h: f64,
pub station_names: BTreeMap<String, String>,
}
#[derive(Debug, Deserialize)]
pub struct FloodConfig {
thresholds: BTreeMap<String, (f64, f64)>,
pub flood_stages: BTreeMap<String, Vec<Stage>>,
upstream_leads: BTreeMap<String, Vec<(String, i64)>>,
pub dam_stations: Vec<String>,
pub monsoon_months: Vec<u32>,
pub ffill_limit_h: usize,
pub basin_anchor: String,
pub forward_rain_h: usize,
pub horizons: Vec<u32>,
pub stale_after_h: f64,
pub heuristic_sigma: f64,
pub heuristic_version: String,
pub latest_window_h: i64,
pub hii_fill_stations: Vec<String>,
pub hii_exact_mirrors: Vec<String>,
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<Item = &str> {
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<FloodConfig> = 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"));
}
}
rust/forecast/src/dam.rs
+98
View File
@@ -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<NaiveDateTime, [f64; 3]>,
}
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<f64>; 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<DamFrame> {
let cfg = &config().dam;
let mut by_time: BTreeMap<NaiveDateTime, [f64; 3]> = 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<NaiveDateTime> = (0..n).map(|i| first + Duration::hours(i as i64)).collect();
let cols: Vec<Vec<f64>> = (0..3)
.map(|k| {
let raw: Vec<f64> = 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());
}
}
rust/forecast/src/features.rs
+166
View File
@@ -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<bool> = 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<usize> = 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<f64> = 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()
}
rust/forecast/src/grid.rs
+86
View File
@@ -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<NaiveDateTime>,
pub level: BTreeMap<String, Series>,
pub discharge: BTreeMap<String, Series>,
}
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<NaiveDateTime> = (0..n).map(|i| first + Duration::hours(i as i64)).collect();
let mut level: BTreeMap<String, Series> = BTreeMap::new();
let mut discharge: BTreeMap<String, Series> = 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<usize> {
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<usize> {
self.level.get(station)?.iter().rposition(|x| !x.is_nan())
}
}
rust/forecast/src/hgb.rs
+104
View File
@@ -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<Vec<f64>, D::Error> {
let raw: Vec<Option<f64>> = Deserialize::deserialize(d)?;
Ok(raw
.into_iter()
.map(|t| t.unwrap_or(f64::INFINITY))
.collect())
}
#[derive(Debug, Deserialize)]
pub struct Tree {
feature: Vec<usize>,
#[serde(deserialize_with = "thresholds")]
threshold: Vec<f64>,
missing_left: Vec<bool>,
left: Vec<usize>,
right: Vec<usize>,
leaf: Vec<bool>,
value: Vec<f64>,
}
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<Tree>,
}
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))
}
}
rust/forecast/src/hii.rs
+224
View File
@@ -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<HiiRow> {
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<String> {
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>) -> 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<Row>, Vec<FillInfo>) {
let cfg = config();
let candidates: BTreeSet<String> = 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<Row> = Vec::new();
let mut info = Vec::new();
for (code, mirror) in mirrors {
let base: BTreeMap<NaiveDateTime, f64> = rid
.iter()
.filter(|r| r.station_code == code)
.map(|r| (r.timestamp, r.water_level))
.collect();
let diffs: Vec<f64> = 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<i64>) -> Vec<Row> {
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<i64>, offset: f64) -> Vec<HiiRow> {
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());
}
}
rust/forecast/src/lib.rs
+24
View File
@@ -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_<station>.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};
rust/forecast/src/model.rs
+165
View File
@@ -0,0 +1,165 @@
//! Trained bundles exported by src/ml/export.py (`flood_<station>.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<String>,
pub regression_target: Option<String>,
pub feature_names: Vec<String>,
pub horizons: Vec<u32>,
/// Validation-residual sigma per horizon (keys are the hours as text).
pub sigma: BTreeMap<String, f64>,
pub thresholds: BundleThresholds,
/// `max_<h>` regressors and `warn_<h>` / `danger_<h>` classifiers; a head
/// skipped at training time is null.
pub heads: BTreeMap<String, Option<Head>>,
}
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<Self> {
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<Self> {
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<String, Arc<Bundle>>;
/// `models/flood_<station>.json.gz`, cached by modification time.
pub struct ModelStore {
dir: PathBuf,
loaded: BTreeMap<String, (SystemTime, Arc<Bundle>)>,
}
fn station_of(path: &Path) -> Option<String> {
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<PathBuf>) -> 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<ModelSet> {
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()
);
}
}
}
rust/forecast/src/predict.rs
+245
View File
@@ -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<S: Serializer>(t: &NaiveDateTime, s: S) -> Result<S::Ok, S::Error> {
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<String>,
pub source: String,
pub threshold_warning: f64,
pub threshold_danger: f64,
#[serde(skip_serializing_if = "Option::is_none")]
pub stages: Option<Vec<StageProbability>>,
}
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<f64>,
}
fn heuristic(c: &Current, horizons: &[u32]) -> Vec<Forecast> {
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<Vec<Option<Forecast>>> {
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<Forecast> {
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<Forecast> {
let grid = Grid::from_rows(rows);
grid.level
.keys()
.flat_map(|station| forecast_station(station, &grid, models, now, rain, dam))
.collect()
}
rust/forecast/src/rain.rs
+27
View File
@@ -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<NaiveDateTime, f64>;
/// 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<f64>]) -> Option<f64> {
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<f64> {
times
.iter()
.map(|t| rain.get(t).copied().unwrap_or(f64::NAN))
.collect()
}
rust/forecast/src/rolling.rs
+206
View File
@@ -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<f64>;
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);
}
}
rust/forecast/testdata/fixture_flood2024.json.gz
Vendored
BIN
View File
Binary file not shown.
rust/forecast/testdata/fixture_peak2024.json.gz
Vendored
BIN
View File
Binary file not shown.
rust/forecast/testdata/fixture_recent2026.json.gz
Vendored
BIN
View File
Binary file not shown.
rust/forecast/testdata/fixture_stale2026.json.gz
Vendored
BIN
View File
Binary file not shown.
rust/forecast/testdata/gen_forecast_fixtures.py
Vendored
+310
View File
@@ -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_<scenario>.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()
rust/forecast/testdata/models/flood_P.1.json.gz
Vendored
BIN
View File
Binary file not shown.
rust/forecast/tests/parity.rs
+322
View File
@@ -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<Row> {
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<Row>,
hii: Vec<HiiRow>,
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<Option<f64>> =
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<DamDay> = 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<Row> {
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::<Vec<_>>(), ["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"]);
}
}