model-selection-rs 0.1.0

//! Stratified K-fold — K-fold that preserves per-class proportions.

use std::collections::HashMap;
use std::hash::Hash;

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

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

/// Stratified K-fold cross-validation.
///
/// Each fold keeps roughly the same class distribution as the full dataset.
/// This ports the previously hand-rolled stratification logic from the guide's
/// Evaluation addendum into a real, tested implementation: samples are grouped
/// by class label, then each class's samples are distributed across the folds so
/// every fold receives a proportional slice of every class.
///
/// The label type is generic (`L: Eq + Hash + Clone`), so string, integer or
/// enum labels all work with no forced mapping step — the same convention as the
/// sibling `imbalance-rs` crate. Labels are supplied at construction and stored,
/// which is what lets a stratified splitter satisfy the plain
/// [`CvSplitter`](crate::splitters::CvSplitter) interface (see the module docs).
///
/// # Small classes
///
/// A class with fewer than `n_splits` samples cannot appear in every fold. Such
/// a class is spread across as many folds as it can fill (one sample each,
/// leading folds first) and a warning is emitted on stderr — matching the
/// "warn and adjust rather than hard-error where reasonable" policy.
///
/// ```
/// use ndarray::array;
/// use model_selection_rs::splitters::{CvSplitter, StratifiedKFold};
///
/// let y = array![0, 0, 0, 0, 1, 1, 1, 1];
/// let skf = StratifiedKFold::new(2, &y).unwrap();
/// let splits = skf.split(y.len()).unwrap();
/// // Each fold's test set holds two 0s and two 1s.
/// assert_eq!(splits.len(), 2);
/// ```
#[derive(Debug, Clone)]
pub struct StratifiedKFold<L> {
    n_splits: usize,
    shuffle: bool,
    seed: u64,
    labels: Vec<L>,
}

impl<L: Eq + Hash + Clone> StratifiedKFold<L> {
    /// Create a `StratifiedKFold` over the class labels `y`.
    ///
    /// # Errors
    ///
    /// Returns [`ModelSelectionError::InvalidSplitCount`] if `n_splits < 2`.
    pub fn new(n_splits: usize, y: &Array1<L>) -> 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,
            labels: y.to_vec(),
        })
    }

    /// Shuffle each class's samples (deterministically, from `seed`) before
    /// distributing them across folds.
    #[must_use]
    pub fn with_shuffle(mut self, seed: u64) -> Self {
        self.shuffle = true;
        self.seed = seed;
        self
    }

    /// Group sample indices by class, preserving first-appearance class order
    /// for deterministic output.
    fn class_indices(&self) -> Vec<Vec<usize>> {
        let mut order: Vec<L> = Vec::new();
        let mut map: HashMap<L, Vec<usize>> = HashMap::new();
        for (i, label) in self.labels.iter().enumerate() {
            map.entry(label.clone()).or_insert_with(|| {
                order.push(label.clone());
                Vec::new()
            });
            map.get_mut(label).unwrap().push(i);
        }
        order.into_iter().map(|c| map.remove(&c).unwrap()).collect()
    }
}

/// Shared stratification core: distribute each class's index list across
/// `n_splits` test folds proportionally. Returns one test-index vector per fold.
///
/// Used by both [`StratifiedKFold`] and
/// [`RepeatedStratifiedKFold`](super::RepeatedStratifiedKFold).
pub(crate) fn stratified_test_folds(
    class_indices: &[Vec<usize>],
    n_splits: usize,
    shuffle: bool,
    seed: u64,
) -> Vec<Vec<usize>> {
    let mut rng = StdRng::seed_from_u64(seed);
    let mut test_folds: Vec<Vec<usize>> = vec![Vec::new(); n_splits];

    for indices in class_indices {
        let mut idx = indices.clone();
        if shuffle {
            idx.shuffle(&mut rng);
        }
        if idx.len() < n_splits {
            eprintln!(
                "model-selection-rs: StratifiedKFold — a class has {} sample(s), \
                 fewer than n_splits={n_splits}; it will be present in only {} of \
                 the {n_splits} folds.",
                idx.len(),
                idx.len()
            );
        }
        // Near-equal chunking of this class across folds; fold j gets chunk j.
        for (fold, (start, end)) in fold_bounds(idx.len(), n_splits).into_iter().enumerate() {
            test_folds[fold].extend_from_slice(&idx[start..end]);
        }
    }

    // Sort each fold so output is order-stable regardless of class iteration.
    for fold in &mut test_folds {
        fold.sort_unstable();
    }
    test_folds
}

/// Turn per-fold test sets into `(train, test)` pairs over `n_samples`.
pub(crate) fn test_folds_to_splits(
    test_folds: Vec<Vec<usize>>,
    n_samples: usize,
) -> Vec<(Vec<usize>, Vec<usize>)> {
    test_folds
        .into_iter()
        .map(|test| {
            let in_test: std::collections::HashSet<usize> = test.iter().copied().collect();
            let train: Vec<usize> = (0..n_samples).filter(|i| !in_test.contains(i)).collect();
            (train, test)
        })
        .collect()
}

impl<L: Eq + Hash + Clone> CvSplitter for StratifiedKFold<L> {
    fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>> {
        if n_samples != self.labels.len() {
            return Err(ModelSelectionError::ShapeMismatch {
                expected: self.labels.len(),
                got: n_samples,
            });
        }
        if self.n_splits > n_samples {
            return Err(ModelSelectionError::NotEnoughSamples {
                needed: self.n_splits,
                got: n_samples,
            });
        }
        let class_indices = self.class_indices();
        let test_folds =
            stratified_test_folds(&class_indices, self.n_splits, self.shuffle, self.seed);
        Ok(test_folds_to_splits(test_folds, n_samples))
    }

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

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

    fn class_proportions<L: Eq + Hash + Clone>(labels: &[L], idx: &[usize]) -> HashMap<L, f64> {
        let mut counts: HashMap<L, usize> = HashMap::new();
        for &i in idx {
            *counts.entry(labels[i].clone()).or_default() += 1;
        }
        let total = idx.len() as f64;
        counts
            .into_iter()
            .map(|(k, v)| (k, v as f64 / total))
            .collect()
    }

    #[test]
    fn balanced_two_class_splits_evenly() {
        let y = array![0, 0, 0, 0, 1, 1, 1, 1];
        let skf = StratifiedKFold::new(2, &y).unwrap();
        for (_, test) in skf.split(y.len()).unwrap() {
            let props = class_proportions(y.as_slice().unwrap(), &test);
            assert_eq!(props[&0], 0.5);
            assert_eq!(props[&1], 0.5);
        }
    }

    #[test]
    fn preserves_proportions_under_heavy_imbalance() {
        // 90 of class 0, 10 of class 1.
        let mut v = vec![0; 90];
        v.extend(std::iter::repeat(1).take(10));
        let y = Array1::from(v);
        let overall = class_proportions(y.as_slice().unwrap(), &(0..100).collect::<Vec<_>>());
        let skf = StratifiedKFold::new(5, &y).unwrap().with_shuffle(1);
        for (_, test) in skf.split(100).unwrap() {
            let props = class_proportions(y.as_slice().unwrap(), &test);
            for class in [0, 1] {
                assert!(
                    (props[&class] - overall[&class]).abs() < 0.05,
                    "class {class}: fold {} vs overall {}",
                    props[&class],
                    overall[&class]
                );
            }
        }
    }

    #[test]
    fn every_sample_tested_once() {
        let y = array![0, 1, 0, 1, 0, 1, 0, 1, 0, 1];
        let skf = StratifiedKFold::new(5, &y).unwrap();
        let mut seen: Vec<usize> = skf
            .split(10)
            .unwrap()
            .iter()
            .flat_map(|(_, te)| te.clone())
            .collect();
        seen.sort_unstable();
        assert_eq!(seen, (0..10).collect::<Vec<_>>());
    }

    #[test]
    fn shape_mismatch_when_n_samples_disagrees() {
        let y = array![0, 1, 0, 1];
        let skf = StratifiedKFold::new(2, &y).unwrap();
        assert!(matches!(
            skf.split(5),
            Err(ModelSelectionError::ShapeMismatch {
                expected: 4,
                got: 5
            })
        ));
    }

    #[test]
    fn string_labels_work() {
        let y = array!["cat", "dog", "cat", "dog"];
        let skf = StratifiedKFold::new(2, &y).unwrap();
        assert_eq!(skf.split(4).unwrap().len(), 2);
    }
}