stats_claw/distributions/discrete/
binomial.rs1use 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 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 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 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 #[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 #[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}