Skip to main content

model_selection_rs/evaluate/
nested_cv.rs

1//! Nested cross-validation: honest performance estimation around an inner
2//! hyperparameter-selection loop.
3
4use ndarray::{Array1, Array2, Axis};
5
6use crate::error::Result;
7use crate::scoring::Scorer;
8use crate::splitters::CvSplitter;
9
10/// Output of [`nested_cross_validate`].
11///
12/// The headline number is [`mean_score`](NestedCvResults::mean_score) over the
13/// outer folds — the non-leaked estimate that nested CV exists to produce. The
14/// per-fold `selected_params` are a diagnostic: wildly different selections
15/// across outer folds signal an unstable model/data pairing, worth surfacing
16/// rather than discarding.
17#[derive(Debug, Clone)]
18pub struct NestedCvResults<P> {
19    /// Outer-fold test scores — the honest performance estimate.
20    pub outer_scores: Vec<f64>,
21    /// Hyperparameters selected by the inner loop on each outer fold.
22    pub selected_params: Vec<P>,
23}
24
25impl<P> NestedCvResults<P> {
26    /// Mean of the outer-fold scores.
27    #[must_use]
28    pub fn mean_score(&self) -> f64 {
29        self.outer_scores.iter().sum::<f64>() / self.outer_scores.len() as f64
30    }
31
32    /// Population standard deviation of the outer-fold scores.
33    #[must_use]
34    pub fn std_score(&self) -> f64 {
35        let m = self.mean_score();
36        let n = self.outer_scores.len();
37        if n < 2 {
38            return 0.0;
39        }
40        (self
41            .outer_scores
42            .iter()
43            .map(|s| (s - m).powi(2))
44            .sum::<f64>()
45            / n as f64)
46            .sqrt()
47    }
48}
49
50/// Run nested cross-validation.
51///
52/// For each **outer** fold:
53/// 1. `tune_fn` receives the outer-training data and the `inner` splitter, runs
54///    whatever hyperparameter search it likes (grid, `tpe`, …) using its own
55///    inner-CV loop, and returns the best hyperparameters `P`.
56/// 2. `fit_fn` refits a final model with those hyperparameters on the *full*
57///    outer-training portion.
58/// 3. that model is scored once on the untouched outer-test portion.
59///
60/// The inner loop never sees the outer-test data, so the outer scores carry none
61/// of the optimistic bias that tuning-and-evaluating on the same data produces.
62/// `tune_fn` is kept agnostic to *how* tuning happens, so this composes with any
63/// search approach rather than reimplementing one.
64///
65/// With the `parallel` feature the outer folds run concurrently over `rayon`.
66///
67/// # Errors
68///
69/// Propagates any error from `outer.split(x.nrows())`.
70pub fn nested_cross_validate<OS, IS, P, Tune, Fit, M>(
71    outer: &OS,
72    inner: &IS,
73    x: &Array2<f64>,
74    y: &Array1<f64>,
75    tune_fn: Tune,
76    fit_fn: Fit,
77    scorer: &(dyn Scorer + Sync),
78) -> Result<NestedCvResults<P>>
79where
80    OS: CvSplitter,
81    IS: CvSplitter + Sync,
82    Tune: Fn(&Array2<f64>, &Array1<f64>, &IS) -> P + Sync,
83    Fit: Fn(&P, &Array2<f64>, &Array1<f64>) -> M + Sync,
84    M: Fn(&Array2<f64>) -> Array1<f64>,
85    P: Send,
86{
87    let outer_splits = outer.split(x.nrows())?;
88
89    let eval = |train: &[usize], test: &[usize]| -> (f64, P) {
90        let x_train = x.select(Axis(0), train);
91        let y_train = y.select(Axis(0), train);
92        let x_test = x.select(Axis(0), test);
93        let y_test = y.select(Axis(0), test);
94
95        let best = tune_fn(&x_train, &y_train, inner);
96        let model = fit_fn(&best, &x_train, &y_train);
97        let score = scorer.score(&y_test, &model(&x_test));
98        (score, best)
99    };
100
101    #[cfg(feature = "parallel")]
102    let results: Vec<(f64, P)> = {
103        use rayon::prelude::*;
104        outer_splits
105            .par_iter()
106            .map(|(tr, te)| eval(tr, te))
107            .collect()
108    };
109    #[cfg(not(feature = "parallel"))]
110    let results: Vec<(f64, P)> = outer_splits.iter().map(|(tr, te)| eval(tr, te)).collect();
111
112    let mut outer_scores = Vec::with_capacity(results.len());
113    let mut selected_params = Vec::with_capacity(results.len());
114    for (score, param) in results {
115        outer_scores.push(score);
116        selected_params.push(param);
117    }
118
119    Ok(NestedCvResults {
120        outer_scores,
121        selected_params,
122    })
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    use crate::evaluate::{cross_validate, BoxedScorer};
129    use crate::scoring::MeanSquaredError;
130    use crate::splitters::KFold;
131
132    /// Ridge regression on one feature with L2 penalty `lambda`, closed form.
133    fn ridge_fit(
134        lambda: &f64,
135        x: &Array2<f64>,
136        y: &Array1<f64>,
137    ) -> impl Fn(&Array2<f64>) -> Array1<f64> {
138        let lambda = *lambda;
139        // Center, solve (xx + lambda) slope = xy, intercept = mean_y - slope*mean_x.
140        let n = x.nrows() as f64;
141        let xs: Vec<f64> = x.column(0).to_vec();
142        let ys: Vec<f64> = y.to_vec();
143        let mx = xs.iter().sum::<f64>() / n;
144        let my = ys.iter().sum::<f64>() / n;
145        let sxx: f64 = xs.iter().map(|v| (v - mx).powi(2)).sum();
146        let sxy: f64 = xs.iter().zip(&ys).map(|(a, b)| (a - mx) * (b - my)).sum();
147        let slope = sxy / (sxx + lambda);
148        let intercept = my - slope * mx;
149        move |xq: &Array2<f64>| xq.column(0).mapv(|v| slope * v + intercept)
150    }
151
152    /// Tune lambda by inner CV, returning the value with the best (lowest) MSE.
153    fn tune_lambda(x: &Array2<f64>, y: &Array1<f64>, inner: &KFold) -> f64 {
154        let candidates = [0.0f64, 0.1, 1.0, 10.0, 100.0];
155        let mut best = (f64::INFINITY, 0.0);
156        for &lam in &candidates {
157            let scorers: Vec<BoxedScorer> = vec![Box::new(MeanSquaredError)];
158            let res = cross_validate(
159                inner,
160                x,
161                y,
162                move |xt, yt| ridge_fit(&lam, xt, yt),
163                &scorers,
164                false,
165            )
166            .unwrap();
167            let mse = res.mean_test_score(0);
168            if mse < best.0 {
169                best = (mse, lam);
170            }
171        }
172        best.1
173    }
174
175    #[test]
176    fn recovers_low_regularization_on_clean_linear_data() {
177        // Clean linear data -> little regularization is best.
178        let x = Array2::from_shape_fn((60, 1), |(i, _)| i as f64 / 10.0);
179        let y = x.column(0).mapv(|v| 2.0 * v + 1.0);
180
181        let outer = KFold::new(5).unwrap().with_shuffle(0);
182        let inner = KFold::new(4).unwrap().with_shuffle(1);
183        let res = nested_cross_validate(
184            &outer,
185            &inner,
186            &x,
187            &y,
188            tune_lambda,
189            ridge_fit,
190            &MeanSquaredError,
191        )
192        .unwrap();
193
194        assert_eq!(res.outer_scores.len(), 5);
195        // On clean linear data the smallest lambdas should win everywhere.
196        assert!(
197            res.selected_params.iter().all(|&l| l <= 1.0),
198            "selected lambdas: {:?}",
199            res.selected_params
200        );
201        // The honest MSE estimate should be small but non-zero.
202        assert!(res.mean_score() < 1.0, "mean MSE {}", res.mean_score());
203    }
204}