Skip to main content

stats_claw/distributions/discrete/
binomial.rs

1//! Binomial distribution numerics, for the
2//! [`BinomialDistribution`].
3//!
4//! Equivalent to `scipy.stats.binom(number_of_trials, success_probability)`. The
5//! PMF uses `ln_gamma` for the binomial coefficient, the CDF is expressed through
6//! the regularized incomplete beta `betai`, the quantile is the smallest `k` with
7//! `cdf(k) ≥ p`, and sampling counts `n` Bernoulli trials.
8//!
9//! # Examples
10//!
11//! ```
12//! use stats_claw::distributions::{Moments, Pmf};
13//! use stats_claw::distributions::BinomialDistribution;
14//!
15//! let d = BinomialDistribution { number_of_trials: 10, success_probability: 0.5, ..Default::default() };
16//! // Mean is n*p = 5; PMF at the mode (5) is the largest.
17//! assert_eq!(d.mean(), Some(5.0));
18//! assert!(d.pmf(5) > d.pmf(3), "mode pmf should exceed flank");
19//! ```
20
21use super::super::{Cdf, Moments, Pmf, Quantile, Sample, count_to_f64};
22use crate::distributions::BinomialDistribution;
23use crate::rng::SplitMix64;
24use crate::special::{betai, ln_gamma};
25
26impl BinomialDistribution {
27    /// Number of trials as a float (`n`).
28    fn n(&self) -> f64 {
29        count_to_f64(self.number_of_trials)
30    }
31}
32
33impl Pmf for BinomialDistribution {
34    fn pmf(&self, k: i64) -> f64 {
35        if k < 0 || k > self.number_of_trials {
36            return 0.0;
37        }
38        let (n, p) = (self.n(), self.success_probability);
39        let kf = count_to_f64(k);
40        // log C(n,k) = ln Γ(n+1) − ln Γ(k+1) − ln Γ(n−k+1)
41        let log_choose = ln_gamma(n + 1.0) - ln_gamma(kf + 1.0) - ln_gamma(n - kf + 1.0);
42        let log_pmf = (n - kf).mul_add((1.0 - p).ln(), kf.mul_add(p.ln(), log_choose));
43        log_pmf.exp()
44    }
45}
46
47impl Cdf for BinomialDistribution {
48    fn cdf(&self, x: f64) -> f64 {
49        if x < 0.0 {
50            return 0.0;
51        }
52        if x >= self.n() {
53            return 1.0;
54        }
55        let k = x.floor();
56        let p = self.success_probability;
57        // P(X ≤ k) = I_{1−p}(n−k, k+1)
58        betai(self.n() - k, k + 1.0, 1.0 - p)
59    }
60}
61
62impl Quantile for BinomialDistribution {
63    fn quantile(&self, p: f64) -> f64 {
64        if p <= 0.0 {
65            return 0.0;
66        }
67        if p >= 1.0 {
68            return self.n();
69        }
70        let mut cumulative = 0.0;
71        for k in 0..=self.number_of_trials {
72            cumulative += self.pmf(k);
73            if cumulative >= p - 1e-12 {
74                return count_to_f64(k);
75            }
76        }
77        self.n()
78    }
79}
80
81impl Moments for BinomialDistribution {
82    fn mean(&self) -> Option<f64> {
83        Some(self.n() * self.success_probability)
84    }
85    fn variance(&self) -> Option<f64> {
86        let p = self.success_probability;
87        Some(self.n() * p * (1.0 - p))
88    }
89}
90
91impl Sample for BinomialDistribution {
92    fn sample(&self, rng: &mut SplitMix64) -> f64 {
93        let mut successes = 0.0;
94        for _ in 0..self.number_of_trials.max(0) {
95            if rng.next_f64() < self.success_probability {
96                successes += 1.0;
97            }
98        }
99        successes
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    /// The mean is `n·p` and the variance `n·p·(1−p)`.
108    #[test]
109    fn moments_are_np_and_npq() {
110        let d = BinomialDistribution {
111            number_of_trials: 20,
112            success_probability: 0.5,
113            ..Default::default()
114        };
115        assert_eq!(d.mean(), Some(10.0));
116        assert_eq!(d.variance(), Some(5.0));
117    }
118
119    /// The PMF sums to one over the full support.
120    #[test]
121    fn pmf_sums_to_one() {
122        let d = BinomialDistribution {
123            number_of_trials: 12,
124            success_probability: 0.3,
125            ..Default::default()
126        };
127        let total: f64 = (0..=12).map(|k| d.pmf(k)).sum();
128        assert!((total - 1.0).abs() < 1e-12, "sum was {total}");
129    }
130}