grabowski
e3694d7591
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.
86 lines
2.3 KiB
Rust
86 lines
2.3 KiB
Rust
//! 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>,
|
|
}
|