rand/distr/bernoulli.rs
1// Copyright 2018 Developers of the Rand project.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9//! The Bernoulli distribution `Bernoulli(p)`.
10
11use crate::distr::Distribution;
12use crate::{Rng, RngExt};
13use core::fmt;
14
15#[cfg(feature = "serde")]
16use serde::{Deserialize, Serialize};
17
18/// The [Bernoulli distribution](https://en.wikipedia.org/wiki/Bernoulli_distribution) `Bernoulli(p)`.
19///
20/// This distribution describes a single boolean random variable, which is true
21/// with probability `p` and false with probability `1 - p`.
22/// It is a special case of the Binomial distribution with `n = 1`.
23///
24/// # Plot
25///
26/// The following plot shows the Bernoulli distribution with `p = 0.1`,
27/// `p = 0.5`, and `p = 0.9`.
28///
29/// 
30///
31/// # Example
32///
33/// ```rust
34/// use rand::distr::{Bernoulli, Distribution};
35///
36/// let d = Bernoulli::new(0.3).unwrap();
37/// let v = d.sample(&mut rand::rng());
38/// println!("{} is from a Bernoulli distribution", v);
39/// ```
40///
41/// # Precision
42///
43/// This `Bernoulli` distribution uses 64 bits from the RNG (a `u64`),
44/// so only probabilities that are multiples of 2<sup>-64</sup> can be
45/// represented.
46#[derive(Clone, Copy, Debug, PartialEq)]
47#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
48pub struct Bernoulli {
49 /// Probability of success, relative to the maximal integer.
50 p_int: u64,
51}
52
53// To sample from the Bernoulli distribution we use a method that compares a
54// random `u64` value `v < (p * 2^64)`.
55//
56// If `p == 1.0`, the integer `v` to compare against can not represented as a
57// `u64`. We manually set it to `u64::MAX` instead (2^64 - 1 instead of 2^64).
58// Note that value of `p < 1.0` can never result in `u64::MAX`, because an
59// `f64` only has 53 bits of precision, and the next largest value of `p` will
60// result in `2^64 - 2048`.
61//
62// Also there is a 100% theoretical concern: if someone consistently wants to
63// generate `true` using the Bernoulli distribution (i.e. by using a probability
64// of `1.0`), just using `u64::MAX` is not enough. On average it would return
65// false once every 2^64 iterations. Some people apparently care about this
66// case.
67//
68// That is why we special-case `u64::MAX` to always return `true`, without using
69// the RNG, and pay the performance price for all uses that *are* reasonable.
70// Luckily, if `new()` and `sample` are close, the compiler can optimize out the
71// extra check.
72const ALWAYS_TRUE: u64 = u64::MAX;
73
74// This is just `2.0.powi(64)`, but written this way because it is not available
75// in `no_std` mode.
76const SCALE: f64 = 2.0 * (1u64 << 63) as f64;
77
78/// Error type returned from [`Bernoulli::new`].
79#[derive(Clone, Copy, Debug, PartialEq, Eq)]
80pub enum BernoulliError {
81 /// `p < 0` or `p > 1`.
82 InvalidProbability,
83}
84
85impl fmt::Display for BernoulliError {
86 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
87 f.write_str(match self {
88 BernoulliError::InvalidProbability => "p is outside [0, 1] in Bernoulli distribution",
89 })
90 }
91}
92
93impl core::error::Error for BernoulliError {}
94
95impl Bernoulli {
96 /// Construct a new `Bernoulli` with the given probability of success `p`.
97 ///
98 /// # Precision
99 ///
100 /// For `p = 1.0`, the resulting distribution will always generate true.
101 /// For `p = 0.0`, the resulting distribution will always generate false.
102 ///
103 /// This method is accurate for any input `p` in the range `[0, 1]` which is
104 /// a multiple of 2<sup>-64</sup>. (Note that not all multiples of
105 /// 2<sup>-64</sup> in `[0, 1]` can be represented as a `f64`.)
106 #[inline]
107 pub fn new(p: f64) -> Result<Bernoulli, BernoulliError> {
108 if !(0.0..1.0).contains(&p) {
109 if p == 1.0 {
110 return Ok(Bernoulli { p_int: ALWAYS_TRUE });
111 }
112 return Err(BernoulliError::InvalidProbability);
113 }
114 Ok(Bernoulli {
115 p_int: (p * SCALE) as u64,
116 })
117 }
118
119 /// Construct a new `Bernoulli` with the probability of success of
120 /// `numerator`-in-`denominator`. I.e. `from_ratio(2, 3)` will return
121 /// a `Bernoulli` with a 2-in-3 chance, or about 67%, of returning `true`.
122 ///
123 /// For `numerator == denominator`, the resulting distribution will always
124 /// return `true`; for `numerator == 0` it will always return `false`.
125 /// For `numerator > denominator` or `denominator == 0`, this returns an
126 /// error.
127 ///
128 /// # Example
129 ///
130 /// ```
131 /// use rand::distr::Bernoulli;
132 ///
133 /// let d = Bernoulli::from_ratio(2, 3).unwrap();
134 /// assert!((d.p() - 2.0 / 3.0).abs() < 1e-9);
135 ///
136 /// // Edge cases:
137 /// assert_eq!(Bernoulli::from_ratio(3, 3).unwrap().p(), 1.0); // always true
138 /// assert_eq!(Bernoulli::from_ratio(0, 3).unwrap().p(), 0.0); // always false
139 /// assert!(Bernoulli::from_ratio(4, 3).is_err()); // numerator > denominator
140 /// assert!(Bernoulli::from_ratio(1, 0).is_err()); // denominator == 0
141 /// ```
142 #[inline]
143 pub fn from_ratio(numerator: u32, denominator: u32) -> Result<Bernoulli, BernoulliError> {
144 if numerator > denominator || denominator == 0 {
145 return Err(BernoulliError::InvalidProbability);
146 }
147 if numerator == denominator {
148 return Ok(Bernoulli { p_int: ALWAYS_TRUE });
149 }
150 let p_int = ((f64::from(numerator) / f64::from(denominator)) * SCALE) as u64;
151 Ok(Bernoulli { p_int })
152 }
153
154 #[inline]
155 /// Returns the probability (`p`) of the distribution.
156 ///
157 /// This value may differ slightly from the input due to loss of precision.
158 pub fn p(&self) -> f64 {
159 if self.p_int == ALWAYS_TRUE {
160 1.0
161 } else {
162 (self.p_int as f64) / SCALE
163 }
164 }
165}
166
167impl Distribution<bool> for Bernoulli {
168 #[inline]
169 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool {
170 // Make sure to always return true for p = 1.0.
171 if self.p_int == ALWAYS_TRUE {
172 return true;
173 }
174 let v: u64 = rng.random();
175 v < self.p_int
176 }
177}
178
179#[cfg(test)]
180mod test {
181 use super::Bernoulli;
182 use crate::RngExt;
183 use crate::distr::Distribution;
184
185 #[test]
186 #[cfg(feature = "serde")]
187 fn test_serializing_deserializing_bernoulli() {
188 let coin_flip = Bernoulli::new(0.5).unwrap();
189 let de_coin_flip: Bernoulli =
190 postcard::from_bytes(&postcard::to_allocvec(&coin_flip).unwrap()).unwrap();
191
192 assert_eq!(coin_flip.p_int, de_coin_flip.p_int);
193 }
194
195 #[test]
196 fn test_trivial() {
197 // We prefer to be explicit here.
198 #![allow(clippy::bool_assert_comparison)]
199
200 let mut r = crate::test::rng(1);
201 let always_false = Bernoulli::new(0.0).unwrap();
202 let always_true = Bernoulli::new(1.0).unwrap();
203 for _ in 0..5 {
204 assert_eq!(r.sample::<bool, _>(&always_false), false);
205 assert_eq!(r.sample::<bool, _>(&always_true), true);
206 assert_eq!(Distribution::<bool>::sample(&always_false, &mut r), false);
207 assert_eq!(Distribution::<bool>::sample(&always_true, &mut r), true);
208 }
209 }
210
211 #[test]
212 #[cfg_attr(miri, ignore)] // Miri is too slow
213 fn test_average() {
214 const P: f64 = 0.3;
215 const NUM: u32 = 3;
216 const DENOM: u32 = 10;
217 let d1 = Bernoulli::new(P).unwrap();
218 let d2 = Bernoulli::from_ratio(NUM, DENOM).unwrap();
219 const N: u32 = 100_000;
220
221 let mut sum1: u32 = 0;
222 let mut sum2: u32 = 0;
223 let mut rng = crate::test::rng(2);
224 for _ in 0..N {
225 if d1.sample(&mut rng) {
226 sum1 += 1;
227 }
228 if d2.sample(&mut rng) {
229 sum2 += 1;
230 }
231 }
232 let avg1 = (sum1 as f64) / (N as f64);
233 assert!((avg1 - P).abs() < 5e-3);
234
235 let avg2 = (sum2 as f64) / (N as f64);
236 assert!((avg2 - (NUM as f64) / (DENOM as f64)).abs() < 5e-3);
237 }
238
239 #[test]
240 fn value_stability() {
241 let mut rng = crate::test::rng(3);
242 let distr = Bernoulli::new(0.4532).unwrap();
243 let mut buf = [false; 10];
244 for x in &mut buf {
245 *x = rng.sample(distr);
246 }
247 assert_eq!(
248 buf,
249 [
250 true, false, false, true, false, false, true, true, true, true
251 ]
252 );
253 }
254
255 #[test]
256 fn bernoulli_distributions_can_be_compared() {
257 assert_eq!(Bernoulli::new(1.0), Bernoulli::new(1.0));
258 }
259}