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}