Add geometric and hypergeometric distributions

rand_distr/CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/)
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+- New `Geometric`, `StandardGeometric` and `Hypergeometric` distributions (#1062)
 - New `Beta` sampling algorithm for improved performance and accuracy (#1000)
 - `Normal` and `LogNormal` now support `from_mean_cv` and `from_zscore` (#1044)
 - Variants of `NormalError` changed (#1044)

rand_distr/benches/distributions.rs

@@ -136,6 +136,8 @@ distr_int!(distr_weighted_alias_method_u32, usize, WeightedAliasIndex::new(vec![
 distr_int!(distr_weighted_alias_method_f64, usize, WeightedAliasIndex::new(vec![1.0f64, 0.001, 1.0/3.0, 4.01, 0.0, 3.3, 22.0, 0.001]).unwrap());
 distr_int!(distr_weighted_alias_method_large_set, usize, WeightedAliasIndex::new((0..10000).rev().chain(1..10001).collect()).unwrap());
 
+distr_int!(distr_geometric, u64, Geometric::new(0.5).unwrap());
+distr_int!(distr_standard_geometric, u64, StandardGeometric);
 
 #[bench]
 fn dist_iter(b: &mut Bencher) {

rand_distr/src/geometric.rs

@@ -0,0 +1,233 @@
+//! The geometric distribution.
+
+use crate::Distribution;
+use rand::Rng;
+use core::fmt;
+
+/// The geometric distribution `Geometric(p)` bounded to `[0, u64::MAX]`.
+/// 
+/// This is the probability distribution of the number of failures before the
+/// first success in a series of Bernoulli trials. It has the density function
+/// `f(k) = (1 - p)^k p` for `k >= 0`, where `p` is the probability of success
+/// on each trial.
+/// 
+/// This is the discrete analogue of the [exponential distribution](crate::Exp).
+/// 
+/// Note that [`StandardGeometric`](crate::StandardGeometric) is an optimised
+/// implementation for `p = 0.5`.
+///
+/// # Example
+///
+/// ```
+/// use rand_distr::{Geometric, Distribution};
+///
+/// let geo = Geometric::new(0.25).unwrap();
+/// let v = geo.sample(&mut rand::thread_rng());
+/// println!("{} is from a Geometric(0.25) distribution", v);
+/// ```
+#[derive(Copy, Clone, Debug)]
+pub struct Geometric
+{
+    p: f64,
+    pi: f64,
+    k: u64
+}
+
+/// Error type returned from `Geometric::new`.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum Error {
+    /// `p < 0 || p > 1` or `nan`
+    InvalidProbability,
+}
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            Error::InvalidProbability => "p is NaN or outside the interval [0, 1] in geometric distribution",
+        })
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for Error {}
+
+impl Geometric {
+    /// Construct a new `Geometric` with the given shape parameter `p`
+    /// (probablity of success on each trial).
+    pub fn new(p: f64) -> Result<Self, Error> {
+        if !p.is_finite() || p < 0.0 || p > 1.0 {
+            Err(Error::InvalidProbability)
+        } else if p == 0.0 || p >= 2.0 / 3.0 {
+            Ok(Geometric { p, pi: p, k: 0 })
+        } else {
+            let (pi, k) = {
+                // choose smallest k such that pi = (1 - p)^(2^k) <= 0.5
+                let mut k = 1;
+                let mut pi = (1.0 - p).powi(2);
+                while pi > 0.5 {
+                    k += 1;
+                    pi = pi * pi;
+                }
+                (pi, k)
+            };
+
+            Ok(Geometric { p, pi, k })
+        }
+    }
+}
+
+impl Distribution<u64> for Geometric
+{
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
+        if self.p >= 2.0 / 3.0 {
+            // use the trivial algorithm:
+            let mut failures = 0;
+            loop {
+                let u = rng.gen::<f64>();
+                if u <= self.p { break; }
+                failures += 1;
+            }
+            return failures;
+        }
+
+        if self.p == 0.0 { return core::u64::MAX; }
+
+        let Geometric { p, pi, k } = *self;
+
+        // Based on the algorithm presented in section 3 of
+        // Karl Bringmann and Tobias Friedrich (July 2013) - Exact and Efficient
+        // Generation of Geometric Random Variates and Random Graphs, published
+        // in International Colloquium on Automata, Languages and Programming
+        // (pp.267-278)
+        // https://people.mpi-inf.mpg.de/~kbringma/paper/2013ICALP-1.pdf
+
+        // Use the trivial algorithm to sample D from Geo(pi) = Geo(p) / 2^k:
+        let d = {
+            let mut failures = 0;
+            while rng.gen::<f64>() < pi {
+                failures += 1;
+            }
+            failures
+        };
+
+        // Use rejection sampling for the remainder M from Geo(p) % 2^k:
+        // choose M uniformly from [0, 2^k), but reject with probability (1 - p)^M
+        // NOTE: The paper suggests using bitwise sampling here, which is 
+        // currently unsupported, but should improve performance by requiring
+        // fewer iterations on average.                 ~ October 28, 2020
+        let m = loop {
+            let m = rng.gen::<u64>() & ((1 << k) - 1);
+            let p_reject = if m <= core::i32::MAX as u64 {
+                (1.0 - p).powi(m as i32)
+            } else {
+                (1.0 - p).powf(m as f64)
+            };
+
+            let u = rng.gen::<f64>();
+            if u < p_reject {
+                break m;
+            }
+        };
+
+        (d << k) + m
+    }
+}
+
+/// Samples integers according to the geometric distribution with success
+/// probability `p = 0.5`. This is equivalent to `Geometeric::new(0.5)`,
+/// but faster.
+/// 
+/// See [`Geometric`](crate::Geometric) for the general geometric distribution.
+/// 
+/// Implemented via iterated [Rng::gen::<u64>().leading_zeros()].
+/// 
+/// # Example
+/// ```
+/// use rand::prelude::*;
+/// use rand_distr::StandardGeometric;
+/// 
+/// let v = StandardGeometric.sample(&mut thread_rng());
+/// println!("{} is from a Geometric(0.5) distribution", v);
+/// ```
+#[derive(Copy, Clone, Debug)]
+pub struct StandardGeometric;
+
+impl Distribution<u64> for StandardGeometric {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
+        let mut result = 0;
+        loop {
+            let x = rng.gen::<u64>().leading_zeros() as u64;
+            result += x;
+            if x < 64 { break; }
+        }
+        result
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_geo_invalid_p() {
+        assert!(Geometric::new(core::f64::NAN).is_err());
+        assert!(Geometric::new(core::f64::INFINITY).is_err());
+        assert!(Geometric::new(core::f64::NEG_INFINITY).is_err());
+
+        assert!(Geometric::new(-0.5).is_err());
+        assert!(Geometric::new(0.0).is_ok());
+        assert!(Geometric::new(1.0).is_ok());
+        assert!(Geometric::new(2.0).is_err());
+    }
+
+    fn test_geo_mean_and_variance<R: Rng>(p: f64, rng: &mut R) {
+        let distr = Geometric::new(p).unwrap();
+
+        let expected_mean = (1.0 - p) / p;
+        let expected_variance = (1.0 - p) / (p * p);
+
+        let mut results = [0.0; 10000];
+        for i in results.iter_mut() {
+            *i = distr.sample(rng) as f64;
+        }
+
+        let mean = results.iter().sum::<f64>() / results.len() as f64;
+        assert!((mean as f64 - expected_mean).abs() < expected_mean / 40.0);
+
+        let variance =
+            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>() / results.len() as f64;
+        assert!((variance - expected_variance).abs() < expected_variance / 10.0);
+    }
+
+    #[test]
+    fn test_geometric() {
+        let mut rng = crate::test::rng(12345);
+
+        test_geo_mean_and_variance(0.10, &mut rng);
+        test_geo_mean_and_variance(0.25, &mut rng);
+        test_geo_mean_and_variance(0.50, &mut rng);
+        test_geo_mean_and_variance(0.75, &mut rng);
+        test_geo_mean_and_variance(0.90, &mut rng);
+    }
+
+    #[test]
+    fn test_standard_geometric() {
+        let mut rng = crate::test::rng(654321);
+
+        let distr = StandardGeometric;
+        let expected_mean = 1.0;
+        let expected_variance = 2.0;
+
+        let mut results = [0.0; 1000];
+        for i in results.iter_mut() {
+            *i = distr.sample(&mut rng) as f64;
+        }
+
+        let mean = results.iter().sum::<f64>() / results.len() as f64;
+        assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0);
+
+        let variance =
+            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>() / results.len() as f64;
+        assert!((variance - expected_variance).abs() < expected_variance / 10.0);
+    }
+}