Add Poisson and binomial distributions

Author: fizyk20

src/distributions/binomial.rs

@@ -0,0 +1,156 @@
+// Copyright 2016-2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// https://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+//! The binomial distribution.
+
+use Rng;
+use distributions::Distribution;
+use distributions::log_gamma::log_gamma;
+use std::f64::consts::PI;
+
+/// The binomial distribution `Binomial(n, p)`.
+///
+/// This distribution has density function: `f(k) = n!/(k! (n-k)!) p^k (1-p)^(n-k)` for `k >= 0`.
+///
+/// # Example
+///
+/// ```rust
+/// use rand::distributions::{Binomial, Distribution};
+///
+/// let bin = Binomial::new(20, 0.3);
+/// let v = bin.sample(&mut rand::thread_rng());
+/// println!("{} is from a binomial distribution", v);
+/// ```
+#[derive(Clone, Copy, Debug)]
+pub struct Binomial {
+    n: u64, // number of trials
+    p: f64, // probability of success
+}
+
+impl Binomial {
+    /// Construct a new `Binomial` with the given shape parameters
+    /// `n`, `p`. Panics if `p <= 0` or `p >= 1`.
+    pub fn new(n: u64, p: f64) -> Binomial {
+        assert!(p > 0.0, "Binomial::new called with `p` <= 0");
+        assert!(p < 1.0, "Binomial::new called with `p` >= 1");
+        Binomial { n: n, p: p }
+    }
+}
+
+impl Distribution<u64> for Binomial {
+    fn sample<R: Rng>(&self, rng: &mut R) -> u64 {
+        // binomial distribution is symmetrical with respect to p -> 1-p, k -> n-k
+        // switch p so that it is less than 0.5 - this allows for lower expected values
+        // we will just invert the result at the end
+        let p = if self.p <= 0.5 {
+            self.p
+        } else {
+            1.0 - self.p
+        };
+
+        // expected value of the sample
+        let expected = self.n as f64 * p;
+
+        let result =
+            // for low expected values we just simulate n drawings
+            if expected < 25.0 {
+                let mut lresult = 0.0;
+                for _ in 0 .. self.n {
+                    if rng.gen::<f64>() < p {
+                        lresult += 1.0;
+                    }
+                }
+                lresult
+            }
+            // high expected value - do the rejection method
+            else {
+                // prepare some cached values
+                let float_n = self.n as f64;
+                let ln_fact_n = log_gamma(float_n + 1.0);
+                let pc = 1.0 - p;
+                let log_p = p.ln();
+                let log_pc = pc.ln();
+                let sq = (expected * (2.0 * pc)).sqrt();
+
+                let mut lresult;
+
+                loop {
+                    let mut comp_dev: f64;
+                    // we use the lorentzian distribution as the comparison distribution
+                    // f(x) ~ 1/(1+x/^2)
+                    loop {
+                        // draw from the lorentzian distribution
+                        comp_dev = (PI*rng.gen::<f64>()).tan();
+                        // shift the peak of the comparison ditribution
+                        lresult = expected + sq * comp_dev;
+                        // repeat the drawing until we are in the range of possible values
+                        if lresult >= 0.0 && lresult < float_n + 1.0 {
+                            break;
+                        }
+                    }
+
+                    // the result should be discrete
+                    lresult = lresult.floor();
+
+                    let log_binomial_dist = ln_fact_n - log_gamma(lresult+1.0) -
+                        log_gamma(float_n - lresult + 1.0) + lresult*log_p + (float_n - lresult)*log_pc;
+                    // this is the binomial probability divided by the comparison probability
+                    // we will generate a uniform random value and if it is larger than this,
+                    // we interpret it as a value falling out of the distribution and repeat
+                    let comparison_coeff = (log_binomial_dist.exp() * sq) * (1.2 * (1.0 + comp_dev*comp_dev));
+
+                    if comparison_coeff >= rng.gen() {
+                        break;
+                    }
+                }
+
+                lresult
+            };
+
+        // invert the result for p < 0.5
+        if p != self.p {
+            self.n - result as u64
+        } else {
+            result as u64
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use distributions::Distribution;
+    use super::Binomial;
+
+    #[test]
+    fn test_binomial() {
+        let binomial = Binomial::new(150, 0.1);
+        let mut rng = ::test::rng(123);
+        let mut sum = 0;
+        for _ in 0..1000 {
+            sum += binomial.sample(&mut rng);
+        }
+        let avg = (sum as f64) / 1000.0;
+        println!("Binomial average: {}", avg);
+        assert!((avg - 15.0).abs() < 0.5); // not 100% certain, but probable enough
+    }
+
+    #[test]
+    #[should_panic]
+    #[cfg_attr(target_env = "msvc", ignore)]
+    fn test_binomial_invalid_lambda_zero() {
+        Binomial::new(20, 0.0);
+    }
+    #[test]
+    #[should_panic]
+    #[cfg_attr(target_env = "msvc", ignore)]
+    fn test_binomial_invalid_lambda_neg() {
+        Binomial::new(20, -10.0);
+    }
+}

src/distributions/exponential.rs

@@ -10,7 +10,7 @@
 
 //! The exponential distribution.
 
-use {Rng};
+use Rng;
 use distributions::{ziggurat, ziggurat_tables, Distribution};
 
 /// Samples floating-point numbers according to the exponential distribution,
@@ -52,10 +52,14 @@ impl Distribution<f64> for Exp1 {
             ziggurat_tables::ZIG_EXP_R - rng.gen::<f64>().ln()
         }
 
-        ziggurat(rng, false,
-                 &ziggurat_tables::ZIG_EXP_X,
-                 &ziggurat_tables::ZIG_EXP_F,
-                 pdf, zero_case)
+        ziggurat(
+            rng,
+            false,
+            &ziggurat_tables::ZIG_EXP_X,
+            &ziggurat_tables::ZIG_EXP_F,
+            pdf,
+            zero_case,
+        )
     }
 }
 
@@ -76,7 +80,7 @@ impl Distribution<f64> for Exp1 {
 #[derive(Clone, Copy, Debug)]
 pub struct Exp {
     /// `lambda` stored as `1/lambda`, since this is what we scale by.
-    lambda_inverse: f64
+    lambda_inverse: f64,
 }
 
 impl Exp {
@@ -85,7 +89,9 @@ impl Exp {
     #[inline]
     pub fn new(lambda: f64) -> Exp {
         assert!(lambda > 0.0, "Exp::new called with `lambda` <= 0");
-        Exp { lambda_inverse: 1.0 / lambda }
+        Exp {
+            lambda_inverse: 1.0 / lambda,
+        }
     }
 }
 

src/distributions/float.rs

@@ -69,7 +69,6 @@ macro_rules! float_impls {
 float_impls! { f32, u32, 23, 127, next_u32 }
 float_impls! { f64, u64, 52, 1023, next_u64 }
 
-
 #[cfg(test)]
 mod tests {
     use Rng;

src/distributions/gamma.rs

@@ -15,7 +15,7 @@
 use self::GammaRepr::*;
 use self::ChiSquaredRepr::*;
 
-use {Rng};
+use Rng;
 use distributions::normal::StandardNormal;
 use distributions::{Distribution, Exp, Uniform};
 
@@ -58,7 +58,7 @@ pub struct Gamma {
 enum GammaRepr {
     Large(GammaLargeShape),
     One(Exp),
-    Small(GammaSmallShape)
+    Small(GammaSmallShape),
 }
 
 // These two helpers could be made public, but saving the
@@ -78,7 +78,7 @@ enum GammaRepr {
 #[derive(Clone, Copy, Debug)]
 struct GammaSmallShape {
     inv_shape: f64,
-    large_shape: GammaLargeShape
+    large_shape: GammaLargeShape,
 }
 
 /// Gamma distribution where the shape parameter is larger than 1.
@@ -89,7 +89,7 @@ struct GammaSmallShape {
 struct GammaLargeShape {
     scale: f64,
     c: f64,
-    d: f64
+    d: f64,
 }
 
 impl Gamma {
@@ -117,7 +117,7 @@ impl GammaSmallShape {
     fn new_raw(shape: f64, scale: f64) -> GammaSmallShape {
         GammaSmallShape {
             inv_shape: 1. / shape,
-            large_shape: GammaLargeShape::new_raw(shape + 1.0, scale)
+            large_shape: GammaLargeShape::new_raw(shape + 1.0, scale),
         }
     }
 }
@@ -128,7 +128,7 @@ impl GammaLargeShape {
         GammaLargeShape {
             scale: scale,
             c: 1. / (9. * d).sqrt(),
-            d: d
+            d: d,
         }
     }
 }
@@ -154,17 +154,19 @@ impl Distribution<f64> for GammaLargeShape {
         loop {
             let x = rng.sample(StandardNormal);
             let v_cbrt = 1.0 + self.c * x;
-            if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0
-                continue
+            if v_cbrt <= 0.0 {
+                // a^3 <= 0 iff a <= 0
+                continue;
             }
 
             let v = v_cbrt * v_cbrt * v_cbrt;
             let u: f64 = rng.sample(Uniform);
 
             let x_sqr = x * x;
-            if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
-                u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
-                return self.d * v * self.scale
+            if u < 1.0 - 0.0331 * x_sqr * x_sqr
+                || u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln())
+            {
+                return self.d * v * self.scale;
             }
         }
     }
@@ -222,7 +224,7 @@ impl Distribution<f64> for ChiSquared {
                 let norm = rng.sample(StandardNormal);
                 norm * norm
             }
-            DoFAnythingElse(ref g) => g.sample(rng)
+            DoFAnythingElse(ref g) => g.sample(rng),
         }
     }
 }
@@ -261,7 +263,7 @@ impl FisherF {
         FisherF {
             numer: ChiSquared::new(m),
             denom: ChiSquared::new(n),
-            dof_ratio: n / m
+            dof_ratio: n / m,
         }
     }
 }
@@ -286,7 +288,7 @@ impl Distribution<f64> for FisherF {
 #[derive(Clone, Copy, Debug)]
 pub struct StudentT {
     chi: ChiSquared,
-    dof: f64
+    dof: f64,
 }
 
 impl StudentT {
@@ -296,7 +298,7 @@ impl StudentT {
         assert!(n > 0.0, "StudentT::new called with `n <= 0`");
         StudentT {
             chi: ChiSquared::new(n),
-            dof: n
+            dof: n,
         }
     }
 }
@@ -310,7 +312,7 @@ impl Distribution<f64> for StudentT {
 #[cfg(test)]
 mod test {
     use distributions::Distribution;
-    use super::{ChiSquared, StudentT, FisherF};
+    use super::{ChiSquared, FisherF, StudentT};
 
     #[test]
     fn test_chi_squared_one() {

src/distributions/integer.rs

@@ -12,7 +12,7 @@
 
 use core::mem;
 
-use {Rng};
+use Rng;
 use distributions::{Distribution, Uniform};
 
 impl Distribution<isize> for Uniform {
@@ -109,24 +109,23 @@ impl Distribution<u128> for Uniform {
     }
 }
 
-
 #[cfg(test)]
 mod tests {
     use Rng;
-    use distributions::{Uniform};
-    
+    use distributions::Uniform;
+
     #[test]
     fn test_integers() {
         let mut rng = ::test::rng(806);
-        
+
         rng.sample::<isize, _>(Uniform);
         rng.sample::<i8, _>(Uniform);
         rng.sample::<i16, _>(Uniform);
         rng.sample::<i32, _>(Uniform);
         rng.sample::<i64, _>(Uniform);
         #[cfg(feature = "i128_support")]
         rng.sample::<i128, _>(Uniform);
-        
+
         rng.sample::<usize, _>(Uniform);
         rng.sample::<u8, _>(Uniform);
         rng.sample::<u16, _>(Uniform);