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
70/// How observations are split for cross-fitting: `n_rep` independent
71/// repetitions of seeded K-fold splitting.
72#[derive(Debug, Clone, Copy, PartialEq, Eq)]
73pub struct SampleSplitting {
74    pub n_folds: usize,
75    pub n_rep: usize,
76    pub seed: u64,
77}
78
79impl Default for SampleSplitting {
80    fn default() -> Self {
81        Self {
82            n_folds: 5,
83            n_rep: 1,
84            seed: 0,
85        }
86    }
87}
88
89impl SampleSplitting {
90    /// Draw the folds for every repetition. Repetition `r` uses its own seed
91    /// derived from `seed`, so repetition 0 equals `KFold::new(n_folds, seed)`.
92    pub fn draw(&self, n: usize) -> Result<Vec<Vec<Fold>>> {
93        if self.n_rep == 0 {
94            return Err(DmlError::InvalidArgument("n_rep must be at least 1".into()));
95        }
96        (0..self.n_rep as u64)
97            .map(|r| {
98                // Golden-ratio increment keeps repetition seeds far apart.
99                let seed = self
100                    .seed
101                    .wrapping_add(r.wrapping_mul(0x9E37_79B9_7F4A_7C15));
102                KFold::new(self.n_folds, seed)?.split(n)
103            })
104            .collect()
105    }
106}
107
108/// Out-of-fold predictions: for each fold, fit a fresh copy of `learner` on the
109/// training rows and predict the held-out rows. Every row is predicted by a
110/// model that never saw it — the core of DML cross-fitting.
111pub fn cross_fit_predict(
112    learner: &dyn Regressor,
113    x: ArrayView2<'_, f64>,
114    target: ArrayView1<'_, f64>,
115    folds: &[Fold],
116) -> Result<Array1<f64>> {
117    if target.len() != x.nrows() {
118        return Err(DmlError::DimensionMismatch(
119            "cross_fit_predict x/target rows".into(),
120        ));
121    }
122    cross_fit_with(x.nrows(), folds, |fold| {
123        let mut model = learner.fresh();
124        model.fit(
125            x.select(Axis(0), &fold.train).view(),
126            gather(target, &fold.train).view(),
127        )?;
128        model.predict(x.select(Axis(0), &fold.test).view())
129    })
130}
131
132/// Like [`cross_fit_predict`], but each model is trained only on the training
133/// rows where `keep[i]` is true (e.g. controls only); predictions still cover
134/// every test row.
135pub fn cross_fit_predict_subset(
136    learner: &dyn Regressor,
137    x: ArrayView2<'_, f64>,
138    target: ArrayView1<'_, f64>,
139    folds: &[Fold],
140    keep: &[bool],
141) -> Result<Array1<f64>> {
142    if keep.len() != x.nrows() {
143        return Err(DmlError::DimensionMismatch(
144            "cross_fit_predict_subset mask length".into(),
145        ));
146    }
147    let subset: Vec<Fold> = folds
148        .iter()
149        .map(|f| Fold {
150            train: f.train.iter().copied().filter(|&i| keep[i]).collect(),
151            test: f.test.clone(),
152        })
153        .collect();
154    if subset.iter().any(|f| f.train.is_empty()) {
155        return Err(DmlError::InvalidArgument(
156            "a training fold has no observations in the requested subgroup".into(),
157        ));
158    }
159    cross_fit_predict(learner, x, target, &subset)
160}
161
162/// Out-of-fold class-1 probabilities from a [`Classifier`].
163pub fn cross_fit_predict_proba(
164    learner: &dyn Classifier,
165    x: ArrayView2<'_, f64>,
166    labels: ArrayView1<'_, f64>,
167    folds: &[Fold],
168) -> Result<Array1<f64>> {
169    if labels.len() != x.nrows() {
170        return Err(DmlError::DimensionMismatch(
171            "cross_fit_predict_proba x/labels rows".into(),
172        ));
173    }
174    cross_fit_with(x.nrows(), folds, |fold| {
175        let mut model = learner.fresh();
176        model.fit(
177            x.select(Axis(0), &fold.train).view(),
178            gather(labels, &fold.train).view(),
179        )?;
180        model.predict_proba(x.select(Axis(0), &fold.test).view())
181    })
182}
183
184/// Run `fit_predict` on every fold and scatter each fold's test-set
185/// predictions into one length-`n` vector.
186///
187/// With the `parallel` feature (on by default) folds run on rayon's thread
188/// pool; results are identical to the sequential order because every fold
189/// writes to disjoint rows.
190pub(crate) fn cross_fit_with<F>(n: usize, folds: &[Fold], fit_predict: F) -> Result<Array1<f64>>
191where
192    F: Fn(&Fold) -> Result<Array1<f64>> + Sync,
193{
194    #[cfg(feature = "parallel")]
195    let preds: Vec<Array1<f64>> = {
196        use rayon::prelude::*;
197        folds.par_iter().map(&fit_predict).collect::<Result<_>>()?
198    };
199    #[cfg(not(feature = "parallel"))]
200    let preds: Vec<Array1<f64>> = folds.iter().map(&fit_predict).collect::<Result<_>>()?;
201
202    let mut out = Array1::<f64>::from_elem(n, f64::NAN);
203    for (fold, pred) in folds.iter().zip(preds) {
204        if pred.len() != fold.test.len() {
205            return Err(DmlError::DimensionMismatch(
206                "learner returned wrong number of predictions".into(),
207            ));
208        }
209        for (&i, &p) in fold.test.iter().zip(pred.iter()) {
210            out[i] = p;
211        }
212    }
213    if out.iter().any(|v| v.is_nan()) {
214        return Err(DmlError::InvalidArgument(
215            "folds do not cover every observation, or a learner returned NaN".into(),
216        ));
217    }
218    Ok(out)
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224    use crate::learner::Ols;
225    use ndarray::Array2;
226
227    #[test]
228    fn folds_partition_all_indices() {
229        let folds = KFold::new(3, 7).unwrap().split(10).unwrap();
230        let mut seen: Vec<usize> = folds.iter().flat_map(|f| f.test.clone()).collect();
231        seen.sort_unstable();
232        assert_eq!(seen, (0..10).collect::<Vec<_>>());
233        for f in &folds {
234            assert!(f.test.iter().all(|i| !f.train.contains(i)));
235            assert_eq!(f.train.len() + f.test.len(), 10);
236        }
237        let sizes: Vec<usize> = folds.iter().map(|f| f.test.len()).collect();
238        assert_eq!(sizes, vec![4, 3, 3]);
239    }
240
241    #[test]
242    fn same_seed_same_split() {
243        let a = KFold::new(5, 42).unwrap().split(50).unwrap();
244        let b = KFold::new(5, 42).unwrap().split(50).unwrap();
245        let c = KFold::new(5, 43).unwrap().split(50).unwrap();
246        assert_eq!(a, b);
247        assert_ne!(a, c);
248    }
249
250    #[test]
251    fn repetitions_differ_and_first_matches_kfold() {
252        let reps = SampleSplitting {
253            n_folds: 4,
254            n_rep: 3,
255            seed: 9,
256        }
257        .draw(40)
258        .unwrap();
259        assert_eq!(reps.len(), 3);
260        assert_eq!(reps[0], KFold::new(4, 9).unwrap().split(40).unwrap());
261        assert_ne!(reps[0], reps[1]);
262        assert_ne!(reps[1], reps[2]);
263    }
264
265    #[test]
266    fn subset_training_uses_only_kept_rows() {
267        // Target is 0 on kept rows and 100 elsewhere: a mean model trained
268        // only on kept rows predicts ~0 everywhere.
269        let x = Array2::<f64>::zeros((20, 0));
270        let keep: Vec<bool> = (0..20).map(|i| i % 2 == 0).collect();
271        let y: Array1<f64> = keep.iter().map(|&k| if k { 0.0 } else { 100.0 }).collect();
272        let folds = KFold::new(4, 1).unwrap().split(20).unwrap();
273        let pred =
274            cross_fit_predict_subset(&Ols::new(), x.view(), y.view(), &folds, &keep).unwrap();
275        assert!(pred.iter().all(|p| p.abs() < 1e-12));
276    }
277
278    #[test]
279    fn rejects_bad_settings() {
280        assert!(KFold::new(1, 0).is_err());
281        assert!(KFold::new(5, 0).unwrap().split(3).is_err());
282    }
283
284    #[test]
285    fn cross_fit_recovers_linear_signal() {
286        let n = 100;
287        let x = Array2::from_shape_fn((n, 2), |(i, j)| ((i * 7 + j * 3) % 11) as f64);
288        let y: Array1<f64> = x
289            .rows()
290            .into_iter()
291            .map(|r| 1.0 + 0.5 * r[0] - 2.0 * r[1])
292            .collect();
293        let folds = KFold::new(5, 1).unwrap().split(n).unwrap();
294        let pred = cross_fit_predict(&Ols::new(), x.view(), y.view(), &folds).unwrap();
295        assert!((&pred - &y).iter().all(|e| e.abs() < 1e-8));
296    }
297}
dml_core/crossfit.rs

dml_core/
crossfit.rs
