Skip to main content

stats_claw/distributions/discrete/
poisson.rs

1//! Poisson distribution numerics, for the
2//! [`PoissonDistribution`].
3//!
4//! Equivalent to `scipy.stats.poisson(rate_parameter)`. The PMF uses `ln_gamma`
5//! for the factorial, the CDF is the regularized upper incomplete gamma
6//! `gamma_q(k+1, λ)`, the quantile is the smallest `k` with `cdf(k) ≥ p`, and
7//! sampling uses Knuth's multiplicative method.
8//!
9//! # Examples
10//!
11//! ```
12//! use stats_claw::distributions::{Moments, Pmf};
13//! use stats_claw::distributions::PoissonDistribution;
14//!
15//! let d = PoissonDistribution { rate_parameter: 3.0, ..Default::default() };
16//! // Mean and variance both equal the rate.
17//! assert_eq!(d.mean(), Some(3.0));
18//! assert_eq!(d.variance(), Some(3.0));
19//! // PMF is positive at a support point.
20//! assert!(d.pmf(3) > 0.0, "pmf at mode was {}", d.pmf(3));
21//! ```
22
23use super::super::{Cdf, Moments, Pmf, Quantile, Sample, count_to_f64};
24use crate::distributions::PoissonDistribution;
25use crate::rng::SplitMix64;
26use crate::special::{gamma_q, ln_gamma};
27
28/// Hard cap on the support scan for the quantile: far beyond the tail of any
29/// realistic Poisson rate, so the loop terminates without a float comparison.
30const MAX_SUPPORT: i64 = 1_000_000;
31
32impl Pmf for PoissonDistribution {
33    fn pmf(&self, k: i64) -> f64 {
34        if k < 0 {
35            return 0.0;
36        }
37        let lambda = self.rate_parameter;
38        let kf = count_to_f64(k);
39        // log pmf = k·ln λ − λ − ln Γ(k+1)
40        let log_pmf = kf.mul_add(lambda.ln(), -lambda) - ln_gamma(kf + 1.0);
41        log_pmf.exp()
42    }
43}
44
45impl Cdf for PoissonDistribution {
46    fn cdf(&self, x: f64) -> f64 {
47        if x < 0.0 {
48            return 0.0;
49        }
50        let k = x.floor();
51        // P(X ≤ k) = Q(k+1, λ) (regularized upper incomplete gamma).
52        gamma_q(k + 1.0, self.rate_parameter)
53    }
54}
55
56impl Quantile for PoissonDistribution {
57    fn quantile(&self, p: f64) -> f64 {
58        if p <= 0.0 {
59            return 0.0;
60        }
61        // Scan the support in order, returning the first `k` whose cumulative
62        // mass reaches `p`. The hard cap guards against a runaway loop for
63        // pathological inputs; the Poisson tail is exhausted long before it for
64        // any realistic rate.
65        let mut cumulative = 0.0;
66        let mut last = 0i64;
67        for k in 0..=MAX_SUPPORT {
68            last = k;
69            cumulative += self.pmf(k);
70            if cumulative >= p - 1e-12 {
71                return count_to_f64(k);
72            }
73        }
74        count_to_f64(last)
75    }
76}
77
78impl Moments for PoissonDistribution {
79    fn mean(&self) -> Option<f64> {
80        Some(self.rate_parameter)
81    }
82    fn variance(&self) -> Option<f64> {
83        Some(self.rate_parameter)
84    }
85}
86
87impl Sample for PoissonDistribution {
88    fn sample(&self, rng: &mut SplitMix64) -> f64 {
89        let threshold = (-self.rate_parameter).exp();
90        let mut product = rng.next_f64();
91        let mut count = 0.0;
92        // Knuth's method: multiply uniforms until the product drops below e^{−λ};
93        // a `loop` keeps the float comparison out of a `while` head.
94        loop {
95            if product <= threshold {
96                return count;
97            }
98            product *= rng.next_f64();
99            count += 1.0;
100        }
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    /// Mean and variance both equal the rate.
109    #[test]
110    fn mean_and_variance_equal_rate() {
111        let d = PoissonDistribution {
112            rate_parameter: 4.0,
113            ..Default::default()
114        };
115        assert_eq!(d.mean(), Some(4.0));
116        assert_eq!(d.variance(), Some(4.0));
117    }
118
119    /// The PMF sums to ~1 over a wide support window.
120    #[test]
121    fn pmf_sums_to_one() {
122        let d = PoissonDistribution {
123            rate_parameter: 4.0,
124            ..Default::default()
125        };
126        let total: f64 = (0..60).map(|k| d.pmf(k)).sum();
127        assert!((total - 1.0).abs() < 1e-12, "sum was {total}");
128    }
129}