model-selection-rs 0.1.0

//! The core cross-validation evaluation loop.

use std::time::{Duration, Instant};

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

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

/// A scorer usable by the evaluation utilities.
///
/// The `Send + Sync` bounds keep the public signatures identical whether or not
/// the `parallel` feature is enabled (feature flags never change the API), and
/// let folds be scored across threads when it is. Every built-in scorer, and any
/// [`make_scorer`](crate::scoring::make_scorer) closure over `Send + Sync` data,
/// satisfies them.
pub type BoxedScorer = Box<dyn Scorer + Send + Sync>;

/// Results of a [`cross_validate`] run.
///
/// Scores are stored as `[scorer][fold]`. Use [`mean_test_score`] /
/// [`std_test_score`] (by index or by name) for summaries.
///
/// [`mean_test_score`]: CvResults::mean_test_score
/// [`std_test_score`]: CvResults::std_test_score
#[derive(Debug, Clone)]
pub struct CvResults {
    /// Metric names, in the order they were supplied.
    pub scorer_names: Vec<String>,
    /// Test scores as `[scorer][fold]`.
    pub test_scores: Vec<Vec<f64>>,
    /// Train scores as `[scorer][fold]`, if `return_train_scores` was set.
    pub train_scores: Option<Vec<Vec<f64>>>,
    /// Wall-clock fit time per fold.
    pub fit_times: Vec<Duration>,
    /// Wall-clock score time per fold.
    pub score_times: Vec<Duration>,
}

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

fn std(xs: &[f64]) -> f64 {
    if xs.len() < 2 {
        return 0.0;
    }
    let m = mean(xs);
    let var = xs.iter().map(|x| (x - m).powi(2)).sum::<f64>() / xs.len() as f64;
    var.sqrt()
}

impl CvResults {
    /// Number of folds evaluated.
    #[must_use]
    pub fn n_splits(&self) -> usize {
        self.fit_times.len()
    }

    /// Mean test score for scorer `idx`.
    #[must_use]
    pub fn mean_test_score(&self, idx: usize) -> f64 {
        mean(&self.test_scores[idx])
    }

    /// Population standard deviation of the test scores for scorer `idx`.
    #[must_use]
    pub fn std_test_score(&self, idx: usize) -> f64 {
        std(&self.test_scores[idx])
    }

    /// Mean test score for the scorer named `name`, if present.
    #[must_use]
    pub fn mean_test_score_by_name(&self, name: &str) -> Option<f64> {
        self.scorer_names
            .iter()
            .position(|n| n == name)
            .map(|i| self.mean_test_score(i))
    }

    /// Total fit time across all folds.
    #[must_use]
    pub fn total_fit_time(&self) -> Duration {
        self.fit_times.iter().sum()
    }
}

/// Compute every scorer's score for a single fold. Shared by the serial and
/// parallel paths so both produce numerically identical results.
struct FoldScores {
    test: Vec<f64>,
    train: Option<Vec<f64>>,
    fit_time: Duration,
    score_time: Duration,
}

fn evaluate_fold<F, M>(
    x: &Array2<f64>,
    y: &Array1<f64>,
    train: &[usize],
    test: &[usize],
    fit_fn: &F,
    scorers: &[BoxedScorer],
    return_train_scores: bool,
) -> FoldScores
where
    F: Fn(&Array2<f64>, &Array1<f64>) -> M,
    M: Fn(&Array2<f64>) -> Array1<f64>,
{
    let x_train = x.select(Axis(0), train);
    let y_train = y.select(Axis(0), train);
    let x_test = x.select(Axis(0), test);
    let y_test = y.select(Axis(0), test);

    let fit_start = Instant::now();
    let model = fit_fn(&x_train, &y_train);
    let fit_time = fit_start.elapsed();

    let score_start = Instant::now();
    let pred_test = model(&x_test);
    let test: Vec<f64> = scorers
        .iter()
        .map(|s| s.score(&y_test, &pred_test))
        .collect();
    let score_time = score_start.elapsed();

    let train = if return_train_scores {
        let pred_train = model(&x_train);
        Some(
            scorers
                .iter()
                .map(|s| s.score(&y_train, &pred_train))
                .collect(),
        )
    } else {
        None
    };

    FoldScores {
        test,
        train,
        fit_time,
        score_time,
    }
}

/// Cross-validate a model-fitting closure with one or more scorers.
///
/// Ties a [`CvSplitter`], a model-fitting closure, and one or more
/// [`Scorer`]s together — the utility you actually call day to day. `fit_fn` is
/// generic (it takes `(x_train, y_train)` and returns a *prediction* closure),
/// so this works with `smartcore`, `linfa`, or a hand-rolled model without this
/// crate depending on any of them.
///
/// Every scorer is evaluated in the **same pass** over the folds (re-fitting per
/// metric would be wasteful), so asking for accuracy *and* F1 together costs one
/// set of fits, not two.
///
/// With the `parallel` feature enabled the folds are fit and scored across a
/// `rayon` thread pool. The results are numerically identical to the serial path
/// — parallelism changes only wall-clock time, and the per-fold order is
/// preserved.
///
/// # Errors
///
/// Propagates any error from `splitter.split(x.nrows())`.
///
/// # Example
///
/// ```
/// use ndarray::{array, Array1, Array2};
/// use model_selection_rs::evaluate::{cross_validate, BoxedScorer};
/// use model_selection_rs::scoring::MeanSquaredError;
/// use model_selection_rs::splitters::KFold;
///
/// // A trivial "model" that predicts the training mean.
/// let x: Array2<f64> = Array2::zeros((10, 1));
/// let y: Array1<f64> = array![1., 2., 3., 4., 5., 6., 7., 8., 9., 10.];
/// let kf = KFold::new(5).unwrap();
/// let scorers: Vec<BoxedScorer> = vec![Box::new(MeanSquaredError)];
/// let res = cross_validate(&kf, &x, &y, |_xt, yt| {
///     let mean = yt.sum() / yt.len() as f64;
///     move |xq: &Array2<f64>| Array1::from_elem(xq.nrows(), mean)
/// }, &scorers, false).unwrap();
/// assert_eq!(res.n_splits(), 5);
/// ```
pub fn cross_validate<S, F, M>(
    splitter: &S,
    x: &Array2<f64>,
    y: &Array1<f64>,
    fit_fn: F,
    scorers: &[BoxedScorer],
    return_train_scores: bool,
) -> Result<CvResults>
where
    S: CvSplitter,
    F: Fn(&Array2<f64>, &Array1<f64>) -> M + Sync,
    M: Fn(&Array2<f64>) -> Array1<f64>,
{
    let splits = splitter.split(x.nrows())?;

    #[cfg(feature = "parallel")]
    let fold_scores: Vec<FoldScores> = {
        use rayon::prelude::*;
        splits
            .par_iter()
            .map(|(train, test)| {
                evaluate_fold(x, y, train, test, &fit_fn, scorers, return_train_scores)
            })
            .collect()
    };

    #[cfg(not(feature = "parallel"))]
    let fold_scores: Vec<FoldScores> = splits
        .iter()
        .map(|(train, test)| {
            evaluate_fold(x, y, train, test, &fit_fn, scorers, return_train_scores)
        })
        .collect();

    Ok(assemble(fold_scores, scorers, return_train_scores))
}

/// Transpose per-fold results into the `[scorer][fold]` layout of `CvResults`.
fn assemble(
    fold_scores: Vec<FoldScores>,
    scorers: &[BoxedScorer],
    return_train_scores: bool,
) -> CvResults {
    let n_scorers = scorers.len();
    let mut test_scores = vec![Vec::with_capacity(fold_scores.len()); n_scorers];
    let mut train_scores = if return_train_scores {
        Some(vec![Vec::with_capacity(fold_scores.len()); n_scorers])
    } else {
        None
    };
    let mut fit_times = Vec::with_capacity(fold_scores.len());
    let mut score_times = Vec::with_capacity(fold_scores.len());

    for fold in fold_scores {
        for (s, v) in fold.test.into_iter().enumerate() {
            test_scores[s].push(v);
        }
        if let (Some(dst), Some(src)) = (train_scores.as_mut(), fold.train) {
            for (s, v) in src.into_iter().enumerate() {
                dst[s].push(v);
            }
        }
        fit_times.push(fold.fit_time);
        score_times.push(fold.score_time);
    }

    CvResults {
        scorer_names: scorers.iter().map(|s| s.name().to_string()).collect(),
        test_scores,
        train_scores,
        fit_times,
        score_times,
    }
}

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

    /// Fit an ordinary least squares line y = a*x + b on one feature.
    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 mean_x = xs.iter().sum::<f64>() / n;
        let mean_y = ys.iter().sum::<f64>() / n;
        let cov: f64 = xs
            .iter()
            .zip(&ys)
            .map(|(a, b)| (a - mean_x) * (b - mean_y))
            .sum();
        let var: f64 = xs.iter().map(|a| (a - mean_x).powi(2)).sum();
        let slope = if var == 0.0 { 0.0 } else { cov / var };
        let intercept = mean_y - slope * mean_x;
        move |xq: &Array2<f64>| xq.column(0).mapv(|v| slope * v + intercept)
    }

    #[test]
    fn multiple_scorers_in_one_pass() {
        let x = Array2::from_shape_fn((20, 1), |(i, _)| i as f64);
        let y = x.column(0).mapv(|v| 2.0 * v + 1.0);
        let kf = KFold::new(4).unwrap();
        let scorers: Vec<BoxedScorer> = vec![Box::new(MeanSquaredError), Box::new(R2Score)];
        let res = cross_validate(&kf, &x, &y, ols_fit, &scorers, true).unwrap();
        assert_eq!(res.scorer_names, vec!["mse", "r2"]);
        assert_eq!(res.n_splits(), 4);
        // Perfect linear data -> near-zero MSE, R2 ~ 1.
        assert!(res.mean_test_score(0) < 1e-9);
        assert!((res.mean_test_score(1) - 1.0).abs() < 1e-9);
        assert!(res.train_scores.is_some());
    }

    #[test]
    fn lookup_by_name() {
        let x = Array2::from_shape_fn((12, 1), |(i, _)| i as f64);
        let y = x.column(0).to_owned();
        let kf = KFold::new(3).unwrap();
        let scorers: Vec<BoxedScorer> = vec![Box::new(MeanSquaredError)];
        let res = cross_validate(&kf, &x, &y, ols_fit, &scorers, false).unwrap();
        assert!(res.mean_test_score_by_name("mse").is_some());
        assert!(res.mean_test_score_by_name("nope").is_none());
    }

    #[test]
    fn array_example_in_docs() {
        let x = array![[0.0], [1.0], [2.0], [3.0]];
        let y = array![0.0, 1.0, 2.0, 3.0];
        let kf = KFold::new(2).unwrap();
        let scorers: Vec<BoxedScorer> = vec![Box::new(MeanSquaredError)];
        let res = cross_validate(&kf, &x, &y, ols_fit, &scorers, false).unwrap();
        assert_eq!(res.n_splits(), 2);
    }
}