Support sampling integers from discrete distributions

src/distribution/bernoulli.rs

@@ -86,9 +86,15 @@ impl std::fmt::Display for 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
+        rng.sample::<bool, _>(self) as u8 as f64
     }
 }
 

src/distribution/binomial.rs

@@ -92,19 +92,25 @@ impl std::fmt::Display for 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`

src/distribution/categorical.rs

@@ -83,9 +83,21 @@ impl std::fmt::Display for 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
     }
 }
 
@@ -281,13 +293,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

src/distribution/discrete_uniform.rs

@@ -57,6 +57,12 @@ impl std::fmt::Display for 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

src/distribution/geometric.rs

@@ -74,17 +74,26 @@ impl std::fmt::Display for 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()
+            // This cast is safe, because the largest finite value this expression can take is when
+            // `x = 1.4e-45` and `1.0 - self.p = 0.9999999999999999`, in which case we get
+            // `930262250532780300`, which when casted to a `u64` is `930262250532780288`.
+            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`

src/distribution/hypergeometric.rs

@@ -120,17 +120,17 @@ impl std::fmt::Display for 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;
@@ -143,6 +143,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`

src/distribution/multinomial.rs

@@ -119,6 +119,23 @@ where
     }
 }
 
+impl<D> ::rand::distributions::Distribution<OVector<u64, D>> for Multinomial<D>
+where
+    D: Dim,
+    nalgebra::DefaultAllocator: nalgebra::allocator::Allocator<f64, D>,
+    nalgebra::DefaultAllocator: nalgebra::allocator::Allocator<u64, D>,
+{
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> OVector<u64, D> {
+        let p_cdf = super::categorical::prob_mass_to_cdf(self.p().as_slice());
+        let mut res = OVector::zeros_generic(self.p.shape_generic().0, Const::<1>);
+        for _ in 0..self.n {
+            let i = super::categorical::sample_unchecked(rng, &p_cdf);
+            res[i] += 1;
+        }
+        res
+    }
+}
+
 impl<D> ::rand::distributions::Distribution<OVector<f64, D>> for Multinomial<D>
 where
     D: Dim,
@@ -129,8 +146,7 @@ where
         let mut res = OVector::zeros_generic(self.p.shape_generic().0, Const::<1>);
         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
     }

src/distribution/negative_binomial.rs

@@ -113,7 +113,7 @@ impl std::fmt::Display for 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) as u64
     }
 }
 

src/distribution/poisson.rs

@@ -72,6 +72,17 @@ impl std::fmt::Display for 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) as u64
+    }
+}
+
 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
@@ -260,6 +271,7 @@ impl Discrete<u64, f64> for Poisson {
         -self.lambda + x as f64 * self.lambda.ln() - factorial::ln_factorial(x)
     }
 }
+
 /// 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