probability_rs/dist/
geometric.rs

1use crate::dist::{Discrete, DistError, Distribution, Moments};
2use crate::rng::RngCore;
3
4/// Geometric(p) over k=1,2,... counts trials until first success.
5#[derive(Debug, Clone, Copy)]
6pub struct Geometric {
7    p: f64,
8}
9
10impl Geometric {
11    pub fn new(p: f64) -> Result<Self, DistError> {
12        if !(0.0..=1.0).contains(&p) || p == 0.0 || !p.is_finite() {
13            return Err(DistError::InvalidParameter);
14        }
15        Ok(Self { p })
16    }
17    #[inline]
18    pub fn p(&self) -> f64 {
19        self.p
20    }
21}
22
23impl Distribution for Geometric {
24    type Value = i64;
25    fn cdf(&self, k: i64) -> f64 {
26        if k < 1 {
27            return 0.0;
28        }
29        1.0 - (1.0 - self.p).powi(k as i32)
30    }
31    fn in_support(&self, k: i64) -> bool {
32        k >= 1
33    }
34    fn sample<R: RngCore>(&self, rng: &mut R) -> i64 {
35        // Inverse CDF: k = ceil(log(1-u)/log(1-p))
36        let u = rng.next_f64();
37        let q = 1.0 - self.p;
38        let k = ((1.0 - u).ln() / q.ln()).ceil() as i64;
39        if k < 1 { 1 } else { k }
40    }
41}
42
43impl Discrete for Geometric {
44    fn pmf(&self, k: i64) -> f64 {
45        if !self.in_support(k) {
46            return 0.0;
47        }
48        self.p * (1.0 - self.p).powi((k - 1) as i32)
49    }
50    fn inv_cdf(&self, p: f64) -> i64 {
51        debug_assert!((0.0..=1.0).contains(&p));
52        if p <= 0.0 {
53            return 1;
54        }
55        if p >= 1.0 {
56            return i64::MAX;
57        }
58        let q = 1.0 - self.p;
59        ((1.0 - p).ln() / q.ln()).ceil() as i64
60    }
61}
62
63impl Moments for Geometric {
64    fn mean(&self) -> f64 {
65        1.0 / self.p
66    }
67    fn variance(&self) -> f64 {
68        (1.0 - self.p) / (self.p * self.p)
69    }
70    fn skewness(&self) -> f64 {
71        let p = self.p;
72        (2.0 - p) / (1.0 - p).sqrt()
73    }
74    fn kurtosis(&self) -> f64 {
75        let p = self.p;
76        6.0 + (p * p) / (1.0 - p)
77    }
78    fn entropy(&self) -> f64 {
79        // Shannon entropy of the geometric distribution on {1,2,...}
80        // H = -(p ln p + (1-p)/p ln(1-p))
81        let p = self.p;
82        let q = 1.0 - p;
83        -(p * p.ln() + (q / p) * q.ln())
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    #[test]
91    fn pmf_values() {
92        let g = Geometric::new(0.25).unwrap();
93        assert!((g.pmf(1) - 0.25).abs() < 1e-15);
94        assert!((g.pmf(2) - 0.25 * 0.75).abs() < 1e-15);
95    }
96    #[test]
97    fn cdf_values() {
98        let g = Geometric::new(0.5).unwrap();
99        assert!((g.cdf(1) - 0.5).abs() < 1e-15);
100        assert!((g.cdf(2) - 0.75).abs() < 1e-15);
101    }
102    #[test]
103    fn moments_higher() {
104        let g = Geometric::new(0.25).unwrap();
105        let p: f64 = 0.25;
106        let skew = (2.0 - p)/((1.0 - p).sqrt());
107        let kurt = 6.0 + (p*p)/(1.0 - p);
108        assert!((g.skewness() - skew).abs() < 1e-12);
109        assert!((g.kurtosis() - kurt).abs() < 1e-12);
110    }
111
112    #[test]
113    fn entropy_geometric() {
114        let p = 0.3;
115        let g = Geometric::new(p).unwrap();
116        let q = 1.0 - p;
117        let expected = -(p * p.ln() + (q / p) * q.ln());
118        assert!((g.entropy() - expected).abs() < 1e-12);
119    }
120}
probability_rs/dist/geometric.rs

probability_rs/dist/
geometric.rs
