prospicio_models/
resample.rs

1//! Resampling: train/test splits, cross-validation, and grid and random
2//! search.
3//!
4//! Splits are row indices, so they work with any [`Design`]. Random splits
5//! draw from [`StreamRng`] with a seed, so they replay exactly.
6
7use prospicio_core::{Error, Result, StreamRng};
8use rayon::prelude::*;
9
10use crate::design::Design;
11use crate::model::{Fitted, Model};
12
13/// One train/test split: row indices into the data.
14#[derive(Debug, Clone, PartialEq)]
15pub struct Split {
16    pub train: Vec<usize>,
17    pub test: Vec<usize>,
18}
19
20/// `k`-fold cross-validation: rows shuffled with `seed`, then cut into `k`
21/// folds of near-equal size; fold `i` is the test set of split `i`.
22///
23/// ```
24/// use prospicio_models::resample::k_fold;
25///
26/// let splits = k_fold(10, 5, 1).unwrap();
27/// assert_eq!(splits.len(), 5);
28/// assert!(splits.iter().all(|s| s.test.len() == 2 && s.train.len() == 8));
29/// ```
30pub fn k_fold(n: usize, k: usize, seed: u64) -> Result<Vec<Split>> {
31    check_k(n, k)?;
32    let order = shuffled(n, seed);
33    Ok(folds(&order, k)
34        .into_iter()
35        .map(|test| complement(n, test))
36        .collect())
37}
38
39/// Grouped `k`-fold: every row of a group lands in the same fold, so a
40/// policy or an insured seen in training never appears in its test set.
41/// Groups are shuffled with `seed` and dealt into `k` folds.
42pub fn group_k_fold(groups: &[String], k: usize, seed: u64) -> Result<Vec<Split>> {
43    let mut names: Vec<&String> = groups.iter().collect();
44    names.sort();
45    names.dedup();
46    check_k(names.len(), k)?;
47    let order = shuffled(names.len(), seed);
48    let fold_of_group: Vec<usize> = {
49        let mut f = vec![0; names.len()];
50        for (fold, members) in folds(&order, k).into_iter().enumerate() {
51            for g in members {
52                f[g] = fold;
53            }
54        }
55        f
56    };
57    let mut tests = vec![Vec::new(); k];
58    for (row, g) in groups.iter().enumerate() {
59        let gi = names.binary_search(&g).expect("every group is listed");
60        tests[fold_of_group[gi]].push(row);
61    }
62    Ok(tests
63        .into_iter()
64        .map(|test| complement(groups.len(), test))
65        .collect())
66}
67
68/// Time-ordered splits for forecasting: for each of the last `n_test`
69/// distinct periods `t`, train on rows with period before `t` and test on
70/// rows with period `t`. For a reserving triangle, the period is the
71/// calendar diagonal (origin + development), and this is the diagonal
72/// backtest of `docs/design/models.md`.
73///
74/// ```
75/// use prospicio_models::resample::time_ordered;
76///
77/// // Diagonals of a 3 × 3 triangle's cells.
78/// let diag = [0, 1, 2, 1, 2, 2];
79/// let splits = time_ordered(&diag, 1).unwrap();
80/// assert_eq!(splits[0].test, [2, 4, 5]);
81/// assert_eq!(splits[0].train, [0, 1, 3]);
82/// ```
83pub fn time_ordered(periods: &[i64], n_test: usize) -> Result<Vec<Split>> {
84    let mut distinct: Vec<i64> = periods.to_vec();
85    distinct.sort_unstable();
86    distinct.dedup();
87    if n_test == 0 || n_test >= distinct.len() {
88        return Err(Error::InvalidParameter {
89            name: "n_test",
90            value: n_test as f64,
91            reason: "must be at least 1 and leave an earlier period to train on",
92        });
93    }
94    Ok(distinct[distinct.len() - n_test..]
95        .iter()
96        .map(|&t| Split {
97            train: (0..periods.len()).filter(|&i| periods[i] < t).collect(),
98            test: (0..periods.len()).filter(|&i| periods[i] == t).collect(),
99        })
100        .collect())
101}
102
103/// Fits `model` on each split's training rows and scores its predictions
104/// on the test rows with `score(y_test, predicted, test_design)`. Returns
105/// one score per split. Splits run in parallel ([`map_splits`]).
106///
107/// ```
108/// use prospicio_models::resample::{cross_validate, k_fold};
109/// # use prospicio_models::{Design, Fitted, Model};
110/// # struct MeanModel;
111/// # struct MeanFit(f64);
112/// # impl Model for MeanModel {
113/// #     type Fitted = MeanFit;
114/// #     fn fit(&self, _: &Design, y: &[f64]) -> prospicio_core::Result<MeanFit> {
115/// #         Ok(MeanFit(y.iter().sum::<f64>() / y.len() as f64))
116/// #     }
117/// # }
118/// # impl Fitted for MeanFit {
119/// #     fn predict(&self, d: &Design) -> prospicio_core::Result<Vec<f64>> { Ok(vec![self.0; d.n_rows()]) }
120/// #     fn predict_distribution(&self, _: &Design, _: usize, _: u64)
121/// #         -> prospicio_core::Result<prospicio_prob::PredictiveDistribution> { unimplemented!() }
122/// # }
123/// let x = Design::new(vec!["x".into()], vec![vec![1.0; 6]]).unwrap();
124/// let y = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
125/// let scores = cross_validate(&MeanModel, &x, &y, &k_fold(6, 3, 1).unwrap(), |t, p, _| {
126///     prospicio_models::metrics::rmse(t, p).unwrap()
127/// })
128/// .unwrap();
129/// assert_eq!(scores.len(), 3);
130/// ```
131pub fn cross_validate<M, S>(
132    model: &M,
133    design: &Design,
134    y: &[f64],
135    splits: &[Split],
136    score: S,
137) -> Result<Vec<f64>>
138where
139    M: Model + Sync,
140    S: Fn(&[f64], &[f64], &Design) -> f64 + Sync,
141{
142    if y.len() != design.n_rows() {
143        return Err(Error::Data(format!(
144            "{} responses for {} design rows",
145            y.len(),
146            design.n_rows()
147        )));
148    }
149    map_splits(splits, |s| {
150        let train = design.select(&s.train);
151        let y_train: Vec<f64> = s.train.iter().map(|&i| y[i]).collect();
152        let fitted = model.fit(&train, &y_train)?;
153        let test = design.select(&s.test);
154        let y_test: Vec<f64> = s.test.iter().map(|&i| y[i]).collect();
155        let pred = fitted.predict(&test)?;
156        Ok(score(&y_test, &pred, &test))
157    })
158}
159
160/// Runs `f` on every split in parallel on the Rayon pool (the outer loop,
161/// as `docs/architecture.md` prescribes) and returns the results in split
162/// order: the first error in split order, if any fails. Each split's work
163/// is independent and seeded on its own, so results do not depend on the
164/// number of threads.
165pub fn map_splits<T, F>(splits: &[Split], f: F) -> Result<Vec<T>>
166where
167    T: Send,
168    F: Fn(&Split) -> Result<T> + Sync + Send,
169{
170    let results: Vec<Result<T>> = splits.par_iter().map(&f).collect();
171    results.into_iter().collect()
172}
173
174/// The result of [`grid_search`]: each candidate's mean score, and the
175/// best (lowest) one.
176#[derive(Debug, Clone, PartialEq)]
177pub struct GridSearch<P> {
178    /// `(candidate, mean score over the splits)`, in the order given.
179    pub scores: Vec<(P, f64)>,
180    /// Index of the candidate with the lowest mean score.
181    pub best: usize,
182}
183
184/// Scores every candidate in `candidates` by [`cross_validate`] on the same
185/// splits, with `make(candidate)` building the model, and picks the lowest
186/// mean score (scores are losses: deviance, RMSE, CRPS).
187pub fn grid_search<P, M, B, S>(
188    candidates: Vec<P>,
189    make: B,
190    design: &Design,
191    y: &[f64],
192    splits: &[Split],
193    score: S,
194) -> Result<GridSearch<P>>
195where
196    M: Model + Sync,
197    B: Fn(&P) -> M,
198    S: Fn(&[f64], &[f64], &Design) -> f64 + Sync,
199{
200    if candidates.is_empty() {
201        return Err(Error::InvalidParameter {
202            name: "candidates",
203            value: 0.0,
204            reason: "must not be empty",
205        });
206    }
207    let mut scores = Vec::with_capacity(candidates.len());
208    for c in candidates {
209        let s = cross_validate(&make(&c), design, y, splits, &score)?;
210        let mean = s.iter().sum::<f64>() / s.len() as f64;
211        scores.push((c, mean));
212    }
213    let best = scores
214        .iter()
215        .enumerate()
216        .min_by(|a, b| a.1.1.total_cmp(&b.1.1))
217        .map(|(i, _)| i)
218        .expect("not empty");
219    Ok(GridSearch { scores, best })
220}
221
222/// Random search: draws `n` candidates with `draw` from stream 0 of
223/// `seed`, then scores them as [`grid_search`] does. With several
224/// hyperparameters, random candidates cover each one's range far better
225/// than a grid of the same size (Bergstra and Bengio, 2012), which is what
226/// neural networks need. [`log_uniform`] draws scale parameters such as
227/// learning rates.
228///
229/// ```
230/// use prospicio_models::resample::{k_fold, log_uniform, random_search};
231/// # use prospicio_models::{Design, Fitted, Model};
232/// # struct Shrunk(f64);
233/// # struct Fit(f64);
234/// # impl Model for Shrunk {
235/// #     type Fitted = Fit;
236/// #     fn fit(&self, _: &Design, y: &[f64]) -> prospicio_core::Result<Fit> {
237/// #         Ok(Fit(y.iter().sum::<f64>() / (y.len() as f64 + self.0)))
238/// #     }
239/// # }
240/// # impl Fitted for Fit {
241/// #     fn predict(&self, d: &Design) -> prospicio_core::Result<Vec<f64>> { Ok(vec![self.0; d.n_rows()]) }
242/// #     fn predict_distribution(&self, _: &Design, _: usize, _: u64)
243/// #         -> prospicio_core::Result<prospicio_prob::PredictiveDistribution> { unimplemented!() }
244/// # }
245/// let x = Design::new(vec!["x".into()], vec![vec![1.0; 8]]).unwrap();
246/// let y = [3.0, 5.0, 4.0, 6.0, 5.0, 4.0, 5.0, 4.0];
247/// let splits = k_fold(8, 4, 1).unwrap();
248/// let found = random_search(
249///     20,
250///     7,
251///     |rng| log_uniform(rng, 1e-3, 10.0),
252///     |&shrink| Shrunk(shrink),
253///     &x,
254///     &y,
255///     &splits,
256///     |t, p, _| prospicio_models::metrics::rmse(t, p).unwrap(),
257/// )
258/// .unwrap();
259/// assert_eq!(found.scores.len(), 20);
260/// let best = found.scores[found.best].1;
261/// assert!(found.scores.iter().all(|(_, s)| *s >= best));
262/// // Heavy shrinkage of the mean towards zero loses.
263/// assert!(found.scores[found.best].0 < 1.0);
264/// ```
265#[allow(clippy::too_many_arguments)]
266pub fn random_search<P, M, D, B, S>(
267    n: usize,
268    seed: u64,
269    mut draw: D,
270    make: B,
271    design: &Design,
272    y: &[f64],
273    splits: &[Split],
274    score: S,
275) -> Result<GridSearch<P>>
276where
277    M: Model + Sync,
278    D: FnMut(&mut StreamRng) -> P,
279    B: Fn(&P) -> M,
280    S: Fn(&[f64], &[f64], &Design) -> f64 + Sync,
281{
282    let mut rng = StreamRng::new(seed, 0);
283    let candidates = (0..n).map(|_| draw(&mut rng)).collect();
284    grid_search(candidates, make, design, y, splits, score)
285}
286
287/// A draw log-uniform between `low` and `high` (both positive): uniform
288/// in orders of magnitude, for learning rates and penalties.
289pub fn log_uniform(rng: &mut StreamRng, low: f64, high: f64) -> f64 {
290    (low.ln() + rng.next_open01() * (high.ln() - low.ln())).exp()
291}
292
293/// A draw uniform over `low..=high`, for layer widths and counts.
294pub fn uniform_int(rng: &mut StreamRng, low: usize, high: usize) -> usize {
295    let span = (high - low + 1) as f64;
296    low + ((rng.next_open01() * span) as usize).min(high - low)
297}
298
299fn check_k(n: usize, k: usize) -> Result<()> {
300    if k < 2 || k > n {
301        return Err(Error::InvalidParameter {
302            name: "k",
303            value: k as f64,
304            reason: "must be at least 2 and at most the number of rows (or groups)",
305        });
306    }
307    Ok(())
308}
309
310/// `0..n` shuffled by Fisher–Yates on stream `stream` of `seed`: the
311/// permutation that the model-agnostic explanations in the bindings apply to
312/// a column, so they shuffle alike in Python and R.
313///
314/// ```
315/// use prospicio_models::resample::permutation;
316///
317/// let mut p = permutation(6, 1, 0);
318/// assert_eq!(p, permutation(6, 1, 0));
319/// p.sort_unstable();
320/// assert_eq!(p, (0..6).collect::<Vec<_>>());
321/// ```
322pub fn permutation(n: usize, seed: u64, stream: u64) -> Vec<usize> {
323    let mut order: Vec<usize> = (0..n).collect();
324    let mut rng = StreamRng::new(seed, stream);
325    for i in (1..n).rev() {
326        let j = ((rng.next_open01() * (i + 1) as f64) as usize).min(i);
327        order.swap(i, j);
328    }
329    order
330}
331
332/// `0..n` shuffled on stream 0 of `seed`.
333fn shuffled(n: usize, seed: u64) -> Vec<usize> {
334    permutation(n, seed, 0)
335}
336
337/// `order` cut into `k` contiguous folds, the first `n mod k` one longer.
338fn folds(order: &[usize], k: usize) -> Vec<Vec<usize>> {
339    let n = order.len();
340    let (base, extra) = (n / k, n % k);
341    let mut out = Vec::with_capacity(k);
342    let mut start = 0;
343    for i in 0..k {
344        let len = base + usize::from(i < extra);
345        let mut fold = order[start..start + len].to_vec();
346        fold.sort_unstable();
347        out.push(fold);
348        start += len;
349    }
350    out
351}
352
353fn complement(n: usize, test: Vec<usize>) -> Split {
354    let mut in_test = vec![false; n];
355    for &i in &test {
356        in_test[i] = true;
357    }
358    Split {
359        train: (0..n).filter(|&i| !in_test[i]).collect(),
360        test,
361    }
362}
363
364#[cfg(test)]
365mod tests {
366    use super::*;
367
368    #[test]
369    fn map_splits_keeps_order_and_the_first_error() {
370        let splits = k_fold(40, 8, 3).unwrap();
371        let sums = map_splits(&splits, |s| Ok(s.test.iter().sum::<usize>())).unwrap();
372        let serial: Vec<usize> = splits.iter().map(|s| s.test.iter().sum()).collect();
373        assert_eq!(sums, serial);
374        let err = map_splits(&splits, |s| {
375            if s.test.contains(&splits[5].test[0]) || s.test.contains(&splits[2].test[0]) {
376                Err(Error::Data(format!("fold with {}", s.test[0])))
377            } else {
378                Ok(())
379            }
380        })
381        .unwrap_err();
382        assert_eq!(
383            err.to_string(),
384            Error::Data(format!("fold with {}", splits[2].test[0])).to_string()
385        );
386    }
387
388    #[test]
389    fn k_fold_partitions_and_replays() {
390        let s = k_fold(23, 4, 9).unwrap();
391        let mut all: Vec<usize> = s.iter().flat_map(|x| x.test.clone()).collect();
392        all.sort_unstable();
393        assert_eq!(all, (0..23).collect::<Vec<_>>());
394        assert_eq!(s, k_fold(23, 4, 9).unwrap());
395        assert_ne!(s, k_fold(23, 4, 10).unwrap());
396        assert!(k_fold(3, 4, 1).is_err());
397    }
398
399    #[test]
400    fn permutations_replay_and_differ_by_stream() {
401        let p = permutation(50, 4, 2);
402        assert_eq!(p, permutation(50, 4, 2));
403        assert_ne!(p, permutation(50, 4, 3));
404        assert_eq!(permutation(50, 4, 0), shuffled(50, 4));
405        let mut sorted = p.clone();
406        sorted.sort_unstable();
407        assert_eq!(sorted, (0..50).collect::<Vec<_>>());
408        assert!(permutation(0, 1, 0).is_empty());
409    }
410
411    #[test]
412    fn groups_stay_together() {
413        let groups: Vec<String> = ["a", "b", "a", "c", "b", "d", "c"]
414            .iter()
415            .map(|s| s.to_string())
416            .collect();
417        for s in group_k_fold(&groups, 2, 3).unwrap() {
418            for &i in &s.test {
419                assert!(s.train.iter().all(|&j| groups[j] != groups[i]));
420            }
421        }
422    }
423}
prospicio_models/resample.rs

prospicio_models/
resample.rs
