zsfm_tabfm/ensemble/
config_gen.rs

1//! Ports `TabFMClassifier`/`TabFMRegressor`'s `_generate_ensemble()` — builds the `n_estimators`
2//! member configs (feature permutation, classification class-shift offset, categorical-value
3//! permutation, row-subsample pattern, normalization method). RNG consumption order matches the
4//! source exactly (see the plan doc / module comments below) so results are bit-identical to the
5//! real wrapper for the same `random_state`.
6
7use super::pyrandom::PyRandom;
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
10pub enum NormMethod {
11    None,
12    Power,
13    Quantile,
14    QuantileRtdl,
15    Robust,
16}
17
18impl NormMethod {
19    pub fn parse(s: &str) -> anyhow::Result<Self> {
20        Ok(match s {
21            "none" => NormMethod::None,
22            "power" => NormMethod::Power,
23            "quantile" => NormMethod::Quantile,
24            "quantile_rtdl" => NormMethod::QuantileRtdl,
25            "robust" => NormMethod::Robust,
26            other => anyhow::bail!("unknown norm_method: {other}"),
27        })
28    }
29}
30
31pub struct MemberConfig {
32    /// Permutation of `0..n_features` (or a subsampled+permuted subset if
33    /// `n_features > max_num_features`).
34    pub feature_permutation: Vec<usize>,
35    /// Amount added (mod `n_classes`) to every training label fed to the model for this member.
36    /// Always `0` for regression or when class-shift is disabled/inapplicable.
37    pub class_shift: usize,
38    /// `cat_col_index -> (original_code -> permuted_code)`, present only when
39    /// `permute_categorical=true` (default: off, so this is `None` for every member).
40    pub cat_permutation: Option<Vec<Vec<usize>>>,
41    /// Row indices to bag from `0..n_train`, present only when row-subsampling is active
42    /// (default: off, so this is `None` for every member — use all rows).
43    pub row_subsample: Option<Vec<usize>>,
44    pub norm_method: NormMethod,
45}
46
47pub struct EnsembleConfigParams {
48    pub n_estimators: usize,
49    pub n_features: usize,
50    pub n_train_rows: usize,
51    pub is_classification: bool,
52    pub n_classes: usize,
53    pub class_shift: bool,
54    pub permute_categorical: bool,
55    /// Number of distinct values per categorical column (only consulted if
56    /// `permute_categorical` is true); empty if there are no categorical columns.
57    pub cat_value_counts: Vec<usize>,
58    pub max_num_rows: Option<usize>,
59    pub norm_methods: Vec<NormMethod>,
60    pub random_state: u64,
61}
62
63/// Feature-permutation generation (`FeatureShuffler`): its own independent `random.Random`
64/// stream, seeded with the same `random_state` but never mixed with the main RNG below.
65fn generate_feature_permutations(n_features: usize, n_estimators: usize, random_state: u64) -> Vec<Vec<usize>> {
66    let mut rng = PyRandom::new(random_state);
67    if n_features <= 5 {
68        let all_perms = permutations(n_features);
69        let k = n_estimators.min(all_perms.len());
70        rng.sample(&all_perms, k)
71    } else {
72        (0..n_estimators).map(|_| rng.sample_indices(n_features, n_features)).collect()
73    }
74}
75
76fn permutations(n: usize) -> Vec<Vec<usize>> {
77    let mut items: Vec<usize> = (0..n).collect();
78    let mut result = Vec::new();
79    permute_rec(&mut items, 0, &mut result);
80    result
81}
82
83fn permute_rec(items: &mut Vec<usize>, k: usize, out: &mut Vec<Vec<usize>>) {
84    if k == items.len() {
85        out.push(items.clone());
86        return;
87    }
88    for i in k..items.len() {
89        items.swap(k, i);
90        permute_rec(items, k + 1, out);
91        items.swap(k, i);
92    }
93}
94
95/// The full `_generate_ensemble()` port. See module docs for the exact RNG call order this
96/// must preserve: (1) feature permutations from the *independent* `FeatureShuffler` stream;
97/// then, from the main stream: (2) class-shift base offsets, (3) categorical permutations
98/// per-member (only if enabled), (4) row-subsample patterns per-member (only if enabled),
99/// (5) one `shuffle()` over the zipped per-member tuples; norm-method assignment happens last,
100/// by position, consuming no RNG.
101pub fn generate_ensemble(p: &EnsembleConfigParams) -> Vec<MemberConfig> {
102    let shuffle_patterns = generate_feature_permutations(p.n_features, p.n_estimators, p.random_state);
103    let n_members_from_shuffle = shuffle_patterns.len();
104
105    let mut rng = PyRandom::new(p.random_state);
106
107    let shift_offsets: Vec<usize> = if p.is_classification && p.class_shift && p.n_estimators > 1 && p.n_classes > 1 {
108        let base_offsets = rng.sample_indices(p.n_classes, p.n_classes);
109        let num_cycles = p.n_estimators.div_ceil(base_offsets.len());
110        base_offsets
111            .iter()
112            .cycle()
113            .take(base_offsets.len() * num_cycles)
114            .take(p.n_estimators)
115            .copied()
116            .collect()
117    } else {
118        vec![0usize; p.n_estimators]
119    };
120
121    let cat_permutations: Vec<Option<Vec<Vec<usize>>>> = if p.permute_categorical && !p.cat_value_counts.is_empty() {
122        (0..p.n_estimators)
123            .map(|_| {
124                Some(
125                    p.cat_value_counts
126                        .iter()
127                        .map(|&n_vals| rng.sample_indices(n_vals, n_vals))
128                        .collect(),
129                )
130            })
131            .collect()
132    } else {
133        vec![None; p.n_estimators]
134    };
135
136    let n_rows_target = p.max_num_rows.map(|m| m.min(p.n_train_rows)).unwrap_or(p.n_train_rows);
137    let row_subsample_patterns: Vec<Option<Vec<usize>>> = if n_rows_target < p.n_train_rows {
138        (0..p.n_estimators).map(|_| Some(rng.sample_indices(p.n_train_rows, n_rows_target))).collect()
139    } else {
140        vec![None; p.n_estimators]
141    };
142
143    let n = n_members_from_shuffle.min(shift_offsets.len());
144    let mut zipped: Vec<(Vec<usize>, usize, Option<Vec<Vec<usize>>>, Option<Vec<usize>>)> = (0..n)
145        .map(|i| {
146            (
147                shuffle_patterns[i].clone(),
148                shift_offsets[i],
149                cat_permutations[i].clone(),
150                row_subsample_patterns[i].clone(),
151            )
152        })
153        .collect();
154    rng.shuffle(&mut zipped);
155
156    let num_cycles = p.n_estimators.div_ceil(p.norm_methods.len());
157    let norm_methods_for_estimators: Vec<NormMethod> =
158        p.norm_methods.iter().cycle().take(p.norm_methods.len() * num_cycles).take(n).copied().collect();
159
160    zipped
161        .into_iter()
162        .zip(norm_methods_for_estimators)
163        .map(|((feature_permutation, class_shift, cat_permutation, row_subsample), norm_method)| MemberConfig {
164            feature_permutation,
165            class_shift,
166            cat_permutation,
167            row_subsample,
168            norm_method,
169        })
170        .collect()
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    #[test]
178    fn test_default_config_shape() {
179        let p = EnsembleConfigParams {
180            n_estimators: 32,
181            n_features: 8,
182            n_train_rows: 10,
183            is_classification: true,
184            n_classes: 3,
185            class_shift: true,
186            permute_categorical: false,
187            cat_value_counts: vec![],
188            max_num_rows: None,
189            norm_methods: vec![NormMethod::None, NormMethod::Power],
190            random_state: 42,
191        };
192        let configs = generate_ensemble(&p);
193        assert_eq!(configs.len(), 32);
194        for (i, c) in configs.iter().enumerate() {
195            assert_eq!(c.feature_permutation.len(), 8);
196            let mut sorted = c.feature_permutation.clone();
197            sorted.sort();
198            assert_eq!(sorted, (0..8).collect::<Vec<_>>(), "member {i} not a valid permutation");
199            assert!(c.class_shift < 3);
200            assert!(c.cat_permutation.is_none());
201            assert!(c.row_subsample.is_none());
202        }
203        // norm methods cycle none/power by final position, deterministically.
204        assert_eq!(configs[0].norm_method, NormMethod::None);
205        assert_eq!(configs[1].norm_method, NormMethod::Power);
206    }
207
208    #[test]
209    fn test_small_feature_count_uses_permutation_enumeration() {
210        let p = EnsembleConfigParams {
211            n_estimators: 32,
212            n_features: 3,
213            n_train_rows: 10,
214            is_classification: false,
215            n_classes: 1,
216            class_shift: true,
217            permute_categorical: false,
218            cat_value_counts: vec![],
219            max_num_rows: None,
220            norm_methods: vec![NormMethod::None],
221            random_state: 42,
222        };
223        let configs = generate_ensemble(&p);
224        // 3! = 6 total permutations, fewer than n_estimators=32 -> truncated ensemble.
225        assert_eq!(configs.len(), 6);
226    }
227
228    #[test]
229    fn test_matches_real_ensemble_generator() {
230        // Generated by directly invoking the real `EnsembleGenerator(n_estimators=32,
231        // norm_methods=None, class_shift=True, random_state=42, task="classification").fit(X, y)`
232        // (8 features, 10 rows, 3 classes) from classifier_and_regressor.py and reading back
233        // `ensemble_configs_` (grouped by norm method: "none" gets even overall positions,
234        // "power" gets odd, per the `norm_methods_for_estimators[i % 2]` cycling below).
235        let p = EnsembleConfigParams {
236            n_estimators: 32,
237            n_features: 8,
238            n_train_rows: 10,
239            is_classification: true,
240            n_classes: 3,
241            class_shift: true,
242            permute_categorical: false,
243            cat_value_counts: vec![],
244            max_num_rows: None,
245            norm_methods: vec![NormMethod::None, NormMethod::Power],
246            random_state: 42,
247        };
248        let configs = generate_ensemble(&p);
249        let expected: [(NormMethod, &[usize], usize); 6] = [
250            (NormMethod::None, &[0, 5, 3, 4, 7, 1, 6, 2], 1),
251            (NormMethod::Power, &[3, 0, 1, 4, 6, 7, 5, 2], 0),
252            (NormMethod::None, &[2, 1, 5, 3, 6, 4, 7, 0], 1),
253            (NormMethod::Power, &[4, 1, 6, 7, 2, 3, 5, 0], 2),
254            (NormMethod::None, &[1, 4, 7, 2, 3, 0, 6, 5], 2),
255            (NormMethod::Power, &[2, 3, 1, 7, 6, 0, 4, 5], 0),
256        ];
257        for (i, (exp_method, exp_perm, exp_shift)) in expected.iter().enumerate() {
258            assert_eq!(configs[i].norm_method, *exp_method, "member {i} norm_method");
259            assert_eq!(&configs[i].feature_permutation, exp_perm, "member {i} permutation");
260            assert_eq!(configs[i].class_shift, *exp_shift, "member {i} shift");
261        }
262    }
263}
zsfm_tabfm/ensemble/config_gen.rs

zsfm_tabfm/ensemble/
config_gen.rs
