#![expect(
clippy::unwrap_used,
reason = "test scope: failed precondition = test failure"
)]
use leto::application::reduction::var_axis;
use leto::application::statistics::{covariance, pearson_correlation};
use leto::{Array2, LetoError, Storage};
#[track_caller]
fn assert_close(actual: f64, expected: f64) {
assert!(
(actual - expected).abs() <= 1.0e-12 * expected.abs().max(1.0),
"actual {actual} expected {expected}"
);
}
fn leto(shape: [usize; 2], data: Vec<f64>) -> Array2<f64> {
Array2::from_shape_vec(shape, data).unwrap()
}
#[test]
fn covariance_closed_form_sample_and_population() {
let a = leto([2, 3], vec![1.0, 2.0, 3.0, 4.0, 6.0, 8.0]);
let c1 = covariance(&a, 1.0).unwrap();
let d1 = c1.storage().as_slice();
assert_close(d1[0], 1.0); assert_close(d1[1], 2.0); assert_close(d1[2], 2.0); assert_close(d1[3], 4.0);
let c0 = covariance(&a, 0.0).unwrap();
let d0 = c0.storage().as_slice();
assert_close(d0[0], 2.0 / 3.0);
assert_close(d0[1], 4.0 / 3.0);
assert_close(d0[3], 8.0 / 3.0);
}
#[test]
fn covariance_diagonal_equals_variance() {
let a = leto(
[3, 4],
vec![1.0, 3.0, 2.0, 5.0, 2.0, 2.0, 2.0, 2.0, 9.0, 1.0, 4.0, 7.0],
);
for &ddof in &[0.0_f64, 1.0] {
let cov = covariance(&a, ddof).unwrap();
let cd = cov.storage().as_slice();
let var = var_axis::<f64, _, 2, 1>(&a, 1, ddof).unwrap();
let vd = var.storage().as_slice();
for i in 0..3usize {
assert_close(cd[i * 3 + i], vd[i]);
}
}
}
#[test]
fn covariance_is_symmetric() {
let a = leto([3, 5], (0..15).map(f64::from).collect());
let cov = covariance(&a, 1.0).unwrap();
let cd = cov.storage().as_slice();
for i in 0..3usize {
for j in 0..3usize {
assert_close(cd[i * 3 + j], cd[j * 3 + i]);
}
}
}
#[test]
fn correlation_diagonal_unit_and_perfect_linear() {
let a = leto([2, 3], vec![1.0, 2.0, 3.0, 4.0, 6.0, 8.0]);
let r = pearson_correlation(&a).unwrap();
let rd = r.storage().as_slice();
assert_close(rd[0], 1.0); assert_close(rd[3], 1.0); assert_close(rd[1], 1.0); assert_close(rd[2], 1.0); }
#[test]
fn correlation_perfect_negative() {
let a = leto([2, 3], vec![1.0, 2.0, 3.0, -1.0, -2.0, -3.0]);
let r = pearson_correlation(&a).unwrap();
let rd = r.storage().as_slice();
assert_close(rd[1], -1.0);
}
#[test]
fn correlation_bounded_and_matches_normalized_covariance() {
let a = leto(
[3, 4],
vec![1.0, 2.0, 4.0, 3.0, 2.0, 1.0, 3.0, 5.0, 7.0, 1.0, 2.0, 4.0],
);
let cov = covariance(&a, 0.0).unwrap();
let cd = cov.storage().as_slice();
let r = pearson_correlation(&a).unwrap();
let rd = r.storage().as_slice();
for i in 0..3usize {
for j in 0..3usize {
let expected = cd[i * 3 + j] / (cd[i * 3 + i].sqrt() * cd[j * 3 + j].sqrt());
assert_close(rd[i * 3 + j], expected);
assert!(rd[i * 3 + j].abs() <= 1.0 + 1.0e-12);
}
}
}
#[test]
fn covariance_rejects_empty_and_excess_ddof() {
let a = leto([2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
match covariance(&a, 3.0) {
Err(LetoError::StorageError { reason }) => {
assert_eq!(
reason,
"variance degrees of freedom (n - ddof) must be positive"
);
}
other => panic!("expected dof error, got {other:?}"),
}
let empty: Array2<f64> = Array2::from_shape_vec([0, 0], vec![]).unwrap();
match covariance(&empty, 0.0) {
Err(LetoError::StorageError { reason }) => {
assert_eq!(reason, "covariance over an empty matrix is undefined");
}
other => panic!("expected empty error, got {other:?}"),
}
match pearson_correlation(&empty) {
Err(LetoError::StorageError { reason }) => {
assert_eq!(reason, "covariance over an empty matrix is undefined");
}
other => panic!("expected propagated covariance error, got {other:?}"),
}
}