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(
66    n_features: usize,
67    n_estimators: usize,
68    random_state: u64,
69) -> Vec<Vec<usize>> {
70    let mut rng = PyRandom::new(random_state);
71    if n_features <= 5 {
72        let all_perms = permutations(n_features);
73        let k = n_estimators.min(all_perms.len());
74        rng.sample(&all_perms, k)
75    } else {
76        (0..n_estimators)
77            .map(|_| rng.sample_indices(n_features, n_features))
78            .collect()
79    }
80}
81
82fn permutations(n: usize) -> Vec<Vec<usize>> {
83    let mut items: Vec<usize> = (0..n).collect();
84    let mut result = Vec::new();
85    permute_rec(&mut items, 0, &mut result);
86    result
87}
88
89fn permute_rec(items: &mut Vec<usize>, k: usize, out: &mut Vec<Vec<usize>>) {
90    if k == items.len() {
91        out.push(items.clone());
92        return;
93    }
94    for i in k..items.len() {
95        items.swap(k, i);
96        permute_rec(items, k + 1, out);
97        items.swap(k, i);
98    }
99}
100
101/// The full `_generate_ensemble()` port. See module docs for the exact RNG call order this
102/// must preserve: (1) feature permutations from the *independent* `FeatureShuffler` stream;
103/// then, from the main stream: (2) class-shift base offsets, (3) categorical permutations
104/// per-member (only if enabled), (4) row-subsample patterns per-member (only if enabled),
105/// (5) one `shuffle()` over the zipped per-member tuples; norm-method assignment happens last,
106/// by position, consuming no RNG.
107pub fn generate_ensemble(p: &EnsembleConfigParams) -> Vec<MemberConfig> {
108    let shuffle_patterns =
109        generate_feature_permutations(p.n_features, p.n_estimators, p.random_state);
110    let n_members_from_shuffle = shuffle_patterns.len();
111
112    let mut rng = PyRandom::new(p.random_state);
113
114    let shift_offsets: Vec<usize> =
115        if p.is_classification && p.class_shift && p.n_estimators > 1 && p.n_classes > 1 {
116            let base_offsets = rng.sample_indices(p.n_classes, p.n_classes);
117            let num_cycles = p.n_estimators.div_ceil(base_offsets.len());
118            base_offsets
119                .iter()
120                .cycle()
121                .take(base_offsets.len() * num_cycles)
122                .take(p.n_estimators)
123                .copied()
124                .collect()
125        } else {
126            vec![0usize; p.n_estimators]
127        };
128
129    let cat_permutations: Vec<Option<Vec<Vec<usize>>>> =
130        if p.permute_categorical && !p.cat_value_counts.is_empty() {
131            (0..p.n_estimators)
132                .map(|_| {
133                    Some(
134                        p.cat_value_counts
135                            .iter()
136                            .map(|&n_vals| rng.sample_indices(n_vals, n_vals))
137                            .collect(),
138                    )
139                })
140                .collect()
141        } else {
142            vec![None; p.n_estimators]
143        };
144
145    let n_rows_target = p
146        .max_num_rows
147        .map(|m| m.min(p.n_train_rows))
148        .unwrap_or(p.n_train_rows);
149    let row_subsample_patterns: Vec<Option<Vec<usize>>> = if n_rows_target < p.n_train_rows {
150        (0..p.n_estimators)
151            .map(|_| Some(rng.sample_indices(p.n_train_rows, n_rows_target)))
152            .collect()
153    } else {
154        vec![None; p.n_estimators]
155    };
156
157    let n = n_members_from_shuffle.min(shift_offsets.len());
158    let mut zipped: Vec<(
159        Vec<usize>,
160        usize,
161        Option<Vec<Vec<usize>>>,
162        Option<Vec<usize>>,
163    )> = (0..n)
164        .map(|i| {
165            (
166                shuffle_patterns[i].clone(),
167                shift_offsets[i],
168                cat_permutations[i].clone(),
169                row_subsample_patterns[i].clone(),
170            )
171        })
172        .collect();
173    rng.shuffle(&mut zipped);
174
175    let num_cycles = p.n_estimators.div_ceil(p.norm_methods.len());
176    let norm_methods_for_estimators: Vec<NormMethod> = p
177        .norm_methods
178        .iter()
179        .cycle()
180        .take(p.norm_methods.len() * num_cycles)
181        .take(n)
182        .copied()
183        .collect();
184
185    zipped
186        .into_iter()
187        .zip(norm_methods_for_estimators)
188        .map(
189            |((feature_permutation, class_shift, cat_permutation, row_subsample), norm_method)| {
190                MemberConfig {
191                    feature_permutation,
192                    class_shift,
193                    cat_permutation,
194                    row_subsample,
195                    norm_method,
196                }
197            },
198        )
199        .collect()
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[test]
207    fn test_default_config_shape() {
208        let p = EnsembleConfigParams {
209            n_estimators: 32,
210            n_features: 8,
211            n_train_rows: 10,
212            is_classification: true,
213            n_classes: 3,
214            class_shift: true,
215            permute_categorical: false,
216            cat_value_counts: vec![],
217            max_num_rows: None,
218            norm_methods: vec![NormMethod::None, NormMethod::Power],
219            random_state: 42,
220        };
221        let configs = generate_ensemble(&p);
222        assert_eq!(configs.len(), 32);
223        for (i, c) in configs.iter().enumerate() {
224            assert_eq!(c.feature_permutation.len(), 8);
225            let mut sorted = c.feature_permutation.clone();
226            sorted.sort();
227            assert_eq!(
228                sorted,
229                (0..8).collect::<Vec<_>>(),
230                "member {i} not a valid permutation"
231            );
232            assert!(c.class_shift < 3);
233            assert!(c.cat_permutation.is_none());
234            assert!(c.row_subsample.is_none());
235        }
236        // norm methods cycle none/power by final position, deterministically.
237        assert_eq!(configs[0].norm_method, NormMethod::None);
238        assert_eq!(configs[1].norm_method, NormMethod::Power);
239    }
240
241    #[test]
242    fn test_small_feature_count_uses_permutation_enumeration() {
243        let p = EnsembleConfigParams {
244            n_estimators: 32,
245            n_features: 3,
246            n_train_rows: 10,
247            is_classification: false,
248            n_classes: 1,
249            class_shift: true,
250            permute_categorical: false,
251            cat_value_counts: vec![],
252            max_num_rows: None,
253            norm_methods: vec![NormMethod::None],
254            random_state: 42,
255        };
256        let configs = generate_ensemble(&p);
257        // 3! = 6 total permutations, fewer than n_estimators=32 -> truncated ensemble.
258        assert_eq!(configs.len(), 6);
259    }
260
261    #[test]
262    fn test_matches_real_ensemble_generator() {
263        // Generated by directly invoking the real `EnsembleGenerator(n_estimators=32,
264        // norm_methods=None, class_shift=True, random_state=42, task="classification").fit(X, y)`
265        // (8 features, 10 rows, 3 classes) from classifier_and_regressor.py and reading back
266        // `ensemble_configs_` (grouped by norm method: "none" gets even overall positions,
267        // "power" gets odd, per the `norm_methods_for_estimators[i % 2]` cycling below).
268        let p = EnsembleConfigParams {
269            n_estimators: 32,
270            n_features: 8,
271            n_train_rows: 10,
272            is_classification: true,
273            n_classes: 3,
274            class_shift: true,
275            permute_categorical: false,
276            cat_value_counts: vec![],
277            max_num_rows: None,
278            norm_methods: vec![NormMethod::None, NormMethod::Power],
279            random_state: 42,
280        };
281        let configs = generate_ensemble(&p);
282        let expected: [(NormMethod, &[usize], usize); 6] = [
283            (NormMethod::None, &[0, 5, 3, 4, 7, 1, 6, 2], 1),
284            (NormMethod::Power, &[3, 0, 1, 4, 6, 7, 5, 2], 0),
285            (NormMethod::None, &[2, 1, 5, 3, 6, 4, 7, 0], 1),
286            (NormMethod::Power, &[4, 1, 6, 7, 2, 3, 5, 0], 2),
287            (NormMethod::None, &[1, 4, 7, 2, 3, 0, 6, 5], 2),
288            (NormMethod::Power, &[2, 3, 1, 7, 6, 0, 4, 5], 0),
289        ];
290        for (i, (exp_method, exp_perm, exp_shift)) in expected.iter().enumerate() {
291            assert_eq!(
292                configs[i].norm_method, *exp_method,
293                "member {i} norm_method"
294            );
295            assert_eq!(
296                &configs[i].feature_permutation, exp_perm,
297                "member {i} permutation"
298            );
299            assert_eq!(configs[i].class_shift, *exp_shift, "member {i} shift");
300        }
301    }
302}
zsfm_tabfm/ensemble/config_gen.rs

zsfm_tabfm/ensemble/
config_gen.rs
