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`. new(lambda: f64) -> Poisson43 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 { sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u6457 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] test_poisson_10()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] test_poisson_15()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] test_poisson_invalid_lambda_zero()148 fn test_poisson_invalid_lambda_zero() { 149 Poisson::new(0.0); 150 } 151 152 #[test] 153 #[should_panic] test_poisson_invalid_lambda_neg()154 fn test_poisson_invalid_lambda_neg() { 155 Poisson::new(-10.0); 156 } 157 } 158