somatize_study/partition/
stratified.rs

1//! `k` folds where every class keeps the share it has in the whole.
2
3use super::dealing::{assemble, checked, deal, grouped_by, ordering};
4use super::{Fold, Partition, PartitionError};
5use crate::Samples;
6use std::fmt;
7
8/// Stratifying is not a different algorithm: it is a [`KFold`](crate::KFold)
9/// applied **inside each class**, the folds concatenated. That is why there is
10/// one scheme here and not sklearn's `KFold` / `StratifiedKFold` pair.
11///
12/// Needs [`by_class`](Samples::by_class).
13#[derive(Debug, Clone, PartialEq, Eq, Hash)]
14pub struct Stratified {
15    /// How many folds.
16    pub k: usize,
17    /// The seed, or `None` for the order they came in.
18    pub shuffle: Option<u64>,
19}
20
21impl Stratified {
22    /// The folds, in order.
23    ///
24    /// A class with fewer members than folds cannot be in all of them, and that
25    /// is an **error**: sklearn warns and carries on, which leaves a result you
26    /// cannot tell from a good one.
27    pub fn folds(&self, samples: &Samples) -> Result<Vec<Fold>, PartitionError> {
28        let n = samples.n();
29        let strata = samples.strata().ok_or(PartitionError::NeedsClasses)?;
30        checked(self.k, n)?;
31
32        let mut tests = vec![Vec::new(); self.k];
33        for (class, members) in grouped_by(strata, &ordering(n, self.shuffle)) {
34            if members.len() < self.k {
35                return Err(PartitionError::ClassTooSmall {
36                    class,
37                    count: members.len(),
38                    k: self.k,
39                });
40            }
41            for (fold, share) in deal(&members, self.k).into_iter().enumerate() {
42                tests[fold].extend(share);
43            }
44        }
45        Ok(assemble(n, tests))
46    }
47}
48
49impl fmt::Display for Stratified {
50    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51        super::scheme(f, "stratified", self.k, self.shuffle)
52    }
53}
54
55impl From<Stratified> for Partition {
56    fn from(cut: Stratified) -> Self {
57        Self::Stratified(cut)
58    }
59}
somatize_study/partition/stratified.rs

somatize_study/partition/
stratified.rs
