dml_core/learner/
search.rs

1//! Hyperparameter tuning by grid search: choose the best of several
2//! candidate learners by K-fold cross-validation (scikit-learn's `KFold` /
3//! `StratifiedKFold` folds without shuffling), then refit it on all data.
4//!
5//! Used as a nuisance learner, the search runs inside every cross-fitting
6//! fold on that fold's training data, which corresponds to DoubleML's
7//! `tune(..., tune_on_folds=True)`.
8
9use ndarray::{Array1, ArrayView1, ArrayView2, Axis};
10
11use super::cv::{check_cv, complement, kfold_test_sets, map_folds, stratified_test_sets};
12use super::{Classifier, Regressor};
13use crate::data::gather;
14use crate::error::{DmlError, Result};
15
16/// Index of the smallest score (the first one on ties).
17fn argmin(scores: &[f64]) -> usize {
18    scores
19        .iter()
20        .enumerate()
21        .fold(
22            (0, f64::INFINITY),
23            |(bi, bv), (i, &v)| if v < bv { (i, v) } else { (bi, bv) },
24        )
25        .0
26}
27
28fn check_candidates(n: usize) -> Result<()> {
29    if n == 0 {
30        return Err(DmlError::InvalidArgument(
31            "a grid search needs at least one candidate".into(),
32        ));
33    }
34    Ok(())
35}
36
37/// Grid search over regression learners, scored by the mean squared error
38/// of the held-out folds.
39pub struct GridSearchRegressor {
40    candidates: Vec<Box<dyn Regressor>>,
41    cv: usize,
42    best: Option<usize>,
43    scores: Vec<f64>,
44    fitted: Option<Box<dyn Regressor>>,
45}
46
47impl GridSearchRegressor {
48    /// `candidates`, e.g. the same learner with different hyperparameters.
49    pub fn new(candidates: Vec<Box<dyn Regressor>>) -> Self {
50        Self {
51            candidates,
52            cv: 5,
53            best: None,
54            scores: Vec::new(),
55            fitted: None,
56        }
57    }
58
59    /// Number of cross-validation folds (default 5).
60    pub fn cv(mut self, cv: usize) -> Self {
61        self.cv = cv;
62        self
63    }
64
65    /// Index of the chosen candidate after fitting.
66    pub fn best_index(&self) -> Option<usize> {
67        self.best
68    }
69
70    /// Cross-validated mean squared error of every candidate.
71    pub fn cv_scores(&self) -> &[f64] {
72        &self.scores
73    }
74
75    fn fit_impl(
76        &mut self,
77        x: ArrayView2<'_, f64>,
78        y: ArrayView1<'_, f64>,
79        w: Option<ArrayView1<'_, f64>>,
80    ) -> Result<()> {
81        check_candidates(self.candidates.len())?;
82        let n = x.nrows();
83        check_cv(n, self.cv)?;
84        let tests = kfold_test_sets(n, self.cv);
85        let fit = |c: &dyn Regressor, rows: &[usize]| -> Result<Box<dyn Regressor>> {
86            let mut m = c.fresh();
87            let (xt, yt) = (x.select(Axis(0), rows), gather(y, rows));
88            match w {
89                Some(w) => m.fit_weighted(xt.view(), yt.view(), gather(w, rows).view())?,
90                None => m.fit(xt.view(), yt.view())?,
91            }
92            Ok(m)
93        };
94        let mut scores = Vec::with_capacity(self.candidates.len());
95        for c in &self.candidates {
96            let fold_mse = map_folds(&tests, |test| {
97                let m = fit(c.as_ref(), &complement(n, test))?;
98                let pred = m.predict(x.select(Axis(0), test).view())?;
99                Ok(test
100                    .iter()
101                    .zip(&pred)
102                    .map(|(&i, p)| (y[i] - p).powi(2))
103                    .sum::<f64>()
104                    / test.len() as f64)
105            })?;
106            scores.push(fold_mse.iter().sum::<f64>() / fold_mse.len() as f64);
107        }
108        let best = argmin(&scores);
109        let all: Vec<usize> = (0..n).collect();
110        self.fitted = Some(fit(self.candidates[best].as_ref(), &all)?);
111        self.best = Some(best);
112        self.scores = scores;
113        Ok(())
114    }
115}
116
117impl Regressor for GridSearchRegressor {
118    fn fit(&mut self, x: ArrayView2<'_, f64>, y: ArrayView1<'_, f64>) -> Result<()> {
119        self.fit_impl(x, y, None)
120    }
121
122    /// Candidates are fitted with the weights; folds are scored unweighted.
123    fn fit_weighted(
124        &mut self,
125        x: ArrayView2<'_, f64>,
126        y: ArrayView1<'_, f64>,
127        w: ArrayView1<'_, f64>,
128    ) -> Result<()> {
129        self.fit_impl(x, y, Some(w))
130    }
131
132    fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>> {
133        self.fitted.as_ref().ok_or(DmlError::NotFitted)?.predict(x)
134    }
135
136    fn fresh(&self) -> Box<dyn Regressor> {
137        Box::new(Self::new(self.candidates.iter().map(|c| c.fresh()).collect()).cv(self.cv))
138    }
139}
140
141impl Clone for GridSearchRegressor {
142    /// An unfitted copy with the same candidates.
143    fn clone(&self) -> Self {
144        Self::new(self.candidates.iter().map(|c| c.fresh()).collect()).cv(self.cv)
145    }
146}
147
148/// Grid search over classifiers, scored by the log loss of the held-out
149/// folds (probabilities clipped to `[1e-15, 1 − 1e-15]`).
150pub struct GridSearchClassifier {
151    candidates: Vec<Box<dyn Classifier>>,
152    cv: usize,
153    best: Option<usize>,
154    scores: Vec<f64>,
155    fitted: Option<Box<dyn Classifier>>,
156}
157
158impl GridSearchClassifier {
159    pub fn new(candidates: Vec<Box<dyn Classifier>>) -> Self {
160        Self {
161            candidates,
162            cv: 5,
163            best: None,
164            scores: Vec::new(),
165            fitted: None,
166        }
167    }
168
169    /// Number of (stratified) cross-validation folds (default 5).
170    pub fn cv(mut self, cv: usize) -> Self {
171        self.cv = cv;
172        self
173    }
174
175    pub fn best_index(&self) -> Option<usize> {
176        self.best
177    }
178
179    /// Cross-validated log loss of every candidate.
180    pub fn cv_scores(&self) -> &[f64] {
181        &self.scores
182    }
183
184    fn fit_impl(
185        &mut self,
186        x: ArrayView2<'_, f64>,
187        y: ArrayView1<'_, f64>,
188        w: Option<ArrayView1<'_, f64>>,
189    ) -> Result<()> {
190        check_candidates(self.candidates.len())?;
191        let n = x.nrows();
192        check_cv(n, self.cv)?;
193        let tests = stratified_test_sets(y, self.cv)?;
194        let fit = |c: &dyn Classifier, rows: &[usize]| -> Result<Box<dyn Classifier>> {
195            let mut m = c.fresh();
196            let (xt, yt) = (x.select(Axis(0), rows), gather(y, rows));
197            match w {
198                Some(w) => m.fit_weighted(xt.view(), yt.view(), gather(w, rows).view())?,
199                None => m.fit(xt.view(), yt.view())?,
200            }
201            Ok(m)
202        };
203        let mut scores = Vec::with_capacity(self.candidates.len());
204        for c in &self.candidates {
205            let fold_loss = map_folds(&tests, |test| {
206                let m = fit(c.as_ref(), &complement(n, test))?;
207                let p = m.predict_proba(x.select(Axis(0), test).view())?;
208                let loss: f64 = test
209                    .iter()
210                    .zip(&p)
211                    .map(|(&i, &q)| {
212                        let q = q.clamp(1e-15, 1.0 - 1e-15);
213                        -(y[i] * q.ln() + (1.0 - y[i]) * (1.0 - q).ln())
214                    })
215                    .sum();
216                Ok(loss / test.len() as f64)
217            })?;
218            scores.push(fold_loss.iter().sum::<f64>() / fold_loss.len() as f64);
219        }
220        let best = argmin(&scores);
221        let all: Vec<usize> = (0..n).collect();
222        self.fitted = Some(fit(self.candidates[best].as_ref(), &all)?);
223        self.best = Some(best);
224        self.scores = scores;
225        Ok(())
226    }
227}
228
229impl Classifier for GridSearchClassifier {
230    fn fit(&mut self, x: ArrayView2<'_, f64>, y: ArrayView1<'_, f64>) -> Result<()> {
231        self.fit_impl(x, y, None)
232    }
233
234    fn fit_weighted(
235        &mut self,
236        x: ArrayView2<'_, f64>,
237        y: ArrayView1<'_, f64>,
238        w: ArrayView1<'_, f64>,
239    ) -> Result<()> {
240        self.fit_impl(x, y, Some(w))
241    }
242
243    fn predict_proba(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>> {
244        self.fitted
245            .as_ref()
246            .ok_or(DmlError::NotFitted)?
247            .predict_proba(x)
248    }
249
250    fn fresh(&self) -> Box<dyn Classifier> {
251        Box::new(Self::new(self.candidates.iter().map(|c| c.fresh()).collect()).cv(self.cv))
252    }
253}
254
255impl Clone for GridSearchClassifier {
256    /// An unfitted copy with the same candidates.
257    fn clone(&self) -> Self {
258        Self::new(self.candidates.iter().map(|c| c.fresh()).collect()).cv(self.cv)
259    }
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use crate::learner::{Lasso, LogisticRegression, Ols};
266    use ndarray::Array2;
267
268    #[test]
269    fn picks_the_candidate_with_the_smallest_cv_error() {
270        // y depends on x0 only; a heavy Lasso penalty underfits.
271        let n = 200;
272        let x = Array2::from_shape_fn((n, 3), |(i, j)| ((i * 7 + j * 13) % 17) as f64 / 17.0 - 0.5);
273        let y: Array1<f64> = (0..n)
274            .map(|i| 3.0 * x[[i, 0]] + 0.01 * ((i % 5) as f64))
275            .collect();
276        let mut gs = GridSearchRegressor::new(vec![
277            Box::new(Lasso::new(1.0)),
278            Box::new(Ols::new()),
279            Box::new(Lasso::new(0.5)),
280        ]);
281        gs.fit(x.view(), y.view()).unwrap();
282        assert_eq!(gs.best_index(), Some(1));
283        assert_eq!(gs.cv_scores().len(), 3);
284        let p = gs.predict(x.view()).unwrap();
285        assert!(p.iter().zip(&y).all(|(a, b)| (a - b).abs() < 0.1));
286    }
287
288    #[test]
289    fn classifier_search_and_errors() {
290        let n = 200;
291        let x = Array2::from_shape_fn((n, 2), |(i, j)| ((i * 3 + j * 11) % 23) as f64 / 23.0 - 0.5);
292        let y: Array1<f64> = (0..n)
293            .map(|i| f64::from(x[[i, 0]] + 0.1 * x[[i, 1]] > 0.0))
294            .collect();
295        let mut gs = GridSearchClassifier::new(vec![
296            Box::new(LogisticRegression::new(0.001)),
297            Box::new(LogisticRegression::new(10.0)),
298        ]);
299        gs.fit(x.view(), y.view()).unwrap();
300        assert_eq!(gs.best_index(), Some(1));
301        assert!(GridSearchClassifier::new(vec![])
302            .fit(x.view(), y.view())
303            .is_err());
304        assert!(gs.fresh().predict_proba(x.view()).is_err());
305    }
306}
dml_core/learner/search.rs

dml_core/learner/
search.rs
