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. 10 11 use crate::distributions::Distribution; 12 use crate::Rng; 13 use core::{fmt, u64}; 14 15 #[cfg(feature = "serde1")] 16 use serde::{Serialize, Deserialize}; 17 /// The Bernoulli distribution. 18 /// 19 /// This is a special case of the Binomial distribution where `n = 1`. 20 /// 21 /// # Example 22 /// 23 /// ```rust 24 /// use rand::distributions::{Bernoulli, Distribution}; 25 /// 26 /// let d = Bernoulli::new(0.3).unwrap(); 27 /// let v = d.sample(&mut rand::thread_rng()); 28 /// println!("{} is from a Bernoulli distribution", v); 29 /// ``` 30 /// 31 /// # Precision 32 /// 33 /// This `Bernoulli` distribution uses 64 bits from the RNG (a `u64`), 34 /// so only probabilities that are multiples of 2<sup>-64</sup> can be 35 /// represented. 36 #[derive(Clone, Copy, Debug)] 37 #[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))] 38 pub struct Bernoulli { 39 /// Probability of success, relative to the maximal integer. 40 p_int: u64, 41 } 42 43 // To sample from the Bernoulli distribution we use a method that compares a 44 // random `u64` value `v < (p * 2^64)`. 45 // 46 // If `p == 1.0`, the integer `v` to compare against can not represented as a 47 // `u64`. We manually set it to `u64::MAX` instead (2^64 - 1 instead of 2^64). 48 // Note that value of `p < 1.0` can never result in `u64::MAX`, because an 49 // `f64` only has 53 bits of precision, and the next largest value of `p` will 50 // result in `2^64 - 2048`. 51 // 52 // Also there is a 100% theoretical concern: if someone consistenly wants to 53 // generate `true` using the Bernoulli distribution (i.e. by using a probability 54 // of `1.0`), just using `u64::MAX` is not enough. On average it would return 55 // false once every 2^64 iterations. Some people apparently care about this 56 // case. 57 // 58 // That is why we special-case `u64::MAX` to always return `true`, without using 59 // the RNG, and pay the performance price for all uses that *are* reasonable. 60 // Luckily, if `new()` and `sample` are close, the compiler can optimize out the 61 // extra check. 62 const ALWAYS_TRUE: u64 = u64::MAX; 63 64 // This is just `2.0.powi(64)`, but written this way because it is not available 65 // in `no_std` mode. 66 const SCALE: f64 = 2.0 * (1u64 << 63) as f64; 67 68 /// Error type returned from `Bernoulli::new`. 69 #[derive(Clone, Copy, Debug, PartialEq, Eq)] 70 pub enum BernoulliError { 71 /// `p < 0` or `p > 1`. 72 InvalidProbability, 73 } 74 75 impl fmt::Display for BernoulliError { fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result76 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 77 f.write_str(match self { 78 BernoulliError::InvalidProbability => "p is outside [0, 1] in Bernoulli distribution", 79 }) 80 } 81 } 82 83 #[cfg(feature = "std")] 84 impl ::std::error::Error for BernoulliError {} 85 86 impl Bernoulli { 87 /// Construct a new `Bernoulli` with the given probability of success `p`. 88 /// 89 /// # Precision 90 /// 91 /// For `p = 1.0`, the resulting distribution will always generate true. 92 /// For `p = 0.0`, the resulting distribution will always generate false. 93 /// 94 /// This method is accurate for any input `p` in the range `[0, 1]` which is 95 /// a multiple of 2<sup>-64</sup>. (Note that not all multiples of 96 /// 2<sup>-64</sup> in `[0, 1]` can be represented as a `f64`.) 97 #[inline] new(p: f64) -> Result<Bernoulli, BernoulliError>98 pub fn new(p: f64) -> Result<Bernoulli, BernoulliError> { 99 if !(p >= 0.0 && p < 1.0) { 100 if p == 1.0 { 101 return Ok(Bernoulli { p_int: ALWAYS_TRUE }); 102 } 103 return Err(BernoulliError::InvalidProbability); 104 } 105 Ok(Bernoulli { 106 p_int: (p * SCALE) as u64, 107 }) 108 } 109 110 /// Construct a new `Bernoulli` with the probability of success of 111 /// `numerator`-in-`denominator`. I.e. `new_ratio(2, 3)` will return 112 /// a `Bernoulli` with a 2-in-3 chance, or about 67%, of returning `true`. 113 /// 114 /// return `true`. If `numerator == 0` it will always return `false`. 115 /// For `numerator > denominator` and `denominator == 0`, this returns an 116 /// error. Otherwise, for `numerator == denominator`, samples are always 117 /// true; for `numerator == 0` samples are always false. 118 #[inline] from_ratio(numerator: u32, denominator: u32) -> Result<Bernoulli, BernoulliError>119 pub fn from_ratio(numerator: u32, denominator: u32) -> Result<Bernoulli, BernoulliError> { 120 if numerator > denominator || denominator == 0 { 121 return Err(BernoulliError::InvalidProbability); 122 } 123 if numerator == denominator { 124 return Ok(Bernoulli { p_int: ALWAYS_TRUE }); 125 } 126 let p_int = ((f64::from(numerator) / f64::from(denominator)) * SCALE) as u64; 127 Ok(Bernoulli { p_int }) 128 } 129 } 130 131 impl Distribution<bool> for Bernoulli { 132 #[inline] sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool133 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool { 134 // Make sure to always return true for p = 1.0. 135 if self.p_int == ALWAYS_TRUE { 136 return true; 137 } 138 let v: u64 = rng.gen(); 139 v < self.p_int 140 } 141 } 142 143 #[cfg(test)] 144 mod test { 145 use super::Bernoulli; 146 use crate::distributions::Distribution; 147 use crate::Rng; 148 149 #[test] 150 #[cfg(feature="serde1")] test_serializing_deserializing_bernoulli()151 fn test_serializing_deserializing_bernoulli() { 152 let coin_flip = Bernoulli::new(0.5).unwrap(); 153 let de_coin_flip : Bernoulli = bincode::deserialize(&bincode::serialize(&coin_flip).unwrap()).unwrap(); 154 155 assert_eq!(coin_flip.p_int, de_coin_flip.p_int); 156 } 157 158 #[test] test_trivial()159 fn test_trivial() { 160 let mut r = crate::test::rng(1); 161 let always_false = Bernoulli::new(0.0).unwrap(); 162 let always_true = Bernoulli::new(1.0).unwrap(); 163 for _ in 0..5 { 164 assert_eq!(r.sample::<bool, _>(&always_false), false); 165 assert_eq!(r.sample::<bool, _>(&always_true), true); 166 assert_eq!(Distribution::<bool>::sample(&always_false, &mut r), false); 167 assert_eq!(Distribution::<bool>::sample(&always_true, &mut r), true); 168 } 169 } 170 171 #[test] 172 #[cfg_attr(miri, ignore)] // Miri is too slow test_average()173 fn test_average() { 174 const P: f64 = 0.3; 175 const NUM: u32 = 3; 176 const DENOM: u32 = 10; 177 let d1 = Bernoulli::new(P).unwrap(); 178 let d2 = Bernoulli::from_ratio(NUM, DENOM).unwrap(); 179 const N: u32 = 100_000; 180 181 let mut sum1: u32 = 0; 182 let mut sum2: u32 = 0; 183 let mut rng = crate::test::rng(2); 184 for _ in 0..N { 185 if d1.sample(&mut rng) { 186 sum1 += 1; 187 } 188 if d2.sample(&mut rng) { 189 sum2 += 1; 190 } 191 } 192 let avg1 = (sum1 as f64) / (N as f64); 193 assert!((avg1 - P).abs() < 5e-3); 194 195 let avg2 = (sum2 as f64) / (N as f64); 196 assert!((avg2 - (NUM as f64) / (DENOM as f64)).abs() < 5e-3); 197 } 198 199 #[test] value_stability()200 fn value_stability() { 201 let mut rng = crate::test::rng(3); 202 let distr = Bernoulli::new(0.4532).unwrap(); 203 let mut buf = [false; 10]; 204 for x in &mut buf { 205 *x = rng.sample(&distr); 206 } 207 assert_eq!(buf, [ 208 true, false, false, true, false, false, true, true, true, true 209 ]); 210 } 211 } 212