1use crate::simd;
2
3#[must_use]
4pub fn mean(x: &[f64]) -> f64 {
5 if x.is_empty() {
6 return 0.0;
7 }
8 simd::sum(x) / x.len() as f64
9}
10
11#[must_use]
12pub fn variance(x: &[f64]) -> f64 {
13 if x.is_empty() {
14 return 0.0;
15 }
16 simd::sq_dev(x, mean(x)) / x.len() as f64
17}
18
19#[must_use]
20pub fn median(x: &[f64]) -> f64 {
21 if x.is_empty() {
22 return f64::NAN;
23 }
24 let mut s = x.to_vec();
25 s.sort_by(f64::total_cmp);
26 let m = s.len() / 2;
27 if s.len().is_multiple_of(2) {
28 (s[m - 1] + s[m]) / 2.0
29 } else {
30 s[m]
31 }
32}
33
34#[must_use]
35pub fn geomean(x: &[f64]) -> f64 {
36 if x.is_empty() {
37 return f64::NAN;
38 }
39 let ln: Vec<f64> = x.iter().map(|v| v.ln()).collect();
40 mean(&ln).exp()
41}
42
43#[must_use]
44pub fn mse_u16(a: &[u16], sa: usize, b: &[u16], sb: usize, w: usize, h: usize) -> f64 {
45 if w == 0 || h == 0 {
46 return 0.0;
47 }
48 let sse: u64 = (0..h).map(|r| simd::sq_diff_u16(&a[r * sa..r * sa + w], &b[r * sb..r * sb + w])).sum();
49 sse as f64 / (w * h) as f64
50}
51
52#[must_use]
53pub fn psnr(mse: f64, max: f64, cap: f64) -> f64 {
54 if mse <= 0.0 {
55 return cap;
56 }
57 (10.0 * (max * max / mse).log10()).min(cap)
58}
59
60const LANCZOS: [f64; 9] = [
61 0.999_999_999_999_809_9,
62 676.520_368_121_885_1,
63 -1_259.139_216_722_402_8,
64 771.323_428_777_653_1,
65 -176.615_029_162_140_6,
66 12.507_343_278_686_905,
67 -0.138_571_095_265_720_12,
68 9.984_369_578_019_572e-6,
69 1.505_632_735_149_311_6e-7,
70];
71
72#[must_use]
73pub fn ln_gamma(x: f64) -> f64 {
74 if x < 0.5 {
75 return (std::f64::consts::PI / (std::f64::consts::PI * x).sin()).ln() - ln_gamma(1.0 - x);
76 }
77 let x = x - 1.0;
78 let t = x + 7.5;
79 let s = LANCZOS[1..].iter().enumerate().fold(LANCZOS[0], |s, (i, c)| s + c / (x + i as f64 + 1.0));
80 0.5 * (2.0 * std::f64::consts::PI).ln() + (x + 0.5) * t.ln() - t + s.ln()
81}
82
83fn beta_fraction(a: f64, b: f64, x: f64) -> f64 {
84 const TINY: f64 = 1e-300;
85 let (qab, qap, qam) = (a + b, a + 1.0, a - 1.0);
86 let mut c = 1.0;
87 let mut d = 1.0 - qab * x / qap;
88 if d.abs() < TINY {
89 d = TINY;
90 }
91 d = 1.0 / d;
92 let mut h = d;
93 for m in 1..=300 {
94 let m = f64::from(m);
95 let m2 = 2.0 * m;
96 let aa = m * (b - m) * x / ((qam + m2) * (a + m2));
97 d = 1.0 + aa * d;
98 d = if d.abs() < TINY { TINY } else { d };
99 c = 1.0 + aa / c;
100 c = if c.abs() < TINY { TINY } else { c };
101 d = 1.0 / d;
102 h *= d * c;
103 let aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2));
104 d = 1.0 + aa * d;
105 d = if d.abs() < TINY { TINY } else { d };
106 c = 1.0 + aa / c;
107 c = if c.abs() < TINY { TINY } else { c };
108 d = 1.0 / d;
109 let del = d * c;
110 h *= del;
111 if (del - 1.0).abs() < 1e-15 {
112 break;
113 }
114 }
115 h
116}
117
118#[must_use]
119pub fn incomplete_beta(a: f64, b: f64, x: f64) -> f64 {
120 if x <= 0.0 {
121 return 0.0;
122 }
123 if x >= 1.0 {
124 return 1.0;
125 }
126 let front = (ln_gamma(a + b) - ln_gamma(a) - ln_gamma(b) + a * x.ln() + b * (1.0 - x).ln()).exp();
127 if x < (a + 1.0) / (a + b + 2.0) {
128 front * beta_fraction(a, b, x) / a
129 } else {
130 1.0 - front * beta_fraction(b, a, 1.0 - x) / b
131 }
132}
133
134#[must_use]
135pub fn student_t_two_sided(t: f64, df: f64) -> f64 {
136 if !t.is_finite() {
137 return 0.0;
138 }
139 incomplete_beta(df / 2.0, 0.5, df / (df + t * t))
140}
141
142#[derive(Clone, Copy, Debug)]
143pub struct Welch {
144 pub t: f64,
145 pub df: f64,
146 pub p: f64,
147}
148
149#[must_use]
150pub fn welch(a: &[f64], b: &[f64]) -> Option<Welch> {
151 if a.len() < 2 || b.len() < 2 {
152 return None;
153 }
154 let (na, nb) = (a.len() as f64, b.len() as f64);
155 let va = variance(a) * na / (na - 1.0) / na;
156 let vb = variance(b) * nb / (nb - 1.0) / nb;
157 let se2 = va + vb;
158 if se2 <= 0.0 {
159 return None;
160 }
161 let t = (mean(a) - mean(b)) / se2.sqrt();
162 let df = se2 * se2 / (va * va / (na - 1.0) + vb * vb / (nb - 1.0));
163 Some(Welch { t, df, p: student_t_two_sided(t, df) })
164}
165
166#[derive(Clone, Copy, Debug)]
167pub struct Line {
168 pub intercept: f64,
169 pub slope: f64,
170 pub se: f64,
171}
172
173#[must_use]
174pub fn line_weighted(x: &[f64], y: &[f64], w: &[f64]) -> Option<Line> {
175 let n = x.len().min(y.len()).min(w.len());
176 if n < 3 {
177 return None;
178 }
179 let (x, y, w) = (&x[..n], &y[..n], &w[..n]);
180 let sw = simd::sum(w);
181 if sw <= 0.0 {
182 return None;
183 }
184 let mx = x.iter().zip(w).map(|(a, b)| a * b).sum::<f64>() / sw;
185 let my = y.iter().zip(w).map(|(a, b)| a * b).sum::<f64>() / sw;
186 let (mut sxx, mut sxy) = (0.0, 0.0);
187 for i in 0..n {
188 sxx += w[i] * (x[i] - mx) * (x[i] - mx);
189 sxy += w[i] * (x[i] - mx) * (y[i] - my);
190 }
191 if sxx <= 0.0 {
192 return None;
193 }
194 let slope = sxy / sxx;
195 let intercept = my - slope * mx;
196 let rss: f64 = (0..n).map(|i| w[i] * (y[i] - intercept - slope * x[i]).powi(2)).sum();
197 let se = (rss / (n as f64 - 2.0) / sxx).sqrt();
198 Some(Line { intercept, slope, se })
199}
200
201#[cfg(test)]
202mod tests {
203 use super::*;
204
205 #[test]
206 fn weighted_line_matches_the_unweighted_formula_and_repeated_points() {
207 let x = [0.0, 1.0, 2.0, 3.0, 4.0];
208 let y = [1.1, 2.9, 5.2, 6.8, 9.1];
209 let l = line_weighted(&x, &y, &[1.0; 5]).expect("five points");
210 assert!((l.slope - 1.99).abs() < 1e-12, "slope {}", l.slope);
211 assert!((l.intercept - 1.04).abs() < 1e-12, "intercept {}", l.intercept);
212 let rss: f64 = x.iter().zip(&y).map(|(a, b)| (b - 1.04 - 1.99 * a).powi(2)).sum();
213 assert!((l.se - (rss / 3.0 / 10.0).sqrt()).abs() < 1e-12, "se {}", l.se);
214 let scaled = line_weighted(&x, &y, &[7.0; 5]).expect("five points");
215 assert!((scaled.slope - l.slope).abs() < 1e-12 && (scaled.se - l.se).abs() < 1e-12);
216 let xw = [0.0, 1.0, 2.0, 3.0];
217 let yw = [0.0, 1.5, 1.8, 3.3];
218 let a = line_weighted(&xw, &yw, &[1.0, 2.0, 1.0, 1.0]).expect("weighted");
219 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");
220 assert!((a.slope - b.slope).abs() < 1e-12 && (a.intercept - b.intercept).abs() < 1e-12);
221 }
222
223 #[test]
224 fn student_t_and_welch_match_tables() {
225 assert!((ln_gamma(5.0) - 24f64.ln()).abs() < 1e-12);
226 assert!((ln_gamma(0.5) - std::f64::consts::PI.sqrt().ln()).abs() < 1e-12);
227 assert!((student_t_two_sided(2.228_138_852, 10.0) - 0.05).abs() < 1e-6);
228 assert!((student_t_two_sided(2.0, 10.0) - 0.073_388_034).abs() < 1e-6);
229 assert!((student_t_two_sided(1.959_963_985, 1e9) - 0.05).abs() < 1e-6);
230 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];
231 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];
232 let w = welch(&a, &b).expect("two samples");
233 assert!((w.t + 2.46).abs() < 5e-3, "t {}", w.t);
234 assert!((w.df - 24.99).abs() < 5e-3, "df {}", w.df);
235 assert!((w.p - 0.021).abs() < 5e-4, "p {}", w.p);
236 }
237
238 #[test]
239 fn moments_and_psnr() {
240 let x = [1.0, 2.0, 3.0, 4.0];
241 assert!((mean(&x) - 2.5).abs() < 1e-12);
242 assert!((variance(&x) - 1.25).abs() < 1e-12);
243 let a = [10u16; 12];
244 let b = [12u16; 12];
245 assert!((mse_u16(&a, 4, &b, 4, 4, 3) - 4.0).abs() < 1e-12);
246 assert!((psnr(4.0, 255.0, 99.0) - 10.0 * (255.0f64 * 255.0 / 4.0).log10()).abs() < 1e-12);
247 assert_eq!(psnr(0.0, 255.0, 99.0), 99.0);
248 assert_eq!(median(&[3.0, 1.0, 2.0]), 2.0);
249 assert_eq!(median(&[4.0, 1.0, 3.0, 2.0]), 2.5);
250 assert!(median(&[]).is_nan());
251 assert!((geomean(&[1.0, 4.0, 16.0]) - 4.0).abs() < 1e-12);
252 assert!((geomean(&[0.5, 2.0]) - 1.0).abs() < 1e-15);
253 assert!(geomean(&[]).is_nan());
254 }
255}