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::TabFMModel;
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: &TabFMModel,
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, &fold_x_train, &fold_y_train, &fold_x_val, cat_mask, n_classes, configs, p.outlier_threshold,
83            p.batch_size,
84        )?;
85        for (m, member_out) in per_member.into_iter().enumerate() {
86            for (local_idx, &orig_idx) in fold.val_idx.iter().enumerate() {
87                oof[m][orig_idx] = member_out[local_idx].clone();
88            }
89        }
90    }
91    Ok(oof)
92}
93
94/// Regression counterpart: `[member][original_train_row]` *scaled* OOF predictions.
95pub fn run_oof_regression(
96    model: &TabFMModel,
97    x_train_raw: &[Vec<Value>],
98    y_train_scaled: &[f64],
99    cat_mask: &[bool],
100    configs: &[MemberConfig],
101    p: &EnsembleParams,
102    num_folds: usize,
103) -> Result<Vec<Vec<f64>>> {
104    let n = x_train_raw.len();
105    let folds = kfold_splits(n, num_folds, p.random_state);
106
107    let mut oof: Vec<Vec<f64>> = vec![vec![0.0; n]; configs.len()];
108    for fold in &folds {
109        let fold_x_train = select_rows(x_train_raw, &fold.train_idx);
110        let fold_y_train: Vec<f64> = fold.train_idx.iter().map(|&i| y_train_scaled[i]).collect();
111        let fold_x_val = select_rows(x_train_raw, &fold.val_idx);
112
113        let per_member = run_members_regression(
114            model, &fold_x_train, &fold_y_train, &fold_x_val, cat_mask, configs, p.outlier_threshold, p.batch_size,
115        )?;
116        for (m, member_out) in per_member.into_iter().enumerate() {
117            for (local_idx, &orig_idx) in fold.val_idx.iter().enumerate() {
118                oof[m][orig_idx] = member_out[local_idx];
119            }
120        }
121    }
122    Ok(oof)
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn test_kfold_covers_every_row_exactly_once() {
131        let folds = kfold_splits(11, 5, 42);
132        let mut seen = vec![0u32; 11];
133        for f in &folds {
134            for &i in &f.val_idx {
135                seen[i] += 1;
136            }
137            assert!(!f.train_idx.is_empty());
138        }
139        assert!(seen.iter().all(|&c| c == 1));
140    }
141
142    #[test]
143    fn test_kfold_fold_count_capped_by_n() {
144        let folds = kfold_splits(3, 5, 42);
145        assert_eq!(folds.len(), 3);
146    }
147}
zsfm_tabfm/ensemble/oof.rs

zsfm_tabfm/ensemble/
oof.rs
