use model_selection_rs::evaluate::{learning_curve, TrainSize};
use model_selection_rs::scoring::R2Score;
use model_selection_rs::splitters::KFold;
use ndarray::{Array1, Array2};
fn constant_mean(_x: &Array2<f64>, y: &Array1<f64>) -> impl Fn(&Array2<f64>) -> Array1<f64> {
let mean = y.sum() / y.len() as f64;
move |xq: &Array2<f64>| Array1::from_elem(xq.nrows(), mean)
}
fn main() {
let x = Array2::from_shape_fn((80, 1), |(i, _)| i as f64);
let y = x.column(0).mapv(|v| 2.0 * v + 1.0);
let kf = KFold::new(4).unwrap().with_shuffle(0);
let sizes = [
TrainSize::Fraction(0.2),
TrainSize::Fraction(0.4),
TrainSize::Fraction(0.6),
TrainSize::Fraction(0.8),
TrainSize::Fraction(1.0),
];
let lc = learning_curve(&kf, &x, &y, constant_mean, &R2Score, &sizes).unwrap();
println!("Learning curve for an under-capacity (constant) model on linear data:\n");
println!(
" {:>10} {:>12} {:>12}",
"train_size", "train R2", "val R2"
);
let train = lc.mean_train_scores();
let val = lc.mean_val_scores();
for i in 0..lc.train_sizes.len() {
println!(
" {:>10} {:>12.3} {:>12.3}",
lc.train_sizes[i], train[i], val[i]
);
}
println!("\nBoth curves sit low and close together — the classic high-bias shape:");
println!("more data does not help a model that cannot represent the signal.");
}