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