Skip to main content

stats_claw/distributions/discrete/
binomial.rs

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