dml_core/
crossfit.rs

1use ndarray::{Array1, ArrayView1, ArrayView2, Axis};
2use rand::seq::SliceRandom;
3use rand::SeedableRng;
4use rand_chacha::ChaCha8Rng;
5
6use crate::data::gather;
7use crate::error::{DmlError, Result};
8use crate::learner::{Classifier, Regressor};
9
10/// One train/test split.
11#[derive(Debug, Clone, PartialEq)]
12pub struct Fold {
13    pub train: Vec<usize>,
14    pub test: Vec<usize>,
15}
16
17/// Shuffled K-fold splitter with a fixed seed, so results are reproducible.
18#[derive(Debug, Clone)]
19pub struct KFold {
20    n_folds: usize,
21    seed: u64,
22}
23
24impl KFold {
25    pub fn new(n_folds: usize, seed: u64) -> Result<Self> {
26        if n_folds < 2 {
27            return Err(DmlError::InvalidArgument(
28                "n_folds must be at least 2".into(),
29            ));
30        }
31        Ok(Self { n_folds, seed })
32    }
33
34    pub fn n_folds(&self) -> usize {
35        self.n_folds
36    }
37
38    /// Split `0..n` into `n_folds` disjoint test sets covering every index once.
39    pub fn split(&self, n: usize) -> Result<Vec<Fold>> {
40        if n < self.n_folds {
41            return Err(DmlError::InvalidArgument(format!(
42                "{n} observations is fewer than {} folds",
43                self.n_folds
44            )));
45        }
46        let mut idx: Vec<usize> = (0..n).collect();
47        idx.shuffle(&mut ChaCha8Rng::seed_from_u64(self.seed));
48
49        let base = n / self.n_folds;
50        let extra = n % self.n_folds;
51        let mut folds = Vec::with_capacity(self.n_folds);
52        let mut start = 0;
53        for k in 0..self.n_folds {
54            let size = base + usize::from(k < extra);
55            let mut test = idx[start..start + size].to_vec();
56            let mut train: Vec<usize> = idx[..start]
57                .iter()
58                .chain(&idx[start + size..])
59                .copied()
60                .collect();
61            test.sort_unstable();
62            train.sort_unstable();
63            folds.push(Fold { train, test });
64            start += size;
65        }
66        Ok(folds)
67    }
68
69    /// Stratified split: the rows of each stratum (e.g. treated / control) are
70    /// shuffled and dealt round-robin over the folds, so every fold holds about
71    /// the same share of each stratum. Fails if a stratum has fewer rows than
72    /// there are folds.
73    pub fn split_stratified(&self, strata: &[usize]) -> Result<Vec<Fold>> {
74        let n = strata.len();
75        if n < self.n_folds {
76            return Err(DmlError::InvalidArgument(format!(
77                "{n} observations is fewer than {} folds",
78                self.n_folds
79            )));
80        }
81        let mut groups: std::collections::BTreeMap<usize, Vec<usize>> = Default::default();
82        for (i, &s) in strata.iter().enumerate() {
83            groups.entry(s).or_default().push(i);
84        }
85        let mut rng = ChaCha8Rng::seed_from_u64(self.seed);
86        let mut tests: Vec<Vec<usize>> = vec![Vec::new(); self.n_folds];
87        let mut next = 0;
88        for (stratum, mut idx) in groups {
89            if idx.len() < self.n_folds {
90                return Err(DmlError::InvalidArgument(format!(
91                    "stratum {stratum} has {} observations, fewer than {} folds",
92                    idx.len(),
93                    self.n_folds
94                )));
95            }
96            idx.shuffle(&mut rng);
97            for i in idx {
98                tests[next].push(i);
99                next = (next + 1) % self.n_folds;
100            }
101        }
102        Ok(tests
103            .into_iter()
104            .map(|mut test| {
105                test.sort_unstable();
106                let mut in_test = vec![false; n];
107                for &i in &test {
108                    in_test[i] = true;
109                }
110                let train = (0..n).filter(|&i| !in_test[i]).collect();
111                Fold { train, test }
112            })
113            .collect())
114    }
115}
116
117/// How observations are split for cross-fitting: `n_rep` independent
118/// repetitions of seeded K-fold splitting.
119#[derive(Debug, Clone, Copy, PartialEq, Eq)]
120pub struct SampleSplitting {
121    pub n_folds: usize,
122    pub n_rep: usize,
123    pub seed: u64,
124}
125
126impl Default for SampleSplitting {
127    fn default() -> Self {
128        Self {
129            n_folds: 5,
130            n_rep: 1,
131            seed: 0,
132        }
133    }
134}
135
136impl SampleSplitting {
137    /// Draw the folds for every repetition. Repetition `r` uses its own seed
138    /// derived from `seed`, so repetition 0 equals `KFold::new(n_folds, seed)`.
139    pub fn draw(&self, n: usize) -> Result<Vec<Vec<Fold>>> {
140        if self.n_rep == 0 {
141            return Err(DmlError::InvalidArgument("n_rep must be at least 1".into()));
142        }
143        (0..self.n_rep as u64)
144            .map(|r| {
145                // Golden-ratio increment keeps repetition seeds far apart.
146                let seed = self
147                    .seed
148                    .wrapping_add(r.wrapping_mul(0x9E37_79B9_7F4A_7C15));
149                KFold::new(self.n_folds, seed)?.split(n)
150            })
151            .collect()
152    }
153
154    /// Like [`draw`](Self::draw), but each fold holds about the same share of
155    /// every stratum (e.g. treatment level), with the same per-repetition seeds.
156    pub fn draw_stratified(&self, strata: &[usize]) -> Result<Vec<Vec<Fold>>> {
157        if self.n_rep == 0 {
158            return Err(DmlError::InvalidArgument("n_rep must be at least 1".into()));
159        }
160        (0..self.n_rep as u64)
161            .map(|r| {
162                let seed = self
163                    .seed
164                    .wrapping_add(r.wrapping_mul(0x9E37_79B9_7F4A_7C15));
165                KFold::new(self.n_folds, seed)?.split_stratified(strata)
166            })
167            .collect()
168    }
169
170    /// Cluster-level splits for every repetition (see [`crate::cluster`]).
171    pub fn draw_clustered(&self, clusters: &[Array1<i64>]) -> Result<Vec<Vec<crate::ClusterFold>>> {
172        if self.n_rep == 0 {
173            return Err(DmlError::InvalidArgument("n_rep must be at least 1".into()));
174        }
175        (0..self.n_rep as u64)
176            .map(|r| {
177                let seed = self
178                    .seed
179                    .wrapping_add(r.wrapping_mul(0x9E37_79B9_7F4A_7C15));
180                crate::cluster::split_clusters(clusters, self.n_folds, seed)
181            })
182            .collect()
183    }
184}
185
186/// Out-of-fold predictions: for each fold, fit a fresh copy of `learner` on the
187/// training rows and predict the held-out rows. Every row is predicted by a
188/// model that never saw it — the core of DML cross-fitting.
189pub fn cross_fit_predict(
190    learner: &dyn Regressor,
191    x: ArrayView2<'_, f64>,
192    target: ArrayView1<'_, f64>,
193    folds: &[Fold],
194) -> Result<Array1<f64>> {
195    if target.len() != x.nrows() {
196        return Err(DmlError::DimensionMismatch(
197            "cross_fit_predict x/target rows".into(),
198        ));
199    }
200    cross_fit_with(x.nrows(), folds, |fold| {
201        let mut model = learner.fresh();
202        model.fit(
203            x.select(Axis(0), &fold.train).view(),
204            gather(target, &fold.train).view(),
205        )?;
206        model.predict(x.select(Axis(0), &fold.test).view())
207    })
208}
209
210/// Like [`cross_fit_predict`], but each model is trained only on the training
211/// rows where `keep[i]` is true (e.g. controls only); predictions still cover
212/// every test row.
213pub fn cross_fit_predict_subset(
214    learner: &dyn Regressor,
215    x: ArrayView2<'_, f64>,
216    target: ArrayView1<'_, f64>,
217    folds: &[Fold],
218    keep: &[bool],
219) -> Result<Array1<f64>> {
220    cross_fit_predict(learner, x, target, &subset_folds(folds, keep, x.nrows())?)
221}
222
223/// Like [`cross_fit_predict_proba`], but each model is trained only on the
224/// training rows where `keep[i]` is true; predictions cover every test row.
225pub fn cross_fit_predict_proba_subset(
226    learner: &dyn Classifier,
227    x: ArrayView2<'_, f64>,
228    labels: ArrayView1<'_, f64>,
229    folds: &[Fold],
230    keep: &[bool],
231) -> Result<Array1<f64>> {
232    cross_fit_predict_proba(learner, x, labels, &subset_folds(folds, keep, x.nrows())?)
233}
234
235/// Restrict every training set to the rows where `keep` is true.
236fn subset_folds(folds: &[Fold], keep: &[bool], n: usize) -> Result<Vec<Fold>> {
237    if keep.len() != n {
238        return Err(DmlError::DimensionMismatch("subgroup mask length".into()));
239    }
240    let subset: Vec<Fold> = folds
241        .iter()
242        .map(|f| Fold {
243            train: f.train.iter().copied().filter(|&i| keep[i]).collect(),
244            test: f.test.clone(),
245        })
246        .collect();
247    if subset.iter().any(|f| f.train.is_empty()) {
248        return Err(DmlError::InvalidArgument(
249            "a training fold has no observations in the requested subgroup".into(),
250        ));
251    }
252    Ok(subset)
253}
254
255/// Out-of-fold class-1 probabilities from a [`Classifier`].
256pub fn cross_fit_predict_proba(
257    learner: &dyn Classifier,
258    x: ArrayView2<'_, f64>,
259    labels: ArrayView1<'_, f64>,
260    folds: &[Fold],
261) -> Result<Array1<f64>> {
262    if labels.len() != x.nrows() {
263        return Err(DmlError::DimensionMismatch(
264            "cross_fit_predict_proba x/labels rows".into(),
265        ));
266    }
267    cross_fit_with(x.nrows(), folds, |fold| {
268        let mut model = learner.fresh();
269        model.fit(
270            x.select(Axis(0), &fold.train).view(),
271            gather(labels, &fold.train).view(),
272        )?;
273        model.predict_proba(x.select(Axis(0), &fold.test).view())
274    })
275}
276
277/// Run `fit_predict` on every fold and scatter each fold's test-set
278/// predictions into one length-`n` vector.
279///
280/// With the `parallel` feature (on by default) folds run on rayon's thread
281/// pool; results are identical to the sequential order because every fold
282/// writes to disjoint rows.
283pub(crate) fn cross_fit_with<F>(n: usize, folds: &[Fold], fit_predict: F) -> Result<Array1<f64>>
284where
285    F: Fn(&Fold) -> Result<Array1<f64>> + Sync,
286{
287    #[cfg(feature = "parallel")]
288    let preds: Vec<Array1<f64>> = {
289        use rayon::prelude::*;
290        folds.par_iter().map(&fit_predict).collect::<Result<_>>()?
291    };
292    #[cfg(not(feature = "parallel"))]
293    let preds: Vec<Array1<f64>> = folds.iter().map(&fit_predict).collect::<Result<_>>()?;
294
295    let mut out = Array1::<f64>::from_elem(n, f64::NAN);
296    for (fold, pred) in folds.iter().zip(preds) {
297        if pred.len() != fold.test.len() {
298            return Err(DmlError::DimensionMismatch(
299                "learner returned wrong number of predictions".into(),
300            ));
301        }
302        for (&i, &p) in fold.test.iter().zip(pred.iter()) {
303            out[i] = p;
304        }
305    }
306    if out.iter().any(|v| v.is_nan()) {
307        return Err(DmlError::InvalidArgument(
308            "folds do not cover every observation, or a learner returned NaN".into(),
309        ));
310    }
311    Ok(out)
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317    use crate::learner::Ols;
318    use ndarray::Array2;
319
320    #[test]
321    fn folds_partition_all_indices() {
322        let folds = KFold::new(3, 7).unwrap().split(10).unwrap();
323        let mut seen: Vec<usize> = folds.iter().flat_map(|f| f.test.clone()).collect();
324        seen.sort_unstable();
325        assert_eq!(seen, (0..10).collect::<Vec<_>>());
326        for f in &folds {
327            assert!(f.test.iter().all(|i| !f.train.contains(i)));
328            assert_eq!(f.train.len() + f.test.len(), 10);
329        }
330        let sizes: Vec<usize> = folds.iter().map(|f| f.test.len()).collect();
331        assert_eq!(sizes, vec![4, 3, 3]);
332    }
333
334    #[test]
335    fn same_seed_same_split() {
336        let a = KFold::new(5, 42).unwrap().split(50).unwrap();
337        let b = KFold::new(5, 42).unwrap().split(50).unwrap();
338        let c = KFold::new(5, 43).unwrap().split(50).unwrap();
339        assert_eq!(a, b);
340        assert_ne!(a, c);
341    }
342
343    #[test]
344    fn repetitions_differ_and_first_matches_kfold() {
345        let reps = SampleSplitting {
346            n_folds: 4,
347            n_rep: 3,
348            seed: 9,
349        }
350        .draw(40)
351        .unwrap();
352        assert_eq!(reps.len(), 3);
353        assert_eq!(reps[0], KFold::new(4, 9).unwrap().split(40).unwrap());
354        assert_ne!(reps[0], reps[1]);
355        assert_ne!(reps[1], reps[2]);
356    }
357
358    #[test]
359    fn subset_training_uses_only_kept_rows() {
360        // Target is 0 on kept rows and 100 elsewhere: a mean model trained
361        // only on kept rows predicts ~0 everywhere.
362        let x = Array2::<f64>::zeros((20, 0));
363        let keep: Vec<bool> = (0..20).map(|i| i % 2 == 0).collect();
364        let y: Array1<f64> = keep.iter().map(|&k| if k { 0.0 } else { 100.0 }).collect();
365        let folds = KFold::new(4, 1).unwrap().split(20).unwrap();
366        let pred =
367            cross_fit_predict_subset(&Ols::new(), x.view(), y.view(), &folds, &keep).unwrap();
368        assert!(pred.iter().all(|p| p.abs() < 1e-12));
369    }
370
371    #[test]
372    fn stratified_folds_balance_each_stratum() {
373        let strata: Vec<usize> = (0..53).map(|i| usize::from(i % 5 == 0)).collect(); // 11 treated
374        let folds = KFold::new(5, 3).unwrap().split_stratified(&strata).unwrap();
375        let mut seen: Vec<usize> = folds.iter().flat_map(|f| f.test.clone()).collect();
376        seen.sort_unstable();
377        assert_eq!(seen, (0..53).collect::<Vec<_>>());
378        for f in &folds {
379            let treated = f.test.iter().filter(|&&i| strata[i] == 1).count();
380            assert!((2..=3).contains(&treated), "treated per fold: {treated}");
381            assert_eq!(f.train.len() + f.test.len(), 53);
382        }
383        assert!(KFold::new(5, 0)
384            .unwrap()
385            .split_stratified(&[0, 0, 0, 0, 0, 1])
386            .is_err());
387    }
388
389    #[test]
390    fn rejects_bad_settings() {
391        assert!(KFold::new(1, 0).is_err());
392        assert!(KFold::new(5, 0).unwrap().split(3).is_err());
393    }
394
395    #[test]
396    fn cross_fit_recovers_linear_signal() {
397        let n = 100;
398        let x = Array2::from_shape_fn((n, 2), |(i, j)| ((i * 7 + j * 3) % 11) as f64);
399        let y: Array1<f64> = x
400            .rows()
401            .into_iter()
402            .map(|r| 1.0 + 0.5 * r[0] - 2.0 * r[1])
403            .collect();
404        let folds = KFold::new(5, 1).unwrap().split(n).unwrap();
405        let pred = cross_fit_predict(&Ols::new(), x.view(), y.view(), &folds).unwrap();
406        assert!((&pred - &y).iter().all(|e| e.abs() < 1e-8));
407    }
408}
dml_core/crossfit.rs

dml_core/
crossfit.rs
