somatize_study/partition/
stratified_grouped.rs

1//! Groups whole, and among the ways of doing that the one that leaves the
2//! classes most even.
3
4use super::dealing::{assemble, checked, grouped_by, heaviest_first, in_order};
5use super::{Fold, Partition, PartitionError};
6use crate::Samples;
7use std::collections::BTreeSet;
8use std::fmt;
9
10/// The two constraints at once, which is where they stop composing cleanly:
11/// with the groups kept whole, exact strata are usually **not reachable at
12/// all**. So this is greedy and approximate, and says so — sklearn's
13/// `StratifiedGroupKFold` is greedy for the same reason, because the exact
14/// problem is a bin packing.
15///
16/// Needs both [`by_class`](Samples::by_class) and
17/// [`in_groups`](Samples::in_groups).
18#[derive(Debug, Clone, PartialEq, Eq, Hash)]
19pub struct StratifiedGrouped {
20    /// How many folds.
21    pub k: usize,
22}
23
24impl StratifiedGrouped {
25    /// The folds, in order.
26    pub fn folds(&self, samples: &Samples) -> Result<Vec<Fold>, PartitionError> {
27        let n = samples.n();
28        let strata = samples.strata().ok_or(PartitionError::NeedsClasses)?;
29        let groups = samples.groups().ok_or(PartitionError::NeedsGroups)?;
30        checked(self.k, n)?;
31        let classes = census(strata, self.k)?;
32
33        let mut tests = vec![Vec::new(); self.k];
34        let mut carried = vec![vec![0usize; classes.len()]; self.k];
35        for (_, indices) in heaviest_first(grouped_by(groups, &in_order(n)), self.k)? {
36            let mine = tally(&indices, strata, &classes);
37            let fold = evenest(&carried, &mine, classes.len());
38            for (class, count) in mine.iter().enumerate() {
39                carried[fold][class] += count;
40            }
41            tests[fold].extend(indices);
42        }
43        Ok(assemble(n, tests))
44    }
45}
46
47/// The classes present, and the check that each reaches every fold.
48fn census(strata: &[u32], k: usize) -> Result<Vec<u32>, PartitionError> {
49    let classes: Vec<u32> = strata
50        .iter()
51        .copied()
52        .collect::<BTreeSet<_>>()
53        .into_iter()
54        .collect();
55    for &class in &classes {
56        let count = strata.iter().filter(|&&c| c == class).count();
57        if count < k {
58            return Err(PartitionError::ClassTooSmall { class, count, k });
59        }
60    }
61    Ok(classes)
62}
63
64/// How many of each class this group carries.
65fn tally(indices: &[usize], strata: &[u32], classes: &[u32]) -> Vec<usize> {
66    let mut counts = vec![0usize; classes.len()];
67    for &index in indices {
68        if let Some(class) = classes.iter().position(|&c| c == strata[index]) {
69            counts[class] += 1;
70        }
71    }
72    counts
73}
74
75/// The fold where putting this group leaves the classes most evenly spread.
76/// Decided once per group and never revisited.
77fn evenest(carried: &[Vec<usize>], mine: &[usize], classes: usize) -> usize {
78    let mut best = 0;
79    let mut best_spread = f64::MAX;
80    for fold in 0..carried.len() {
81        let mut spread = 0.0;
82        for class in 0..classes {
83            let shares: Vec<f64> = carried
84                .iter()
85                .enumerate()
86                .map(|(other, counts)| {
87                    let extra = if other == fold { mine[class] } else { 0 };
88                    (counts[class] + extra) as f64
89                })
90                .collect();
91            spread += deviation(&shares);
92        }
93        if spread < best_spread {
94            best_spread = spread;
95            best = fold;
96        }
97    }
98    best
99}
100
101/// How far from equal a spread is. Standard deviation, unnormalised: only its
102/// ordering is used.
103fn deviation(values: &[f64]) -> f64 {
104    let mean = values.iter().sum::<f64>() / values.len() as f64;
105    (values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / values.len() as f64).sqrt()
106}
107
108impl fmt::Display for StratifiedGrouped {
109    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
110        write!(f, "stratified-grouped:{}", self.k)
111    }
112}
113
114impl From<StratifiedGrouped> for Partition {
115    fn from(cut: StratifiedGrouped) -> Self {
116        Self::StratifiedGrouped(cut)
117    }
118}
somatize_study/partition/stratified_grouped.rs

somatize_study/partition/
stratified_grouped.rs
