model-selection-rs 0.1.0

//! Standard K-fold cross-validation.

use rand::rngs::StdRng;
use rand::seq::SliceRandom;
use rand::SeedableRng;

use super::{fold_bounds, CvSplitter};
use crate::error::{ModelSelectionError, Result};

/// Plain K-fold cross-validation.
///
/// The sample indices are partitioned into `k` folds; each fold serves as the
/// test set once while the remaining `k - 1` folds form the training set. With
/// [`shuffle`](KFold::with_shuffle) the indices are permuted (deterministically,
/// from the given seed) before partitioning.
///
/// This is a from-scratch implementation, so the crate has **zero** required
/// dependency on `smartcore`; it is functionally equivalent to
/// `smartcore::model_selection::KFold` for anyone who would rather use that.
///
/// ```
/// use model_selection_rs::splitters::{CvSplitter, KFold};
///
/// let kf = KFold::new(5).unwrap().with_shuffle(42);
/// let splits = kf.split(50).unwrap();
/// assert_eq!(splits.len(), 5);
/// ```
#[derive(Debug, Clone)]
pub struct KFold {
    n_splits: usize,
    shuffle: bool,
    seed: u64,
}

impl KFold {
    /// Create a `KFold` with `n_splits` folds and no shuffling.
    ///
    /// # Errors
    ///
    /// Returns [`ModelSelectionError::InvalidSplitCount`] if `n_splits < 2`.
    pub fn new(n_splits: usize) -> Result<Self> {
        if n_splits < 2 {
            return Err(ModelSelectionError::InvalidSplitCount {
                msg: format!("n_splits must be >= 2, got {n_splits}"),
            });
        }
        Ok(Self {
            n_splits,
            shuffle: false,
            seed: 0,
        })
    }

    /// Enable shuffling of the sample order before folding, using `seed`.
    #[must_use]
    pub fn with_shuffle(mut self, seed: u64) -> Self {
        self.shuffle = true;
        self.seed = seed;
        self
    }

    /// Build the (possibly shuffled) index order the folds are carved from.
    pub(crate) fn ordered_indices(&self, n_samples: usize) -> Vec<usize> {
        let mut indices: Vec<usize> = (0..n_samples).collect();
        if self.shuffle {
            let mut rng = StdRng::seed_from_u64(self.seed);
            indices.shuffle(&mut rng);
        }
        indices
    }
}

impl CvSplitter for KFold {
    fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>> {
        if self.n_splits > n_samples {
            return Err(ModelSelectionError::NotEnoughSamples {
                needed: self.n_splits,
                got: n_samples,
            });
        }
        let indices = self.ordered_indices(n_samples);
        let bounds = fold_bounds(n_samples, self.n_splits);

        let splits = bounds
            .into_iter()
            .map(|(start, end)| {
                let test: Vec<usize> = indices[start..end].to_vec();
                let mut train: Vec<usize> = Vec::with_capacity(n_samples - test.len());
                train.extend_from_slice(&indices[..start]);
                train.extend_from_slice(&indices[end..]);
                (train, test)
            })
            .collect();
        Ok(splits)
    }

    fn n_splits(&self) -> usize {
        self.n_splits
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashSet;

    #[test]
    fn rejects_fewer_than_two_folds() {
        assert!(KFold::new(1).is_err());
        assert!(KFold::new(0).is_err());
    }

    #[test]
    fn errors_when_more_folds_than_samples() {
        let kf = KFold::new(5).unwrap();
        assert!(matches!(
            kf.split(3),
            Err(ModelSelectionError::NotEnoughSamples { needed: 5, got: 3 })
        ));
    }

    #[test]
    fn every_sample_tested_exactly_once() {
        let kf = KFold::new(4).unwrap();
        let splits = kf.split(23).unwrap();
        let mut seen: Vec<usize> = splits.iter().flat_map(|(_, te)| te.clone()).collect();
        seen.sort_unstable();
        assert_eq!(seen, (0..23).collect::<Vec<_>>());
    }

    #[test]
    fn train_and_test_are_disjoint_and_cover_all() {
        let kf = KFold::new(3).unwrap().with_shuffle(7);
        for (train, test) in kf.split(20).unwrap() {
            let tr: HashSet<_> = train.iter().collect();
            let te: HashSet<_> = test.iter().collect();
            assert!(tr.is_disjoint(&te));
            assert_eq!(tr.len() + te.len(), 20);
        }
    }

    #[test]
    fn fold_sizes_differ_by_at_most_one() {
        let kf = KFold::new(4).unwrap();
        let sizes: Vec<usize> = kf
            .split(23)
            .unwrap()
            .iter()
            .map(|(_, te)| te.len())
            .collect();
        let max = *sizes.iter().max().unwrap();
        let min = *sizes.iter().min().unwrap();
        assert!(max - min <= 1);
    }

    #[test]
    fn shuffle_is_deterministic_for_a_seed() {
        let a = KFold::new(3).unwrap().with_shuffle(99).split(15).unwrap();
        let b = KFold::new(3).unwrap().with_shuffle(99).split(15).unwrap();
        assert_eq!(a, b);
    }
}