diff --git a/crates/ml-alpha/src/calibration.rs b/crates/ml-alpha/src/calibration.rs new file mode 100644 index 000000000..b24b57286 --- /dev/null +++ b/crates/ml-alpha/src/calibration.rs @@ -0,0 +1,202 @@ +//! Phase 1d.0 — post-hoc probability calibration. +//! +//! Platt scaling: logistic regression `P_calib = σ(a·logit + b)` fit on a +//! held-out calibration set. Preserves rank order (AUC unchanged); fixes +//! the sigmoid threshold and prediction-distribution shape. +//! +//! Isotonic regression: monotonic non-parametric calibrator via PAV +//! (pool-adjacent-violators). Better when miscalibration isn't sigmoidal. + +/// A fitted calibrator that maps raw logits → calibrated probabilities. +pub trait Calibrator { + fn transform(&self, logits: &[f32]) -> Vec; +} + +/// Platt scaling — fit a logistic regression on (logit, label) pairs via +/// gradient descent on BCE loss. Two parameters: `a` (slope) and `b` (intercept). +/// +/// `P_calibrated = σ(a · logit + b)` +#[derive(Debug, Clone)] +pub struct PlattScaler { + pub a: f32, + pub b: f32, +} + +impl PlattScaler { + /// Fit `(a, b)` by minimising mean BCE over the calibration set. + /// Returns the fitted scaler. `max_iters` is the gradient-descent budget; + /// `lr` is the learning rate. + pub fn fit( + logits: &[f32], + labels: &[f32], + max_iters: usize, + lr: f32, + ) -> Result { + if logits.len() != labels.len() { + return Err("Platt: logits and labels length mismatch"); + } + if logits.is_empty() { + return Err("Platt: empty calibration set"); + } + let n = logits.len() as f32; + let mut a = 1.0_f32; + let mut b = 0.0_f32; + for _ in 0..max_iters { + let mut grad_a = 0.0_f32; + let mut grad_b = 0.0_f32; + for (&l, &y) in logits.iter().zip(labels.iter()) { + let z = (a * l + b).clamp(-50.0, 50.0); + let p = 1.0 / (1.0 + (-z).exp()); + let err = p - y; + grad_a += err * l; + grad_b += err; + } + a -= lr * grad_a / n; + b -= lr * grad_b / n; + } + Ok(Self { a, b }) + } +} + +impl Calibrator for PlattScaler { + fn transform(&self, logits: &[f32]) -> Vec { + logits + .iter() + .map(|&l| { + let z = (self.a * l + self.b).clamp(-50.0, 50.0); + 1.0 / (1.0 + (-z).exp()) + }) + .collect() + } +} + +/// Isotonic regression calibrator (pool-adjacent-violators algorithm). +/// +/// Fits a monotone-non-decreasing step function from sorted-logit positions +/// to observed-positive-rate. For new inputs, linear interpolation between +/// the nearest two cut points. +#[derive(Debug, Clone)] +pub struct IsotonicCalibrator { + cuts: Vec, + values: Vec, +} + +impl IsotonicCalibrator { + pub fn fit(logits: &[f32], labels: &[f32]) -> Result { + if logits.len() != labels.len() { + return Err("isotonic: logit/label length mismatch"); + } + if logits.is_empty() { + return Err("isotonic: empty calibration set"); + } + let mut paired: Vec<(f32, f32)> = + logits.iter().copied().zip(labels.iter().copied()).collect(); + paired.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal)); + + let n = paired.len(); + let mut cuts: Vec = paired.iter().map(|p| p.0).collect(); + let mut values: Vec = paired.iter().map(|p| p.1).collect(); + let mut weights: Vec = vec![1.0; n]; + + let mut i = 0; + while i + 1 < values.len() { + if values[i] > values[i + 1] { + let new_val = (values[i] * weights[i] + values[i + 1] * weights[i + 1]) + / (weights[i] + weights[i + 1]); + let new_w = weights[i] + weights[i + 1]; + values[i] = new_val; + weights[i] = new_w; + values.remove(i + 1); + weights.remove(i + 1); + cuts.remove(i + 1); + if i > 0 { + i -= 1; + } + } else { + i += 1; + } + } + Ok(Self { cuts, values }) + } +} + +impl Calibrator for IsotonicCalibrator { + fn transform(&self, logits: &[f32]) -> Vec { + logits + .iter() + .map(|&l| { + if self.cuts.is_empty() { + return 0.5; + } + if l <= self.cuts[0] { + return self.values[0]; + } + if l >= *self.cuts.last().unwrap() { + return *self.values.last().unwrap(); + } + let idx = self.cuts.partition_point(|&c| c <= l).saturating_sub(1); + if idx + 1 >= self.cuts.len() { + return self.values[idx]; + } + let lo = self.cuts[idx]; + let hi = self.cuts[idx + 1]; + let lov = self.values[idx]; + let hiv = self.values[idx + 1]; + let frac = if hi > lo { (l - lo) / (hi - lo) } else { 0.0 }; + lov + frac * (hiv - lov) + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_module_compiles() { + let _ = std::any::type_name::(); + } + + #[test] + fn test_platt_perfect_separable() { + let logits: Vec = (0..200) + .map(|i| if i % 2 == 0 { 2.0 } else { -2.0 }) + .collect(); + let labels: Vec = (0..200) + .map(|i| if i % 2 == 0 { 1.0 } else { 0.0 }) + .collect(); + // 2000 iters @ lr=1e-2: plenty of budget for the 2-param convex problem. + let cal = PlattScaler::fit(&logits, &labels, 2000, 1e-2).expect("fit"); + let probs = cal.transform(&logits); + // Invariant: positive samples have HIGHER probability than negative ones. + // Don't lock a specific threshold (per pearl_tests_must_prove_not_lock). + let pos_min = probs.iter().enumerate() + .filter(|(i, _)| i % 2 == 0) + .map(|(_, p)| *p) + .fold(f32::INFINITY, f32::min); + let neg_max = probs.iter().enumerate() + .filter(|(i, _)| i % 2 == 1) + .map(|(_, p)| *p) + .fold(f32::NEG_INFINITY, f32::max); + assert!(pos_min > neg_max, + "expected clean separation: min(pos_probs)={pos_min} > max(neg_probs)={neg_max}"); + assert!(pos_min > 0.5, "positive class mean prob must exceed 0.5, got pos_min={pos_min}"); + assert!(neg_max < 0.5, "negative class mean prob must drop below 0.5, got neg_max={neg_max}"); + } + + #[test] + fn test_isotonic_monotone_output() { + let logits: Vec = (0..100).map(|i| i as f32 * 0.1).collect(); + let labels: Vec = (0..100) + .map(|i| if i >= 50 { 1.0 } else { 0.0 }) + .collect(); + let cal = IsotonicCalibrator::fit(&logits, &labels).expect("fit"); + let probs = cal.transform(&logits); + for w in probs.windows(2) { + assert!(w[1] >= w[0] - 1e-6, "monotonicity violated: {} → {}", w[0], w[1]); + } + assert!(probs.last().unwrap() >= &0.5); + assert!(probs.first().unwrap() <= &0.5); + } +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index 625691149..d2137bec2 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -51,6 +51,7 @@ pub mod mlp; pub mod training; pub mod eval; pub mod metrics_detail; +pub mod calibration; pub use fxcache_reader::{FxCacheReader, FxCacheRecord, FxCacheMetadata}; pub use purged_split::{PurgedSplit, SplitIndices};