Support sampling integers from discrete distributions

statrs-dev · Nov 5, 2021 · b7efa3a · b7efa3a
1 parent 435739b
commit b7efa3a
Show file tree

Hide file tree

Showing 9 changed files with 88 additions and 27 deletions.
diff --git a/src/distribution/bernoulli.rs b/src/distribution/bernoulli.rs
@@ -80,6 +80,12 @@ impl Bernoulli {
     }
 }
 
+impl ::rand::distributions::Distribution<bool> for Bernoulli {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool {
+        rng.gen_bool(self.p())
+    }
+}
+
 impl ::rand::distributions::Distribution<f64> for Bernoulli {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
         rng.gen_bool(self.p()) as u8 as f64

diff --git a/src/distribution/binomial.rs b/src/distribution/binomial.rs
@@ -87,19 +87,25 @@ impl Binomial {
     }
 }
 
-impl ::rand::distributions::Distribution<f64> for Binomial {
-    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
-        (0..self.n).fold(0.0, |acc, _| {
+impl ::rand::distributions::Distribution<u64> for Binomial {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
+        (0..self.n).fold(0, |acc, _| {
             let n: f64 = rng.gen();
             if n < self.p {
-                acc + 1.0
+                acc + 1
             } else {
                 acc
             }
         })
     }
 }
 
+impl ::rand::distributions::Distribution<f64> for Binomial {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
+        rng.sample::<u64, _>(self) as f64
+    }
+}
+
 impl DiscreteCDF<u64, f64> for Binomial {
     /// Calculates the cumulative distribution function for the
     /// binomial distribution at `x`

diff --git a/src/distribution/categorical.rs b/src/distribution/categorical.rs
@@ -74,9 +74,21 @@ impl Categorical {
     }
 }
 
+impl ::rand::distributions::Distribution<usize> for Categorical {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
+        sample_unchecked(rng, &self.cdf)
+    }
+}
+
+impl ::rand::distributions::Distribution<u64> for Categorical {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
+        sample_unchecked(rng, &self.cdf) as u64
+    }
+}
+
 impl ::rand::distributions::Distribution<f64> for Categorical {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
-        sample_unchecked(rng, &self.cdf)
+        sample_unchecked(rng, &self.cdf) as f64
     }
 }
 
@@ -253,13 +265,9 @@ impl Discrete<u64, f64> for Categorical {
 
 /// Draws a sample from the categorical distribution described by `cdf`
 /// without doing any bounds checking
-pub fn sample_unchecked<R: Rng + ?Sized>(rng: &mut R, cdf: &[f64]) -> f64 {
+pub fn sample_unchecked<R: Rng + ?Sized>(rng: &mut R, cdf: &[f64]) -> usize {
     let draw = rng.gen::<f64>() * cdf.last().unwrap();
-    cdf.iter()
-        .enumerate()
-        .find(|(_, val)| **val >= draw)
-        .map(|(i, _)| i)
-        .unwrap() as f64
+    cdf.iter().position(|val| *val >= draw).unwrap()
 }
 
 /// Computes the cdf from the given probability masses. Performs

diff --git a/src/distribution/discrete_uniform.rs b/src/distribution/discrete_uniform.rs
@@ -51,6 +51,12 @@ impl DiscreteUniform {
     }
 }
 
+impl ::rand::distributions::Distribution<i64> for DiscreteUniform {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i64 {
+        rng.gen_range(self.min..=self.max)
+    }
+}
+
 impl ::rand::distributions::Distribution<f64> for DiscreteUniform {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
         rng.gen_range(self.min..=self.max) as f64

diff --git a/src/distribution/geometric.rs b/src/distribution/geometric.rs
@@ -68,17 +68,23 @@ impl Geometric {
     }
 }
 
-impl ::rand::distributions::Distribution<f64> for Geometric {
-    fn sample<R: Rng + ?Sized>(&self, r: &mut R) -> f64 {
+impl ::rand::distributions::Distribution<u64> for Geometric {
+    fn sample<R: Rng + ?Sized>(&self, r: &mut R) -> u64 {
         if ulps_eq!(self.p, 1.0) {
-            1.0
+            1
         } else {
             let x: f64 = r.sample(OpenClosed01);
-            x.log(1.0 - self.p).ceil()
+            x.log(1.0 - self.p).ceil() as u64
         }
     }
 }
 
+impl ::rand::distributions::Distribution<f64> for Geometric {
+    fn sample<R: Rng + ?Sized>(&self, r: &mut R) -> f64 {
+        r.sample::<u64, _>(self) as f64
+    }
+}
+
 impl DiscreteCDF<u64, f64> for Geometric {
     /// Calculates the cumulative distribution function for the geometric
     /// distribution at `x`

diff --git a/src/distribution/hypergeometric.rs b/src/distribution/hypergeometric.rs
@@ -110,17 +110,17 @@ impl Hypergeometric {
     }
 }
 
-impl ::rand::distributions::Distribution<f64> for Hypergeometric {
-    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
+impl ::rand::distributions::Distribution<u64> for Hypergeometric {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
         let mut population = self.population as f64;
         let mut successes = self.successes as f64;
         let mut draws = self.draws;
-        let mut x = 0.0;
+        let mut x = 0;
         loop {
             let p = successes / population;
             let next: f64 = rng.gen();
             if next < p {
-                x += 1.0;
+                x += 1;
                 successes -= 1.0;
             }
             population -= 1.0;
@@ -133,6 +133,12 @@ impl ::rand::distributions::Distribution<f64> for Hypergeometric {
     }
 }
 
+impl ::rand::distributions::Distribution<f64> for Hypergeometric {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
+        rng.sample::<u64, _>(self) as f64
+    }
+}
+
 impl DiscreteCDF<u64, f64> for Hypergeometric {
     /// Calculates the cumulative distribution function for the hypergeometric
     /// distribution at `x`

diff --git a/src/distribution/multinomial.rs b/src/distribution/multinomial.rs
@@ -90,14 +90,25 @@ impl Multinomial {
     }
 }
 
+impl ::rand::distributions::Distribution<Vec<u64>> for Multinomial {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Vec<u64> {
+        let p_cdf = super::categorical::prob_mass_to_cdf(self.p());
+        let mut res = vec![0; self.p.len()];
+        for _ in 0..self.n {
+            let i = super::categorical::sample_unchecked(rng, &p_cdf);
+            res[i] += 1;
+        }
+        res
+    }
+}
+
 impl ::rand::distributions::Distribution<Vec<f64>> for Multinomial {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Vec<f64> {
         let p_cdf = super::categorical::prob_mass_to_cdf(self.p());
         let mut res = vec![0.0; self.p.len()];
         for _ in 0..self.n {
             let i = super::categorical::sample_unchecked(rng, &p_cdf);
-            let el = res.get_mut(i as usize).unwrap();
-            *el += 1.0;
+            res[i] += 1.0;
         }
         res
     }

diff --git a/src/distribution/negative_binomial.rs b/src/distribution/negative_binomial.rs
@@ -89,7 +89,7 @@ impl NegativeBinomial {
 impl ::rand::distributions::Distribution<u64> for NegativeBinomial {
     fn sample<R: Rng + ?Sized>(&self, r: &mut R) -> u64 {
         let lambda = distribution::gamma::sample_unchecked(r, self.r, (1.0 - self.p) / self.p);
-        poisson::sample_unchecked(r, lambda).floor() as u64
+        poisson::sample_unchecked(r, lambda)
     }
 }
 

diff --git a/src/distribution/poisson.rs b/src/distribution/poisson.rs
@@ -67,14 +67,25 @@ impl Poisson {
     }
 }
 
+impl ::rand::distributions::Distribution<u64> for Poisson {
+    /// Generates one sample from the Poisson distribution either by
+    /// Knuth's method if lambda < 30.0 or Rejection method PA by
+    /// A. C. Atkinson from the Journal of the Royal Statistical Society
+    /// Series C (Applied Statistics) Vol. 28 No. 1. (1979) pp. 29 - 35
+    /// otherwise
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
+        sample_unchecked(rng, self.lambda)
+    }
+}
+
 impl ::rand::distributions::Distribution<f64> for Poisson {
     /// Generates one sample from the Poisson distribution either by
     /// Knuth's method if lambda < 30.0 or Rejection method PA by
     /// A. C. Atkinson from the Journal of the Royal Statistical Society
     /// Series C (Applied Statistics) Vol. 28 No. 1. (1979) pp. 29 - 35
     /// otherwise
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
-        sample_unchecked(rng, self.lambda)
+        sample_unchecked(rng, self.lambda) as f64
     }
 }
 
@@ -238,18 +249,19 @@ impl Discrete<u64, f64> for Poisson {
         -self.lambda + x as f64 * self.lambda.ln() - factorial::ln_factorial(x as u64)
     }
 }
+
 /// Generates one sample from the Poisson distribution either by
 /// Knuth's method if lambda < 30.0 or Rejection method PA by
 /// A. C. Atkinson from the Journal of the Royal Statistical Society
 /// Series C (Applied Statistics) Vol. 28 No. 1. (1979) pp. 29 - 35
 /// otherwise
-pub fn sample_unchecked<R: Rng + ?Sized>(rng: &mut R, lambda: f64) -> f64 {
+pub fn sample_unchecked<R: Rng + ?Sized>(rng: &mut R, lambda: f64) -> u64 {
     if lambda < 30.0 {
         let limit = (-lambda).exp();
-        let mut count = 0.0;
+        let mut count = 0;
         let mut product: f64 = rng.gen();
         while product >= limit {
-            count += 1.0;
+            count += 1;
             product *= rng.gen::<f64>();
         }
         count
@@ -273,7 +285,7 @@ pub fn sample_unchecked<R: Rng + ?Sized>(rng: &mut R, lambda: f64) -> f64 {
             let lhs = y + (v / (temp * temp)).ln();
             let rhs = k + n * lambda.ln() - factorial::ln_factorial(n as u64);
             if lhs <= rhs {
-                return n;
+                return n as u64;
             }
         }
     }