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}
128
129impl From<Array2<f64>> for PhiScaledCovariance {
130    #[inline]
131    fn from(cov: Array2<f64>) -> Self {
132        Self(cov)
133    }
134}
135
136impl From<PhiScaledCovariance> for Array2<f64> {
137    #[inline]
138    fn from(cov: PhiScaledCovariance) -> Self {
139        cov.0
140    }
141}
142
143impl Deref for PhiScaledCovariance {
144    type Target = Array2<f64>;
145    #[inline]
146    fn deref(&self) -> &Array2<f64> {
147        &self.0
148    }
149}
150
151impl DerefMut for PhiScaledCovariance {
152    #[inline]
153    fn deref_mut(&mut self) -> &mut Array2<f64> {
154        &mut self.0
155    }
156}
157
158/// Raw penalised Hessian `H = X' W_H X + S(lambda)` with NO dispersion
159/// scaling. Equivalent to `phi * Vb^{-1}` only when `phi == 1`. Use this
160/// for whitening / precision-matrix paths, and pair it with a
161/// [`Dispersion`] at the boundary if the consumer cares about `phi`.
162#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
163#[serde(transparent)]
164pub struct UnscaledPrecision(pub Array2<f64>);
165
166impl UnscaledPrecision {
167    /// Wrap an `Array2` that is already on the unscaled
168    /// `H = XᵀW_H X + S(λ)` scale (no `φ` factor).  Caller is responsible
169    /// for ensuring the matrix actually represents the penalised Hessian.
170    #[inline]
171    pub fn wrap(hessian: Array2<f64>) -> Self {
172        Self(hessian)
173    }
174
175    /// Borrow the underlying penalised Hessian `H` without taking ownership.
176    #[inline]
177    pub fn as_array(&self) -> &Array2<f64> {
178        &self.0
179    }
180
181}
182
183impl From<Array2<f64>> for UnscaledPrecision {
184    #[inline]
185    fn from(h: Array2<f64>) -> Self {
186        Self(h)
187    }
188}
189
190impl From<UnscaledPrecision> for Array2<f64> {
191    #[inline]
192    fn from(h: UnscaledPrecision) -> Self {
193        h.0
194    }
195}
196
197impl Deref for UnscaledPrecision {
198    type Target = Array2<f64>;
199    #[inline]
200    fn deref(&self) -> &Array2<f64> {
201        &self.0
202    }
203}
204
205impl DerefMut for UnscaledPrecision {
206    #[inline]
207    fn deref_mut(&mut self) -> &mut Array2<f64> {
208        &mut self.0
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215    use ndarray::array;
216
217    // ── se_from_covariance ────────────────────────────────────────────────────
218
219    #[test]
220    fn se_from_diagonal_matrix_is_sqrt_of_diagonal() {
221        // cov = diag(4, 9) → se = [2, 3]
222        let cov = array![[4.0_f64, 0.0], [0.0, 9.0]];
223        let se = se_from_covariance(&cov).unwrap();
224        assert_eq!(se.len(), 2);
225        assert!((se[0] - 2.0).abs() < 1e-14);
226        assert!((se[1] - 3.0).abs() < 1e-14);
227    }
228
229    #[test]
230    fn se_snaps_only_backward_error_scale_negative_diagonal() {
231        let cov = array![[1.0_f64, 0.0], [0.0, -4.0 * f64::EPSILON]];
232        let se = se_from_covariance(&cov).unwrap();
233        assert_eq!(se[1], 0.0);
234        let materially_indefinite = array![[1.0_f64, 0.0], [0.0, -1e-8]];
235        assert!(matches!(
236            se_from_covariance(&materially_indefinite),
237            Err(CovarianceStandardErrorError::NegativeDiagonal { index: 1, .. })
238        ));
239    }
240
241    // ── PhiScaledCovariance ───────────────────────────────────────────────────
242
243    #[test]
244    fn phi_scaled_covariance_wrap_and_as_array_round_trip() {
245        let m = array![[1.0_f64, 2.0], [3.0, 4.0]];
246        let wrapped = PhiScaledCovariance::wrap(m.clone());
247        assert_eq!(*wrapped.as_array(), m);
248    }
249
250    #[test]
251    fn phi_scaled_covariance_deref_gives_array2() {
252        let m = array![[5.0_f64]];
253        let wrapped = PhiScaledCovariance::wrap(m.clone());
254        assert_eq!(wrapped.nrows(), 1);
255        assert_eq!(wrapped[[0, 0]], 5.0);
256    }
257
258    // ── UnscaledPrecision ─────────────────────────────────────────────────────
259
260    #[test]
261    fn unscaled_precision_wrap_and_as_array_round_trip() {
262        let h = array![[2.0_f64, 0.0], [0.0, 3.0]];
263        let wrapped = UnscaledPrecision::wrap(h.clone());
264        assert_eq!(*wrapped.as_array(), h);
265    }
266}
gam_problem/dispersion_cov.rs

gam_problem/
dispersion_cov.rs
