model_selection_rs/splitters/
repeated.rs

1//! Repeated K-fold variants for more robust performance estimates.
2
3use std::hash::Hash;
4
5use ndarray::Array1;
6
7use super::stratified_kfold::{stratified_test_folds, test_folds_to_splits};
8use super::{CvSplitter, KFold};
9use crate::error::{ModelSelectionError, Result};
10
11/// Repeated plain K-fold.
12///
13/// Runs [`KFold`] `n_repeats` times, each with a different shuffle seed, and
14/// concatenates the resulting splits. Repeated evaluation smooths out the
15/// variance that a single random fold assignment introduces, giving a more
16/// robust performance estimate. Total splits = `n_repeats * n_splits`.
17///
18/// ```
19/// use model_selection_rs::splitters::{CvSplitter, RepeatedKFold};
20///
21/// let rkf = RepeatedKFold::new(5, 3, 0).unwrap();
22/// assert_eq!(rkf.n_splits(), 15);
23/// assert_eq!(rkf.split(50).unwrap().len(), 15);
24/// ```
25#[derive(Debug, Clone)]
26pub struct RepeatedKFold {
27    n_splits: usize,
28    n_repeats: usize,
29    base_seed: u64,
30}
31
32impl RepeatedKFold {
33    /// Create a `RepeatedKFold` with `n_splits` folds repeated `n_repeats` times.
34    ///
35    /// # Errors
36    ///
37    /// Returns [`ModelSelectionError::InvalidSplitCount`] if `n_splits < 2` or
38    /// `n_repeats < 1`.
39    pub fn new(n_splits: usize, n_repeats: usize, base_seed: u64) -> Result<Self> {
40        if n_splits < 2 {
41            return Err(ModelSelectionError::InvalidSplitCount {
42                msg: format!("n_splits must be >= 2, got {n_splits}"),
43            });
44        }
45        if n_repeats < 1 {
46            return Err(ModelSelectionError::InvalidSplitCount {
47                msg: format!("n_repeats must be >= 1, got {n_repeats}"),
48            });
49        }
50        Ok(Self {
51            n_splits,
52            n_repeats,
53            base_seed,
54        })
55    }
56}
57
58impl CvSplitter for RepeatedKFold {
59    fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>> {
60        let mut all = Vec::with_capacity(self.n_splits * self.n_repeats);
61        for r in 0..self.n_repeats {
62            let kf = KFold::new(self.n_splits)?.with_shuffle(self.base_seed.wrapping_add(r as u64));
63            all.extend(kf.split(n_samples)?);
64        }
65        Ok(all)
66    }
67
68    fn n_splits(&self) -> usize {
69        self.n_splits * self.n_repeats
70    }
71}
72
73/// Repeated stratified K-fold.
74///
75/// The stratified analogue of [`RepeatedKFold`]: runs stratified K-fold
76/// `n_repeats` times with different seeds, preserving class proportions in every
77/// fold of every repeat. Total splits = `n_repeats * n_splits`.
78///
79/// ```
80/// use ndarray::Array1;
81/// use model_selection_rs::splitters::{CvSplitter, RepeatedStratifiedKFold};
82///
83/// let y = Array1::from(vec![0, 1, 0, 1, 0, 1, 0, 1, 0, 1]);
84/// let rskf = RepeatedStratifiedKFold::new(2, 3, 0, &y).unwrap();
85/// assert_eq!(rskf.n_splits(), 6);
86/// ```
87#[derive(Debug, Clone)]
88pub struct RepeatedStratifiedKFold<L> {
89    n_splits: usize,
90    n_repeats: usize,
91    base_seed: u64,
92    labels: Vec<L>,
93}
94
95impl<L: Eq + Hash + Clone> RepeatedStratifiedKFold<L> {
96    /// Create a `RepeatedStratifiedKFold` over labels `y`.
97    ///
98    /// # Errors
99    ///
100    /// Returns [`ModelSelectionError::InvalidSplitCount`] if `n_splits < 2` or
101    /// `n_repeats < 1`.
102    pub fn new(n_splits: usize, n_repeats: usize, base_seed: u64, y: &Array1<L>) -> Result<Self> {
103        if n_splits < 2 {
104            return Err(ModelSelectionError::InvalidSplitCount {
105                msg: format!("n_splits must be >= 2, got {n_splits}"),
106            });
107        }
108        if n_repeats < 1 {
109            return Err(ModelSelectionError::InvalidSplitCount {
110                msg: format!("n_repeats must be >= 1, got {n_repeats}"),
111            });
112        }
113        Ok(Self {
114            n_splits,
115            n_repeats,
116            base_seed,
117            labels: y.to_vec(),
118        })
119    }
120
121    fn class_indices(&self) -> Vec<Vec<usize>> {
122        use std::collections::HashMap;
123        let mut order: Vec<L> = Vec::new();
124        let mut map: HashMap<L, Vec<usize>> = HashMap::new();
125        for (i, label) in self.labels.iter().enumerate() {
126            map.entry(label.clone()).or_insert_with(|| {
127                order.push(label.clone());
128                Vec::new()
129            });
130            map.get_mut(label).unwrap().push(i);
131        }
132        order.into_iter().map(|c| map.remove(&c).unwrap()).collect()
133    }
134}
135
136impl<L: Eq + Hash + Clone> CvSplitter for RepeatedStratifiedKFold<L> {
137    fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>> {
138        if n_samples != self.labels.len() {
139            return Err(ModelSelectionError::ShapeMismatch {
140                expected: self.labels.len(),
141                got: n_samples,
142            });
143        }
144        if self.n_splits > n_samples {
145            return Err(ModelSelectionError::NotEnoughSamples {
146                needed: self.n_splits,
147                got: n_samples,
148            });
149        }
150        let class_indices = self.class_indices();
151        let mut all = Vec::with_capacity(self.n_splits * self.n_repeats);
152        for r in 0..self.n_repeats {
153            let test_folds = stratified_test_folds(
154                &class_indices,
155                self.n_splits,
156                true,
157                self.base_seed.wrapping_add(r as u64),
158            );
159            all.extend(test_folds_to_splits(test_folds, n_samples));
160        }
161        Ok(all)
162    }
163
164    fn n_splits(&self) -> usize {
165        self.n_splits * self.n_repeats
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    #[test]
174    fn repeated_kfold_yields_product_of_splits() {
175        let rkf = RepeatedKFold::new(5, 3, 42).unwrap();
176        assert_eq!(rkf.n_splits(), 15);
177        assert_eq!(rkf.split(50).unwrap().len(), 15);
178    }
179
180    #[test]
181    fn repeats_differ_from_each_other() {
182        let rkf = RepeatedKFold::new(2, 2, 1).unwrap();
183        let splits = rkf.split(20).unwrap();
184        // The first repeat's first-fold test set should differ from the second
185        // repeat's first-fold test set (different seed).
186        assert_ne!(splits[0].1, splits[2].1);
187    }
188
189    #[test]
190    fn repeated_stratified_counts() {
191        let y = Array1::from(vec![0, 1, 0, 1, 0, 1, 0, 1, 0, 1]);
192        let rskf = RepeatedStratifiedKFold::new(2, 4, 0, &y).unwrap();
193        assert_eq!(rskf.n_splits(), 8);
194        assert_eq!(rskf.split(10).unwrap().len(), 8);
195    }
196}
model_selection_rs/splitters/repeated.rs

model_selection_rs/splitters/
repeated.rs
