Skip to content

Commit e707f75

Browse files
Poisson u64 sampling (#1498)
This addresses rust-random/rand#1497 by adding `Distribution<u64>` It also solves rust-random/rand#1312 by not allowing `lambda` bigger than `1.844e19` (this also makes them always fit into `u64`)
1 parent ce7b3c4 commit e707f75

File tree

1 file changed

+29
-15
lines changed

1 file changed

+29
-15
lines changed

rand_distr/src/poisson.rs

Lines changed: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,6 @@ use rand::Rng;
2323
/// This distribution has density function:
2424
/// `f(k) = λ^k * exp(-λ) / k!` for `k >= 0`.
2525
///
26-
/// # Known issues
27-
///
28-
/// See documentation of [`Poisson::new`].
29-
///
3026
/// # Plot
3127
///
3228
/// The following plot shows the Poisson distribution with various values of `λ`.
@@ -40,7 +36,7 @@ use rand::Rng;
4036
/// use rand_distr::{Poisson, Distribution};
4137
///
4238
/// let poi = Poisson::new(2.0).unwrap();
43-
/// let v = poi.sample(&mut rand::thread_rng());
39+
/// let v: f64 = poi.sample(&mut rand::thread_rng());
4440
/// println!("{} is from a Poisson(2) distribution", v);
4541
/// ```
4642
#[derive(Clone, Copy, Debug, PartialEq)]
@@ -52,20 +48,23 @@ where
5248

5349
/// Error type returned from [`Poisson::new`].
5450
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
55-
// Marked non_exhaustive to allow a new error code in the solution to #1312.
56-
#[non_exhaustive]
5751
pub enum Error {
5852
/// `lambda <= 0`
5953
ShapeTooSmall,
6054
/// `lambda = ∞` or `lambda = nan`
6155
NonFinite,
56+
/// `lambda` is too large, see [Poisson::MAX_LAMBDA]
57+
ShapeTooLarge,
6258
}
6359

6460
impl fmt::Display for Error {
6561
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
6662
f.write_str(match self {
6763
Error::ShapeTooSmall => "lambda is not positive in Poisson distribution",
6864
Error::NonFinite => "lambda is infinite or nan in Poisson distribution",
65+
Error::ShapeTooLarge => {
66+
"lambda is too large in Poisson distribution, see Poisson::MAX_LAMBDA"
67+
}
6968
})
7069
}
7170
}
@@ -125,14 +124,7 @@ where
125124
/// Construct a new `Poisson` with the given shape parameter
126125
/// `lambda`.
127126
///
128-
/// # Known issues
129-
///
130-
/// Although this method should return an [`Error`] on invalid parameters,
131-
/// some (extreme) values of `lambda` are known to return a [`Poisson`]
132-
/// object which hangs when [sampled](Distribution::sample).
133-
/// Large (less extreme) values of `lambda` may result in successful
134-
/// sampling but with reduced precision.
135-
/// See [#1312](https://github.com/rust-random/rand/issues/1312).
127+
/// The maximum allowed lambda is [MAX_LAMBDA](Self::MAX_LAMBDA).
136128
pub fn new(lambda: F) -> Result<Poisson<F>, Error> {
137129
if !lambda.is_finite() {
138130
return Err(Error::NonFinite);
@@ -145,11 +137,25 @@ where
145137
let method = if lambda < F::from(12.0).unwrap() {
146138
Method::Knuth(KnuthMethod::new(lambda))
147139
} else {
140+
if lambda > F::from(Self::MAX_LAMBDA).unwrap() {
141+
return Err(Error::ShapeTooLarge);
142+
}
148143
Method::Rejection(RejectionMethod::new(lambda))
149144
};
150145

151146
Ok(Poisson(method))
152147
}
148+
149+
/// The maximum supported value of `lambda`
150+
///
151+
/// This value was selected such that
152+
/// `MAX_LAMBDA + 1e6 * sqrt(MAX_LAMBDA) < 2^64 - 1`,
153+
/// thus ensuring that the probability of sampling a value larger than
154+
/// `u64::MAX` is less than 1e-1000.
155+
///
156+
/// Applying this limit also solves
157+
/// [#1312](https://github.com/rust-random/rand/issues/1312).
158+
pub const MAX_LAMBDA: f64 = 1.844e19;
153159
}
154160

155161
impl<F> Distribution<F> for KnuthMethod<F>
@@ -232,6 +238,14 @@ where
232238
}
233239
}
234240

241+
impl Distribution<u64> for Poisson<f64> {
242+
#[inline]
243+
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
244+
// `as` from float to int saturates
245+
<Poisson<f64> as Distribution<f64>>::sample(self, rng) as u64
246+
}
247+
}
248+
235249
#[cfg(test)]
236250
mod test {
237251
use super::*;

0 commit comments

Comments
 (0)