use crate::simd;
#[must_use]
pub fn mean(x: &[f64]) -> f64 {
if x.is_empty() {
return 0.0;
}
simd::sum(x) / x.len() as f64
}
#[must_use]
pub fn variance(x: &[f64]) -> f64 {
if x.is_empty() {
return 0.0;
}
simd::sq_dev(x, mean(x)) / x.len() as f64
}
#[must_use]
pub fn median(x: &[f64]) -> f64 {
if x.is_empty() {
return f64::NAN;
}
let mut s = x.to_vec();
s.sort_by(f64::total_cmp);
let m = s.len() / 2;
if s.len() % 2 == 0 {
(s[m - 1] + s[m]) / 2.0
} else {
s[m]
}
}
#[must_use]
pub fn geomean(x: &[f64]) -> f64 {
if x.is_empty() {
return f64::NAN;
}
let ln: Vec<f64> = x.iter().map(|v| v.ln()).collect();
mean(&ln).exp()
}
#[must_use]
pub fn mse_u16(a: &[u16], sa: usize, b: &[u16], sb: usize, w: usize, h: usize) -> f64 {
if w == 0 || h == 0 {
return 0.0;
}
let sse: u64 = (0..h).map(|r| simd::sq_diff_u16(&a[r * sa..r * sa + w], &b[r * sb..r * sb + w])).sum();
sse as f64 / (w * h) as f64
}
#[must_use]
pub fn psnr(mse: f64, max: f64, cap: f64) -> f64 {
if mse <= 0.0 {
return cap;
}
(10.0 * (max * max / mse).log10()).min(cap)
}
const LANCZOS: [f64; 9] = [
0.999_999_999_999_809_9,
676.520_368_121_885_1,
-1_259.139_216_722_402_8,
771.323_428_777_653_1,
-176.615_029_162_140_6,
12.507_343_278_686_905,
-0.138_571_095_265_720_12,
9.984_369_578_019_572e-6,
1.505_632_735_149_311_6e-7,
];
#[must_use]
pub fn ln_gamma(x: f64) -> f64 {
if x < 0.5 {
return (std::f64::consts::PI / (std::f64::consts::PI * x).sin()).ln() - ln_gamma(1.0 - x);
}
let x = x - 1.0;
let t = x + 7.5;
let s = LANCZOS[1..].iter().enumerate().fold(LANCZOS[0], |s, (i, c)| s + c / (x + i as f64 + 1.0));
0.5 * (2.0 * std::f64::consts::PI).ln() + (x + 0.5) * t.ln() - t + s.ln()
}
fn beta_fraction(a: f64, b: f64, x: f64) -> f64 {
const TINY: f64 = 1e-300;
let (qab, qap, qam) = (a + b, a + 1.0, a - 1.0);
let mut c = 1.0;
let mut d = 1.0 - qab * x / qap;
if d.abs() < TINY {
d = TINY;
}
d = 1.0 / d;
let mut h = d;
for m in 1..=300 {
let m = f64::from(m);
let m2 = 2.0 * m;
let aa = m * (b - m) * x / ((qam + m2) * (a + m2));
d = 1.0 + aa * d;
d = if d.abs() < TINY { TINY } else { d };
c = 1.0 + aa / c;
c = if c.abs() < TINY { TINY } else { c };
d = 1.0 / d;
h *= d * c;
let aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2));
d = 1.0 + aa * d;
d = if d.abs() < TINY { TINY } else { d };
c = 1.0 + aa / c;
c = if c.abs() < TINY { TINY } else { c };
d = 1.0 / d;
let del = d * c;
h *= del;
if (del - 1.0).abs() < 1e-15 {
break;
}
}
h
}
#[must_use]
pub fn incomplete_beta(a: f64, b: f64, x: f64) -> f64 {
if x <= 0.0 {
return 0.0;
}
if x >= 1.0 {
return 1.0;
}
let front = (ln_gamma(a + b) - ln_gamma(a) - ln_gamma(b) + a * x.ln() + b * (1.0 - x).ln()).exp();
if x < (a + 1.0) / (a + b + 2.0) {
front * beta_fraction(a, b, x) / a
} else {
1.0 - front * beta_fraction(b, a, 1.0 - x) / b
}
}
#[must_use]
pub fn student_t_two_sided(t: f64, df: f64) -> f64 {
if !t.is_finite() {
return 0.0;
}
incomplete_beta(df / 2.0, 0.5, df / (df + t * t))
}
#[derive(Clone, Copy, Debug)]
pub struct Welch {
pub t: f64,
pub df: f64,
pub p: f64,
}
#[must_use]
pub fn welch(a: &[f64], b: &[f64]) -> Option<Welch> {
if a.len() < 2 || b.len() < 2 {
return None;
}
let (na, nb) = (a.len() as f64, b.len() as f64);
let va = variance(a) * na / (na - 1.0) / na;
let vb = variance(b) * nb / (nb - 1.0) / nb;
let se2 = va + vb;
if se2 <= 0.0 {
return None;
}
let t = (mean(a) - mean(b)) / se2.sqrt();
let df = se2 * se2 / (va * va / (na - 1.0) + vb * vb / (nb - 1.0));
Some(Welch { t, df, p: student_t_two_sided(t, df) })
}
#[derive(Clone, Copy, Debug)]
pub struct Line {
pub intercept: f64,
pub slope: f64,
pub se: f64,
}
#[must_use]
pub fn line_weighted(x: &[f64], y: &[f64], w: &[f64]) -> Option<Line> {
let n = x.len().min(y.len()).min(w.len());
if n < 3 {
return None;
}
let (x, y, w) = (&x[..n], &y[..n], &w[..n]);
let sw = simd::sum(w);
if sw <= 0.0 {
return None;
}
let mx = x.iter().zip(w).map(|(a, b)| a * b).sum::<f64>() / sw;
let my = y.iter().zip(w).map(|(a, b)| a * b).sum::<f64>() / sw;
let (mut sxx, mut sxy) = (0.0, 0.0);
for i in 0..n {
sxx += w[i] * (x[i] - mx) * (x[i] - mx);
sxy += w[i] * (x[i] - mx) * (y[i] - my);
}
if sxx <= 0.0 {
return None;
}
let slope = sxy / sxx;
let intercept = my - slope * mx;
let rss: f64 = (0..n).map(|i| w[i] * (y[i] - intercept - slope * x[i]).powi(2)).sum();
let se = (rss / (n as f64 - 2.0) / sxx).sqrt();
Some(Line { intercept, slope, se })
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn weighted_line_matches_the_unweighted_formula_and_repeated_points() {
let x = [0.0, 1.0, 2.0, 3.0, 4.0];
let y = [1.1, 2.9, 5.2, 6.8, 9.1];
let l = line_weighted(&x, &y, &[1.0; 5]).expect("five points");
assert!((l.slope - 1.99).abs() < 1e-12, "slope {}", l.slope);
assert!((l.intercept - 1.04).abs() < 1e-12, "intercept {}", l.intercept);
let rss: f64 = x.iter().zip(&y).map(|(a, b)| (b - 1.04 - 1.99 * a).powi(2)).sum();
assert!((l.se - (rss / 3.0 / 10.0).sqrt()).abs() < 1e-12, "se {}", l.se);
let scaled = line_weighted(&x, &y, &[7.0; 5]).expect("five points");
assert!((scaled.slope - l.slope).abs() < 1e-12 && (scaled.se - l.se).abs() < 1e-12);
let xw = [0.0, 1.0, 2.0, 3.0];
let yw = [0.0, 1.5, 1.8, 3.3];
let a = line_weighted(&xw, &yw, &[1.0, 2.0, 1.0, 1.0]).expect("weighted");
let b = line_weighted(&[0.0, 1.0, 1.0, 2.0, 3.0], &[0.0, 1.5, 1.5, 1.8, 3.3], &[1.0; 5]).expect("repeated");
assert!((a.slope - b.slope).abs() < 1e-12 && (a.intercept - b.intercept).abs() < 1e-12);
}
#[test]
fn student_t_and_welch_match_tables() {
assert!((ln_gamma(5.0) - 24f64.ln()).abs() < 1e-12);
assert!((ln_gamma(0.5) - std::f64::consts::PI.sqrt().ln()).abs() < 1e-12);
assert!((student_t_two_sided(2.228_138_852, 10.0) - 0.05).abs() < 1e-6);
assert!((student_t_two_sided(2.0, 10.0) - 0.073_388_034).abs() < 1e-6);
assert!((student_t_two_sided(1.959_963_985, 1e9) - 0.05).abs() < 1e-6);
let a = [27.5, 21.0, 19.0, 23.6, 17.0, 17.9, 16.9, 20.1, 21.9, 22.6, 23.1, 19.6, 19.0, 21.7, 21.4];
let b = [27.1, 22.0, 20.8, 23.4, 23.4, 23.5, 25.8, 22.0, 24.8, 20.2, 21.9, 22.1, 22.9, 20.5, 24.4];
let w = welch(&a, &b).expect("two samples");
assert!((w.t + 2.46).abs() < 5e-3, "t {}", w.t);
assert!((w.df - 24.99).abs() < 5e-3, "df {}", w.df);
assert!((w.p - 0.021).abs() < 5e-4, "p {}", w.p);
}
#[test]
fn moments_and_psnr() {
let x = [1.0, 2.0, 3.0, 4.0];
assert!((mean(&x) - 2.5).abs() < 1e-12);
assert!((variance(&x) - 1.25).abs() < 1e-12);
let a = [10u16; 12];
let b = [12u16; 12];
assert!((mse_u16(&a, 4, &b, 4, 4, 3) - 4.0).abs() < 1e-12);
assert!((psnr(4.0, 255.0, 99.0) - 10.0 * (255.0f64 * 255.0 / 4.0).log10()).abs() < 1e-12);
assert_eq!(psnr(0.0, 255.0, 99.0), 99.0);
assert_eq!(median(&[3.0, 1.0, 2.0]), 2.0);
assert_eq!(median(&[4.0, 1.0, 3.0, 2.0]), 2.5);
assert!(median(&[]).is_nan());
assert!((geomean(&[1.0, 4.0, 16.0]) - 4.0).abs() < 1e-12);
assert!((geomean(&[0.5, 2.0]) - 1.0).abs() < 1e-15);
assert!(geomean(&[]).is_nan());
}
}