model-selection-rs 0.1.0

//! Learning curves: score vs. training-set size.

use ndarray::{Array1, Array2, Axis};

use crate::error::Result;
use crate::scoring::Scorer;
use crate::splitters::CvSplitter;

/// A training-set size, given either absolutely or as a fraction of the largest
/// usable training set.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TrainSize {
    /// A fixed number of training samples.
    Count(usize),
    /// A fraction in `(0.0, 1.0]` of the largest usable training set.
    Fraction(f64),
}

/// Output of [`learning_curve`].
///
/// Both train and validation scores are returned as `[size][fold]`: the gap
/// between them at each size is the actual diagnostic (high bias vs. high
/// variance), so neither can be dropped.
#[derive(Debug, Clone)]
pub struct LearningCurve {
    /// Absolute training-set sizes actually used, ascending.
    pub train_sizes: Vec<usize>,
    /// Training scores as `[size][fold]`.
    pub train_scores: Vec<Vec<f64>>,
    /// Validation scores as `[size][fold]`.
    pub val_scores: Vec<Vec<f64>>,
}

impl LearningCurve {
    /// Mean training score at each size (one value per size).
    #[must_use]
    pub fn mean_train_scores(&self) -> Vec<f64> {
        self.train_scores.iter().map(|row| mean(row)).collect()
    }

    /// Mean validation score at each size (one value per size).
    #[must_use]
    pub fn mean_val_scores(&self) -> Vec<f64> {
        self.val_scores.iter().map(|row| mean(row)).collect()
    }
}

fn mean(xs: &[f64]) -> f64 {
    xs.iter().sum::<f64>() / xs.len() as f64
}

/// One (size, fold) evaluation.
struct Job {
    size_idx: usize,
    fold_idx: usize,
    train_score: f64,
    val_score: f64,
}

/// Compute a learning curve: for each training size and each fold, fit on a
/// prefix of the fold's training set and score on both that prefix and the
/// held-out validation set.
///
/// `train_sizes` are resolved against the *smallest* training set across folds
/// (so every fold can supply every size), then clamped to `1..=min_train` and
/// sorted ascending.
///
/// With the `parallel` feature the `size × fold` grid — which can be much larger
/// than a plain cross-validation — is fanned out over `rayon`.
///
/// # Errors
///
/// Propagates any error from `splitter.split(x.nrows())`.
pub fn learning_curve<S, F, M>(
    splitter: &S,
    x: &Array2<f64>,
    y: &Array1<f64>,
    fit_fn: F,
    scorer: &(dyn Scorer + Sync),
    train_sizes: &[TrainSize],
) -> Result<LearningCurve>
where
    S: CvSplitter,
    F: Fn(&Array2<f64>, &Array1<f64>) -> M + Sync,
    M: Fn(&Array2<f64>) -> Array1<f64>,
{
    let splits = splitter.split(x.nrows())?;
    let min_train = splits.iter().map(|(tr, _)| tr.len()).min().unwrap_or(0);

    // Resolve, clamp, sort, dedup.
    let mut abs_sizes: Vec<usize> = train_sizes
        .iter()
        .map(|ts| match ts {
            TrainSize::Count(c) => (*c).clamp(1, min_train.max(1)),
            TrainSize::Fraction(f) => {
                ((f * min_train as f64).round() as usize).clamp(1, min_train.max(1))
            }
        })
        .collect();
    abs_sizes.sort_unstable();
    abs_sizes.dedup();

    let eval = |size_idx: usize, fold_idx: usize| -> Job {
        let (train, val) = &splits[fold_idx];
        let size = abs_sizes[size_idx];
        let sub = &train[..size];

        let x_sub = x.select(Axis(0), sub);
        let y_sub = y.select(Axis(0), sub);
        let x_val = x.select(Axis(0), val);
        let y_val = y.select(Axis(0), val);

        let model = fit_fn(&x_sub, &y_sub);
        let train_score = scorer.score(&y_sub, &model(&x_sub));
        let val_score = scorer.score(&y_val, &model(&x_val));
        Job {
            size_idx,
            fold_idx,
            train_score,
            val_score,
        }
    };

    let coords: Vec<(usize, usize)> = (0..abs_sizes.len())
        .flat_map(|s| (0..splits.len()).map(move |f| (s, f)))
        .collect();

    #[cfg(feature = "parallel")]
    let jobs: Vec<Job> = {
        use rayon::prelude::*;
        coords.par_iter().map(|&(s, f)| eval(s, f)).collect()
    };
    #[cfg(not(feature = "parallel"))]
    let jobs: Vec<Job> = coords.iter().map(|&(s, f)| eval(s, f)).collect();

    let mut train_scores = vec![vec![0.0; splits.len()]; abs_sizes.len()];
    let mut val_scores = vec![vec![0.0; splits.len()]; abs_sizes.len()];
    for job in jobs {
        train_scores[job.size_idx][job.fold_idx] = job.train_score;
        val_scores[job.size_idx][job.fold_idx] = job.val_score;
    }

    Ok(LearningCurve {
        train_sizes: abs_sizes,
        train_scores,
        val_scores,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::scoring::R2Score;
    use crate::splitters::KFold;

    /// OLS on one feature (see cross_validate tests).
    fn ols_fit(x: &Array2<f64>, y: &Array1<f64>) -> impl Fn(&Array2<f64>) -> Array1<f64> {
        let n = x.nrows() as f64;
        let xs: Vec<f64> = x.column(0).to_vec();
        let ys: Vec<f64> = y.to_vec();
        let mx = xs.iter().sum::<f64>() / n;
        let my = ys.iter().sum::<f64>() / n;
        let cov: f64 = xs.iter().zip(&ys).map(|(a, b)| (a - mx) * (b - my)).sum();
        let var: f64 = xs.iter().map(|a| (a - mx).powi(2)).sum();
        let slope = if var == 0.0 { 0.0 } else { cov / var };
        let intercept = my - slope * mx;
        move |xq: &Array2<f64>| xq.column(0).mapv(|v| slope * v + intercept)
    }

    /// A high-bias (under-capacity) model: always predicts a constant 0.
    fn constant_zero(_x: &Array2<f64>, _y: &Array1<f64>) -> impl Fn(&Array2<f64>) -> Array1<f64> {
        |xq: &Array2<f64>| Array1::zeros(xq.nrows())
    }

    #[test]
    fn shapes_are_correct() {
        let x = Array2::from_shape_fn((30, 1), |(i, _)| i as f64);
        let y = x.column(0).mapv(|v| 2.0 * v + 1.0);
        let kf = KFold::new(3).unwrap();
        let lc = learning_curve(
            &kf,
            &x,
            &y,
            ols_fit,
            &R2Score,
            &[
                TrainSize::Fraction(0.3),
                TrainSize::Fraction(0.6),
                TrainSize::Fraction(1.0),
            ],
        )
        .unwrap();
        assert_eq!(lc.train_sizes.len(), 3);
        assert_eq!(lc.train_scores.len(), 3);
        assert_eq!(lc.train_scores[0].len(), 3); // folds
    }

    #[test]
    fn high_bias_curves_are_both_mediocre() {
        // Linear data, but a constant model can't capture it: both train and
        // validation R2 stay low and close together (high-bias signature).
        let x = Array2::from_shape_fn((40, 1), |(i, _)| i as f64);
        let y = x.column(0).mapv(|v| 3.0 * v + 2.0);
        let kf = KFold::new(4).unwrap();
        let lc = learning_curve(
            &kf,
            &x,
            &y,
            constant_zero,
            &R2Score,
            &[TrainSize::Fraction(0.5), TrainSize::Fraction(1.0)],
        )
        .unwrap();
        let train = lc.mean_train_scores();
        let val = lc.mean_val_scores();
        for (t, v) in train.iter().zip(&val) {
            assert!(*t < 0.5, "train R2 should be poor, got {t}");
            assert!(*v < 0.5, "val R2 should be poor, got {v}");
        }
    }
}