gam_problem/
dispersion_cov.rs

1//! Newtype wrappers that disambiguate the two coefficient-space second-order
2//! quantities used throughout inference.
3//!
4//! Background ("dispersion ownership"):
5//!
6//! The fitter stores two related matrices for a fitted model.
7//!
8//! * `FitInference::beta_covariance` is the posterior coefficient covariance
9//!   `Vb = phi * H^{-1}`, with `H = X' W_H X + S(lambda)` and `phi` the
10//!   dispersion parameter. This matrix is *already* multiplied by `phi`
11//!   (see `solver/estimate.rs`'s `scaled_covariance` call).
12//! * `FitInference::penalized_hessian` is the raw penalised Hessian `H`,
13//!   with NO dispersion scaling.
14//!
15//! Several downstream consumers (HMC whitening, Laplace sampling, smooth
16//! tests, etc.) have to know which of these representations they hold so
17//! they apply `phi` exactly once. Passing both as bare `Array2<f64>` makes
18//! that easy to get wrong: the same matrix shape can mean either thing,
19//! and the compiler will not catch a missing — or duplicated —
20//! `phi` factor.
21//!
22//! The lightweight newtypes below give us a way to label the convention at
23//! API boundaries without changing the storage type of the existing
24//! `FitInference` fields. Storage stays `Array2<f64>` to avoid cascading
25//! changes into modules outside the dispersion-ownership refactor's scope
26//! (pirls, families, GPU paths, main, etc.); callers that want to be
27//! explicit can wrap with `PhiScaledCovariance::wrap` /
28//! `UnscaledPrecision::wrap` at the boundary.
29//!
30//! `Dispersion` lives in `gam-problem` as the neutral validated scale contract.
31
32use ndarray::{Array1, Array2};
33use serde::{Deserialize, Serialize};
34use std::ops::{Deref, DerefMut};
35use thiserror::Error;
36
37pub use crate::Dispersion;
38
39#[derive(Clone, Copy, Debug, Error, PartialEq)]
40pub enum CovarianceStandardErrorError {
41    #[error("covariance must be square, got {rows}x{cols}")]
42    NotSquare { rows: usize, cols: usize },
43    #[error("covariance entry ({row}, {col}) is non-finite: {value}")]
44    NonFinite { row: usize, col: usize, value: f64 },
45    #[error("covariance dimension {dimension} is too large for a sub-unit rounding-error bound")]
46    DimensionTooLarge { dimension: usize },
47    #[error(
48        "covariance diagonal {index} is materially negative: {value} (backward-error tolerance {tolerance})"
49    )]
50    NegativeDiagonal {
51        index: usize,
52        value: f64,
53        tolerance: f64,
54    },
55}
56
57/// Compute standard errors from a finite square covariance matrix.
58///
59/// A negative diagonal is snapped to zero only when its magnitude is within a
60/// dimension-scaled floating-point backward-error bound relative to the matrix
61/// scale. Material negativity is rejected instead of being projected onto the
62/// PSD cone one diagonal entry at a time.
63pub fn se_from_covariance(cov: &Array2<f64>) -> Result<Array1<f64>, CovarianceStandardErrorError> {
64    let (rows, cols) = cov.dim();
65    if rows != cols {
66        return Err(CovarianceStandardErrorError::NotSquare { rows, cols });
67    }
68    let mut max_abs = 0.0_f64;
69    for ((row, col), &value) in cov.indexed_iter() {
70        if !value.is_finite() {
71            return Err(CovarianceStandardErrorError::NonFinite { row, col, value });
72        }
73        max_abs = max_abs.max(value.abs());
74    }
75    let relative_error = 16.0 * (rows.max(1) as f64) * f64::EPSILON;
76    let tolerance = if relative_error < 1.0 {
77        max_abs / relative_error.recip()
78    } else {
79        return Err(CovarianceStandardErrorError::DimensionTooLarge { dimension: rows });
80    };
81    let mut standard_errors = Array1::zeros(rows);
82    for (index, &value) in cov.diag().iter().enumerate() {
83        standard_errors[index] = if value == 0.0 {
84            0.0
85        } else if value > 0.0 {
86            value.sqrt()
87        } else if -value <= tolerance {
88            0.0
89        } else {
90            return Err(CovarianceStandardErrorError::NegativeDiagonal {
91                index,
92                value,
93                tolerance,
94            });
95        };
96    }
97    Ok(standard_errors)
98}
99
100/// Posterior coefficient covariance `Vb = phi * H^{-1}` — the matrix users
101/// see as `Cov(beta_hat)`. This newtype documents that `phi` has already
102/// been multiplied in.
103///
104/// `#[serde(transparent)]` keeps the on-disk wire format identical to the
105/// pre-newtype `Array2<f64>` storage so saved models round-trip cleanly.
106/// `Deref<Target = Array2<f64>>` lets out-of-scope read sites continue
107/// calling `Array2` methods (`.iter()`, `.nrows()`, `.dim()`, …) on the
108/// wrapper without modification.
109#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
110#[serde(transparent)]
111pub struct PhiScaledCovariance(pub Array2<f64>);
112
113impl PhiScaledCovariance {
114    /// Wrap an array that is known to already be on the `phi * H^{-1}`
115    /// scale.
116    #[inline]
117    pub fn wrap(cov: Array2<f64>) -> Self {
118        Self(cov)
119    }
120
121    /// Borrow the underlying `φ · H⁻¹` matrix without taking ownership.
122    #[inline]
123    pub fn as_array(&self) -> &Array2<f64> {
124        &self.0
125    }
126
127    /// Consume the wrapper and return the raw `φ · H⁻¹` matrix.
128    #[inline]
129    pub fn into_array(self) -> Array2<f64> {
130        self.0
131    }
132}
133
134impl From<Array2<f64>> for PhiScaledCovariance {
135    #[inline]
136    fn from(cov: Array2<f64>) -> Self {
137        Self(cov)
138    }
139}
140
141impl From<PhiScaledCovariance> for Array2<f64> {
142    #[inline]
143    fn from(cov: PhiScaledCovariance) -> Self {
144        cov.0
145    }
146}
147
148impl Deref for PhiScaledCovariance {
149    type Target = Array2<f64>;
150    #[inline]
151    fn deref(&self) -> &Array2<f64> {
152        &self.0
153    }
154}
155
156impl DerefMut for PhiScaledCovariance {
157    #[inline]
158    fn deref_mut(&mut self) -> &mut Array2<f64> {
159        &mut self.0
160    }
161}
162
163/// Raw penalised Hessian `H = X' W_H X + S(lambda)` with NO dispersion
164/// scaling. Equivalent to `phi * Vb^{-1}` only when `phi == 1`. Use this
165/// for whitening / precision-matrix paths, and pair it with a
166/// [`Dispersion`] at the boundary if the consumer cares about `phi`.
167#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
168#[serde(transparent)]
169pub struct UnscaledPrecision(pub Array2<f64>);
170
171impl UnscaledPrecision {
172    /// Wrap an `Array2` that is already on the unscaled
173    /// `H = XᵀW_H X + S(λ)` scale (no `φ` factor).  Caller is responsible
174    /// for ensuring the matrix actually represents the penalised Hessian.
175    #[inline]
176    pub fn wrap(hessian: Array2<f64>) -> Self {
177        Self(hessian)
178    }
179
180    /// Borrow the underlying penalised Hessian `H` without taking ownership.
181    #[inline]
182    pub fn as_array(&self) -> &Array2<f64> {
183        &self.0
184    }
185
186    /// Consume the wrapper and return the raw `H` matrix.
187    #[inline]
188    pub fn into_array(self) -> Array2<f64> {
189        self.0
190    }
191}
192
193impl From<Array2<f64>> for UnscaledPrecision {
194    #[inline]
195    fn from(h: Array2<f64>) -> Self {
196        Self(h)
197    }
198}
199
200impl From<UnscaledPrecision> for Array2<f64> {
201    #[inline]
202    fn from(h: UnscaledPrecision) -> Self {
203        h.0
204    }
205}
206
207impl Deref for UnscaledPrecision {
208    type Target = Array2<f64>;
209    #[inline]
210    fn deref(&self) -> &Array2<f64> {
211        &self.0
212    }
213}
214
215impl DerefMut for UnscaledPrecision {
216    #[inline]
217    fn deref_mut(&mut self) -> &mut Array2<f64> {
218        &mut self.0
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225    use ndarray::array;
226
227    // ── se_from_covariance ────────────────────────────────────────────────────
228
229    #[test]
230    fn se_from_diagonal_matrix_is_sqrt_of_diagonal() {
231        // cov = diag(4, 9) → se = [2, 3]
232        let cov = array![[4.0_f64, 0.0], [0.0, 9.0]];
233        let se = se_from_covariance(&cov).unwrap();
234        assert_eq!(se.len(), 2);
235        assert!((se[0] - 2.0).abs() < 1e-14);
236        assert!((se[1] - 3.0).abs() < 1e-14);
237    }
238
239    #[test]
240    fn se_snaps_only_backward_error_scale_negative_diagonal() {
241        let cov = array![[1.0_f64, 0.0], [0.0, -4.0 * f64::EPSILON]];
242        let se = se_from_covariance(&cov).unwrap();
243        assert_eq!(se[1], 0.0);
244        let materially_indefinite = array![[1.0_f64, 0.0], [0.0, -1e-8]];
245        assert!(matches!(
246            se_from_covariance(&materially_indefinite),
247            Err(CovarianceStandardErrorError::NegativeDiagonal { index: 1, .. })
248        ));
249    }
250
251    // ── PhiScaledCovariance ───────────────────────────────────────────────────
252
253    #[test]
254    fn phi_scaled_covariance_wrap_and_as_array_round_trip() {
255        let m = array![[1.0_f64, 2.0], [3.0, 4.0]];
256        let wrapped = PhiScaledCovariance::wrap(m.clone());
257        assert_eq!(*wrapped.as_array(), m);
258    }
259
260    #[test]
261    fn phi_scaled_covariance_deref_gives_array2() {
262        let m = array![[5.0_f64]];
263        let wrapped = PhiScaledCovariance::wrap(m.clone());
264        assert_eq!(wrapped.nrows(), 1);
265        assert_eq!(wrapped[[0, 0]], 5.0);
266    }
267
268    #[test]
269    fn phi_scaled_covariance_into_array_consumes() {
270        let m = array![[7.0_f64]];
271        let wrapped = PhiScaledCovariance::wrap(m.clone());
272        assert_eq!(wrapped.into_array(), m);
273    }
274
275    // ── UnscaledPrecision ─────────────────────────────────────────────────────
276
277    #[test]
278    fn unscaled_precision_wrap_and_as_array_round_trip() {
279        let h = array![[2.0_f64, 0.0], [0.0, 3.0]];
280        let wrapped = UnscaledPrecision::wrap(h.clone());
281        assert_eq!(*wrapped.as_array(), h);
282    }
283}
gam_problem/dispersion_cov.rs

gam_problem/
dispersion_cov.rs
