model_selection_rs/splitters/

stratified_group_kfold.rs

//! Stratified **and** group-aware K-fold.

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

use ndarray::Array1;

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

/// K-fold that tries to satisfy **two** constraints at once: keep class
/// proportions balanced across folds *and* never let a group straddle the
/// train/test boundary.
///
/// # Approximation
///
/// Perfectly satisfying both constraints simultaneously is not always possible —
/// a group is indivisible, so its whole class makeup lands in one fold. This
/// implementation therefore uses the same documented greedy heuristic as
/// scikit-learn's `StratifiedGroupKFold`, and does not claim exactness:
///
/// 1. Count each class within each group.
/// 2. Visit groups in order of decreasing spread (standard deviation) of their
///    per-class counts, so the "lumpiest" groups are placed first while there is
///    still freedom to balance around them.
/// 3. Place each group in whichever fold keeps the per-fold class distribution
///    closest to uniform (minimising the mean, over classes, of the standard
///    deviation of each class's per-fold share); ties break toward the smaller
///    fold.
///
/// Group integrity is exact (groups are never split); only class balance is
/// approximate. `n_splits` must not exceed the number of distinct groups.
///
/// ```
/// use ndarray::array;
/// use model_selection_rs::splitters::{CvSplitter, StratifiedGroupKFold};
///
/// let y      = array![0, 0, 1, 1, 0, 1, 0, 1];
/// let groups = array![1, 1, 2, 2, 3, 3, 4, 4];
/// let sgkf = StratifiedGroupKFold::new(2, &y, &groups).unwrap();
/// let splits = sgkf.split(y.len()).unwrap();
/// assert_eq!(splits.len(), 2);
/// ```
#[derive(Debug, Clone)]
pub struct StratifiedGroupKFold<L, G> {
    n_splits: usize,
    labels: Vec<L>,
    groups: Vec<G>,
}

impl<L: Eq + Hash + Clone, G: Eq + Hash + Clone> StratifiedGroupKFold<L, G> {
    /// Create a `StratifiedGroupKFold` over class labels `y` and `groups`.
    ///
    /// # Errors
    ///
    /// Returns [`ModelSelectionError::InvalidSplitCount`] if `n_splits < 2`, or
    /// [`ModelSelectionError::ShapeMismatch`] if `y` and `groups` differ in
    /// length.
    pub fn new(n_splits: usize, y: &Array1<L>, groups: &Array1<G>) -> Result<Self> {
        if n_splits < 2 {
            return Err(ModelSelectionError::InvalidSplitCount {
                msg: format!("n_splits must be >= 2, got {n_splits}"),
            });
        }
        if y.len() != groups.len() {
            return Err(ModelSelectionError::ShapeMismatch {
                expected: y.len(),
                got: groups.len(),
            });
        }
        Ok(Self {
            n_splits,
            labels: y.to_vec(),
            groups: groups.to_vec(),
        })
    }
}

/// Population standard deviation of a slice (0.0 for length < 2).
fn std_dev(values: &[f64]) -> f64 {
    let n = values.len();
    if n < 2 {
        return 0.0;
    }
    let mean = values.iter().sum::<f64>() / n as f64;
    let var = values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / n as f64;
    var.sqrt()
}

impl<L: Eq + Hash + Clone, G: Eq + Hash + Clone> CvSplitter for StratifiedGroupKFold<L, G> {
    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,
            });
        }

        // Index the classes 0..n_classes by first appearance.
        let mut class_index: HashMap<L, usize> = HashMap::new();
        for label in &self.labels {
            let next = class_index.len();
            class_index.entry(label.clone()).or_insert(next);
        }
        let n_classes = class_index.len();

        let group_members = collect_groups(&self.groups);
        if group_members.len() < self.n_splits {
            return Err(ModelSelectionError::InvalidSplitCount {
                msg: format!(
                    "n_splits={} exceeds the number of distinct groups ({})",
                    self.n_splits,
                    group_members.len()
                ),
            });
        }

        // Per-group class counts, and overall per-class totals.
        let mut group_class_counts: Vec<Vec<f64>> = Vec::with_capacity(group_members.len());
        let mut class_totals = vec![0.0f64; n_classes];
        for members in &group_members {
            let mut counts = vec![0.0f64; n_classes];
            for &i in members {
                let c = class_index[&self.labels[i]];
                counts[c] += 1.0;
                class_totals[c] += 1.0;
            }
            group_class_counts.push(counts);
        }
        // Guard against a zero divisor for classes with no samples.
        for total in &mut class_totals {
            if *total == 0.0 {
                *total = 1.0;
            }
        }

        // Order groups by decreasing spread of their class counts.
        let mut order: Vec<usize> = (0..group_members.len()).collect();
        order.sort_by(|&a, &b| {
            std_dev(&group_class_counts[b])
                .partial_cmp(&std_dev(&group_class_counts[a]))
                .unwrap_or(std::cmp::Ordering::Equal)
                .then(a.cmp(&b))
        });

        // Greedy assignment.
        let mut fold_class_counts = vec![vec![0.0f64; n_classes]; self.n_splits];
        let mut fold_sizes = vec![0usize; self.n_splits];
        let mut assignment = vec![0usize; group_members.len()];

        for &g in &order {
            let counts = &group_class_counts[g];
            let mut best_fold = 0usize;
            let mut best_std = f64::INFINITY;
            let mut best_size = usize::MAX;

            for fold in 0..self.n_splits {
                // Tentatively add this group's counts to `fold`.
                for c in 0..n_classes {
                    fold_class_counts[fold][c] += counts[c];
                }
                // Mean over classes of the std (over folds) of each class share.
                let mut std_sum = 0.0;
                for c in 0..n_classes {
                    let shares: Vec<f64> = (0..self.n_splits)
                        .map(|f| fold_class_counts[f][c] / class_totals[c])
                        .collect();
                    std_sum += std_dev(&shares);
                }
                let mean_std = std_sum / n_classes as f64;
                // Undo.
                for c in 0..n_classes {
                    fold_class_counts[fold][c] -= counts[c];
                }

                let size = fold_sizes[fold];
                if mean_std < best_std - 1e-12
                    || ((mean_std - best_std).abs() <= 1e-12 && size < best_size)
                {
                    best_std = mean_std;
                    best_fold = fold;
                    best_size = size;
                }
            }

            for c in 0..n_classes {
                fold_class_counts[best_fold][c] += counts[c];
            }
            fold_sizes[best_fold] += group_members[g].len();
            assignment[g] = best_fold;
        }

        // Materialise per-fold test indices.
        let mut test_folds: Vec<Vec<usize>> = vec![Vec::new(); self.n_splits];
        for (g, members) in group_members.into_iter().enumerate() {
            test_folds[assignment[g]].extend(members);
        }
        for fold in &mut test_folds {
            fold.sort_unstable();
        }

        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;

    fn assert_no_leakage(groups: &[i32], splits: &[(Vec<usize>, Vec<usize>)]) {
        for (train, test) in splits {
            let tr: HashSet<i32> = train.iter().map(|&i| groups[i]).collect();
            let te: HashSet<i32> = test.iter().map(|&i| groups[i]).collect();
            assert!(tr.is_disjoint(&te), "group leaked across boundary");
        }
    }

    #[test]
    fn no_group_leaks_and_all_tested_once() {
        let y = array![0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 0, 1];
        let groups = array![1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6];
        let sgkf = StratifiedGroupKFold::new(3, &y, &groups).unwrap();
        let splits = sgkf.split(y.len()).unwrap();
        assert_no_leakage(groups.as_slice().unwrap(), &splits);

        let mut seen: Vec<usize> = splits.iter().flat_map(|(_, te)| te.clone()).collect();
        seen.sort_unstable();
        assert_eq!(seen, (0..12).collect::<Vec<_>>());
    }

    #[test]
    fn keeps_class_balance_reasonably() {
        // 12 groups, each purely one class, 6 of each class.
        let mut y = Vec::new();
        let mut groups = Vec::new();
        for g in 0..12 {
            let class = g % 2; // alternate pure-class groups
            for _ in 0..3 {
                y.push(class);
                groups.push(g);
            }
        }
        let y = Array1::from(y);
        let groups = Array1::from(groups);
        let sgkf = StratifiedGroupKFold::new(3, &y, &groups).unwrap();
        let splits = sgkf.split(y.len()).unwrap();
        for (_, test) in &splits {
            let ones = test.iter().filter(|&&i| y[i] == 1).count();
            let frac = ones as f64 / test.len() as f64;
            assert!(
                (frac - 0.5).abs() < 0.2,
                "fold class-1 share {frac} off balance"
            );
        }
    }

    #[test]
    fn shape_mismatch_on_unequal_lengths() {
        let y = array![0, 1, 0];
        let groups = array![1, 2, 3, 4];
        assert!(matches!(
            StratifiedGroupKFold::new(2, &y, &groups),
            Err(ModelSelectionError::ShapeMismatch { .. })
        ));
    }
}