qualia_core_db/solvers/learning/resampling/
folds.rs1struct 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
17fn 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#[derive(Debug, Clone)]
28pub struct Fold {
29 pub train: Vec<usize>,
30 pub test: Vec<usize>,
31}
32
33pub 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 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
62pub 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
72pub 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 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 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}