Revision control
Copy as Markdown
Other Tools
#![cfg(feature = "ndarray")]
use ndarray::{Array1, Array2};
use rand::SeedableRng;
use rand_chacha::ChaCha20Rng;
use smawk::online_column_minima;
mod random_monge;
use random_monge::random_monge_matrix;
#[derive(Debug)]
struct LinRegression {
alpha: f64,
beta: f64,
r_squared: f64,
}
/// Square an expression. Works equally well for floats and matrices.
macro_rules! squared {
($x:expr) => {
$x * $x
};
}
/// Compute the mean of a 1-dimensional array.
macro_rules! mean {
($a:expr) => {
$a.mean().expect("Mean of empty array")
};
}
/// Compute a simple linear regression from the list of values.
///
fn linear_regression(values: &[(usize, i32)]) -> LinRegression {
let xs = values.iter().map(|&(x, _)| x as f64).collect::<Array1<_>>();
let ys = values.iter().map(|&(_, y)| y as f64).collect::<Array1<_>>();
let xs_mean = mean!(&xs);
let ys_mean = mean!(&ys);
let xs_ys_mean = mean!(&xs * &ys);
let cov_xs_ys = ((&xs - xs_mean) * (&ys - ys_mean)).sum();
let var_xs = squared!(&xs - xs_mean).sum();
let beta = cov_xs_ys / var_xs;
let alpha = ys_mean - beta * xs_mean;
let r_squared = squared!(xs_ys_mean - xs_mean * ys_mean)
/ ((mean!(&xs * &xs) - squared!(xs_mean)) * (mean!(&ys * &ys) - squared!(ys_mean)));
LinRegression {
alpha: alpha,
beta: beta,
r_squared: r_squared,
}
}
/// Check that the number of matrix accesses in `online_column_minima`
/// grows as O(*n*) for *n* ✕ *n* matrix.
#[test]
fn online_linear_complexity() {
let mut rng = ChaCha20Rng::seed_from_u64(0);
let mut data = vec![];
for &size in &[1, 2, 3, 4, 5, 10, 15, 20, 30, 40, 50, 60, 70, 80, 90, 100] {
let matrix: Array2<i32> = random_monge_matrix(size, size, &mut rng);
let count = std::cell::RefCell::new(0);
online_column_minima(0, size, |_, i, j| {
*count.borrow_mut() += 1;
matrix[[i, j]]
});
data.push((size, count.into_inner()));
}
let lin_reg = linear_regression(&data);
assert!(
lin_reg.r_squared > 0.95,
"r² = {:.4} is lower than expected for a linear fit\nData points: {:?}\n{:?}",
lin_reg.r_squared,
data,
lin_reg
);
}