zsfm_tabfm/ensemble/
calibration.rs

1//! Platt scaling (binary) and vector scaling (multiclass) output calibration, fit on
2//! out-of-fold probabilities (see `oof.rs`). The wrapper fits these via `scipy.optimize.minimize`
3//! (L-BFGS-B) on a regularized negative-log-likelihood; this ports the same objective but
4//! optimizes it with an in-house box-constrained coordinate descent (golden-section line search
5//! per coordinate) — expect looser numerical agreement than the always-on ensembling path, which
6//! matches the real optimizer's exact trajectory less closely by construction.
7
8const EPS: f64 = 1e-12;
9
10/// Minimizes `f` over `[lo, hi]` via golden-section search (safe under box constraints, unlike
11/// Brent's method's unbounded bracket growth — used here instead of `scalers::brent_minimize`).
12fn golden_section_bounded(f: impl Fn(f64) -> f64, lo: f64, hi: f64) -> f64 {
13    const GR: f64 = 0.618_033_988_749_895; // 1/phi
14    let (mut a, mut b) = (lo, hi);
15    let mut c = b - GR * (b - a);
16    let mut d = a + GR * (b - a);
17    let mut fc = f(c);
18    let mut fd = f(d);
19    for _ in 0..100 {
20        if (b - a).abs() < 1e-10 {
21            break;
22        }
23        if fc < fd {
24            b = d;
25            d = c;
26            fd = fc;
27            c = b - GR * (b - a);
28            fc = f(c);
29        } else {
30            a = c;
31            c = d;
32            fc = fd;
33            d = a + GR * (b - a);
34            fd = f(d);
35        }
36    }
37    0.5 * (a + b)
38}
39
40pub struct PlattParams {
41    pub a: f64,
42    pub b: f64,
43}
44
45impl PlattParams {
46    /// `p_all`: OOF probabilities `[N][2]`; `y`: true class indices (0 or 1).
47    pub fn fit(p_all: &[Vec<f64>], y: &[usize], lambda: f64) -> Self {
48        let z: Vec<f64> = p_all
49            .iter()
50            .map(|p| ((p[1] + EPS) / (p[0] + EPS)).ln())
51            .collect();
52        let loss = |a: f64, b: f64| -> f64 {
53            let n = z.len() as f64;
54            let mut nll = 0.0;
55            for i in 0..z.len() {
56                let p1 = sigmoid(a * z[i] + b);
57                let p_correct = if y[i] == 1 { p1 } else { 1.0 - p1 };
58                nll -= (p_correct + EPS).ln();
59            }
60            nll / n + lambda * ((a - 1.0).powi(2) + b.powi(2))
61        };
62
63        let mut a = 1.0f64;
64        let mut b = 0.0f64;
65        for _ in 0..20 {
66            a = golden_section_bounded(|av| loss(av, b), 0.8, 1.2);
67            b = golden_section_bounded(|bv| loss(a, bv), -1.0, 1.0);
68        }
69        PlattParams { a, b }
70    }
71
72    pub fn apply(&self, p: &[f64]) -> Vec<f64> {
73        let z = ((p[1] + EPS) / (p[0] + EPS)).ln();
74        let p1 = sigmoid(self.a * z + self.b);
75        vec![1.0 - p1, p1]
76    }
77}
78
79pub struct VectorScalingParams {
80    pub w: Vec<f64>,
81    pub b: Vec<f64>,
82}
83
84impl VectorScalingParams {
85    /// `p_all`: OOF probabilities `[N][K]`; `y`: true class indices.
86    pub fn fit(p_all: &[Vec<f64>], y: &[usize], lambda: f64) -> Self {
87        let k = p_all[0].len();
88        let z: Vec<Vec<f64>> = p_all
89            .iter()
90            .map(|p| p.iter().map(|&v| (v + EPS).ln()).collect())
91            .collect();
92
93        let loss = |w: &[f64], b: &[f64]| -> f64 {
94            let n = z.len() as f64;
95            let mut nll = 0.0;
96            for (i, zi) in z.iter().enumerate() {
97                let logits: Vec<f64> = (0..k).map(|c| w[c] * zi[c] + b[c]).collect();
98                let probs = softmax(&logits);
99                nll -= (probs[y[i]] + EPS).ln();
100            }
101            let reg: f64 = w.iter().map(|&wv| (wv - 1.0).powi(2)).sum::<f64>()
102                + b.iter().map(|&bv| bv.powi(2)).sum::<f64>();
103            nll / n + lambda * reg
104        };
105
106        let mut w = vec![1.0f64; k];
107        let mut b = vec![0.0f64; k];
108        for _ in 0..20 {
109            for c in 0..k {
110                let (w2, b2) = (w.clone(), b.clone());
111                w[c] = golden_section_bounded(
112                    |wc| {
113                        let mut wt = w2.clone();
114                        wt[c] = wc;
115                        loss(&wt, &b2)
116                    },
117                    0.8,
118                    1.2,
119                );
120            }
121            for c in 0..k {
122                let (w2, b2) = (w.clone(), b.clone());
123                b[c] = golden_section_bounded(
124                    |bc| {
125                        let mut bt = b2.clone();
126                        bt[c] = bc;
127                        loss(&w2, &bt)
128                    },
129                    -1.0,
130                    1.0,
131                );
132            }
133        }
134        VectorScalingParams { w, b }
135    }
136
137    pub fn apply(&self, p: &[f64]) -> Vec<f64> {
138        let k = p.len();
139        let logits: Vec<f64> = (0..k)
140            .map(|c| self.w[c] * (p[c] + EPS).ln() + self.b[c])
141            .collect();
142        softmax(&logits)
143    }
144}
145
146fn sigmoid(x: f64) -> f64 {
147    1.0 / (1.0 + (-x).exp())
148}
149
150fn softmax(logits: &[f64]) -> Vec<f64> {
151    let max = logits.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
152    let exp: Vec<f64> = logits.iter().map(|&v| (v - max).exp()).collect();
153    let sum: f64 = exp.iter().sum();
154    exp.into_iter().map(|v| v / sum).collect()
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    #[test]
162    fn test_platt_improves_calibration_of_overconfident_probs() {
163        // Overconfident-but-correct predictions should shrink toward the true labels less
164        // aggressively than raw probs after Platt scaling; here we just check the fit doesn't
165        // diverge and produces valid probabilities.
166        let p_all = vec![
167            vec![0.99, 0.01],
168            vec![0.02, 0.98],
169            vec![0.6, 0.4],
170            vec![0.4, 0.6],
171        ];
172        let y = vec![0, 1, 0, 1];
173        let params = PlattParams::fit(&p_all, &y, 1e-2);
174        for p in &p_all {
175            let out = params.apply(p);
176            assert!((out[0] + out[1] - 1.0).abs() < 1e-6);
177            assert!(out[0] >= 0.0 && out[1] >= 0.0);
178        }
179    }
180
181    #[test]
182    fn test_vector_scaling_valid_probabilities() {
183        let p_all = vec![
184            vec![0.7, 0.2, 0.1],
185            vec![0.1, 0.8, 0.1],
186            vec![0.2, 0.2, 0.6],
187        ];
188        let y = vec![0, 1, 2];
189        let params = VectorScalingParams::fit(&p_all, &y, 1e-2);
190        for p in &p_all {
191            let out = params.apply(p);
192            let sum: f64 = out.iter().sum();
193            assert!((sum - 1.0).abs() < 1e-6);
194        }
195    }
196}
zsfm_tabfm/ensemble/calibration.rs

zsfm_tabfm/ensemble/
calibration.rs
