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    /// Rescale an estimated dispersion in place.
164    ///
165    /// Fixed dispersions are left unchanged and return `Ok(false)`. The product
166    /// is checked before mutation so this operation is atomic on error.
167    pub fn rescale_estimate(&mut self, multiplier: f64) -> Result<bool, DispersionError> {
168        if !(multiplier.is_finite() && multiplier > 0.0) {
169            return Err(DispersionError::InvalidMultiplier { multiplier });
170        }
171        if !self.is_estimated() {
172            return Ok(false);
173        }
174        let phi = self.phi * multiplier;
175        if !phi.is_finite() || (self.phi > 0.0 && phi == 0.0) {
176            return Err(DispersionError::RescaleNotRepresentable {
177                phi: self.phi,
178                multiplier,
179            });
180        }
181        self.phi = phi;
182        Ok(true)
183    }
184}
185
186impl From<Dispersion> for DispersionWire {
187    fn from(value: Dispersion) -> Self {
188        Self {
189            source: value.source,
190            phi: value.phi,
191        }
192    }
193}
194
195impl TryFrom<DispersionWire> for Dispersion {
196    type Error = DispersionError;
197
198    fn try_from(value: DispersionWire) -> Result<Self, Self::Error> {
199        match value.source {
200            DispersionSource::Known => Self::known(value.phi),
201            DispersionSource::Estimated => Self::estimated(value.phi),
202        }
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209
210    #[test]
211    fn constructors_enforce_distinct_domains() {
212        assert_eq!(Dispersion::known(2.5).unwrap().phi(), 2.5);
213        assert_eq!(
214            Dispersion::estimated(0.0).unwrap(),
215            Dispersion::ZERO_ESTIMATE
216        );
217        assert!(Dispersion::known(0.0).is_err());
218        assert_eq!(
219            Dispersion::estimated(-0.0).unwrap().phi().to_bits(),
220            0.0_f64.to_bits()
221        );
222        assert!(Dispersion::estimated(-1.0).is_err());
223        assert!(Dispersion::known(f64::NAN).is_err());
224    }
225
226    #[test]
227    fn reciprocal_is_exactly_fallible_at_the_boundary() {
228        assert_eq!(Dispersion::known(4.0).unwrap().reciprocal(), Ok(0.25));
229        assert_eq!(
230            Dispersion::ZERO_ESTIMATE.reciprocal(),
231            Err(DispersionError::ZeroHasNoReciprocal)
232        );
233    }
234
235    #[test]
236    fn sqrt_preserves_zero_boundary() {
237        assert_eq!(Dispersion::ZERO_ESTIMATE.sqrt(), 0.0);
238        assert_eq!(Dispersion::estimated(9.0).unwrap().sqrt(), 3.0);
239    }
240
241    #[test]
242    fn source_is_part_of_identity() {
243        assert_ne!(
244            Dispersion::known(1.0).unwrap(),
245            Dispersion::estimated(1.0).unwrap()
246        );
247        assert!(!Dispersion::UNIT.is_estimated());
248    }
249}