Skip to main content

probability_rs/dist/
binomial.rs

1use crate::dist::{Discrete, DistError, Distribution, Moments};
2use crate::rng::RngCore;
3
4/// Binomial(n, p) over k=0..n.
5#[derive(Debug, Clone, Copy)]
6pub struct Binomial {
7    n: u64,
8    p: f64,
9}
10
11impl Binomial {
12    pub fn new(n: u64, p: f64) -> Result<Self, DistError> {
13        if !(0.0..=1.0).contains(&p) || !p.is_finite() {
14            return Err(DistError::InvalidParameter);
15        }
16        Ok(Self { n, p })
17    }
18    #[inline]
19    pub fn n(&self) -> u64 {
20        self.n
21    }
22    #[inline]
23    pub fn p(&self) -> f64 {
24        self.p
25    }
26
27    fn pmf_recurrence(&self, k: u64) -> f64 {
28        if k > self.n {
29            return 0.0;
30        }
31        // Start at k=0: (1-p)^n
32        let mut p0 = (1.0 - self.p).powi(self.n as i32);
33        if k == 0 {
34            return p0;
35        }
36        // p(k) = p(k-1) * (n-k+1)/k * (p/(1-p))
37        let odds = self.p / (1.0 - self.p);
38        for i in 1..=k {
39            p0 *= (self.n - i + 1) as f64 / (i as f64) * odds;
40        }
41        p0
42    }
43
44    fn cdf_sum(&self, k: u64) -> f64 {
45        if k >= self.n {
46            return 1.0;
47        }
48        let mut sum = 0.0;
49        for i in 0..=k {
50            sum += self.pmf_recurrence(i);
51        }
52        sum
53    }
54}
55
56impl Distribution for Binomial {
57    type Value = i64;
58    fn cdf(&self, x: i64) -> f64 {
59        if x < 0 { 0.0 } else { self.cdf_sum(x as u64) }
60    }
61    fn in_support(&self, x: i64) -> bool {
62        x >= 0 && (x as u64) <= self.n
63    }
64    fn sample<R: RngCore>(&self, rng: &mut R) -> i64 {
65        // Inversion by summing pmf from 0.. until exceeds u
66        let u = rng.next_f64();
67        let mut acc = 0.0;
68        for k in 0..=self.n {
69            acc += self.pmf_recurrence(k);
70            if u <= acc {
71                return k as i64;
72            }
73        }
74        self.n as i64
75    }
76}
77
78impl Discrete for Binomial {
79    fn pmf(&self, x: i64) -> f64 {
80        if x < 0 {
81            return 0.0;
82        }
83        self.pmf_recurrence(x as u64)
84    }
85    fn inv_cdf(&self, p: f64) -> i64 {
86        debug_assert!((0.0..=1.0).contains(&p));
87        if p <= 0.0 {
88            return 0;
89        }
90        if p >= 1.0 {
91            return self.n as i64;
92        }
93        let mut acc = 0.0;
94        for k in 0..=self.n {
95            acc += self.pmf_recurrence(k);
96            if p <= acc {
97                return k as i64;
98            }
99        }
100        self.n as i64
101    }
102}
103
104impl Moments for Binomial {
105    fn mean(&self) -> f64 {
106        (self.n as f64) * self.p
107    }
108    fn variance(&self) -> f64 {
109        (self.n as f64) * self.p * (1.0 - self.p)
110    }
111    fn skewness(&self) -> f64 {
112        let n = self.n as f64;
113        let p = self.p;
114        (1.0 - 2.0 * p) / (n * p * (1.0 - p)).sqrt()
115    }
116    fn kurtosis(&self) -> f64 {
117        let n = self.n as f64;
118        let p = self.p;
119        (1.0 - 6.0 * p * (1.0 - p)) / (n * p * (1.0 - p))
120    }
121    fn entropy(&self) -> f64 {
122        // Exact sum: H = - sum_{k=0}^n pmf(k) ln pmf(k)
123        let mut h = 0.0;
124        for k in 0..=self.n {
125            let pk = self.pmf_recurrence(k);
126            if pk > 0.0 { h -= pk * pk.ln(); }
127        }
128        h
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    #[test]
136    fn pmf_values() {
137        let b = Binomial::new(5, 0.4).unwrap();
138        // pmf(0) = (1-p)^n
139        assert!((b.pmf(0) - 0.6f64.powi(5)).abs() < 1e-15);
140        // pmf(5) = p^n
141        assert!((b.pmf(5) - 0.4f64.powi(5)).abs() < 1e-15);
142    }
143    #[test]
144    fn cdf_monotone() {
145        let b = Binomial::new(10, 0.3).unwrap();
146        assert!(b.cdf(3) <= b.cdf(4));
147        assert!(b.cdf(9) <= 1.0);
148    }
149    #[test]
150    fn moments_higher() {
151        let b = Binomial::new(10, 0.3).unwrap();
152        let n: f64 = 10.0; let p: f64 = 0.3; let q: f64 = 0.7;
153        let skew = (1.0 - 2.0*p)/((n*p*q).sqrt());
154        let kurt = (1.0 - 6.0*p*q)/(n*p*q);
155        assert!((b.skewness() - skew).abs() < 1e-12);
156        assert!((b.kurtosis() - kurt).abs() < 1e-12);
157    }
158}