Skip to main content

qualia_core_db/solvers/learning/resampling/
folds.rs

1//! Resampling index generators — k-fold / LOOCV splits and train/test partition,
2//! with a deterministic shuffle so results are reproducible.
3
4/// Deterministic LCG (Numerical Recipes constants) for reproducible shuffles —
5/// no RNG dependency.
6struct Lcg(u64);
7impl Lcg {
8    fn next_below(&mut self, bound: usize) -> usize {
9        self.0 = self
10            .0
11            .wrapping_mul(6364136223846793005)
12            .wrapping_add(1442695040888963407);
13        ((self.0 >> 33) as usize) % bound.max(1)
14    }
15}
16
17/// In-place Fisher–Yates shuffle of `idx` seeded by `seed`.
18fn shuffle(idx: &mut [usize], seed: u64) {
19    let mut rng = Lcg(seed ^ 0x9E3779B97F4A7C15);
20    for i in (1..idx.len()).rev() {
21        let j = rng.next_below(i + 1);
22        idx.swap(i, j);
23    }
24}
25
26/// One train/test split by row index.
27#[derive(Debug, Clone)]
28pub struct Fold {
29    pub train: Vec<usize>,
30    pub test: Vec<usize>,
31}
32
33/// `k`-fold splits over `n` rows. `shuffle_rows` randomizes fold membership
34/// (seeded); otherwise folds are contiguous blocks. Empty if `k < 2` or `k > n`.
35pub fn k_fold(n: usize, k: usize, shuffle_rows: bool, seed: u64) -> Vec<Fold> {
36    if k < 2 || k > n {
37        return Vec::new();
38    }
39    let mut order: Vec<usize> = (0..n).collect();
40    if shuffle_rows {
41        shuffle(&mut order, seed);
42    }
43    // Fold sizes: the first `n % k` folds get one extra element.
44    let base = n / k;
45    let rem = n % k;
46    let mut folds = Vec::with_capacity(k);
47    let mut start = 0;
48    for f in 0..k {
49        let len = base + usize::from(f < rem);
50        let test: Vec<usize> = order[start..start + len].to_vec();
51        let train: Vec<usize> = order[..start]
52            .iter()
53            .chain(order[start + len..].iter())
54            .copied()
55            .collect();
56        folds.push(Fold { train, test });
57        start += len;
58    }
59    folds
60}
61
62/// Leave-one-out cross-validation = `n`-fold (each test set is a single row).
63pub fn loocv(n: usize) -> Vec<Fold> {
64    (0..n)
65        .map(|i| Fold {
66            test: vec![i],
67            train: (0..n).filter(|&j| j != i).collect(),
68        })
69        .collect()
70}
71
72/// A single train/test partition: `test_fraction` of the (shuffled) rows go to
73/// test. `None` if `test_fraction` is not in `(0,1)` or `n < 2`.
74pub fn train_test_split(
75    n: usize,
76    test_fraction: f64,
77    seed: u64,
78) -> Option<(Vec<usize>, Vec<usize>)> {
79    if n < 2 || !(0.0..1.0).contains(&test_fraction) || test_fraction <= 0.0 {
80        return None;
81    }
82    let n_test = ((n as f64 * test_fraction).round() as usize).clamp(1, n - 1);
83    let mut order: Vec<usize> = (0..n).collect();
84    shuffle(&mut order, seed);
85    let test = order[..n_test].to_vec();
86    let train = order[n_test..].to_vec();
87    Some((train, test))
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use std::collections::HashSet;
94
95    #[test]
96    fn k_fold_partitions_every_row_once_as_test() {
97        let folds = k_fold(10, 3, true, 42);
98        assert_eq!(folds.len(), 3);
99        let mut all_test: Vec<usize> = folds.iter().flat_map(|f| f.test.iter().copied()).collect();
100        all_test.sort_unstable();
101        assert_eq!(
102            all_test,
103            (0..10).collect::<Vec<_>>(),
104            "every row is a test exactly once"
105        );
106        // train and test are disjoint and together cover all rows.
107        for f in &folds {
108            let t: HashSet<usize> = f.test.iter().copied().collect();
109            assert!(
110                f.train.iter().all(|i| !t.contains(i)),
111                "train/test disjoint"
112            );
113            assert_eq!(f.train.len() + f.test.len(), 10);
114        }
115    }
116
117    #[test]
118    fn k_fold_sizes_balanced() {
119        // 10 rows, 3 folds → sizes 4,3,3.
120        let folds = k_fold(10, 3, false, 0);
121        let mut sizes: Vec<usize> = folds.iter().map(|f| f.test.len()).collect();
122        sizes.sort_unstable();
123        assert_eq!(sizes, vec![3, 3, 4]);
124    }
125
126    #[test]
127    fn loocv_has_n_folds_of_one() {
128        let folds = loocv(5);
129        assert_eq!(folds.len(), 5);
130        assert!(folds
131            .iter()
132            .all(|f| f.test.len() == 1 && f.train.len() == 4));
133    }
134
135    #[test]
136    fn train_test_split_sizes_and_disjoint() {
137        let (train, test) = train_test_split(100, 0.25, 7).unwrap();
138        assert_eq!(test.len(), 25);
139        assert_eq!(train.len(), 75);
140        let ts: HashSet<usize> = test.iter().copied().collect();
141        assert!(train.iter().all(|i| !ts.contains(i)));
142        assert!(train_test_split(1, 0.5, 0).is_none());
143        assert!(train_test_split(10, 1.5, 0).is_none());
144    }
145
146    #[test]
147    fn guards() {
148        assert!(k_fold(5, 1, false, 0).is_empty());
149        assert!(k_fold(3, 5, false, 0).is_empty());
150    }
151}