model_selection_rs/splitters/
mod.rs

1//! Cross-validation splitters.
2//!
3//! Every splitter implements the single [`CvSplitter`] trait, which yields
4//! `(train_indices, test_indices)` pairs — **indices only**, never materialized
5//! data copies, matching scikit-learn's memory-efficient convention. Callers
6//! apply the indices to their data with ndarray fancy indexing
7//! ([`ndarray::ArrayBase::select`]).
8//!
9//! # Design note: how label- and group-aware splitters fit one trait
10//!
11//! Milestone 1 of the project plan left one question open: should stratified /
12//! group splitters need a *separate* trait taking `y` / `groups`, or can a
13//! single trait serve everything? This crate resolves it in favour of **one
14//! trait**. Label- and group-aware splitters
15//! ([`StratifiedKFold`], [`GroupKFold`], [`StratifiedGroupKFold`],
16//! [`StratifiedShuffleSplit`], [`RepeatedStratifiedKFold`]) take their labels /
17//! groups at **construction time** and store them, then implement the same
18//! [`CvSplitter::split`] as everything else.
19//!
20//! The upside is uniformity: [`cross_validate`](crate::evaluate::cross_validate),
21//! [`learning_curve`](crate::evaluate::learning_curve) and friends accept *any*
22//! `S: CvSplitter` with no special cases. The cost is that a stratified splitter
23//! owns a copy of its label array — cheap in practice, since you always have `y`
24//! in scope when you set up a CV loop, and labels are one small column.
25//!
26//! # Fallibility
27//!
28//! The plan sketched `split` as infallible (`-> Vec<..>`). It is promoted to
29//! `-> Result<Vec<..>>` here because validation (too few samples, an impossible
30//! time-series window, a label array whose length disagrees with `n_samples`)
31//! genuinely can fail and a library should surface that rather than panic. The
32//! per-split index math itself never fails once validation passes.
33
34use crate::error::Result;
35
36mod group_kfold;
37mod kfold;
38mod leave_one_out;
39mod repeated;
40mod shuffle_split;
41mod stratified_group_kfold;
42mod stratified_kfold;
43mod stratified_shuffle_split;
44mod time_series_split;
45
46pub use group_kfold::GroupKFold;
47pub use kfold::KFold;
48pub use leave_one_out::LeaveOneOut;
49pub use repeated::{RepeatedKFold, RepeatedStratifiedKFold};
50pub use shuffle_split::{ShuffleSplit, SubsetSize};
51pub use stratified_group_kfold::StratifiedGroupKFold;
52pub use stratified_kfold::StratifiedKFold;
53pub use stratified_shuffle_split::StratifiedShuffleSplit;
54pub use time_series_split::TimeSeriesSplit;
55
56/// A cross-validation splitting strategy.
57///
58/// Implementors return `(train, test)` index pairs for a dataset of
59/// `n_samples` rows. The indices are into the original row order; apply them
60/// with [`ndarray::ArrayBase::select`].
61///
62/// ```
63/// use model_selection_rs::splitters::{CvSplitter, KFold};
64///
65/// let kf = KFold::new(3).unwrap();
66/// let splits = kf.split(6).unwrap();
67/// assert_eq!(splits.len(), 3);
68/// for (train, test) in &splits {
69///     assert_eq!(train.len() + test.len(), 6);
70/// }
71/// ```
72pub trait CvSplitter {
73    /// Produce every `(train_indices, test_indices)` pair for `n_samples` rows.
74    ///
75    /// # Errors
76    ///
77    /// Returns [`ModelSelectionError`](crate::error::ModelSelectionError) if the
78    /// configuration cannot produce valid splits for `n_samples` (for example
79    /// more folds than samples, or — for stored-label splitters — an
80    /// `n_samples` that disagrees with the stored label array length).
81    fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>>;
82
83    /// Number of splits this strategy yields.
84    fn n_splits(&self) -> usize;
85}
86
87/// Split `n_samples` indices into `k` contiguous, near-equal folds.
88///
89/// The first `n % k` folds receive one extra element, matching the fold-size
90/// convention used by scikit-learn's `KFold`. Returns the fold *boundaries* as
91/// `(start, end)` half-open ranges over a `0..n_samples` index space; callers
92/// map those ranges onto whatever (possibly shuffled) index order they hold.
93pub(crate) fn fold_bounds(n_samples: usize, k: usize) -> Vec<(usize, usize)> {
94    let base = n_samples / k;
95    let remainder = n_samples % k;
96    let mut bounds = Vec::with_capacity(k);
97    let mut start = 0;
98    for fold in 0..k {
99        let size = base + usize::from(fold < remainder);
100        bounds.push((start, start + size));
101        start += size;
102    }
103    bounds
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    /// A trivial splitter used to prove the trait's bounds work end-to-end
111    /// before any real splitter is layered on top (plan Milestone 1 DoD).
112    struct PassthroughSplitter;
113
114    impl CvSplitter for PassthroughSplitter {
115        fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>> {
116            // Everything is "train", nothing is "test" — the point is only to
117            // exercise the trait object / generic bounds, not to be useful.
118            Ok(vec![((0..n_samples).collect(), Vec::new())])
119        }
120        fn n_splits(&self) -> usize {
121            1
122        }
123    }
124
125    #[test]
126    fn passthrough_works_as_trait_object() {
127        let s: &dyn CvSplitter = &PassthroughSplitter;
128        let splits = s.split(5).unwrap();
129        assert_eq!(s.n_splits(), 1);
130        assert_eq!(splits[0].0, vec![0, 1, 2, 3, 4]);
131        assert!(splits[0].1.is_empty());
132    }
133
134    #[test]
135    fn fold_bounds_distributes_remainder_to_leading_folds() {
136        // 7 into 3 -> sizes 3, 2, 2
137        assert_eq!(fold_bounds(7, 3), vec![(0, 3), (3, 5), (5, 7)]);
138        // exact division -> equal sizes
139        assert_eq!(fold_bounds(6, 3), vec![(0, 2), (2, 4), (4, 6)]);
140    }
141}
model_selection_rs/splitters/mod.rs

model_selection_rs/splitters/
mod.rs
