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