Skip to main content

nereids_fitting/
statistics.rs

1//! Consistency of a fit with its prior, and agreement between two
2//! estimates.
3
4use 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/// A χ² test: the statistic `q`, its degrees of freedom `dof`, and the
15/// probability `p` of a `q` at least as large under the hypothesis tested.
16#[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    /// The test of `q` against `χ²` with `dof` degrees of freedom.
25    ///
26    /// # Errors
27    /// `FittingError::InvalidConfig` if `q` is negative or not a number, or
28    /// `dof` is 0.
29    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
65/// The consistency of the fitted values `estimate` of a `prior`'s
66/// parameters with it, in the prior's order, given the fit's covariance
67/// `posterior` of them.
68///
69/// With `C = LLᵀ` the prior covariance, `z = L⁻¹(estimate − mean)` and
70/// `I − L⁻¹·posterior·L⁻ᵀ = Σᵢ rᵢuᵢuᵢᵀ`, `rᵢ` is the share of the prior's
71/// variance along `uᵢ` that the counts remove, and `q = Σᵢ (uᵢᵀz)²/rᵢ` over
72/// the `dof` directions with `rᵢ ≥ 8e-4`.  If the prior and the counts' model
73/// are right and `estimate` and `posterior` are the fit's without its bounds
74/// ([`Unbounded`](crate::poisson::Unbounded)), `q` follows `χ²` with `dof`
75/// degrees of freedom, and `p = P(χ²_dof ≥ q)`.
76///
77/// `None` when `posterior` is not finite or no direction has `rᵢ ≥ 8e-4`.
78///
79/// # Errors
80/// `FittingError::LengthMismatch` if `estimate` or `posterior` does not
81/// match the prior's parameters; `FittingError::InvalidConfig` if the prior's
82/// mean is not finite or a measured sd not finite and positive, `estimate` is
83/// not finite while `posterior` is, or `posterior` is not symmetric to 1e-12
84/// of `√(Σᵢᵢ Σⱼⱼ)` or not within the prior, with some `rᵢ` below `−8e-4` or
85/// above `1 + 8e-4`, or `q` is not a number, as overflow in whitening can
86/// give; `FittingError::EvaluationFailed` if a decomposition fails.
87pub 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
174/// The agreement of two independent estimates `a` and `b` of the same
175/// parameters, each a mean with its covariance:
176/// `q = dᵀ(C_a + C_b)⁻¹d`, with `d` the difference of the means, follows `χ²`
177/// with as many degrees of freedom as parameters when both estimate the same
178/// values.
179///
180/// # Errors
181/// `FittingError::InvalidConfig` if `a` and `b` are over different
182/// parameters, either has a mean that is not finite or a measured sd not
183/// finite and positive, or `q` is not a number, as overflow in whitening can
184/// give; `FittingError::EvaluationFailed` if `C_a + C_b` is
185/// not positive definite in floating point.
186pub 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}