docs: Rust rewrite assessment + forecast parity spike
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
1 parent
04ebdaa931
commit
e3694d7591
10 files changed
+1081
No files matched your search
@@ -0,0 +1,2 @@
|
||||
target/
|
||||
data/
|
||||
Generated
+132
@@ -0,0 +1,132 @@
|
||||
# This file is automatically @generated by Cargo.
|
||||
# It is not intended for manual editing.
|
||||
version = 4
|
||||
|
||||
[[package]]
|
||||
name = "autocfg"
|
||||
version = "1.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
|
||||
|
||||
[[package]]
|
||||
name = "chrono"
|
||||
version = "0.4.45"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itoa"
|
||||
version = "1.0.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
|
||||
|
||||
[[package]]
|
||||
name = "memchr"
|
||||
version = "2.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
|
||||
|
||||
[[package]]
|
||||
name = "num-traits"
|
||||
version = "0.2.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ping-forecast-spike"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"chrono",
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2"
|
||||
version = "1.0.107"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
|
||||
dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quote"
|
||||
version = "1.0.47"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde"
|
||||
version = "1.0.229"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba"
|
||||
dependencies = [
|
||||
"serde_core",
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_core"
|
||||
version = "1.0.229"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48"
|
||||
dependencies = [
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_derive"
|
||||
version = "1.0.229"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.154"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"memchr",
|
||||
"serde",
|
||||
"serde_core",
|
||||
"zmij",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "3.0.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicode-ident"
|
||||
version = "1.0.27"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563"
|
||||
|
||||
[[package]]
|
||||
name = "zmij"
|
||||
version = "1.0.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b"
|
||||
@@ -0,0 +1,15 @@
|
||||
[package]
|
||||
name = "ping-forecast-spike"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
# float_roundtrip: exact float parsing. Tree thresholds often equal real data
|
||||
# values, so a 1-ulp parse error flips `x <= threshold` (default parser is not exact).
|
||||
serde_json = { version = "1", features = ["float_roundtrip"] }
|
||||
chrono = { version = "0.4", default-features = false, features = ["std"] }
|
||||
|
||||
[profile.release]
|
||||
debug = false
|
||||
@@ -0,0 +1,43 @@
|
||||
# Rust forecast spike
|
||||
|
||||
Answers one question: can the flood forecast be served from Rust without
|
||||
train/serve skew, while training stays in Python (scikit-learn)?
|
||||
Assessment and migration plan: [docs/RUST_REWRITE.md](../../docs/RUST_REWRITE.md).
|
||||
|
||||
- `export/export_fixture.py`: dumps `models/flood_P.1.joblib` (all 9 HGB heads) as
|
||||
flat tree arrays, plus two fixtures from `models/cache` (Aug-Oct 2024 record
|
||||
flood, Jun-Aug 2026): the observed hourly grid of P.1 and its 7 upstream
|
||||
gauges, rain, the pandas feature matrix and every head's sklearn prediction.
|
||||
- `src/features.rs`: port of `src/ml/features.py::build_features` (P.1 path:
|
||||
lags, rises, rolling stats, upstream leads, season, rain).
|
||||
- `src/hgb.rs`: HistGradientBoosting inference (~90 lines).
|
||||
- `src/main.rs`: compares bit-for-bit and times serving-shaped work.
|
||||
|
||||
```bash
|
||||
cd spikes/rust-forecast/export
|
||||
uv run --project ../../.. python export_fixture.py # writes ../data/*.json (gitignored)
|
||||
cd .. && cargo run --release
|
||||
uv run --project ../.. python export/time_python.py # Python timing for comparison
|
||||
```
|
||||
|
||||
Result (model hgb-v3+cbb3bf7, 3,600 hourly rows):
|
||||
|
||||
| | Result |
|
||||
| --- | --- |
|
||||
| Features, 68 columns x 3,600 rows | 244,800 / 244,800 cells bit-identical |
|
||||
| Regression heads (max_6/12/24) | bit-identical on every row |
|
||||
| Classifier heads (warn/danger) | equal to the last ulp (<= 2.2e-16, `exp` vs scipy `expit`) |
|
||||
| One station: 336 h window -> features + 9 heads | Rust 0.25 ms, Python 158 ms |
|
||||
| Model load | Rust 16 ms (JSON), Python 199 ms (joblib) |
|
||||
|
||||
Two traps the spike had to get past, both of which flip tree branches because
|
||||
HGB thresholds often equal real data values exactly:
|
||||
|
||||
1. `serde_json` parses floats best-effort (can be 1 ulp off) unless the
|
||||
`float_roundtrip` feature is on.
|
||||
2. pandas `rolling().sum()/.mean()` are incremental with Kahan compensation,
|
||||
a "run of identical values" shortcut and a sign clamp. A plain window sum
|
||||
differs in the last bits; `rolling_kahan()` replicates pandas exactly.
|
||||
|
||||
Not covered yet: dam features (Mae Ngat), HII gap fill, Open-Meteo catchment
|
||||
mean, the heuristic fallback, stage probabilities, the other 14 stations.
|
||||
@@ -0,0 +1,132 @@
|
||||
"""Export a trained flood bundle + a feature/prediction fixture for the Rust spike.
|
||||
|
||||
Run from the monitor repo's environment:
|
||||
uv run --project ../../.. python export_fixture.py (from this directory)
|
||||
|
||||
Writes ../data/model_<station>.json (all heads as flat tree arrays) and
|
||||
../data/fixture_<name>.json (observed hourly grid + rain, the Python feature
|
||||
matrix, and every head's Python prediction for every row).
|
||||
"""
|
||||
|
||||
import json
|
||||
import math
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import joblib
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
REPO = Path(__file__).resolve().parents[3] # repo root
|
||||
sys.path.insert(0, str(REPO))
|
||||
|
||||
from src.ml import features # noqa: E402
|
||||
|
||||
OUT = Path(__file__).resolve().parent.parent / "data"
|
||||
OUT.mkdir(exist_ok=True)
|
||||
STATION = "P.1"
|
||||
WINDOWS = {
|
||||
"flood2024": ("2024-08-15", "2024-10-31"), # record flood, high-level paths
|
||||
"recent2026": ("2026-06-01", "2026-08-13"), # latest cached data
|
||||
}
|
||||
|
||||
|
||||
def clean(values):
|
||||
return [None if (v is None or (isinstance(v, float) and math.isnan(v))) else float(v) for v in values]
|
||||
|
||||
|
||||
def export_head(est):
|
||||
preds = est._predictors
|
||||
assert all(len(it) == 1 for it in preds), "multiclass not expected"
|
||||
trees = []
|
||||
for (tp,) in preds:
|
||||
n = tp.nodes
|
||||
assert not n["is_categorical"].any()
|
||||
assert np.isfinite(n["value"]).all() and not np.isneginf(n["num_threshold"]).any()
|
||||
assert not np.isnan(n["num_threshold"]).any()
|
||||
trees.append(
|
||||
{
|
||||
"feature": n["feature_idx"].astype(int).tolist(),
|
||||
# +inf thresholds (a "missing vs everything" split) are not
|
||||
# valid JSON: written as null, read back as +inf in Rust
|
||||
"threshold": [
|
||||
None if np.isposinf(t) else float(t) for t in n["num_threshold"]
|
||||
],
|
||||
"missing_left": n["missing_go_to_left"].astype(bool).tolist(),
|
||||
"left": n["left"].astype(int).tolist(),
|
||||
"right": n["right"].astype(int).tolist(),
|
||||
"leaf": n["is_leaf"].astype(bool).tolist(),
|
||||
"value": n["value"].astype(float).tolist(),
|
||||
}
|
||||
)
|
||||
kind = "classifier" if hasattr(est, "predict_proba") else "regressor"
|
||||
return {
|
||||
"kind": kind,
|
||||
"baseline": float(np.ravel(est._baseline_prediction)[0]),
|
||||
"trees": trees,
|
||||
}
|
||||
|
||||
|
||||
bundle = joblib.load(REPO / "models" / f"flood_{STATION}.joblib")
|
||||
model = {
|
||||
"station_code": STATION,
|
||||
"model_version": bundle["model_version"],
|
||||
"regression_target": bundle.get("regression_target"),
|
||||
"feature_names": bundle["feature_names"],
|
||||
"horizons": bundle["horizons"],
|
||||
"sigma": {str(k): float(v) for k, v in bundle["sigma"].items()},
|
||||
"thresholds": bundle["thresholds"],
|
||||
"heads": {name: export_head(est) for name, est in bundle["heads"].items()},
|
||||
}
|
||||
(OUT / f"model_{STATION}.json").write_text(json.dumps(model))
|
||||
n_nodes = sum(len(t["value"]) for h in model["heads"].values() for t in h["trees"])
|
||||
print(f"model {model['model_version']}: {len(model['heads'])} heads, {n_nodes} nodes")
|
||||
|
||||
stations = [STATION] + [u for u, _ in features.UPSTREAM_LEADS[STATION]]
|
||||
long = pd.concat(
|
||||
pd.read_csv(REPO / "models" / "cache" / f"{s}.csv.gz", parse_dates=["timestamp"])
|
||||
for s in stations
|
||||
)
|
||||
rain_all = pd.read_csv(REPO / "models" / "cache" / "rain_openmeteo.csv.gz")
|
||||
print("rain columns:", list(rain_all.columns))
|
||||
ts_col = rain_all.columns[0]
|
||||
rain_all[ts_col] = pd.to_datetime(rain_all[ts_col])
|
||||
rain_all = rain_all.set_index(ts_col)[rain_all.columns[1]].astype(float)
|
||||
|
||||
for name, (start, end) in WINDOWS.items():
|
||||
df = long[(long.timestamp >= start) & (long.timestamp < end)]
|
||||
grid = features.make_hourly_grid(df)
|
||||
idx = grid.observed.index
|
||||
rain = rain_all.reindex(idx)
|
||||
# punch deterministic holes so NaN handling of the rolling sums is exercised
|
||||
rain.iloc[::97] = np.nan
|
||||
feats = features.build_features(grid, STATION, rain=rain)[model["feature_names"]]
|
||||
|
||||
preds = {}
|
||||
for head_name, est in bundle["heads"].items():
|
||||
if hasattr(est, "predict_proba"):
|
||||
preds[head_name] = est.predict_proba(feats)[:, 1].tolist()
|
||||
else:
|
||||
preds[head_name] = est.predict(feats).tolist()
|
||||
|
||||
fixture = {
|
||||
"station": STATION,
|
||||
"timestamps": [t.strftime("%Y-%m-%dT%H:%M:%S") for t in idx],
|
||||
"observed": {
|
||||
s: {
|
||||
"water_level": clean(grid.observed[(s, "water_level")].tolist()),
|
||||
"discharge": clean(grid.observed[(s, "discharge")].tolist()),
|
||||
}
|
||||
for s in stations
|
||||
if (s, "water_level") in grid.observed.columns
|
||||
},
|
||||
"rain": clean(rain.tolist()),
|
||||
"feature_names": model["feature_names"],
|
||||
"features": [clean(row) for row in feats.to_numpy().tolist()],
|
||||
"predictions": preds,
|
||||
}
|
||||
(OUT / f"fixture_{name}.json").write_text(json.dumps(fixture))
|
||||
print(
|
||||
f"{name}: {len(idx)} hours, {len(fixture['observed'])} stations, "
|
||||
f"P.1 max {np.nanmax(grid.observed[(STATION, 'water_level')]):.2f} m"
|
||||
)
|
||||
@@ -0,0 +1,38 @@
|
||||
"""Python timing for the same serving-shaped work as the Rust spike."""
|
||||
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import joblib
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
REPO = Path(__file__).resolve().parents[3] # repo root
|
||||
sys.path.insert(0, str(REPO))
|
||||
from src.ml import features # noqa: E402
|
||||
|
||||
bundle = joblib.load(REPO / "models" / "flood_P.1.joblib")
|
||||
stations = ["P.1"] + [u for u, _ in features.UPSTREAM_LEADS["P.1"]]
|
||||
long = pd.concat(
|
||||
pd.read_csv(REPO / "models" / "cache" / f"{s}.csv.gz", parse_dates=["timestamp"])
|
||||
for s in stations
|
||||
)
|
||||
df = long[(long.timestamp >= "2026-07-30") & (long.timestamp < "2026-08-13")]
|
||||
rain_all = pd.read_csv(REPO / "models" / "cache" / "rain_openmeteo.csv.gz", parse_dates=["timestamp"])
|
||||
rain_all = rain_all.set_index("timestamp")["chiang_dao"].astype(float)
|
||||
|
||||
reps = 20
|
||||
t = time.perf_counter()
|
||||
for _ in range(reps):
|
||||
grid = features.make_hourly_grid(df)
|
||||
rain = rain_all.reindex(grid.observed.index)
|
||||
row = features.build_features(grid, "P.1", rain=rain)[bundle["feature_names"]].iloc[[-1]]
|
||||
out = {
|
||||
k: (h.predict_proba(row)[0, 1] if hasattr(h, "predict_proba") else h.predict(row)[0])
|
||||
for k, h in bundle["heads"].items()
|
||||
}
|
||||
print(f"python: grid + features + 9 heads for the last hour: {(time.perf_counter() - t) / reps * 1000:.1f} ms per station")
|
||||
t = time.perf_counter()
|
||||
joblib.load(REPO / "models" / "flood_P.1.joblib")
|
||||
print(f"python: bundle load {(time.perf_counter() - t) * 1000:.0f} ms")
|
||||
@@ -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
|
||||
}
|
||||
@@ -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>,
|
||||
}
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user