model-selection-rs 0.1.0

//! Cross-splitter consistency and edge-case tests (plan Milestone 10).
//!
//! These sit above the per-splitter unit tests: they assert properties that
//! must hold *across* the whole splitter family, and probe the degenerate
//! inputs each splitter promises to handle a particular way.

use std::collections::HashSet;

use model_selection_rs::splitters::{
    CvSplitter, GroupKFold, KFold, LeaveOneOut, RepeatedKFold, RepeatedStratifiedKFold,
    ShuffleSplit, StratifiedGroupKFold, StratifiedKFold, StratifiedShuffleSplit, SubsetSize,
    TimeSeriesSplit,
};
use ndarray::Array1;

/// Every splitter: train and test indices of each fold must be disjoint and
/// within range.
#[test]
fn all_splitters_produce_disjoint_in_range_folds() {
    let n = 30;
    let y = Array1::from((0..n).map(|i| (i % 3) as i32).collect::<Vec<_>>());
    let groups = Array1::from((0..n).map(|i| (i / 3) as i32).collect::<Vec<_>>());

    let splitters: Vec<Box<dyn CvSplitter>> = vec![
        Box::new(KFold::new(5).unwrap()),
        Box::new(KFold::new(5).unwrap().with_shuffle(0)),
        Box::new(StratifiedKFold::new(3, &y).unwrap()),
        Box::new(GroupKFold::new(4, &groups).unwrap()),
        Box::new(StratifiedGroupKFold::new(3, &y, &groups).unwrap()),
        Box::new(TimeSeriesSplit::new(4).unwrap()),
        Box::new(ShuffleSplit::new(5).with_test_size(SubsetSize::Fraction(0.2))),
        Box::new(StratifiedShuffleSplit::new(5, &y).with_test_size(SubsetSize::Fraction(0.2))),
        Box::new(RepeatedKFold::new(3, 2, 0).unwrap()),
        Box::new(RepeatedStratifiedKFold::new(3, 2, 0, &y).unwrap()),
        Box::new(LeaveOneOut),
    ];

    for splitter in &splitters {
        for (train, test) in splitter.split(n).unwrap() {
            let tr: HashSet<usize> = train.iter().copied().collect();
            let te: HashSet<usize> = test.iter().copied().collect();
            assert!(tr.is_disjoint(&te), "train/test overlap");
            assert!(train.iter().all(|&i| i < n), "train index out of range");
            assert!(test.iter().all(|&i| i < n), "test index out of range");
            assert!(!test.is_empty(), "empty test fold");
        }
    }
}

/// KFold-family splitters partition the dataset: every sample is tested exactly
/// once across the full set of splits. (ShuffleSplit is deliberately excluded —
/// its splits are independent draws, not a partition.)
#[test]
fn kfold_family_partitions_test_sets() {
    let n = 24;
    let y = Array1::from((0..n).map(|i| (i % 2) as i32).collect::<Vec<_>>());
    let groups = Array1::from((0..n).map(|i| (i / 2) as i32).collect::<Vec<_>>());

    let partitioners: Vec<(&str, Box<dyn CvSplitter>)> = vec![
        ("kfold", Box::new(KFold::new(6).unwrap())),
        ("stratified", Box::new(StratifiedKFold::new(6, &y).unwrap())),
        ("group", Box::new(GroupKFold::new(6, &groups).unwrap())),
        (
            "stratified_group",
            Box::new(StratifiedGroupKFold::new(6, &y, &groups).unwrap()),
        ),
        ("loo", Box::new(LeaveOneOut)),
    ];

    for (name, splitter) in &partitioners {
        let mut seen: Vec<usize> = splitter
            .split(n)
            .unwrap()
            .iter()
            .flat_map(|(_, te)| te.clone())
            .collect();
        seen.sort_unstable();
        assert_eq!(
            seen,
            (0..n).collect::<Vec<_>>(),
            "{name} did not partition test sets"
        );
    }
}

#[test]
fn kfold_errors_when_more_folds_than_samples() {
    assert!(KFold::new(10).unwrap().split(5).is_err());
}

#[test]
fn stratified_single_class_still_splits() {
    // A single class is degenerate but valid — it should behave like plain
    // KFold and not error.
    let y = Array1::from(vec![7i32; 10]);
    let skf = StratifiedKFold::new(5, &y).unwrap();
    let splits = skf.split(10).unwrap();
    assert_eq!(splits.len(), 5);
    let mut seen: Vec<usize> = splits.iter().flat_map(|(_, te)| te.clone()).collect();
    seen.sort_unstable();
    assert_eq!(seen, (0..10).collect::<Vec<_>>());
}

#[test]
fn stratified_tiny_class_is_adjusted_not_errored() {
    // Class 1 has a single sample but n_splits is 3: must warn+adjust, not error.
    let mut v = vec![0i32; 9];
    v.push(1);
    let y = Array1::from(v);
    let skf = StratifiedKFold::new(3, &y).unwrap();
    let splits = skf.split(10).unwrap();
    assert_eq!(splits.len(), 3);
    // The lone class-1 sample (index 9) is tested in exactly one fold.
    let appearances = splits.iter().filter(|(_, te)| te.contains(&9)).count();
    assert_eq!(appearances, 1);
}

#[test]
fn group_single_group_cannot_make_two_folds() {
    let groups = Array1::from(vec![42i32; 8]);
    let gkf = GroupKFold::new(2, &groups).unwrap();
    assert!(gkf.split(8).is_err(), "one group can't fill two folds");
}

#[test]
fn time_series_all_identical_order_still_chronological() {
    // TimeSeriesSplit works on row order, so "identical timestamps" == plain
    // positional order; the chronological property is about position.
    let tss = TimeSeriesSplit::new(3).unwrap();
    for (train, test) in tss.split(12).unwrap() {
        assert!(train.iter().max().unwrap() < test.iter().min().unwrap());
    }
}

#[test]
fn n_splits_reported_matches_produced_count() {
    let n = 20;
    let y = Array1::from((0..n).map(|i| (i % 2) as i32).collect::<Vec<_>>());
    let cases: Vec<Box<dyn CvSplitter>> = vec![
        Box::new(KFold::new(5).unwrap()),
        Box::new(StratifiedKFold::new(4, &y).unwrap()),
        Box::new(TimeSeriesSplit::new(3).unwrap()),
        Box::new(ShuffleSplit::new(7)),
        Box::new(RepeatedKFold::new(5, 2, 0).unwrap()),
    ];
    for splitter in &cases {
        assert_eq!(splitter.n_splits(), splitter.split(n).unwrap().len());
    }
}