model-selection-rs 0.1.0

//! Validation curves: score vs. a single hyperparameter.

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

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

/// Output of [`validation_curve`].
///
/// The companion diagnostic to a [learning curve](super::learning_curve): here
/// the training-set size is fixed and a single hyperparameter varies. Kept
/// deliberately distinct from the learning curve, which the two are easy to
/// conflate. Scores are `[param][fold]`.
#[derive(Debug, Clone)]
pub struct ValidationCurve<P> {
    /// The parameter values, in the order supplied.
    pub param_values: Vec<P>,
    /// Training scores as `[param][fold]`.
    pub train_scores: Vec<Vec<f64>>,
    /// Validation scores as `[param][fold]`.
    pub val_scores: Vec<Vec<f64>>,
}

impl<P> ValidationCurve<P> {
    /// Mean training score at each parameter value.
    #[must_use]
    pub fn mean_train_scores(&self) -> Vec<f64> {
        self.train_scores.iter().map(|r| mean(r)).collect()
    }

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

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

struct Job {
    param_idx: usize,
    fold_idx: usize,
    train_score: f64,
    val_score: f64,
}

/// Compute a validation curve: for each hyperparameter value and each fold, fit
/// with that value and score on train and validation.
///
/// `fit_fn` here takes the parameter value as its first argument (e.g. a closure
/// over "fit a decision tree with this `max_depth`").
///
/// With the `parallel` feature the `param × fold` grid is fanned out over
/// `rayon`.
///
/// # Errors
///
/// Propagates any error from `splitter.split(x.nrows())`.
pub fn validation_curve<S, F, M, P>(
    splitter: &S,
    x: &Array2<f64>,
    y: &Array1<f64>,
    fit_fn: F,
    scorer: &(dyn Scorer + Sync),
    param_range: &[P],
) -> Result<ValidationCurve<P>>
where
    S: CvSplitter,
    F: Fn(&P, &Array2<f64>, &Array1<f64>) -> M + Sync,
    M: Fn(&Array2<f64>) -> Array1<f64>,
    P: Clone + Sync,
{
    let splits = splitter.split(x.nrows())?;

    let eval = |param_idx: usize, fold_idx: usize| -> Job {
        let (train, val) = &splits[fold_idx];
        let x_train = x.select(Axis(0), train);
        let y_train = y.select(Axis(0), train);
        let x_val = x.select(Axis(0), val);
        let y_val = y.select(Axis(0), val);

        let model = fit_fn(&param_range[param_idx], &x_train, &y_train);
        let train_score = scorer.score(&y_train, &model(&x_train));
        let val_score = scorer.score(&y_val, &model(&x_val));
        Job {
            param_idx,
            fold_idx,
            train_score,
            val_score,
        }
    };

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

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

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

    Ok(ValidationCurve {
        param_values: param_range.to_vec(),
        train_scores,
        val_scores,
    })
}

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

    /// Polynomial-ish toy: fit y with a ridge-regularised constant+slope where
    /// `lambda` shrinks the slope. Too much shrinkage underfits; none overfits
    /// noise. Here we just check the mechanics + a sweet-spot shape using a
    /// hyperparameter `degree` on clearly nonlinear data.
    ///
    /// Model: predict with a polynomial of the given `degree` on one feature,
    /// fit by least squares via normal equations on a Vandermonde matrix.
    fn poly_fit(
        degree: &usize,
        x: &Array2<f64>,
        y: &Array1<f64>,
    ) -> impl Fn(&Array2<f64>) -> Array1<f64> {
        let d = *degree;
        let xs: Vec<f64> = x.column(0).to_vec();
        let ys: Vec<f64> = y.to_vec();
        let coeffs = fit_poly(&xs, &ys, d);
        move |xq: &Array2<f64>| xq.column(0).mapv(|v| eval_poly(&coeffs, v))
    }

    fn eval_poly(coeffs: &[f64], x: f64) -> f64 {
        coeffs
            .iter()
            .enumerate()
            .map(|(i, c)| c * x.powi(i as i32))
            .sum()
    }

    /// Least-squares polynomial fit via normal equations (small, dense).
    fn fit_poly(xs: &[f64], ys: &[f64], degree: usize) -> Vec<f64> {
        let n = xs.len();
        let m = degree + 1;
        // Vandermonde X (n x m).
        let x: Vec<Vec<f64>> = xs
            .iter()
            .map(|&v| (0..m).map(|p| v.powi(p as i32)).collect())
            .collect();
        // Normal equations: (XtX) c = Xt y.
        let mut xtx = vec![vec![0.0; m]; m];
        let mut xty = vec![0.0; m];
        for i in 0..n {
            for a in 0..m {
                xty[a] += x[i][a] * ys[i];
                for b in 0..m {
                    xtx[a][b] += x[i][a] * x[i][b];
                }
            }
        }
        solve(xtx, xty)
    }

    /// Gaussian elimination with partial pivoting.
    #[allow(clippy::needless_range_loop)] // explicit column indices read clearer here
    fn solve(mut a: Vec<Vec<f64>>, mut b: Vec<f64>) -> Vec<f64> {
        let n = b.len();
        for col in 0..n {
            let pivot = (col..n)
                .max_by(|&r1, &r2| a[r1][col].abs().partial_cmp(&a[r2][col].abs()).unwrap())
                .unwrap();
            a.swap(col, pivot);
            b.swap(col, pivot);
            let d = a[col][col];
            if d.abs() < 1e-12 {
                continue;
            }
            for row in (col + 1)..n {
                let f = a[row][col] / d;
                for k in col..n {
                    a[row][k] -= f * a[col][k];
                }
                b[row] -= f * b[col];
            }
        }
        let mut sol = vec![0.0; n];
        for row in (0..n).rev() {
            let mut s = b[row];
            for k in (row + 1)..n {
                s -= a[row][k] * sol[k];
            }
            sol[row] = if a[row][row].abs() < 1e-12 {
                0.0
            } else {
                s / a[row][row]
            };
        }
        sol
    }

    #[test]
    fn degree_sweet_spot_shows_up() {
        // Quadratic truth: degree 1 underfits, degree 2 is ideal.
        let x = Array2::from_shape_fn((40, 1), |(i, _)| (i as f64) / 10.0 - 2.0);
        let y = x.column(0).mapv(|v| v * v - 0.5 * v + 1.0);
        let kf = KFold::new(4).unwrap().with_shuffle(0);
        let vc =
            validation_curve(&kf, &x, &y, poly_fit, &MeanSquaredError, &[1usize, 2, 3]).unwrap();
        let val = vc.mean_val_scores(); // MSE, lower is better
                                        // degree 2 should beat degree 1 clearly.
        assert!(
            val[1] < val[0],
            "degree 2 MSE {} should beat degree 1 {}",
            val[1],
            val[0]
        );
    }
}