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