model_selection_rs/splitters/
stratified_kfold.rs

1//! Stratified K-fold — K-fold that preserves per-class proportions.
2
3use std::collections::HashMap;
4use std::hash::Hash;
5
6use ndarray::Array1;
7use rand::rngs::StdRng;
8use rand::seq::SliceRandom;
9use rand::SeedableRng;
10
11use super::{fold_bounds, CvSplitter};
12use crate::error::{ModelSelectionError, Result};
13
14/// Stratified K-fold cross-validation.
15///
16/// Each fold keeps roughly the same class distribution as the full dataset.
17/// This ports the previously hand-rolled stratification logic from the guide's
18/// Evaluation addendum into a real, tested implementation: samples are grouped
19/// by class label, then each class's samples are distributed across the folds so
20/// every fold receives a proportional slice of every class.
21///
22/// The label type is generic (`L: Eq + Hash + Clone`), so string, integer or
23/// enum labels all work with no forced mapping step — the same convention as the
24/// sibling `imbalance-rs` crate. Labels are supplied at construction and stored,
25/// which is what lets a stratified splitter satisfy the plain
26/// [`CvSplitter`](crate::splitters::CvSplitter) interface (see the module docs).
27///
28/// # Small classes
29///
30/// A class with fewer than `n_splits` samples cannot appear in every fold. Such
31/// a class is spread across as many folds as it can fill (one sample each,
32/// leading folds first) and a warning is emitted on stderr — matching the
33/// "warn and adjust rather than hard-error where reasonable" policy.
34///
35/// ```
36/// use ndarray::array;
37/// use model_selection_rs::splitters::{CvSplitter, StratifiedKFold};
38///
39/// let y = array![0, 0, 0, 0, 1, 1, 1, 1];
40/// let skf = StratifiedKFold::new(2, &y).unwrap();
41/// let splits = skf.split(y.len()).unwrap();
42/// // Each fold's test set holds two 0s and two 1s.
43/// assert_eq!(splits.len(), 2);
44/// ```
45#[derive(Debug, Clone)]
46pub struct StratifiedKFold<L> {
47    n_splits: usize,
48    shuffle: bool,
49    seed: u64,
50    labels: Vec<L>,
51}
52
53impl<L: Eq + Hash + Clone> StratifiedKFold<L> {
54    /// Create a `StratifiedKFold` over the class labels `y`.
55    ///
56    /// # Errors
57    ///
58    /// Returns [`ModelSelectionError::InvalidSplitCount`] if `n_splits < 2`.
59    pub fn new(n_splits: usize, y: &Array1<L>) -> Result<Self> {
60        if n_splits < 2 {
61            return Err(ModelSelectionError::InvalidSplitCount {
62                msg: format!("n_splits must be >= 2, got {n_splits}"),
63            });
64        }
65        Ok(Self {
66            n_splits,
67            shuffle: false,
68            seed: 0,
69            labels: y.to_vec(),
70        })
71    }
72
73    /// Shuffle each class's samples (deterministically, from `seed`) before
74    /// distributing them across folds.
75    #[must_use]
76    pub fn with_shuffle(mut self, seed: u64) -> Self {
77        self.shuffle = true;
78        self.seed = seed;
79        self
80    }
81
82    /// Group sample indices by class, preserving first-appearance class order
83    /// for deterministic output.
84    fn class_indices(&self) -> Vec<Vec<usize>> {
85        let mut order: Vec<L> = Vec::new();
86        let mut map: HashMap<L, Vec<usize>> = HashMap::new();
87        for (i, label) in self.labels.iter().enumerate() {
88            map.entry(label.clone()).or_insert_with(|| {
89                order.push(label.clone());
90                Vec::new()
91            });
92            map.get_mut(label).unwrap().push(i);
93        }
94        order.into_iter().map(|c| map.remove(&c).unwrap()).collect()
95    }
96}
97
98/// Shared stratification core: distribute each class's index list across
99/// `n_splits` test folds proportionally. Returns one test-index vector per fold.
100///
101/// Used by both [`StratifiedKFold`] and
102/// [`RepeatedStratifiedKFold`](super::RepeatedStratifiedKFold).
103pub(crate) fn stratified_test_folds(
104    class_indices: &[Vec<usize>],
105    n_splits: usize,
106    shuffle: bool,
107    seed: u64,
108) -> Vec<Vec<usize>> {
109    let mut rng = StdRng::seed_from_u64(seed);
110    let mut test_folds: Vec<Vec<usize>> = vec![Vec::new(); n_splits];
111
112    for indices in class_indices {
113        let mut idx = indices.clone();
114        if shuffle {
115            idx.shuffle(&mut rng);
116        }
117        if idx.len() < n_splits {
118            eprintln!(
119                "model-selection-rs: StratifiedKFold — a class has {} sample(s), \
120                 fewer than n_splits={n_splits}; it will be present in only {} of \
121                 the {n_splits} folds.",
122                idx.len(),
123                idx.len()
124            );
125        }
126        // Near-equal chunking of this class across folds; fold j gets chunk j.
127        for (fold, (start, end)) in fold_bounds(idx.len(), n_splits).into_iter().enumerate() {
128            test_folds[fold].extend_from_slice(&idx[start..end]);
129        }
130    }
131
132    // Sort each fold so output is order-stable regardless of class iteration.
133    for fold in &mut test_folds {
134        fold.sort_unstable();
135    }
136    test_folds
137}
138
139/// Turn per-fold test sets into `(train, test)` pairs over `n_samples`.
140pub(crate) fn test_folds_to_splits(
141    test_folds: Vec<Vec<usize>>,
142    n_samples: usize,
143) -> Vec<(Vec<usize>, Vec<usize>)> {
144    test_folds
145        .into_iter()
146        .map(|test| {
147            let in_test: std::collections::HashSet<usize> = test.iter().copied().collect();
148            let train: Vec<usize> = (0..n_samples).filter(|i| !in_test.contains(i)).collect();
149            (train, test)
150        })
151        .collect()
152}
153
154impl<L: Eq + Hash + Clone> CvSplitter for StratifiedKFold<L> {
155    fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>> {
156        if n_samples != self.labels.len() {
157            return Err(ModelSelectionError::ShapeMismatch {
158                expected: self.labels.len(),
159                got: n_samples,
160            });
161        }
162        if self.n_splits > n_samples {
163            return Err(ModelSelectionError::NotEnoughSamples {
164                needed: self.n_splits,
165                got: n_samples,
166            });
167        }
168        let class_indices = self.class_indices();
169        let test_folds =
170            stratified_test_folds(&class_indices, self.n_splits, self.shuffle, self.seed);
171        Ok(test_folds_to_splits(test_folds, n_samples))
172    }
173
174    fn n_splits(&self) -> usize {
175        self.n_splits
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182    use ndarray::array;
183
184    fn class_proportions<L: Eq + Hash + Clone>(labels: &[L], idx: &[usize]) -> HashMap<L, f64> {
185        let mut counts: HashMap<L, usize> = HashMap::new();
186        for &i in idx {
187            *counts.entry(labels[i].clone()).or_default() += 1;
188        }
189        let total = idx.len() as f64;
190        counts
191            .into_iter()
192            .map(|(k, v)| (k, v as f64 / total))
193            .collect()
194    }
195
196    #[test]
197    fn balanced_two_class_splits_evenly() {
198        let y = array![0, 0, 0, 0, 1, 1, 1, 1];
199        let skf = StratifiedKFold::new(2, &y).unwrap();
200        for (_, test) in skf.split(y.len()).unwrap() {
201            let props = class_proportions(y.as_slice().unwrap(), &test);
202            assert_eq!(props[&0], 0.5);
203            assert_eq!(props[&1], 0.5);
204        }
205    }
206
207    #[test]
208    fn preserves_proportions_under_heavy_imbalance() {
209        // 90 of class 0, 10 of class 1.
210        let mut v = vec![0; 90];
211        v.extend(std::iter::repeat(1).take(10));
212        let y = Array1::from(v);
213        let overall = class_proportions(y.as_slice().unwrap(), &(0..100).collect::<Vec<_>>());
214        let skf = StratifiedKFold::new(5, &y).unwrap().with_shuffle(1);
215        for (_, test) in skf.split(100).unwrap() {
216            let props = class_proportions(y.as_slice().unwrap(), &test);
217            for class in [0, 1] {
218                assert!(
219                    (props[&class] - overall[&class]).abs() < 0.05,
220                    "class {class}: fold {} vs overall {}",
221                    props[&class],
222                    overall[&class]
223                );
224            }
225        }
226    }
227
228    #[test]
229    fn every_sample_tested_once() {
230        let y = array![0, 1, 0, 1, 0, 1, 0, 1, 0, 1];
231        let skf = StratifiedKFold::new(5, &y).unwrap();
232        let mut seen: Vec<usize> = skf
233            .split(10)
234            .unwrap()
235            .iter()
236            .flat_map(|(_, te)| te.clone())
237            .collect();
238        seen.sort_unstable();
239        assert_eq!(seen, (0..10).collect::<Vec<_>>());
240    }
241
242    #[test]
243    fn shape_mismatch_when_n_samples_disagrees() {
244        let y = array![0, 1, 0, 1];
245        let skf = StratifiedKFold::new(2, &y).unwrap();
246        assert!(matches!(
247            skf.split(5),
248            Err(ModelSelectionError::ShapeMismatch {
249                expected: 4,
250                got: 5
251            })
252        ));
253    }
254
255    #[test]
256    fn string_labels_work() {
257        let y = array!["cat", "dog", "cat", "dog"];
258        let skf = StratifiedKFold::new(2, &y).unwrap();
259        assert_eq!(skf.split(4).unwrap().len(), 2);
260    }
261}
model_selection_rs/splitters/stratified_kfold.rs

model_selection_rs/splitters/
stratified_kfold.rs
