zsfm_tabfm/ensemble/
oof.rs

1//! Out-of-fold (OOF) prediction generation, used to fit calibration/NNLS. Fold splitting uses
2//! our own seeded shuffle (`PyRandom`) rather than sklearn's `KFold(shuffle=True)`, which draws
3//! from NumPy's legacy `RandomState` — a related but distinct RNG family we haven't ported (see
4//! the plan's scoping note). Statistically equivalent, not bit-identical.
5
6use anyhow::Result;
7use serde_json::Value;
8
9use crate::infer::BurnTabFMModel;
10
11use super::config_gen::MemberConfig;
12use super::orchestrate::{run_members_classification, run_members_regression, EnsembleParams};
13use super::pyrandom::PyRandom;
14
15pub struct KFoldSplit {
16    pub train_idx: Vec<usize>,
17    pub val_idx: Vec<usize>,
18}
19
20/// A shuffled K-fold split of `0..n`.
21pub fn kfold_splits(n: usize, k: usize, random_state: u64) -> Vec<KFoldSplit> {
22    let mut indices: Vec<usize> = (0..n).collect();
23    PyRandom::new(random_state).shuffle(&mut indices);
24
25    let k = k.min(n).max(1);
26    let base = n / k;
27    let remainder = n % k;
28    let mut folds = Vec::with_capacity(k);
29    let mut start = 0;
30    for i in 0..k {
31        let size = base + if i < remainder { 1 } else { 0 };
32        let val_idx: Vec<usize> = indices[start..start + size].to_vec();
33        let train_idx: Vec<usize> = indices
34            .iter()
35            .enumerate()
36            .filter(|(pos, _)| *pos < start || *pos >= start + size)
37            .map(|(_, &v)| v)
38            .collect();
39        folds.push(KFoldSplit { train_idx, val_idx });
40        start += size;
41    }
42    folds
43}
44
45fn select_rows<T: Clone>(rows: &[T], idx: &[usize]) -> Vec<T> {
46    idx.iter().map(|&i| rows[i].clone()).collect()
47}
48
49/// Runs the full `n_estimators`-member ensemble on each of `num_folds` folds (fold's validation
50/// rows as query, the rest as context), assembling `[member][original_train_row][class]`
51/// un-shifted OOF logits (every training row appears in exactly one fold's validation set).
52///
53/// A flattened-across-folds version (grouping folds by `(train_size, val_size)` shape and running
54/// all fold×member tasks through one `rayon` pass per shape group) was tried here, on the theory
55/// that 5 sequential per-fold rounds (default `num_folds_for_cv=5`) waste parallelism when each
56/// round only chunks `n_estimators` members into a handful of `rayon` chunks. Measured worse in
57/// every configuration tried (same or ~25% slower with default chunking, ~75% slower forcing one
58/// batch per shape group) — each fold's own `predict_batch` call already saturates Accelerate's
59/// internal BLAS threading on this machine, so adding a `rayon` layer across folds/shape-groups on
60/// top oversubscribes rather than helping (the same class of issue Round 1 documented for thread
61/// count). Reverted; kept simple.
62pub fn run_oof_classification(
63    model: &BurnTabFMModel,
64    x_train_raw: &[Vec<Value>],
65    y_train_codes: &[f64],
66    cat_mask: &[bool],
67    n_classes: usize,
68    configs: &[MemberConfig],
69    p: &EnsembleParams,
70    num_folds: usize,
71) -> Result<Vec<Vec<Vec<f64>>>> {
72    let n = x_train_raw.len();
73    let folds = kfold_splits(n, num_folds, p.random_state);
74
75    let mut oof: Vec<Vec<Vec<f64>>> = vec![vec![Vec::new(); n]; configs.len()];
76    for fold in &folds {
77        let fold_x_train = select_rows(x_train_raw, &fold.train_idx);
78        let fold_y_train: Vec<f64> = fold.train_idx.iter().map(|&i| y_train_codes[i]).collect();
79        let fold_x_val = select_rows(x_train_raw, &fold.val_idx);
80
81        let per_member = run_members_classification(
82            model,
83            &fold_x_train,
84            &fold_y_train,
85            &fold_x_val,
86            cat_mask,
87            n_classes,
88            configs,
89            p.outlier_threshold,
90            p.batch_size,
91        )?;
92        for (m, member_out) in per_member.into_iter().enumerate() {
93            for (local_idx, &orig_idx) in fold.val_idx.iter().enumerate() {
94                oof[m][orig_idx] = member_out[local_idx].clone();
95            }
96        }
97    }
98    Ok(oof)
99}
100
101/// Regression counterpart: `[member][original_train_row]` *scaled* OOF predictions.
102pub fn run_oof_regression(
103    model: &BurnTabFMModel,
104    x_train_raw: &[Vec<Value>],
105    y_train_scaled: &[f64],
106    cat_mask: &[bool],
107    configs: &[MemberConfig],
108    p: &EnsembleParams,
109    num_folds: usize,
110) -> Result<Vec<Vec<f64>>> {
111    let n = x_train_raw.len();
112    let folds = kfold_splits(n, num_folds, p.random_state);
113
114    let mut oof: Vec<Vec<f64>> = vec![vec![0.0; n]; configs.len()];
115    for fold in &folds {
116        let fold_x_train = select_rows(x_train_raw, &fold.train_idx);
117        let fold_y_train: Vec<f64> = fold.train_idx.iter().map(|&i| y_train_scaled[i]).collect();
118        let fold_x_val = select_rows(x_train_raw, &fold.val_idx);
119
120        let per_member = run_members_regression(
121            model,
122            &fold_x_train,
123            &fold_y_train,
124            &fold_x_val,
125            cat_mask,
126            configs,
127            p.outlier_threshold,
128            p.batch_size,
129        )?;
130        for (m, member_out) in per_member.into_iter().enumerate() {
131            for (local_idx, &orig_idx) in fold.val_idx.iter().enumerate() {
132                oof[m][orig_idx] = member_out[local_idx];
133            }
134        }
135    }
136    Ok(oof)
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn test_kfold_covers_every_row_exactly_once() {
145        let folds = kfold_splits(11, 5, 42);
146        let mut seen = vec![0u32; 11];
147        for f in &folds {
148            for &i in &f.val_idx {
149                seen[i] += 1;
150            }
151            assert!(!f.train_idx.is_empty());
152        }
153        assert!(seen.iter().all(|&c| c == 1));
154    }
155
156    #[test]
157    fn test_kfold_fold_count_capped_by_n() {
158        let folds = kfold_splits(3, 5, 42);
159        assert_eq!(folds.len(), 3);
160    }
161}
zsfm_tabfm/ensemble/oof.rs

zsfm_tabfm/ensemble/
oof.rs
