1use std::f64::consts::PI;
5
6use faer::{Mat, Side};
7
8use crate::error::FittingError;
9use crate::lm::FlatMatrix;
10use crate::poisson::{NEWTON_DECREMENT_TOL, Prior, symmetric};
11
12const TERM_SHIFT: f64 = 0.1;
13
14#[derive(Debug, Clone, Copy, PartialEq)]
17pub struct Consistency {
18 pub q: f64,
19 pub dof: usize,
20 pub p: f64,
21}
22
23impl Consistency {
24 pub fn new(q: f64, dof: usize) -> Result<Self, FittingError> {
30 if q.is_nan() || q < 0.0 || dof == 0 {
31 return Err(FittingError::InvalidConfig(format!(
32 "a χ² test has a statistic of 0 or more on 1 or more degrees of freedom; got \
33 {q} on {dof}"
34 )));
35 }
36 Ok(Self {
37 q,
38 dof,
39 p: chi_squared_survival(q, dof),
40 })
41 }
42}
43
44fn chi_squared_survival(q: f64, dof: usize) -> f64 {
45 if q == f64::INFINITY {
46 return 0.0;
47 }
48 let x = q / 2.0;
49 let (mut sum, mut log_term, offset) = if dof.is_multiple_of(2) {
50 (0.0, -x, 1.0)
51 } else {
52 (
53 libm::erfc(x.sqrt()),
54 std::f64::consts::LN_2 - x + 0.5 * (x / PI).ln(),
55 1.5,
56 )
57 };
58 for n in 0..dof / 2 {
59 sum += log_term.exp();
60 log_term += x.ln() - (n as f64 + offset).ln();
61 }
62 sum
63}
64
65pub fn consistency(
88 prior: &Prior,
89 estimate: &[f64],
90 posterior: &FlatMatrix,
91) -> Result<Option<Consistency>, FittingError> {
92 let n = prior.parameters.len();
93 for (expected, actual, field) in [
94 (n, estimate.len(), "estimate"),
95 (n, posterior.nrows, "posterior rows"),
96 (n, posterior.ncols, "posterior columns"),
97 (n * n, posterior.data.len(), "posterior entries"),
98 ] {
99 if actual != expected {
100 return Err(FittingError::LengthMismatch {
101 expected,
102 actual,
103 field,
104 });
105 }
106 }
107 if !prior.is_valid() {
108 return Err(FittingError::InvalidConfig(format!(
109 "the consistency of a prior needs a finite mean and finite positive sds; got \
110 {prior:?}"
111 )));
112 }
113 let posterior = Mat::from_fn(n, n, |a, b| posterior.get(a, b));
114 if !(0..n).all(|a| (0..n).all(|b| posterior[(a, b)].is_finite())) {
115 return Ok(None);
116 }
117 if !estimate.iter().all(|x| x.is_finite()) {
118 return Err(FittingError::InvalidConfig(format!(
119 "the consistency of a prior needs finite estimates; got {estimate:?}"
120 )));
121 }
122 if !symmetric(n, |a, b| posterior[(a, b)]) {
123 return Err(FittingError::InvalidConfig(format!(
124 "a posterior covariance must be symmetric; got {posterior:?}"
125 )));
126 }
127 let mut z: Vec<f64> = estimate
128 .iter()
129 .zip(&prior.mean)
130 .map(|(x, m)| x - m)
131 .collect();
132 prior.whiten(&mut z);
133 let half: Vec<Vec<f64>> = (0..n)
134 .map(|b| {
135 let mut column: Vec<f64> = (0..n).map(|a| posterior[(a, b)]).collect();
136 prior.whiten(&mut column);
137 column
138 })
139 .collect();
140 let whitened: Vec<Vec<f64>> = (0..n)
141 .map(|a| {
142 let mut row: Vec<f64> = half.iter().map(|column| column[a]).collect();
143 prior.whiten(&mut row);
144 row
145 })
146 .collect();
147 let shrinkage = Mat::from_fn(n, n, |a, b| {
148 f64::from(u8::from(a == b)) - 0.5 * (whitened[a][b] + whitened[b][a])
149 });
150 let eigen = shrinkage
151 .self_adjoint_eigen(Side::Lower)
152 .map_err(|e| FittingError::EvaluationFailed(format!("{e:?}")))?;
153 let floor = (2.0 * (2.0 * NEWTON_DECREMENT_TOL).sqrt() / TERM_SHIFT).powi(2);
154 let shares = eigen.S().column_vector();
155 if (0..n).any(|d| !(-floor..=1.0 + floor).contains(&shares[d])) {
156 return Err(FittingError::InvalidConfig(format!(
157 "a posterior covariance must lie within the prior's; the shares of the prior's \
158 variance it removes are {:?}",
159 (0..n).map(|d| shares[d]).collect::<Vec<f64>>()
160 )));
161 }
162 let (mut q, mut dof) = (0.0, 0);
163 for d in 0..n {
164 let share = shares[d];
165 if share >= floor {
166 let projection: f64 = (0..n).map(|a| eigen.U()[(a, d)] * z[a]).sum();
167 q += projection.powi(2) / share;
168 dof += 1;
169 }
170 }
171 (dof > 0).then(|| Consistency::new(q, dof)).transpose()
172}
173
174pub fn agreement(a: &Prior, b: &Prior) -> Result<Consistency, FittingError> {
187 if a.parameters != b.parameters || !a.is_valid() || !b.is_valid() {
188 return Err(FittingError::InvalidConfig(format!(
189 "two valid estimates of the same parameters are compared; got {a:?} and {b:?}"
190 )));
191 }
192 let k = a.parameters.len();
193 let mut sum = FlatMatrix::zeros(k, k);
194 for i in 0..k {
195 for j in 0..k {
196 *sum.get_mut(i, j) = a.covariance(i, j) + b.covariance(i, j);
197 }
198 }
199 let joint = Prior::factored(&a.parameters, &b.mean, &sum);
200 if !joint.is_valid() {
201 return Err(FittingError::EvaluationFailed(format!(
202 "the sum of the two covariances is not positive definite: {sum:?}"
203 )));
204 }
205 let mut d: Vec<f64> = a.mean.iter().zip(&b.mean).map(|(x, y)| x - y).collect();
206 joint.whiten(&mut d);
207 Consistency::new(d.iter().map(|w| w * w).sum(), k)
208}
209
210#[cfg(test)]
211mod tests {
212 use super::*;
213
214 #[test]
215 fn the_chi_squared_survival_matches_scipy() {
216 for (dof, q, p) in [
217 (1, 3.841_458_820_694_120_5, 0.05),
218 (2, 5.991_464_547_107_98, 0.05),
219 (3, 7.814_727_903_251_176_5, 0.05),
220 (4, 9.487_729_036_781_154, 0.05),
221 (5, 11.070_497_693_516_351, 0.05),
222 (6, 12.591_587_243_743_977, 0.05),
223 (1, 2.0, 0.157_299_207_050_281_05),
224 (5, 40.0, 1.493_367_900_050_396e-7),
225 (4, 0.3, 0.989_814_172_888_816_5),
226 ] {
227 let survival = chi_squared_survival(q, dof);
228 assert!((survival / p - 1.0).abs() <= 1e-13, "{dof} {q}: {survival}");
229 }
230 for (dof, q, p) in [
231 (1500, 1500.0, 0.495_144_193_335_767_9),
232 (1501, 1501.0, 0.495_145_811_152_950_14),
233 (1500, 1700.0, 2.217_079_970_428_596_7e-4),
234 ] {
235 let survival = chi_squared_survival(q, dof);
236 assert!((survival / p - 1.0).abs() <= 1e-10, "{dof} {q}: {survival}");
237 }
238 for dof in 1..=6 {
239 assert_eq!(chi_squared_survival(0.0, dof), 1.0);
240 assert_eq!(chi_squared_survival(2000.0, dof), 0.0);
241 assert_eq!(chi_squared_survival(f64::INFINITY, dof), 0.0);
242 }
243 }
244
245 fn two_by_two(data: [f64; 4]) -> FlatMatrix {
246 FlatMatrix {
247 data: data.to_vec(),
248 nrows: 2,
249 ncols: 2,
250 }
251 }
252
253 fn rotated(shares: [f64; 2], angle: f64) -> [[f64; 2]; 2] {
254 let (s, c) = angle.sin_cos();
255 let u = [[c, -s], [s, c]];
256 let mut m = [[0.0; 2]; 2];
257 for (i, row) in m.iter_mut().enumerate() {
258 for (j, value) in row.iter_mut().enumerate() {
259 *value = (0..2).map(|d| u[i][d] * shares[d] * u[j][d]).sum();
260 }
261 }
262 m
263 }
264
265 #[test]
266 fn the_statistic_weighs_each_direction_by_the_variance_the_counts_remove() {
267 let factor = [[2.0, 0.0], [1.0, 2.0]];
268 let covariance = FlatMatrix {
269 data: vec![4.0, 2.0, 2.0, 5.0],
270 nrows: 2,
271 ncols: 2,
272 };
273 let prior = Prior::correlated(&[0, 1], &[3.0, -1.0], &covariance).unwrap();
274 let angle: f64 = 0.5;
275 let (s, c) = angle.sin_cos();
276 let z = [1.0, 3.0];
277 let estimate: Vec<f64> = (0..2)
278 .map(|i| [3.0, -1.0][i] + (0..2).map(|j| factor[i][j] * z[j]).sum::<f64>())
279 .collect();
280 let posterior = |shares: [f64; 2]| {
281 let m = rotated(shares, angle);
282 let mut data = vec![0.0; 4];
283 for i in 0..2 {
284 for j in 0..2 {
285 data[i * 2 + j] = (0..2)
286 .flat_map(|a| (0..2).map(move |b| (a, b)))
287 .map(|(a, b)| {
288 factor[i][a] * (f64::from(u8::from(a == b)) - m[a][b]) * factor[j][b]
289 })
290 .sum();
291 }
292 }
293 FlatMatrix {
294 data,
295 nrows: 2,
296 ncols: 2,
297 }
298 };
299 let along_first = c * z[0] + s * z[1];
300 let both = consistency(&prior, &estimate, &posterior([0.5, 0.25]))
301 .unwrap()
302 .unwrap();
303 let along_second = -s * z[0] + c * z[1];
304 let q = along_first.powi(2) / 0.5 + along_second.powi(2) / 0.25;
305 assert_eq!(both.dof, 2);
306 assert!((both.q / q - 1.0).abs() <= 1e-12, "{} vs {q}", both.q);
307 assert!((both.p / (-q / 2.0).exp() - 1.0).abs() <= 1e-12);
308 let one = consistency(&prior, &estimate, &posterior([0.5, 1e-4]))
309 .unwrap()
310 .unwrap();
311 assert_eq!(one.dof, 1);
312 let q = along_first.powi(2) / 0.5;
313 assert!((one.q / q - 1.0).abs() <= 1e-12, "{} vs {q}", one.q);
314 assert!(
315 consistency(&prior, &estimate, &posterior([1e-4, 1e-4]))
316 .unwrap()
317 .is_none()
318 );
319 let mut undetermined = posterior([0.5, 0.25]);
320 undetermined.data[3] = f64::NAN;
321 assert!(
322 consistency(&prior, &estimate, &undetermined)
323 .unwrap()
324 .is_none()
325 );
326 assert!(consistency(&prior, &estimate[..1], &posterior([0.5, 0.25])).is_err());
327 let mut truncated = posterior([0.5, 0.25]);
328 truncated.data.pop();
329 assert!(consistency(&prior, &estimate, &truncated).is_err());
330 let nan = [estimate[0], f64::NAN];
331 assert!(consistency(&prior, &nan, &posterior([0.5, 0.25])).is_err());
332 assert!(consistency(&prior, &nan, &undetermined).unwrap().is_none());
333 let negative = Prior::measured(0, 0.0, -1.0);
334 let unit = FlatMatrix {
335 data: vec![0.5],
336 nrows: 1,
337 ncols: 1,
338 };
339 assert!(consistency(&negative, &[1.0], &unit).is_err());
340 let unknown = FlatMatrix {
341 data: vec![f64::NAN],
342 ..unit
343 };
344 assert!(consistency(&negative, &[1.0], &unknown).is_err());
345 let identity =
346 Prior::correlated(&[0, 1], &[0.0, 0.0], &two_by_two([1.0, 0.0, 0.0, 1.0])).unwrap();
347 for impossible in [
348 [0.5, 5.0, -5.0, 0.5],
349 [1.5, 0.0, 0.0, 0.5],
350 [-1.0, 0.0, 0.0, 0.5],
351 ] {
352 let result = consistency(&identity, &[2.0, 1.0], &two_by_two(impossible));
353 assert!(result.is_err(), "{impossible:?}: {result:?}");
354 }
355 let far = consistency(
356 &identity,
357 &[1e155, 1e155],
358 &two_by_two([0.5, 0.0, 0.0, 0.5]),
359 )
360 .unwrap()
361 .unwrap();
362 assert_eq!((far.q, far.p), (f64::INFINITY, 0.0));
363 }
364
365 #[test]
366 fn two_estimates_agree_by_their_difference_over_the_sum_of_their_covariances() {
367 let a = Prior::correlated(&[3, 7], &[1.0, 2.0], &two_by_two([4.0, 1.0, 1.0, 2.0])).unwrap();
368 let b =
369 Prior::correlated(&[3, 7], &[2.0, 0.5], &two_by_two([1.0, -0.5, -0.5, 3.0])).unwrap();
370 let (sum, d) = ([5.0, 0.5, 0.5, 5.0], [-1.0, 1.5]);
371 let determinant = sum[0] * sum[3] - sum[1] * sum[2];
372 let q = (d[0] * d[0] * sum[3] - 2.0 * d[0] * d[1] * sum[1] + d[1] * d[1] * sum[0])
373 / determinant;
374 let result = agreement(&a, &b).unwrap();
375 assert_eq!(result.dof, 2);
376 assert!((result.q / q - 1.0).abs() <= 1e-14, "{} vs {q}", result.q);
377 assert!((result.p / (-q / 2.0).exp() - 1.0).abs() <= 1e-14);
378 let other =
379 Prior::correlated(&[3, 8], &[2.0, 0.5], &two_by_two([1.0, 0.0, 0.0, 3.0])).unwrap();
380 assert!(agreement(&a, &other).is_err());
381 let invalid = Prior::measured(3, 0.0, -1.0);
382 let one = Prior::measured(3, 0.0, 1.0);
383 assert!(agreement(&invalid, &one).is_err());
384 let far =
385 |mean| Prior::correlated(&[3, 7], &[mean, 0.0], &two_by_two([1.0, 0.0, 0.0, 1.0]));
386 assert!(agreement(&far(f64::MAX).unwrap(), &far(-f64::MAX).unwrap()).is_err());
387 for (q, dof) in [(-1.0, 2), (f64::NAN, 2), (1.0, 0)] {
388 assert!(Consistency::new(q, dof).is_err(), "{q} {dof}");
389 }
390 }
391}