feat(forecast): Rust forecast crate, bit-for-bit with the Python serving path
CI / Test suite (push) Successful in 1m2s
CI / Format & lint (push) Successful in 22s
Rust / fmt, clippy, test, build (push) Successful in 9m53s
Security / Dependency vulnerabilities (push) Successful in 1m52s
Security / Static analysis (push) Successful in 20s
Security / License report (push) Successful in 1m41s
CI / Test suite (push) Successful in 1m2s
CI / Format & lint (push) Successful in 22s
Rust / fmt, clippy, test, build (push) Successful in 9m53s
Security / Dependency vulnerabilities (push) Successful in 1m52s
Security / Static analysis (push) Successful in 20s
Security / License report (push) Successful in 1m41s
Phase 2a of docs/RUST_REWRITE.md. rust/forecast (ping-forecast) ports the serving half of src/ml: the HII mirror gap fill, the hourly grid, every feature (pandas' Kahan rolling sums included), the Mae Ngat dam frame, HistGradientBoosting tree inference, the heuristic fallback, stale-threshold handling and the P.1 flood-stage probabilities. Hand-over from Python (src/ml/export.py): - models/flood_<station>.json.gz next to each joblib bundle; train_all() writes it, retrain.sh promotes it, `python -m src.ml.export` exports existing ones. - rust/forecast/flood_config.json: thresholds, stages, leads and notify constants, compiled in. tests/test_model_export.py fails while it differs from the Python source. Parity (rust/forecast/tests/parity.rs, fixtures from testdata/gen_forecast_fixtures.py on the production code path): four scenarios incl. the 2024 flood with RID lagging behind HII. Merged readings, all features and all levels identical; with all 15 local models 377/384 probabilities bit-identical, the rest within 2 ulp (exp). Python refactor without behaviour change: predict.forecasts_from_long, and get_latest_forecasts/load_latest/serving_frame take now/rain for reproducible comparisons; the HII query orders its rows.
This commit is contained in:
grabowski committed
1 parent
6d4beae65a
commit
5c4f464383
31 files changed
+2889
-29
No files matched your search
@@ -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
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
@@ -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()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
+310
@@ -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()
|
||||
BIN
Binary file not shown.
@@ -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"]);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user