rand_distr/
geometric.rs

1//! The geometric distribution `Geometric(p)`.
2
3use crate::Distribution;
4use core::fmt;
5#[allow(unused_imports)]
6use num_traits::Float;
7use rand::{Rng, RngExt};
8
9/// The [geometric distribution](https://en.wikipedia.org/wiki/Geometric_distribution) `Geometric(p)`.
10///
11/// This is the probability distribution of the number of failures
12/// (bounded to `[0, u64::MAX]`) before the first success in a
13/// series of [`Bernoulli`](crate::Bernoulli) trials, where the
14/// probability of success on each trial is `p`.
15///
16/// This is the discrete analogue of the [exponential distribution](crate::Exp).
17///
18/// See [`StandardGeometric`](crate::StandardGeometric) for an optimised
19/// implementation for `p = 0.5`.
20///
21/// # Density function
22///
23/// `f(k) = (1 - p)^k p` for `k >= 0`.
24///
25/// # Plot
26///
27/// The following plot illustrates the geometric distribution for various
28/// values of `p`. Note how higher `p` values shift the distribution to
29/// the left, and the mean of the distribution is `1/p`.
30///
31/// ![Geometric distribution](https://raw.githubusercontent.com/rust-random/charts/main/charts/geometric.svg)
32///
33/// # Example
34/// ```
35/// use rand_distr::{Geometric, Distribution};
36///
37/// let geo = Geometric::new(0.25).unwrap();
38/// let v = geo.sample(&mut rand::rng());
39/// println!("{} is from a Geometric(0.25) distribution", v);
40/// ```
41#[derive(Copy, Clone, Debug, PartialEq)]
42#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
43pub struct Geometric {
44    p: f64,
45    pi: f64,
46    k: u64,
47}
48
49/// Error type returned from [`Geometric::new`].
50#[derive(Clone, Copy, Debug, PartialEq, Eq)]
51pub enum Error {
52    /// `p < 0 || p > 1` or `nan`
53    InvalidProbability,
54}
55
56impl fmt::Display for Error {
57    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
58        f.write_str(match self {
59            Error::InvalidProbability => {
60                "p is NaN or outside the interval [0, 1] in geometric distribution"
61            }
62        })
63    }
64}
65
66#[cfg(feature = "std")]
67impl std::error::Error for Error {}
68
69impl Geometric {
70    /// Construct a new `Geometric` distribution
71    ///
72    /// The shape parameter `p` is the probability of success on each trial.
73    ///
74    /// ### Edge cases
75    ///
76    /// If `p == 0.0` or `1.0 - p` rounds to `1.0` then sampling returns
77    /// `u64::MAX`.
78    pub fn new(p: f64) -> Result<Self, Error> {
79        let mut pi = 1.0 - p;
80        if !p.is_finite() || !(0.0..=1.0).contains(&p) {
81            Err(Error::InvalidProbability)
82        } else if pi == 1.0 || p >= 2.0 / 3.0 {
83            Ok(Geometric { p, pi, k: 0 })
84        } else {
85            let (pi, k) = {
86                // choose smallest k such that pi = (1 - p)^(2^k) <= 0.5
87                let mut k = 1;
88                pi = pi * pi;
89                while pi > 0.5 {
90                    k += 1;
91                    pi = pi * pi;
92                }
93                (pi, k)
94            };
95
96            Ok(Geometric { p, pi, k })
97        }
98    }
99}
100
101impl Distribution<u64> for Geometric {
102    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
103        if self.p >= 2.0 / 3.0 {
104            // use the trivial algorithm:
105            let mut failures = 0;
106            loop {
107                let u = rng.random::<f64>();
108                if u <= self.p {
109                    break;
110                }
111                failures += 1;
112            }
113            return failures;
114        }
115
116        if self.pi == 1.0 {
117            return u64::MAX;
118        }
119
120        let Geometric { p, pi, k } = *self;
121
122        // Based on the algorithm presented in section 3 of
123        // Karl Bringmann and Tobias Friedrich (July 2013) - Exact and Efficient
124        // Generation of Geometric Random Variates and Random Graphs, published
125        // in International Colloquium on Automata, Languages and Programming
126        // (pp.267-278)
127        // https://people.mpi-inf.mpg.de/~kbringma/paper/2013ICALP-1.pdf
128
129        // Use the trivial algorithm to sample D from Geo(pi) = Geo(p) / 2^k:
130        let d = {
131            let mut failures = 0;
132            while rng.random::<f64>() < pi {
133                failures += 1;
134            }
135            failures
136        };
137
138        // Use rejection sampling for the remainder M from Geo(p) % 2^k:
139        // choose M uniformly from [0, 2^k), but reject with probability (1 - p)^M
140        // NOTE: The paper suggests using bitwise sampling here, which is
141        // currently unsupported, but should improve performance by requiring
142        // fewer iterations on average.                 ~ October 28, 2020
143        let m = loop {
144            let m = rng.random::<u64>() & ((1 << k) - 1);
145            let p_reject = if m <= i32::MAX as u64 {
146                (1.0 - p).powi(m as i32)
147            } else {
148                (1.0 - p).powf(m as f64)
149            };
150
151            let u = rng.random::<f64>();
152            if u < p_reject {
153                break m;
154            }
155        };
156
157        (d << k) + m
158    }
159}
160
161/// The standard geometric distribution `Geometric(0.5)`.
162///
163/// This is equivalent to `Geometric::new(0.5)`, but faster.
164///
165/// See [`Geometric`](crate::Geometric) for the general geometric distribution.
166///
167/// # Plot
168///
169/// The following plot illustrates the standard geometric distribution.
170///
171/// ![Standard Geometric distribution](https://raw.githubusercontent.com/rust-random/charts/main/charts/standard_geometric.svg)
172///
173/// # Example
174/// ```
175/// use rand::prelude::*;
176/// use rand_distr::StandardGeometric;
177///
178/// let v = StandardGeometric.sample(&mut rand::rng());
179/// println!("{} is from a Geometric(0.5) distribution", v);
180/// ```
181///
182/// # Notes
183/// Implemented via iterated
184/// [`Rng::gen::<u64>().leading_zeros()`](Rng::gen::<u64>().leading_zeros()).
185#[derive(Copy, Clone, Debug)]
186#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
187pub struct StandardGeometric;
188
189impl Distribution<u64> for StandardGeometric {
190    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
191        let mut result = 0;
192        loop {
193            let x = rng.random::<u64>().leading_zeros() as u64;
194            result += x;
195            if x < 64 {
196                break;
197            }
198        }
199        result
200    }
201}
202
203#[cfg(test)]
204mod test {
205    use super::*;
206
207    #[test]
208    fn test_geo_invalid_p() {
209        assert!(Geometric::new(f64::NAN).is_err());
210        assert!(Geometric::new(f64::INFINITY).is_err());
211        assert!(Geometric::new(f64::NEG_INFINITY).is_err());
212
213        assert!(Geometric::new(-0.5).is_err());
214        assert!(Geometric::new(0.0).is_ok());
215        assert!(Geometric::new(1.0).is_ok());
216        assert!(Geometric::new(2.0).is_err());
217    }
218
219    fn test_geo_mean_and_variance<R: Rng>(p: f64, rng: &mut R) {
220        let distr = Geometric::new(p).unwrap();
221
222        let expected_mean = (1.0 - p) / p;
223        let expected_variance = (1.0 - p) / (p * p);
224
225        let mut results = [0.0; 10000];
226        for i in results.iter_mut() {
227            *i = distr.sample(rng) as f64;
228        }
229
230        let mean = results.iter().sum::<f64>() / results.len() as f64;
231        assert!((mean - expected_mean).abs() < expected_mean / 40.0);
232
233        let variance =
234            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>() / results.len() as f64;
235        assert!((variance - expected_variance).abs() < expected_variance / 10.0);
236    }
237
238    #[test]
239    fn test_geometric() {
240        let mut rng = crate::test::rng(12345);
241
242        test_geo_mean_and_variance(0.10, &mut rng);
243        test_geo_mean_and_variance(0.25, &mut rng);
244        test_geo_mean_and_variance(0.50, &mut rng);
245        test_geo_mean_and_variance(0.75, &mut rng);
246        test_geo_mean_and_variance(0.90, &mut rng);
247    }
248
249    #[test]
250    fn test_standard_geometric() {
251        let mut rng = crate::test::rng(654321);
252
253        let distr = StandardGeometric;
254        let expected_mean = 1.0;
255        let expected_variance = 2.0;
256
257        let mut results = [0.0; 1000];
258        for i in results.iter_mut() {
259            *i = distr.sample(&mut rng) as f64;
260        }
261
262        let mean = results.iter().sum::<f64>() / results.len() as f64;
263        assert!((mean - expected_mean).abs() < expected_mean / 50.0);
264
265        let variance =
266            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>() / results.len() as f64;
267        assert!((variance - expected_variance).abs() < expected_variance / 10.0);
268    }
269
270    #[test]
271    fn geometric_distributions_can_be_compared() {
272        assert_eq!(Geometric::new(1.0), Geometric::new(1.0));
273    }
274
275    #[test]
276    fn small_p() {
277        let a = f64::EPSILON / 2.0;
278        assert!(1.0 - a < 1.0); // largest repr. value < 1
279        assert!(Geometric::new(a).is_ok());
280
281        let b = f64::EPSILON / 4.0;
282        assert!(b > 0.0);
283        assert!(1.0 - b == 1.0); // rounds to 1
284        let d = Geometric::new(b).unwrap();
285        let mut rng = crate::test::VoidRng;
286        assert_eq!(d.sample(&mut rng), u64::MAX);
287    }
288}
rand_distr/geometric.rs

rand_distr/
geometric.rs
