model-selection-rs 0.1.0

//! Cross-validation splitters.
//!
//! Every splitter implements the single [`CvSplitter`] trait, which yields
//! `(train_indices, test_indices)` pairs — **indices only**, never materialized
//! data copies, matching scikit-learn's memory-efficient convention. Callers
//! apply the indices to their data with ndarray fancy indexing
//! ([`ndarray::ArrayBase::select`]).
//!
//! # Design note: how label- and group-aware splitters fit one trait
//!
//! Milestone 1 of the project plan left one question open: should stratified /
//! group splitters need a *separate* trait taking `y` / `groups`, or can a
//! single trait serve everything? This crate resolves it in favour of **one
//! trait**. Label- and group-aware splitters
//! ([`StratifiedKFold`], [`GroupKFold`], [`StratifiedGroupKFold`],
//! [`StratifiedShuffleSplit`], [`RepeatedStratifiedKFold`]) take their labels /
//! groups at **construction time** and store them, then implement the same
//! [`CvSplitter::split`] as everything else.
//!
//! The upside is uniformity: [`cross_validate`](crate::evaluate::cross_validate),
//! [`learning_curve`](crate::evaluate::learning_curve) and friends accept *any*
//! `S: CvSplitter` with no special cases. The cost is that a stratified splitter
//! owns a copy of its label array — cheap in practice, since you always have `y`
//! in scope when you set up a CV loop, and labels are one small column.
//!
//! # Fallibility
//!
//! The plan sketched `split` as infallible (`-> Vec<..>`). It is promoted to
//! `-> Result<Vec<..>>` here because validation (too few samples, an impossible
//! time-series window, a label array whose length disagrees with `n_samples`)
//! genuinely can fail and a library should surface that rather than panic. The
//! per-split index math itself never fails once validation passes.

use crate::error::Result;

mod group_kfold;
mod kfold;
mod leave_one_out;
mod repeated;
mod shuffle_split;
mod stratified_group_kfold;
mod stratified_kfold;
mod stratified_shuffle_split;
mod time_series_split;

pub use group_kfold::GroupKFold;
pub use kfold::KFold;
pub use leave_one_out::LeaveOneOut;
pub use repeated::{RepeatedKFold, RepeatedStratifiedKFold};
pub use shuffle_split::{ShuffleSplit, SubsetSize};
pub use stratified_group_kfold::StratifiedGroupKFold;
pub use stratified_kfold::StratifiedKFold;
pub use stratified_shuffle_split::StratifiedShuffleSplit;
pub use time_series_split::TimeSeriesSplit;

/// A cross-validation splitting strategy.
///
/// Implementors return `(train, test)` index pairs for a dataset of
/// `n_samples` rows. The indices are into the original row order; apply them
/// with [`ndarray::ArrayBase::select`].
///
/// ```
/// use model_selection_rs::splitters::{CvSplitter, KFold};
///
/// let kf = KFold::new(3).unwrap();
/// let splits = kf.split(6).unwrap();
/// assert_eq!(splits.len(), 3);
/// for (train, test) in &splits {
///     assert_eq!(train.len() + test.len(), 6);
/// }
/// ```
pub trait CvSplitter {
    /// Produce every `(train_indices, test_indices)` pair for `n_samples` rows.
    ///
    /// # Errors
    ///
    /// Returns [`ModelSelectionError`](crate::error::ModelSelectionError) if the
    /// configuration cannot produce valid splits for `n_samples` (for example
    /// more folds than samples, or — for stored-label splitters — an
    /// `n_samples` that disagrees with the stored label array length).
    fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>>;

    /// Number of splits this strategy yields.
    fn n_splits(&self) -> usize;
}

/// Split `n_samples` indices into `k` contiguous, near-equal folds.
///
/// The first `n % k` folds receive one extra element, matching the fold-size
/// convention used by scikit-learn's `KFold`. Returns the fold *boundaries* as
/// `(start, end)` half-open ranges over a `0..n_samples` index space; callers
/// map those ranges onto whatever (possibly shuffled) index order they hold.
pub(crate) fn fold_bounds(n_samples: usize, k: usize) -> Vec<(usize, usize)> {
    let base = n_samples / k;
    let remainder = n_samples % k;
    let mut bounds = Vec::with_capacity(k);
    let mut start = 0;
    for fold in 0..k {
        let size = base + usize::from(fold < remainder);
        bounds.push((start, start + size));
        start += size;
    }
    bounds
}

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

    /// A trivial splitter used to prove the trait's bounds work end-to-end
    /// before any real splitter is layered on top (plan Milestone 1 DoD).
    struct PassthroughSplitter;

    impl CvSplitter for PassthroughSplitter {
        fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>> {
            // Everything is "train", nothing is "test" — the point is only to
            // exercise the trait object / generic bounds, not to be useful.
            Ok(vec![((0..n_samples).collect(), Vec::new())])
        }
        fn n_splits(&self) -> usize {
            1
        }
    }

    #[test]
    fn passthrough_works_as_trait_object() {
        let s: &dyn CvSplitter = &PassthroughSplitter;
        let splits = s.split(5).unwrap();
        assert_eq!(s.n_splits(), 1);
        assert_eq!(splits[0].0, vec![0, 1, 2, 3, 4]);
        assert!(splits[0].1.is_empty());
    }

    #[test]
    fn fold_bounds_distributes_remainder_to_leading_folds() {
        // 7 into 3 -> sizes 3, 2, 2
        assert_eq!(fold_bounds(7, 3), vec![(0, 3), (3, 5), (5, 7)]);
        // exact division -> equal sizes
        assert_eq!(fold_bounds(6, 3), vec![(0, 2), (2, 4), (4, 6)]);
    }
}