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 binomial distribution. 11 12 use Rng; 13 use distributions::{Distribution, Bernoulli, Cauchy}; 14 use distributions::utils::log_gamma; 15 16 /// The binomial distribution `Binomial(n, p)`. 17 /// 18 /// This distribution has density function: 19 /// `f(k) = n!/(k! (n-k)!) p^k (1-p)^(n-k)` for `k >= 0`. 20 /// 21 /// # Example 22 /// 23 /// ``` 24 /// use rand::distributions::{Binomial, Distribution}; 25 /// 26 /// let bin = Binomial::new(20, 0.3); 27 /// let v = bin.sample(&mut rand::thread_rng()); 28 /// println!("{} is from a binomial distribution", v); 29 /// ``` 30 #[derive(Clone, Copy, Debug)] 31 pub struct Binomial { 32 /// Number of trials. 33 n: u64, 34 /// Probability of success. 35 p: f64, 36 } 37 38 impl Binomial { 39 /// Construct a new `Binomial` with the given shape parameters `n` (number 40 /// of trials) and `p` (probability of success). 41 /// 42 /// Panics if `p < 0` or `p > 1`. new(n: u64, p: f64) -> Binomial43 pub fn new(n: u64, p: f64) -> Binomial { 44 assert!(p >= 0.0, "Binomial::new called with p < 0"); 45 assert!(p <= 1.0, "Binomial::new called with p > 1"); 46 Binomial { n, p } 47 } 48 } 49 50 impl Distribution<u64> for Binomial { sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u6451 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 { 52 // Handle these values directly. 53 if self.p == 0.0 { 54 return 0; 55 } else if self.p == 1.0 { 56 return self.n; 57 } 58 59 // For low n, it is faster to sample directly. For both methods, 60 // performance is independent of p. On Intel Haswell CPU this method 61 // appears to be faster for approx n < 300. 62 if self.n < 300 { 63 let mut result = 0; 64 let d = Bernoulli::new(self.p); 65 for _ in 0 .. self.n { 66 result += rng.sample(d) as u32; 67 } 68 return result as u64; 69 } 70 71 // binomial distribution is symmetrical with respect to p -> 1-p, k -> n-k 72 // switch p so that it is less than 0.5 - this allows for lower expected values 73 // we will just invert the result at the end 74 let p = if self.p <= 0.5 { 75 self.p 76 } else { 77 1.0 - self.p 78 }; 79 80 // prepare some cached values 81 let float_n = self.n as f64; 82 let ln_fact_n = log_gamma(float_n + 1.0); 83 let pc = 1.0 - p; 84 let log_p = p.ln(); 85 let log_pc = pc.ln(); 86 let expected = self.n as f64 * p; 87 let sq = (expected * (2.0 * pc)).sqrt(); 88 89 let mut lresult; 90 91 // we use the Cauchy distribution as the comparison distribution 92 // f(x) ~ 1/(1+x^2) 93 let cauchy = Cauchy::new(0.0, 1.0); 94 loop { 95 let mut comp_dev: f64; 96 loop { 97 // draw from the Cauchy distribution 98 comp_dev = rng.sample(cauchy); 99 // shift the peak of the comparison ditribution 100 lresult = expected + sq * comp_dev; 101 // repeat the drawing until we are in the range of possible values 102 if lresult >= 0.0 && lresult < float_n + 1.0 { 103 break; 104 } 105 } 106 107 // the result should be discrete 108 lresult = lresult.floor(); 109 110 let log_binomial_dist = ln_fact_n - log_gamma(lresult+1.0) - 111 log_gamma(float_n - lresult + 1.0) + lresult*log_p + (float_n - lresult)*log_pc; 112 // this is the binomial probability divided by the comparison probability 113 // we will generate a uniform random value and if it is larger than this, 114 // we interpret it as a value falling out of the distribution and repeat 115 let comparison_coeff = (log_binomial_dist.exp() * sq) * (1.2 * (1.0 + comp_dev*comp_dev)); 116 117 if comparison_coeff >= rng.gen() { 118 break; 119 } 120 } 121 122 // invert the result for p < 0.5 123 if p != self.p { 124 self.n - lresult as u64 125 } else { 126 lresult as u64 127 } 128 } 129 } 130 131 #[cfg(test)] 132 mod test { 133 use Rng; 134 use distributions::Distribution; 135 use super::Binomial; 136 test_binomial_mean_and_variance<R: Rng>(n: u64, p: f64, rng: &mut R)137 fn test_binomial_mean_and_variance<R: Rng>(n: u64, p: f64, rng: &mut R) { 138 let binomial = Binomial::new(n, p); 139 140 let expected_mean = n as f64 * p; 141 let expected_variance = n as f64 * p * (1.0 - p); 142 143 let mut results = [0.0; 1000]; 144 for i in results.iter_mut() { *i = binomial.sample(rng) as f64; } 145 146 let mean = results.iter().sum::<f64>() / results.len() as f64; 147 assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0); 148 149 let variance = 150 results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>() 151 / results.len() as f64; 152 assert!((variance - expected_variance).abs() < expected_variance / 10.0); 153 } 154 155 #[test] test_binomial()156 fn test_binomial() { 157 let mut rng = ::test::rng(351); 158 test_binomial_mean_and_variance(150, 0.1, &mut rng); 159 test_binomial_mean_and_variance(70, 0.6, &mut rng); 160 test_binomial_mean_and_variance(40, 0.5, &mut rng); 161 test_binomial_mean_and_variance(20, 0.7, &mut rng); 162 test_binomial_mean_and_variance(20, 0.5, &mut rng); 163 } 164 165 #[test] test_binomial_end_points()166 fn test_binomial_end_points() { 167 let mut rng = ::test::rng(352); 168 assert_eq!(rng.sample(Binomial::new(20, 0.0)), 0); 169 assert_eq!(rng.sample(Binomial::new(20, 1.0)), 20); 170 } 171 172 #[test] 173 #[should_panic] test_binomial_invalid_lambda_neg()174 fn test_binomial_invalid_lambda_neg() { 175 Binomial::new(20, -10.0); 176 } 177 } 178