docs: Rust rewrite assessment + forecast parity spike
CI / Format & lint (push) Successful in 20s
CI / Test suite (push) Successful in 43s
Docs / Validate documentation (push) Successful in 18s

spikes/rust-forecast ports build_features (P.1 path) and HistGradientBoosting
inference to Rust and checks them against the pandas/sklearn output on two
fixtures (2024 record flood, summer 2026): 244,800 feature cells and every
regression prediction bit-identical, classifier heads to the last ulp;
0.25 ms vs 158 ms per station. docs/RUST_REWRITE.md: what a rewrite would
mean (Rust service, Python training), costs, phased plan, Python alternative.
This commit is contained in:
grabowski committed 2026-10-17 21:53:05 +07:00
1 parent 04ebdaa931
commit e3694d7591
10 files changed
+1081

No files matched your search

+318
View File
@@ -0,0 +1,318 @@
//! Port of `src/ml/features.py::build_features` (pandas) to plain Rust.
//!
//! Pandas semantics reproduced: `ffill(limit)`, `shift(k)` (NaN-padded),
//! `rolling(w, min_periods=1)` aggregations that skip NaN and yield NaN for an
//! all-NaN window, NaN propagation through subtraction.
use chrono::{Datelike, NaiveDateTime};
use std::collections::HashMap;
use std::f64::consts::PI;
pub const FFILL_LIMIT_H: usize = 3;
pub const BASIN_ANCHOR: &str = "P.1";
const MONSOON_MONTHS: [u32; 5] = [6, 7, 8, 9, 10];
pub fn upstream_leads(station: &str) -> &'static [(&'static str, i64)] {
match station {
"P.1" => &[
("P.103", 1),
("P.67", 7),
("P.21", 9),
("P.75", 12),
("P.4A", 12),
("P.92", 15),
("P.20", 17),
],
"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.4A" => &[("P.92", 3), ("P.20", 5)],
"P.92" => &[("P.20", 2)],
"P.5" => &[("P.1", 12), ("P.103", 13)],
"P.81" => &[("P.1", 4), ("P.103", 5)],
_ => &[], // spike: P.103's table omitted
}
}
pub type Series = Vec<f64>;
pub struct Grid {
pub times: Vec<NaiveDateTime>,
/// station -> observed (unfilled) hourly series, NaN where missing
pub level: HashMap<String, Series>,
pub discharge: HashMap<String, Series>,
}
fn nan_series(n: usize) -> Series {
vec![f64::NAN; n]
}
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
}
/// pandas `shift(k)`: positive k looks back, negative k looks ahead.
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()
}
fn sub(a: &[f64], b: &[f64]) -> Series {
a.iter().zip(b).map(|(x, y)| x - y).collect()
}
#[derive(Clone, Copy, PartialEq)]
enum Agg {
Mean,
Sum,
Max,
Min,
}
/// pandas `rolling(w, min_periods=1).<agg>()` over the trailing window [i-w+1, i].
///
/// Sum/Mean replicate pandas' `roll_sum`/`roll_mean` (window/aggregations.pyx)
/// operation for operation: an incremental window with separate Kahan
/// compensations for adds and removes, the "run of identical values" shortcut
/// and the sign clamp of the mean. A plain per-window sum differs in the last
/// bits, and HGB thresholds often sit exactly on real data values, so last-bit
/// differences flip tree branches. Max/Min involve no arithmetic: exact as is.
fn rolling(v: &[f64], w: usize, agg: Agg) -> Series {
match agg {
Agg::Sum | Agg::Mean => rolling_kahan(v, w, agg == Agg::Mean),
Agg::Max | Agg::Min => (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 if agg == Agg::Max {
acc.max(x)
} else {
acc.min(x)
}
})
})
.collect(),
}
}
fn rolling_kahan(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 = if v.is_empty() { f64::NAN } else { v[0] };
let mut out = Vec::with_capacity(v.len());
for i in 0..v.len() {
// window [s, e) for row i; previous window was [i-w, i) clipped
if 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 {
if same_run >= nobs {
prev
} else {
let r = sum / nobs as f64;
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
}
/// Hours since the last observed value (0 when observed; NaN before the first).
fn hours_since_observed(mask: &[bool]) -> Series {
let mut last: Option<usize> = None;
mask.iter()
.enumerate()
.map(|(i, &m)| {
if m {
last = Some(i);
}
last.map_or(f64::NAN, |l| (i - l) as f64)
})
.collect()
}
/// Feature columns by name, in pandas' construction order.
pub fn build_features(grid: &Grid, station: &str, rain: Option<&[f64]>) -> Vec<(String, Series)> {
let n = grid.times.len();
let filled = |s: &str| -> Series {
grid.level
.get(s)
.map_or_else(|| nan_series(n), |v| ffill(v, FFILL_LIMIT_H))
};
let mut cols: Vec<(String, Series)> = Vec::new();
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(&level, 6, Agg::Mean));
put("roll_mean_24".into(), rolling(&level, 24, Agg::Mean));
put("roll_max_6".into(), rolling(&level, 6, Agg::Max));
put("roll_max_24".into(), rolling(&level, 24, Agg::Max));
put("roll_max_72".into(), rolling(&level, 72, Agg::Max));
put("roll_min_24".into(), rolling(&level, 24, Agg::Min));
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)),
);
let age = hours_since_observed(&mask);
put(
"obs_age_h".into(),
age.iter()
.map(|&a| {
if a <= FFILL_LIMIT_H as f64 {
a
} else {
f64::NAN
}
})
.collect(),
);
let mask_f: Series = mask.iter().map(|&m| if m { 1.0 } else { 0.0 }).collect();
put("cov_24h".into(), rolling(&mask_f, 24, Agg::Mean));
for &(up, lead) in upstream_leads(station) {
let u = filled(up);
let u_rise_6 = sub(&u, &shift(&u, 6));
let u_rollmax_24 = rolling(&u, 24, Agg::Max);
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 != BASIN_ANCHOR {
let p1 = filled(BASIN_ANCHOR);
put("P1_level".into(), p1.clone());
put("P1_rollmax_24".into(), rolling(&p1, 24, Agg::Max));
put("P1_rise_24".into(), sub(&p1, &shift(&p1, 24)));
}
let doy: Vec<f64> = grid.times.iter().map(|t| t.ordinal() as f64).collect();
put(
"doy_sin".into(),
doy.iter().map(|d| (2.0 * PI * d / 365.25).sin()).collect(),
);
put(
"doy_cos".into(),
doy.iter().map(|d| (2.0 * PI * d / 365.25).cos()).collect(),
);
put(
"is_monsoon".into(),
grid.times
.iter()
.map(|t| {
if MONSOON_MONTHS.contains(&t.month()) {
1.0
} else {
0.0
}
})
.collect(),
);
if let Some(r) = rain {
put("rain_6h".into(), rolling(r, 6, Agg::Sum));
put("rain_24h".into(), rolling(r, 24, Agg::Sum));
put("rain_72h".into(), rolling(r, 72, Agg::Sum));
// forward sum over (t, t+24]
let mut ahead = shift(r, -1);
ahead.reverse();
let mut fc = rolling(&ahead, 24, Agg::Sum);
fc.reverse();
put("rain_fc24".into(), fc);
}
cols
}
+85
View File
@@ -0,0 +1,85 @@
//! Inference for scikit-learn HistGradientBoosting models exported as flat
//! tree arrays (see export/export_fixture.py). Mirrors sklearn's
//! `_predictor.pyx`: NaN follows `missing_go_to_left`, otherwise
//! `x <= threshold` goes left; raw = baseline + sum of leaf values, summed in
//! tree order (same order as sklearn, so results are bit-identical).
use serde::Deserialize;
use std::collections::BTreeMap;
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(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]
}
}
#[derive(Deserialize, PartialEq, Clone, Copy)]
#[serde(rename_all = "lowercase")]
pub enum Kind {
Regressor,
Classifier,
}
#[derive(Deserialize)]
pub struct Head {
pub kind: Kind,
baseline: f64,
trees: Vec<Tree>,
}
impl Head {
/// Regressor: the prediction. Classifier: P(class 1), like predict_proba[:, 1].
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()),
}
}
}
#[derive(Deserialize)]
pub struct Bundle {
pub station_code: String,
pub model_version: String,
pub regression_target: Option<String>,
pub feature_names: Vec<String>,
pub horizons: Vec<u32>,
pub sigma: BTreeMap<String, f64>,
pub thresholds: BTreeMap<String, f64>,
pub heads: BTreeMap<String, Head>,
}
+211
View File
@@ -0,0 +1,211 @@
//! Parity + timing check: Rust features and HGB inference vs the Python output
//! recorded in data/fixture_*.json.
mod features;
mod hgb;
use chrono::NaiveDateTime;
use serde::Deserialize;
use std::collections::HashMap;
use std::time::Instant;
#[derive(Deserialize)]
struct StationObs {
water_level: Vec<Option<f64>>,
discharge: Vec<Option<f64>>,
}
#[derive(Deserialize)]
struct Fixture {
station: String,
timestamps: Vec<String>,
observed: HashMap<String, StationObs>,
rain: Vec<Option<f64>>,
feature_names: Vec<String>,
features: Vec<Vec<Option<f64>>>,
predictions: HashMap<String, Vec<f64>>,
}
fn nanify(v: &[Option<f64>]) -> Vec<f64> {
v.iter().map(|x| x.unwrap_or(f64::NAN)).collect()
}
fn same(a: f64, b: f64, tol: f64) -> bool {
(a.is_nan() && b.is_nan()) || (a - b).abs() <= tol * (1.0 + b.abs())
}
fn matrix(cols: &[(String, Vec<f64>)], names: &[String]) -> Vec<Vec<f64>> {
let by_name: HashMap<&str, &Vec<f64>> = cols.iter().map(|(k, v)| (k.as_str(), v)).collect();
let picked: Vec<&Vec<f64>> = names
.iter()
.map(|n| {
*by_name
.get(n.as_str())
.unwrap_or_else(|| panic!("missing feature {n}"))
})
.collect();
let rows = picked[0].len();
(0..rows)
.map(|i| picked.iter().map(|c| c[i]).collect())
.collect()
}
fn main() {
let data = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("data");
let t = Instant::now();
let bundle: hgb::Bundle =
serde_json::from_slice(&std::fs::read(data.join("model_P.1.json")).unwrap()).unwrap();
println!(
"model {} ({}): {} heads, {} features, target={:?}, horizons {:?}, sigma {:?}, thresholds {:?}; loaded in {:?}",
bundle.station_code,
bundle.model_version,
bundle.heads.len(),
bundle.feature_names.len(),
bundle.regression_target,
bundle.horizons,
bundle.sigma,
bundle.thresholds,
t.elapsed()
);
for name in ["flood2024", "recent2026"] {
let fx: Fixture = serde_json::from_slice(
&std::fs::read(data.join(format!("fixture_{name}.json"))).unwrap(),
)
.unwrap();
assert_eq!(fx.feature_names, bundle.feature_names);
let grid = features::Grid {
times: fx
.timestamps
.iter()
.map(|s| NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S").unwrap())
.collect(),
level: fx
.observed
.iter()
.map(|(k, v)| (k.clone(), nanify(&v.water_level)))
.collect(),
discharge: fx
.observed
.iter()
.map(|(k, v)| (k.clone(), nanify(&v.discharge)))
.collect(),
};
let rain = nanify(&fx.rain);
// 1. feature parity
let t = Instant::now();
let cols = features::build_features(&grid, &fx.station, Some(&rain));
let feat_time = t.elapsed();
let rust_x = matrix(&cols, &bundle.feature_names);
let py_x: Vec<Vec<f64>> = fx.features.iter().map(|r| nanify(r)).collect();
let (mut cells, mut exact, mut close, mut worst) =
(0usize, 0usize, 0usize, (0.0f64, String::new()));
for (r, p) in rust_x.iter().zip(&py_x) {
for (j, (a, b)) in r.iter().zip(p).enumerate() {
cells += 1;
if a.to_bits() == b.to_bits() || (a.is_nan() && b.is_nan()) {
exact += 1;
} else if same(*a, *b, 1e-9) {
close += 1;
let d = (a - b).abs();
if d > worst.0 {
worst = (d, bundle.feature_names[j].clone());
}
} else {
panic!(
"{name}: feature {} differs: rust {a} vs python {b}",
bundle.feature_names[j]
);
}
}
}
println!(
"\n[{name}] {} rows x {} features: {exact} bit-identical, {close} within 1e-9 (max |diff| {:.1e} in {}), 0 mismatches; Rust feature build {:?}",
rust_x.len(),
bundle.feature_names.len(),
worst.0,
if worst.1.is_empty() { "-" } else { &worst.1 },
feat_time
);
let _ = cells;
// 2. inference parity, on Python's features and on Rust's own features
let t = Instant::now();
let mut report = Vec::new();
for (head_name, head) in &bundle.heads {
let want = &fx.predictions[head_name];
let (mut max_py, mut max_rs, mut ident) = (0.0f64, 0.0f64, 0usize);
for i in 0..rust_x.len() {
let on_py = head.predict(&py_x[i]);
let on_rs = head.predict(&rust_x[i]);
if on_py.to_bits() == want[i].to_bits() {
ident += 1;
}
max_py = max_py.max((on_py - want[i]).abs());
max_rs = max_rs.max((on_rs - want[i]).abs());
}
report.push(format!(
" {head_name:<10} bit-identical {ident}/{} max|diff| on py-features {max_py:.1e}, on rust-features {max_rs:.1e}",
rust_x.len()
));
}
let infer_time = t.elapsed();
println!("{}", report.join("\n"));
println!(
" {} heads x {} rows x 2 inference passes in {:?}",
bundle.heads.len(),
rust_x.len(),
infer_time
);
// 3. serving-shaped work: 14-day window, forecast for the last hour
let tail = 336.min(grid.times.len());
let start = grid.times.len() - tail;
let window = features::Grid {
times: grid.times[start..].to_vec(),
level: grid
.level
.iter()
.map(|(k, v)| (k.clone(), v[start..].to_vec()))
.collect(),
discharge: grid
.discharge
.iter()
.map(|(k, v)| (k.clone(), v[start..].to_vec()))
.collect(),
};
let w_rain = rain[start..].to_vec();
let reps = 200;
let t = Instant::now();
let mut last = Vec::new();
for _ in 0..reps {
let cols = features::build_features(&window, &fx.station, Some(&w_rain));
let x = matrix(&cols, &bundle.feature_names);
let row = x.last().unwrap();
last = bundle
.heads
.iter()
.map(|(k, h)| (k.clone(), h.predict(row)))
.collect();
}
let per = t.elapsed() / reps;
let level_now = window.level[&fx.station]
.iter()
.rev()
.find(|x| !x.is_nan())
.copied()
.unwrap();
let rise24 = last.iter().find(|(k, _)| k == "max_24").unwrap().1;
let warn24 = last.iter().find(|(k, _)| k == "warn_24").unwrap().1;
println!(
" serving: 336 h window -> features + {} heads for the last hour: {:?} per station; at {} level {:.2} m, predicted 24 h max {:.2} m, warn_24 head p={:.3}",
bundle.heads.len(),
per,
window.times.last().unwrap(),
level_now,
(rise24 + level_now).max(level_now),
warn24
);
}
}