Skip to main content

stats_claw/distributions/discrete/
poisson.rs

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