Skip to main content

gam_problem/
dispersion.rs

1//! Validated dispersion/scale contract used by covariance and sampling code.
2
3use serde::{Deserialize, Serialize};
4use thiserror::Error;
5
6/// Why a dispersion value is part of a fitted model.
7///
8/// This tag is deliberately private: callers may inspect it through
9/// [`Dispersion::is_estimated`], but cannot construct an unchecked value by
10/// naming an enum variant.
11#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
12enum DispersionSource {
13    Known,
14    Estimated,
15}
16
17/// Serde representation for [`Dispersion`].
18///
19/// Deserialization goes through `TryFrom`, so persisted bytes cannot bypass the
20/// same numerical invariant as an in-process constructor.
21#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
22struct DispersionWire {
23    source: DispersionSource,
24    phi: f64,
25}
26
27/// A validated response-level dispersion `phi`.
28///
29/// A known/fixed dispersion is finite and strictly positive. An estimated
30/// dispersion is finite and non-negative: zero is a meaningful boundary result
31/// for an exact fit, but any operation that divides by it must explicitly fail.
32/// The distinction prevents a zero estimate from being silently promoted to a
33/// tiny positive fixed scale.
34#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
35#[serde(try_from = "DispersionWire", into = "DispersionWire")]
36pub struct Dispersion {
37    source: DispersionSource,
38    phi: f64,
39}
40
41#[derive(Clone, Copy, Debug, Error, PartialEq)]
42pub enum DispersionError {
43    #[error("dispersion phi must be finite, got {phi}")]
44    NonFinite { phi: f64 },
45    #[error("a known dispersion must be strictly positive, got {phi}")]
46    NonPositiveKnown { phi: f64 },
47    #[error("an estimated dispersion must be non-negative, got {phi}")]
48    NegativeEstimate { phi: f64 },
49    #[error("zero estimated dispersion has no finite reciprocal")]
50    ZeroHasNoReciprocal,
51    #[error("the reciprocal of dispersion phi={phi} is not representable as a finite f64")]
52    ReciprocalNotRepresentable { phi: f64 },
53    #[error("dispersion multiplier must be finite and strictly positive, got {multiplier}")]
54    InvalidMultiplier { multiplier: f64 },
55    #[error(
56        "rescaling dispersion phi={phi} by multiplier={multiplier} is not representable as a finite f64"
57    )]
58    RescaleNotRepresentable { phi: f64, multiplier: f64 },
59}
60
61impl Dispersion {
62    /// Exact unit fixed dispersion for fixed-scale likelihoods.
63    pub const UNIT: Self = Self {
64        source: DispersionSource::Known,
65        phi: 1.0,
66    };
67
68    /// Exact boundary estimate produced by a zero-residual fit.
69    pub const ZERO_ESTIMATE: Self = Self {
70        source: DispersionSource::Estimated,
71        phi: 0.0,
72    };
73
74    /// Construct a finite, strictly-positive fixed dispersion.
75    #[inline]
76    pub fn known(phi: f64) -> Result<Self, DispersionError> {
77        if !phi.is_finite() {
78            return Err(DispersionError::NonFinite { phi });
79        }
80        if phi <= 0.0 {
81            return Err(DispersionError::NonPositiveKnown { phi });
82        }
83        Ok(Self {
84            source: DispersionSource::Known,
85            phi,
86        })
87    }
88
89    /// Construct a finite, non-negative estimated dispersion.
90    #[inline]
91    pub fn estimated(phi: f64) -> Result<Self, DispersionError> {
92        if !phi.is_finite() {
93            return Err(DispersionError::NonFinite { phi });
94        }
95        if phi < 0.0 {
96            return Err(DispersionError::NegativeEstimate { phi });
97        }
98        let phi = if phi == 0.0 { 0.0 } else { phi };
99        Ok(Self {
100            source: DispersionSource::Estimated,
101            phi,
102        })
103    }
104
105    /// Construct a dispersion from a finite positive precision/shape.
106    ///
107    /// This checks the division itself; a positive subnormal denominator whose
108    /// reciprocal overflows is rejected instead of being floored.
109    #[inline]
110    pub fn from_reciprocal(value: f64, estimated: bool) -> Result<Self, DispersionError> {
111        if !value.is_finite() {
112            return Err(DispersionError::NonFinite { phi: value });
113        }
114        if value <= 0.0 {
115            return Err(DispersionError::NonPositiveKnown { phi: value });
116        }
117        let phi = 1.0 / value;
118        if !phi.is_finite() || phi == 0.0 {
119            return Err(DispersionError::ReciprocalNotRepresentable { phi: value });
120        }
121        if estimated {
122            Self::estimated(phi)
123        } else {
124            Self::known(phi)
125        }
126    }
127
128    #[inline]
129    pub const fn phi(self) -> f64 {
130        self.phi
131    }
132
133    #[inline]
134    pub const fn is_estimated(self) -> bool {
135        matches!(self.source, DispersionSource::Estimated)
136    }
137
138    /// Whether this is the exact, validated boundary estimate `phi = 0`.
139    #[inline]
140    pub const fn is_zero_estimate(self) -> bool {
141        self.is_estimated() && self.phi == 0.0
142    }
143
144    /// Return `1 / phi`, rejecting the exact boundary and overflow.
145    #[inline]
146    pub fn reciprocal(self) -> Result<f64, DispersionError> {
147        if self.phi == 0.0 {
148            return Err(DispersionError::ZeroHasNoReciprocal);
149        }
150        let reciprocal = 1.0 / self.phi;
151        if !reciprocal.is_finite() {
152            return Err(DispersionError::ReciprocalNotRepresentable { phi: self.phi });
153        }
154        Ok(reciprocal)
155    }
156
157    /// Return `sqrt(phi)`. Zero is represented exactly.
158    #[inline]
159    pub fn sqrt(self) -> f64 {
160        self.phi.sqrt()
161    }
162
163}
164
165impl From<Dispersion> for DispersionWire {
166    fn from(value: Dispersion) -> Self {
167        Self {
168            source: value.source,
169            phi: value.phi,
170        }
171    }
172}
173
174impl TryFrom<DispersionWire> for Dispersion {
175    type Error = DispersionError;
176
177    fn try_from(value: DispersionWire) -> Result<Self, Self::Error> {
178        match value.source {
179            DispersionSource::Known => Self::known(value.phi),
180            DispersionSource::Estimated => Self::estimated(value.phi),
181        }
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn constructors_enforce_distinct_domains() {
191        assert_eq!(Dispersion::known(2.5).unwrap().phi(), 2.5);
192        assert_eq!(
193            Dispersion::estimated(0.0).unwrap(),
194            Dispersion::ZERO_ESTIMATE
195        );
196        assert!(Dispersion::known(0.0).is_err());
197        assert_eq!(
198            Dispersion::estimated(-0.0).unwrap().phi().to_bits(),
199            0.0_f64.to_bits()
200        );
201        assert!(Dispersion::estimated(-1.0).is_err());
202        assert!(Dispersion::known(f64::NAN).is_err());
203    }
204
205    #[test]
206    fn reciprocal_is_exactly_fallible_at_the_boundary() {
207        assert_eq!(Dispersion::known(4.0).unwrap().reciprocal(), Ok(0.25));
208        assert_eq!(
209            Dispersion::ZERO_ESTIMATE.reciprocal(),
210            Err(DispersionError::ZeroHasNoReciprocal)
211        );
212    }
213
214    #[test]
215    fn sqrt_preserves_zero_boundary() {
216        assert_eq!(Dispersion::ZERO_ESTIMATE.sqrt(), 0.0);
217        assert_eq!(Dispersion::estimated(9.0).unwrap().sqrt(), 3.0);
218    }
219
220    #[test]
221    fn source_is_part_of_identity() {
222        assert_ne!(
223            Dispersion::known(1.0).unwrap(),
224            Dispersion::estimated(1.0).unwrap()
225        );
226        assert!(!Dispersion::UNIT.is_estimated());
227    }
228}