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