rand_distr/
poisson.rs

Help
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
12use num_traits::{Float, FloatConst};
13use crate::{Cauchy, Distribution, Standard};
14use rand::Rng;
15use core::fmt;
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_distr::{Poisson, Distribution};
26///
27/// let poi = Poisson::new(2.0).unwrap();
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#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
33pub struct Poisson<F>
34where F: Float + FloatConst, Standard: Distribution<F>
35{
36    lambda: F,
37    // precalculated values
38    exp_lambda: F,
39    log_lambda: F,
40    sqrt_2lambda: F,
41    magic_val: F,
42}
43
44/// Error type returned from `Poisson::new`.
45#[derive(Clone, Copy, Debug, PartialEq, Eq)]
46pub enum Error {
47    /// `lambda <= 0` or `nan`.
48    ShapeTooSmall,
49}
50
51impl fmt::Display for Error {
52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53        f.write_str(match self {
54            Error::ShapeTooSmall => "lambda is not positive in Poisson distribution",
55        })
56    }
57}
58
59#[cfg(feature = "std")]
60#[cfg_attr(doc_cfg, doc(cfg(feature = "std")))]
61impl std::error::Error for Error {}
62
63impl<F> Poisson<F>
64where F: Float + FloatConst, Standard: Distribution<F>
65{
66    /// Construct a new `Poisson` with the given shape parameter
67    /// `lambda`.
68    pub fn new(lambda: F) -> Result<Poisson<F>, Error> {
69        if !(lambda > F::zero()) {
70            return Err(Error::ShapeTooSmall);
71        }
72        let log_lambda = lambda.ln();
73        Ok(Poisson {
74            lambda,
75            exp_lambda: (-lambda).exp(),
76            log_lambda,
77            sqrt_2lambda: (F::from(2.0).unwrap() * lambda).sqrt(),
78            magic_val: lambda * log_lambda - crate::utils::log_gamma(F::one() + lambda),
79        })
80    }
81}
82
83impl<F> Distribution<F> for Poisson<F>
84where F: Float + FloatConst, Standard: Distribution<F>
85{
86    #[inline]
87    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
88        // using the algorithm from Numerical Recipes in C
89
90        // for low expected values use the Knuth method
91        if self.lambda < F::from(12.0).unwrap() {
92            let mut result = F::zero();
93            let mut p = F::one();
94            while p > self.exp_lambda {
95                p = p*rng.gen::<F>();
96                result = result + F::one();
97            }
98            result - F::one()
99        }
100        // high expected values - rejection method
101        else {
102            // we use the Cauchy distribution as the comparison distribution
103            // f(x) ~ 1/(1+x^2)
104            let cauchy = Cauchy::new(F::zero(), F::one()).unwrap();
105            let mut result;
106
107            loop {
108                let mut comp_dev;
109
110                loop {
111                    // draw from the Cauchy distribution
112                    comp_dev = rng.sample(cauchy);
113                    // shift the peak of the comparison distribution
114                    result = self.sqrt_2lambda * comp_dev + self.lambda;
115                    // repeat the drawing until we are in the range of possible values
116                    if result >= F::zero() {
117                        break;
118                    }
119                }
120                // now the result is a random variable greater than 0 with Cauchy distribution
121                // the result should be an integer value
122                result = result.floor();
123
124                // this is the ratio of the Poisson distribution to the comparison distribution
125                // the magic value scales the distribution function to a range of approximately 0-1
126                // since it is not exact, we multiply the ratio by 0.9 to avoid ratios greater than 1
127                // this doesn't change the resulting distribution, only increases the rate of failed drawings
128                let check = F::from(0.9).unwrap()
129                    * (F::one() + comp_dev * comp_dev)
130                    * (result * self.log_lambda
131                        - crate::utils::log_gamma(F::one() + result)
132                        - self.magic_val)
133                        .exp();
134
135                // check with uniform random value - if below the threshold, we are within the target distribution
136                if rng.gen::<F>() <= check {
137                    break;
138                }
139            }
140            result
141        }
142    }
143}
144
145#[cfg(test)]
146mod test {
147    use super::*;
148
149    fn test_poisson_avg_gen<F: Float + FloatConst>(lambda: F, tol: F)
150        where Standard: Distribution<F>
151    {
152        let poisson = Poisson::new(lambda).unwrap();
153        let mut rng = crate::test::rng(123);
154        let mut sum = F::zero();
155        for _ in 0..1000 {
156            sum = sum + poisson.sample(&mut rng);
157        }
158        let avg = sum / F::from(1000.0).unwrap();
159        assert!((avg - lambda).abs() < tol);
160    }
161
162    #[test]
163    fn test_poisson_avg() {
164        test_poisson_avg_gen::<f64>(10.0, 0.5);
165        test_poisson_avg_gen::<f64>(15.0, 0.5);
166        test_poisson_avg_gen::<f32>(10.0, 0.5);
167        test_poisson_avg_gen::<f32>(15.0, 0.5);
168    }
169
170    #[test]
171    #[should_panic]
172    fn test_poisson_invalid_lambda_zero() {
173        Poisson::new(0.0).unwrap();
174    }
175
176    #[test]
177    #[should_panic]
178    fn test_poisson_invalid_lambda_neg() {
179        Poisson::new(-10.0).unwrap();
180    }
181}
rand_distr/poisson.rs

rand_distr/
poisson.rs
