Skip to main content

stats_claw/distributions/positive/
gamma.rs

1//! Gamma distribution numerics, for the [`GammaDistribution`] parameter struct.
2//!
3//! Equivalent to `scipy.stats.gamma(shape_parameter, scale = scale_parameter)`.
4//! The density is closed form; the CDF is the regularized lower incomplete gamma
5//! `gamma_p`; the quantile inverts the CDF by bracketed bisection; and sampling
6//! uses the Marsaglia–Tsang method (with the shape<1 boosting trick).
7//!
8//! # Examples
9//!
10//! ```
11//! use stats_claw::distributions::{Cdf, Moments};
12//! use stats_claw::distributions::GammaDistribution;
13//!
14//! let d = GammaDistribution { shape_parameter: 2.0, scale_parameter: 3.0, ..Default::default() };
15//! // Mean is shape * scale = 6.
16//! assert_eq!(d.mean(), Some(6.0));
17//! // With shape=1 the Gamma is exponential; cdf(1) with rate 0.5 is 1 - e^{-0.5}.
18//! let exp = GammaDistribution { shape_parameter: 1.0, scale_parameter: 2.0, ..Default::default() };
19//! assert!((exp.cdf(2.0) - (1.0 - (-1.0f64).exp())).abs() < 1e-10);
20//! ```
21
22use super::super::{Cdf, Moments, Pdf, Quantile, Sample, bisection_quantile};
23use crate::distributions::GammaDistribution;
24use crate::rng::SplitMix64;
25use crate::special::{gamma_p, ln_gamma};
26
27impl Pdf for GammaDistribution {
28    fn pdf(&self, x: f64) -> f64 {
29        if x < 0.0 {
30            return 0.0;
31        }
32        let (k, theta) = (self.shape_parameter, self.scale_parameter);
33        // log-space for stability: (k-1)·ln x − x/θ − k·ln θ − ln Γ(k)
34        let log_norm = k.mul_add(theta.ln(), ln_gamma(k));
35        let ln_density = (k - 1.0).mul_add(x.ln(), -(x / theta) - log_norm);
36        ln_density.exp()
37    }
38}
39
40impl Cdf for GammaDistribution {
41    fn cdf(&self, x: f64) -> f64 {
42        if x <= 0.0 {
43            0.0
44        } else {
45            gamma_p(self.shape_parameter, x / self.scale_parameter)
46        }
47    }
48}
49
50impl Quantile for GammaDistribution {
51    fn quantile(&self, p: f64) -> f64 {
52        let mean = self.shape_parameter * self.scale_parameter;
53        let sd = self.scale_parameter * self.shape_parameter.sqrt();
54        let hi = 20.0f64.mul_add(sd, mean).max(self.scale_parameter);
55        bisection_quantile(p, 0.0, hi, |x| self.cdf(x))
56    }
57}
58
59impl Moments for GammaDistribution {
60    fn mean(&self) -> Option<f64> {
61        Some(self.shape_parameter * self.scale_parameter)
62    }
63    fn variance(&self) -> Option<f64> {
64        Some(self.shape_parameter * self.scale_parameter * self.scale_parameter)
65    }
66}
67
68impl Sample for GammaDistribution {
69    fn sample(&self, rng: &mut SplitMix64) -> f64 {
70        self.scale_parameter * marsaglia_tsang(self.shape_parameter, rng)
71    }
72}
73
74/// Draws a standard Gamma(`shape`, 1) variate via Marsaglia–Tsang (2000).
75///
76/// For `shape ≥ 1` it uses the squeeze acceptance directly; for `shape < 1` it
77/// draws at `shape + 1` and rescales by `U^(1/shape)`, which is exact. Shared
78/// with the Beta sampler (`X/(X+Y)` of two unit-scale gammas).
79///
80/// # Arguments
81///
82/// * `shape` — the gamma shape parameter (`> 0`).
83/// * `rng` — the deterministic generator to draw from.
84///
85/// # Returns
86///
87/// One Gamma(`shape`, 1) variate.
88pub(super) fn marsaglia_tsang(shape: f64, rng: &mut SplitMix64) -> f64 {
89    if shape < 1.0 {
90        let boosted = marsaglia_tsang(shape + 1.0, rng);
91        return boosted * rng.next_f64().powf(1.0 / shape);
92    }
93    let depth = shape - 1.0 / 3.0;
94    let coeff = 1.0 / (9.0 * depth).sqrt();
95    loop {
96        let normal = rng.standard_normal();
97        let base = coeff.mul_add(normal, 1.0);
98        if base <= 0.0 {
99            continue;
100        }
101        let cubed = base * base * base;
102        let uniform = rng.next_f64();
103        let normal_sq = normal * normal;
104        // Marsaglia–Tsang acceptance in log space.
105        let accept = 0.5f64.mul_add(normal_sq, depth.mul_add(-cubed, depth) + depth * cubed.ln());
106        if uniform.ln() < accept {
107            return depth * cubed;
108        }
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    /// With shape 1 the Gamma reduces to an exponential with rate `1/scale`.
117    #[test]
118    fn shape_one_cdf_is_exponential() {
119        let d = GammaDistribution {
120            shape_parameter: 1.0,
121            scale_parameter: 2.0,
122            ..Default::default()
123        };
124        assert!((d.cdf(2.0) - (1.0 - (-1.0f64).exp())).abs() < 1e-10);
125    }
126
127    /// The mean is `shape · scale`.
128    #[test]
129    fn mean_is_shape_times_scale() {
130        let d = GammaDistribution {
131            shape_parameter: 2.5,
132            scale_parameter: 1.5,
133            ..Default::default()
134        };
135        assert_eq!(d.mean(), Some(3.75));
136    }
137}