vendor/rand-0.6.5/src/distributions/poisson.rs

   1 // Copyright 2018 Developers of the Rand project.
   2 // Copyright 2016-2017 The Rust Project Developers.
   3 //
   4 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   5 // https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   6 // <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
   7 // option. This file may not be copied, modified, or distributed
   8 // except according to those terms.
   9
  10 //! The Poisson distribution.
  11
  12 use Rng;
  13 use distributions::{Distribution, Cauchy};
  14 use distributions::utils::log_gamma;
  15
  16 /// The Poisson distribution `Poisson(lambda)`.
  17 ///
  18 /// This distribution has a density function:
  19 /// `f(k) = lambda^k * exp(-lambda) / k!` for `k >= 0`.
  20 ///
  21 /// # Example
  22 ///
  23 /// ```
  24 /// use rand::distributions::{Poisson, Distribution};
  25 ///
  26 /// let poi = Poisson::new(2.0);
  27 /// let v = poi.sample(&mut rand::thread_rng());
  28 /// println!("{} is from a Poisson(2) distribution", v);
  29 /// ```
  30 #[derive(Clone, Copy, Debug)]
  31 pub struct Poisson {
  32     lambda: f64,
  33     // precalculated values
  34     exp_lambda: f64,
  35     log_lambda: f64,
  36     sqrt_2lambda: f64,
  37     magic_val: f64,
  38 }
  39
  40 impl Poisson {
  41     /// Construct a new `Poisson` with the given shape parameter
  42     /// `lambda`. Panics if `lambda <= 0`.
  43     pub fn new(lambda: f64) -> Poisson {
  44         assert!(lambda > 0.0, "Poisson::new called with lambda <= 0");
  45         let log_lambda = lambda.ln();
  46         Poisson {
  47             lambda,
  48             exp_lambda: (-lambda).exp(),
  49             log_lambda,
  50             sqrt_2lambda: (2.0 * lambda).sqrt(),
  51             magic_val: lambda * log_lambda - log_gamma(1.0 + lambda),
  52         }
  53     }
  54 }
  55
  56 impl Distribution<u64> for Poisson {
  57     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
  58         // using the algorithm from Numerical Recipes in C
  59
  60         // for low expected values use the Knuth method
  61         if self.lambda < 12.0 {
  62             let mut result = 0;
  63             let mut p = 1.0;
  64             while p > self.exp_lambda {
  65                 p *= rng.gen::<f64>();
  66                 result += 1;
  67             }
  68             result - 1
  69         }
  70         // high expected values - rejection method
  71         else {
  72             let mut int_result: u64;
  73
  74             // we use the Cauchy distribution as the comparison distribution
  75             // f(x) ~ 1/(1+x^2)
  76             let cauchy = Cauchy::new(0.0, 1.0);
  77
  78             loop {
  79                 let mut result;
  80                 let mut comp_dev;
  81
  82                 loop {
  83                     // draw from the Cauchy distribution
  84                     comp_dev = rng.sample(cauchy);
  85                     // shift the peak of the comparison ditribution
  86                     result = self.sqrt_2lambda * comp_dev + self.lambda;
  87                     // repeat the drawing until we are in the range of possible values
  88                     if result >= 0.0 {
  89                         break;
  90                     }
  91                 }
  92                 // now the result is a random variable greater than 0 with Cauchy distribution
  93                 // the result should be an integer value
  94                 result = result.floor();
  95                 int_result = result as u64;
  96
  97                 // this is the ratio of the Poisson distribution to the comparison distribution
  98                 // the magic value scales the distribution function to a range of approximately 0-1
  99                 // since it is not exact, we multiply the ratio by 0.9 to avoid ratios greater than 1
 100                 // this doesn't change the resulting distribution, only increases the rate of failed drawings
 101                 let check = 0.9 * (1.0 + comp_dev * comp_dev)
 102                     * (result * self.log_lambda - log_gamma(1.0 + result) - self.magic_val).exp();
 103
 104                 // check with uniform random value - if below the threshold, we are within the target distribution
 105                 if rng.gen::<f64>() <= check {
 106                     break;
 107                 }
 108             }
 109             int_result
 110         }
 111     }
 112 }
 113
 114 #[cfg(test)]
 115 mod test {
 116     use distributions::Distribution;
 117     use super::Poisson;
 118
 119     #[test]
 120     fn test_poisson_10() {
 121         let poisson = Poisson::new(10.0);
 122         let mut rng = ::test::rng(123);
 123         let mut sum = 0;
 124         for _ in 0..1000 {
 125             sum += poisson.sample(&mut rng);
 126         }
 127         let avg = (sum as f64) / 1000.0;
 128         println!("Poisson average: {}", avg);
 129         assert!((avg - 10.0).abs() < 0.5); // not 100% certain, but probable enough
 130     }
 131
 132     #[test]
 133     fn test_poisson_15() {
 134         // Take the 'high expected values' path
 135         let poisson = Poisson::new(15.0);
 136         let mut rng = ::test::rng(123);
 137         let mut sum = 0;
 138         for _ in 0..1000 {
 139             sum += poisson.sample(&mut rng);
 140         }
 141         let avg = (sum as f64) / 1000.0;
 142         println!("Poisson average: {}", avg);
 143         assert!((avg - 15.0).abs() < 0.5); // not 100% certain, but probable enough
 144     }
 145
 146     #[test]
 147     #[should_panic]
 148     fn test_poisson_invalid_lambda_zero() {
 149         Poisson::new(0.0);
 150     }
 151
 152     #[test]
 153     #[should_panic]
 154     fn test_poisson_invalid_lambda_neg() {
 155         Poisson::new(-10.0);
 156     }
 157 }