model_selection_rs/evaluate/
nested_cv.rs1use ndarray::{Array1, Array2, Axis};
5
6use crate::error::Result;
7use crate::scoring::Scorer;
8use crate::splitters::CvSplitter;
9
10#[derive(Debug, Clone)]
18pub struct NestedCvResults<P> {
19 pub outer_scores: Vec<f64>,
21 pub selected_params: Vec<P>,
23}
24
25impl<P> NestedCvResults<P> {
26 #[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 #[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
50pub 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 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 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 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 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 assert!(
197 res.selected_params.iter().all(|&l| l <= 1.0),
198 "selected lambdas: {:?}",
199 res.selected_params
200 );
201 assert!(res.mean_score() < 1.0, "mean MSE {}", res.mean_score());
203 }
204}