zsfm_tabfm/ensemble/
aggregate.rs

1//! Temperature-scaled softmax and the classification/regression ensemble-combination paths from
2//! `TabFMClassifier._process_logits` / `TabFMRegressor._combine_predictions`.
3
4/// `TabFMClassifier.softmax` (temperature divides logits *before* the max-subtracted softmax).
5pub fn softmax_temperature(logits: &[f64], temperature: f64) -> Vec<f64> {
6    let scaled: Vec<f64> = logits.iter().map(|&v| v / temperature).collect();
7    let max = scaled.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
8    let exp: Vec<f64> = scaled.iter().map(|&v| (v - max).exp()).collect();
9    let sum: f64 = exp.iter().sum();
10    exp.into_iter().map(|v| v / sum).collect()
11}
12
13/// Undoes a member's class-label shift: that member was fed `(y_true + shift) % n_classes` as
14/// training labels, so its output logit at position `(c + shift) % n_classes` is the prediction
15/// for original class `c`.
16pub fn unshift_logits(logits: &[f64], shift: usize, n_classes: usize) -> Vec<f64> {
17    (0..n_classes)
18        .map(|c| logits[(c + shift) % n_classes])
19        .collect()
20}
21
22fn average_vecs(vecs: &[Vec<f64>]) -> Vec<f64> {
23    let n = vecs.len() as f64;
24    let dim = vecs[0].len();
25    let mut out = vec![0.0; dim];
26    for v in vecs {
27        for (o, x) in out.iter_mut().zip(v.iter()) {
28            *o += x / n;
29        }
30    }
31    out
32}
33
34fn weighted_average_vecs(vecs: &[Vec<f64>], weights: &[f64]) -> Vec<f64> {
35    let dim = vecs[0].len();
36    let mut out = vec![0.0; dim];
37    for (v, &w) in vecs.iter().zip(weights.iter()) {
38        for (o, x) in out.iter_mut().zip(v.iter()) {
39            *o += x * w;
40        }
41    }
42    out
43}
44
45pub enum ClassAggMode<'a> {
46    /// Weighted average of per-member (temperature-softmax'd) probabilities.
47    NnlsWeighted(&'a [f64]),
48    /// Average logits first, then a single temperature-softmax. The wrapper's default
49    /// (`average_logits=True`).
50    AverageLogits,
51    /// Temperature-softmax each member, then plain-average the probabilities.
52    AverageProbs,
53}
54
55/// `_process_logits`: `logits_all` is `[n_estimators][n_classes]`, already un-shifted back to
56/// original class order. Returns final `[n_classes]` probabilities.
57pub fn aggregate_classification(
58    logits_all: &[Vec<f64>],
59    temperature: f64,
60    mode: &ClassAggMode,
61) -> Vec<f64> {
62    match mode {
63        ClassAggMode::NnlsWeighted(weights) => {
64            let probs_all: Vec<Vec<f64>> = logits_all
65                .iter()
66                .map(|l| softmax_temperature(l, temperature))
67                .collect();
68            weighted_average_vecs(&probs_all, weights)
69        }
70        ClassAggMode::AverageLogits => {
71            let avg = average_vecs(logits_all);
72            softmax_temperature(&avg, temperature)
73        }
74        ClassAggMode::AverageProbs => {
75            let probs_all: Vec<Vec<f64>> = logits_all
76                .iter()
77                .map(|l| softmax_temperature(l, temperature))
78                .collect();
79            average_vecs(&probs_all)
80        }
81    }
82}
83
84/// `_combine_predictions`, NNLS-off (default) path: average the still-*scaled* per-member
85/// predictions first; the caller inverse-transforms the single averaged result afterward.
86pub fn average_scaled_predictions(scaled: &[f64]) -> f64 {
87    scaled.iter().sum::<f64>() / scaled.len() as f64
88}
89
90/// `_combine_predictions`, NNLS-on path: each member's prediction has already been
91/// inverse-transformed by the caller; combine via the fitted weights.
92pub fn weighted_unscaled_predictions(unscaled: &[f64], weights: &[f64]) -> f64 {
93    unscaled
94        .iter()
95        .zip(weights.iter())
96        .map(|(p, w)| p * w)
97        .sum()
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn test_softmax_sums_to_one() {
106        let p = softmax_temperature(&[1.0, 2.0, 3.0], 0.9);
107        let sum: f64 = p.iter().sum();
108        assert!((sum - 1.0).abs() < 1e-9);
109    }
110
111    #[test]
112    fn test_unshift_roundtrip() {
113        // member fed y'=(y+shift)%3; its logit position for original class c is (c+shift)%3.
114        let n_classes = 3;
115        let shift = 1;
116        // suppose the model's raw (shifted-space) confidence peaks at shifted-position 0,
117        // meaning it's confident about shifted-class 0 = original class (0 - shift) mod 3 = 2.
118        let shifted_logits = vec![9.0, 0.1, 0.2];
119        let unshifted = unshift_logits(&shifted_logits, shift, n_classes);
120        // unshifted[c] = shifted_logits[(c+shift)%3]; unshifted[2] = shifted_logits[(2+1)%3=0] = 9.0
121        assert_eq!(unshifted[2], 9.0);
122    }
123
124    #[test]
125    fn test_average_logits_vs_average_probs_differ() {
126        let logits_all = vec![vec![10.0, 0.0], vec![0.0, 10.0]];
127        let via_logits = aggregate_classification(&logits_all, 0.9, &ClassAggMode::AverageLogits);
128        let via_probs = aggregate_classification(&logits_all, 0.9, &ClassAggMode::AverageProbs);
129        // average-logits of [10,0] and [0,10] -> [5,5] -> softmax -> [0.5,0.5]
130        assert!((via_logits[0] - 0.5).abs() < 1e-6);
131        // average-probs: softmax([10,0]/0.9)~[~1,~0], softmax([0,10]/0.9)~[~0,~1] -> avg ~[0.5,0.5] too here
132        // (symmetric case coincides) — use an asymmetric case to actually differentiate:
133        let logits_all2 = vec![vec![10.0, 0.0], vec![1.0, 0.0]];
134        let l2 = aggregate_classification(&logits_all2, 0.9, &ClassAggMode::AverageLogits);
135        let p2 = aggregate_classification(&logits_all2, 0.9, &ClassAggMode::AverageProbs);
136        assert!(
137            (l2[0] - p2[0]).abs() > 1e-3,
138            "expected the two paths to diverge on asymmetric input"
139        );
140        let _ = via_probs;
141    }
142}
zsfm_tabfm/ensemble/aggregate.rs

zsfm_tabfm/ensemble/
aggregate.rs
