model_selection_rs/splitters/

group_kfold.rs

//! Group-aware K-fold — no group appears in both train and test of a fold.

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

use ndarray::Array1;

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

/// Group K-fold cross-validation.
///
/// Guarantees that the samples of any one group (e.g. a patient id, a user id)
/// never straddle the train/test boundary within a fold. Group leakage — the
/// same entity contributing rows to both training and evaluation — is an
/// easy-to-miss correctness bug in real ML work; this splitter makes it
/// structurally impossible, complementing what
/// [`StratifiedKFold`](super::StratifiedKFold) does for class balance.
///
/// Groups are assigned whole to folds using a greedy largest-group-first
/// heuristic (the same idea scikit-learn uses): repeatedly place the largest
/// remaining group into the fold that currently holds the fewest samples. This
/// keeps fold sizes close without ever splitting a group.
///
/// The number of splits must not exceed the number of distinct groups.
///
/// ```
/// use ndarray::array;
/// use model_selection_rs::splitters::{CvSplitter, GroupKFold};
///
/// let groups = array![1, 1, 2, 2, 3, 3, 4, 4];
/// let gkf = GroupKFold::new(2, &groups).unwrap();
/// let splits = gkf.split(groups.len()).unwrap();
/// assert_eq!(splits.len(), 2);
/// ```
#[derive(Debug, Clone)]
pub struct GroupKFold<G> {
    n_splits: usize,
    groups: Vec<G>,
}

impl<G: Eq + Hash + Clone> GroupKFold<G> {
    /// Create a `GroupKFold` over the per-sample `groups`.
    ///
    /// # Errors
    ///
    /// Returns [`ModelSelectionError::InvalidSplitCount`] if `n_splits < 2`.
    pub fn new(n_splits: usize, groups: &Array1<G>) -> Result<Self> {
        if n_splits < 2 {
            return Err(ModelSelectionError::InvalidSplitCount {
                msg: format!("n_splits must be >= 2, got {n_splits}"),
            });
        }
        Ok(Self {
            n_splits,
            groups: groups.to_vec(),
        })
    }
}

/// Assign whole groups to `n_splits` folds, largest group first into the
/// currently-smallest fold. Returns one test-index vector per fold.
///
/// Shared with [`StratifiedGroupKFold`](super::StratifiedGroupKFold)'s tests via
/// the group-collection helper below.
pub(crate) fn group_test_folds<G: Eq + Hash + Clone>(
    groups: &[G],
    n_splits: usize,
) -> Result<Vec<Vec<usize>>> {
    let grouped = collect_groups(groups);
    if grouped.len() < n_splits {
        return Err(ModelSelectionError::InvalidSplitCount {
            msg: format!(
                "n_splits={n_splits} exceeds the number of distinct groups ({})",
                grouped.len()
            ),
        });
    }

    // Sort groups by size, descending (ties broken by first-appearance order,
    // which `collect_groups` already established, for determinism).
    let mut members: Vec<Vec<usize>> = grouped;
    members.sort_by_key(|m| std::cmp::Reverse(m.len()));

    let mut test_folds: Vec<Vec<usize>> = vec![Vec::new(); n_splits];
    let mut fold_sizes = vec![0usize; n_splits];
    for group_members in members {
        // Fold with the fewest samples so far (lowest index breaks ties).
        let target = fold_sizes
            .iter()
            .enumerate()
            .min_by_key(|(_, &size)| size)
            .map(|(i, _)| i)
            .unwrap();
        fold_sizes[target] += group_members.len();
        test_folds[target].extend(group_members);
    }
    for fold in &mut test_folds {
        fold.sort_unstable();
    }
    Ok(test_folds)
}

/// Group sample indices by group key, preserving first-appearance order.
pub(crate) fn collect_groups<G: Eq + Hash + Clone>(groups: &[G]) -> Vec<Vec<usize>> {
    let mut order: Vec<G> = Vec::new();
    let mut map: HashMap<G, Vec<usize>> = HashMap::new();
    for (i, g) in groups.iter().enumerate() {
        map.entry(g.clone()).or_insert_with(|| {
            order.push(g.clone());
            Vec::new()
        });
        map.get_mut(g).unwrap().push(i);
    }
    order.into_iter().map(|g| map.remove(&g).unwrap()).collect()
}

impl<G: Eq + Hash + Clone> CvSplitter for GroupKFold<G> {
    fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>> {
        if n_samples != self.groups.len() {
            return Err(ModelSelectionError::ShapeMismatch {
                expected: self.groups.len(),
                got: n_samples,
            });
        }
        let test_folds = group_test_folds(&self.groups, self.n_splits)?;
        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;
    use std::collections::HashSet;

    /// Every group must sit entirely on one side of every fold's split.
    fn assert_no_leakage<G: Eq + Hash + Clone>(groups: &[G], splits: &[(Vec<usize>, Vec<usize>)]) {
        for (train, test) in splits {
            let train_groups: HashSet<G> = train.iter().map(|&i| groups[i].clone()).collect();
            let test_groups: HashSet<G> = test.iter().map(|&i| groups[i].clone()).collect();
            assert!(
                train_groups.is_disjoint(&test_groups),
                "a group leaked across the train/test boundary"
            );
        }
    }

    #[test]
    fn no_group_leaks() {
        let groups = array![1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4, 5];
        let gkf = GroupKFold::new(3, &groups).unwrap();
        let splits = gkf.split(groups.len()).unwrap();
        assert_no_leakage(groups.as_slice().unwrap(), &splits);
    }

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

    #[test]
    fn errors_when_more_folds_than_groups() {
        let groups = array![1, 1, 2, 2];
        let gkf = GroupKFold::new(3, &groups).unwrap();
        assert!(matches!(
            gkf.split(4),
            Err(ModelSelectionError::InvalidSplitCount { .. })
        ));
    }

    #[test]
    fn string_group_ids_work() {
        let groups = array!["a", "a", "b", "c"];
        let gkf = GroupKFold::new(3, &groups).unwrap();
        let splits = gkf.split(4).unwrap();
        assert_no_leakage(groups.as_slice().unwrap(), &splits);
    }
}