model_selection_rs/splitters/
kfold.rs

1//! Standard K-fold cross-validation.
2
3use rand::rngs::StdRng;
4use rand::seq::SliceRandom;
5use rand::SeedableRng;
6
7use super::{fold_bounds, CvSplitter};
8use crate::error::{ModelSelectionError, Result};
9
10/// Plain K-fold cross-validation.
11///
12/// The sample indices are partitioned into `k` folds; each fold serves as the
13/// test set once while the remaining `k - 1` folds form the training set. With
14/// [`shuffle`](KFold::with_shuffle) the indices are permuted (deterministically,
15/// from the given seed) before partitioning.
16///
17/// This is a from-scratch implementation, so the crate has **zero** required
18/// dependency on `smartcore`; it is functionally equivalent to
19/// `smartcore::model_selection::KFold` for anyone who would rather use that.
20///
21/// ```
22/// use model_selection_rs::splitters::{CvSplitter, KFold};
23///
24/// let kf = KFold::new(5).unwrap().with_shuffle(42);
25/// let splits = kf.split(50).unwrap();
26/// assert_eq!(splits.len(), 5);
27/// ```
28#[derive(Debug, Clone)]
29pub struct KFold {
30    n_splits: usize,
31    shuffle: bool,
32    seed: u64,
33}
34
35impl KFold {
36    /// Create a `KFold` with `n_splits` folds and no shuffling.
37    ///
38    /// # Errors
39    ///
40    /// Returns [`ModelSelectionError::InvalidSplitCount`] if `n_splits < 2`.
41    pub fn new(n_splits: usize) -> Result<Self> {
42        if n_splits < 2 {
43            return Err(ModelSelectionError::InvalidSplitCount {
44                msg: format!("n_splits must be >= 2, got {n_splits}"),
45            });
46        }
47        Ok(Self {
48            n_splits,
49            shuffle: false,
50            seed: 0,
51        })
52    }
53
54    /// Enable shuffling of the sample order before folding, using `seed`.
55    #[must_use]
56    pub fn with_shuffle(mut self, seed: u64) -> Self {
57        self.shuffle = true;
58        self.seed = seed;
59        self
60    }
61
62    /// Build the (possibly shuffled) index order the folds are carved from.
63    pub(crate) fn ordered_indices(&self, n_samples: usize) -> Vec<usize> {
64        let mut indices: Vec<usize> = (0..n_samples).collect();
65        if self.shuffle {
66            let mut rng = StdRng::seed_from_u64(self.seed);
67            indices.shuffle(&mut rng);
68        }
69        indices
70    }
71}
72
73impl CvSplitter for KFold {
74    fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>> {
75        if self.n_splits > n_samples {
76            return Err(ModelSelectionError::NotEnoughSamples {
77                needed: self.n_splits,
78                got: n_samples,
79            });
80        }
81        let indices = self.ordered_indices(n_samples);
82        let bounds = fold_bounds(n_samples, self.n_splits);
83
84        let splits = bounds
85            .into_iter()
86            .map(|(start, end)| {
87                let test: Vec<usize> = indices[start..end].to_vec();
88                let mut train: Vec<usize> = Vec::with_capacity(n_samples - test.len());
89                train.extend_from_slice(&indices[..start]);
90                train.extend_from_slice(&indices[end..]);
91                (train, test)
92            })
93            .collect();
94        Ok(splits)
95    }
96
97    fn n_splits(&self) -> usize {
98        self.n_splits
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use std::collections::HashSet;
106
107    #[test]
108    fn rejects_fewer_than_two_folds() {
109        assert!(KFold::new(1).is_err());
110        assert!(KFold::new(0).is_err());
111    }
112
113    #[test]
114    fn errors_when_more_folds_than_samples() {
115        let kf = KFold::new(5).unwrap();
116        assert!(matches!(
117            kf.split(3),
118            Err(ModelSelectionError::NotEnoughSamples { needed: 5, got: 3 })
119        ));
120    }
121
122    #[test]
123    fn every_sample_tested_exactly_once() {
124        let kf = KFold::new(4).unwrap();
125        let splits = kf.split(23).unwrap();
126        let mut seen: Vec<usize> = splits.iter().flat_map(|(_, te)| te.clone()).collect();
127        seen.sort_unstable();
128        assert_eq!(seen, (0..23).collect::<Vec<_>>());
129    }
130
131    #[test]
132    fn train_and_test_are_disjoint_and_cover_all() {
133        let kf = KFold::new(3).unwrap().with_shuffle(7);
134        for (train, test) in kf.split(20).unwrap() {
135            let tr: HashSet<_> = train.iter().collect();
136            let te: HashSet<_> = test.iter().collect();
137            assert!(tr.is_disjoint(&te));
138            assert_eq!(tr.len() + te.len(), 20);
139        }
140    }
141
142    #[test]
143    fn fold_sizes_differ_by_at_most_one() {
144        let kf = KFold::new(4).unwrap();
145        let sizes: Vec<usize> = kf
146            .split(23)
147            .unwrap()
148            .iter()
149            .map(|(_, te)| te.len())
150            .collect();
151        let max = *sizes.iter().max().unwrap();
152        let min = *sizes.iter().min().unwrap();
153        assert!(max - min <= 1);
154    }
155
156    #[test]
157    fn shuffle_is_deterministic_for_a_seed() {
158        let a = KFold::new(3).unwrap().with_shuffle(99).split(15).unwrap();
159        let b = KFold::new(3).unwrap().with_shuffle(99).split(15).unwrap();
160        assert_eq!(a, b);
161    }
162}
model_selection_rs/splitters/kfold.rs

model_selection_rs/splitters/
kfold.rs
