From daa02bd5700173fc41de1f6117c80e293323c1af Mon Sep 17 00:00:00 2001 From: Vinzent Steinberg Date: Tue, 19 Feb 2019 19:05:24 +0100 Subject: [PATCH 1/7] Implement the BTPE for binomial sampling This should be more efficient than our previous algorithm, see [1]. The implementation is close to the algorithm as specified in the paper, where many goto statements are used. It can still be refactored to be more readable. [1] Voratas Kachitvichyanukul and Bruce W. Schmeiser. 1988. Binomial random variate generation. Commun. ACM 31, 2 (February 1988), 216-222. http://dx.doi.org/10.1145/42372.42381 --- src/distributions/binomial.rs | 205 +++++++++++++++++++++++++++------- 1 file changed, 163 insertions(+), 42 deletions(-) diff --git a/src/distributions/binomial.rs b/src/distributions/binomial.rs index b1ad873df68..793b5b2d869 100644 --- a/src/distributions/binomial.rs +++ b/src/distributions/binomial.rs @@ -10,8 +10,7 @@ //! The binomial distribution. use Rng; -use distributions::{Distribution, Cauchy}; -use distributions::utils::log_gamma; +use distributions::Distribution; /// The binomial distribution `Binomial(n, p)`. /// @@ -56,9 +55,9 @@ impl Distribution for Binomial { return self.n; } - // binomial distribution is symmetrical with respect to p -> 1-p, k -> n-k - // switch p so that it is less than 0.5 - this allows for lower expected values - // we will just invert the result at the end + // The binomial distribution is symmetrical with respect to p -> 1-p, + // k -> n-k switch p so that it is less than 0.5 - this allows for lower + // expected values we will just invert the result at the end let p = if self.p <= 0.5 { self.p } else { @@ -68,12 +67,20 @@ impl Distribution for Binomial { let result; // For small n * min(p, 1 - p), the BINV algorithm based on the inverse - // transformation of the binomial distribution is more efficient: + // transformation of the binomial distribution is efficient. Otherwise, + // the BTPE algorithm is used. // // Voratas Kachitvichyanukul and Bruce W. Schmeiser. 1988. Binomial // random variate generation. Commun. ACM 31, 2 (February 1988), // 216-222. http://dx.doi.org/10.1145/42372.42381 - if (self.n as f64) * p < 10. && self.n <= (::std::i32::MAX as u64) { + + // Threshold for prefering the BINV algorithm. The paper suggests 10, + // Ranlib uses 30, and GSL uses 14. + const BINV_THRESHOLD: f64 = 10.; + + if (self.n as f64) * p < BINV_THRESHOLD && + self.n <= (::std::i32::MAX as u64) { + // Use the BINV algorithm. let q = 1. - p; let s = p / q; let a = ((self.n + 1) as f64) * s; @@ -87,52 +94,166 @@ impl Distribution for Binomial { } result = x; } else { - // FIXME: Using the BTPE algorithm is probably faster. - - // prepare some cached values - let float_n = self.n as f64; - let ln_fact_n = log_gamma(float_n + 1.0); - let pc = 1.0 - p; - let log_p = p.ln(); - let log_pc = pc.ln(); - let expected = self.n as f64 * p; - let sq = (expected * (2.0 * pc)).sqrt(); - let mut lresult; - - // we use the Cauchy distribution as the comparison distribution - // f(x) ~ 1/(1+x^2) - let cauchy = Cauchy::new(0.0, 1.0); - loop { - let mut comp_dev: f64; + // Use the BTPE algorithm. + + // Threshold for using the squeeze algorithm. This can be freely + // chosen based on performance. Ranlib and GSL use 20. + const SQUEEZE_THRESHOLD: i64 = 20; + + // Step 0: Calculate constants as functions of `n` and `p`. + let n = self.n as f64; + let r = p; + let q = 1. - r; + let f_m = n*r + r; + let m = f_m as i64; + // radius of triangle region, since height=1 also area of region + let p1 = (2.195 * (n*r*q).sqrt() - 4.6 * q).floor() + 0.5; + // tip of triangle + let x_m = (m as f64) + 0.5; + // left edge of triangle + let x_l = x_m - p1; + // right edge of triangle + let x_r = x_m + p1; + let c = 0.134 + 20.5 / (15.3 + (m as f64)); + // p1 + area of parallelogram region + let p2 = p1 * (1. + 2. * c); + + fn lambda(a: f64) -> f64 { + a * (1. + 0.5 * a) + } + + let lambda_l = lambda((f_m - x_l) / (f_m - x_l * r)); + let lambda_r = lambda((x_r - f_m) / (x_r * q)); + // p1 + area of left tail + let p3 = p2 + c / lambda_l; + // p1 + area of right tail + let p4 = p3 + c / lambda_r; + + // return value + let mut y: i64; + + 'outer: loop { + // Step 1: Generate `u` for selecting the region. If region 1 is + // selected, generate a triangularly distributed variate. + let u = rng.gen_range(0., p4); + let mut v = rng.gen_range(0., 1.); + if !(u > p1) { + y = (x_m - p1 * v + u) as i64; + break; + } + loop { - // draw from the Cauchy distribution - comp_dev = rng.sample(cauchy); - // shift the peak of the comparison ditribution - lresult = expected + sq * comp_dev; - // repeat the drawing until we are in the range of possible values - if lresult >= 0.0 && lresult < float_n + 1.0 { - break; + // Step 2: Region 2, parallelograms. Check if region 2 is used. If + // so, generate `y`. + if !(u > p2) { + let x = x_l + (u - p1) / c; + v = v * c + 1.0 - (x - x_m).abs() / p1; + if v > 1. { + continue 'outer; + } else { + y = x as i64; + break; + } + } + + // Step 3: Region 3, left exponential tail. + if !(u > p3) { + y = (x_l + v.ln() / lambda_l) as i64; + if y < 0 { + continue 'outer; + } else { + v *= (u - p2) * lambda_l; + break; + } } + + // Step 4: Region 4, right exponential tail. + y = (x_r - v.ln() / lambda_r) as i64; + if y > 0 && (y as u64) > self.n { + continue 'outer; + } else { + v *= (u - p3) * lambda_r; + } + + break; } - // the result should be discrete - lresult = lresult.floor(); - let log_binomial_dist = ln_fact_n - log_gamma(lresult+1.0) - - log_gamma(float_n - lresult + 1.0) + lresult*log_p + (float_n - lresult)*log_pc; - // this is the binomial probability divided by the comparison probability - // we will generate a uniform random value and if it is larger than this, - // we interpret it as a value falling out of the distribution and repeat - let comparison_coeff = (log_binomial_dist.exp() * sq) * (1.2 * (1.0 + comp_dev*comp_dev)); + // Step 5: Acceptance/rejection comparison. + + // Step 5.0: Test for appropriate method of evaluating f(y). + let k = (y - m).abs(); + if !(k > SQUEEZE_THRESHOLD && (k as f64) < 0.5 * n*r*q - 1.) { + // Step 5.1: Evaluate f(y) via the recursive relationship. Start the + // search from the mode. + let s = r / q; + let a = s * (n + 1.); + let mut f = 1.0; + if m < y { + let mut i = m; + loop { + i += 1; + f *= a / (i as f64) - s; + if i == y { + break; + } + } + } else if m > y { + let mut i = y; + loop { + i += 1; + f /= a / (i as f64) - s; + if i == m { + break; + } + } + } + if v > f { + continue; + } else { + break; + } + } - if comparison_coeff >= rng.gen() { + // Step 5.2: Squeezing. Check the value of ln(v) againts upper and + // lower bound of ln(f(y)). + let k = k as f64; + let rho = (k / (n*r*q)) * ((k * (k / 3. + 0.625) + 1./6.) / (n*r*q) + 0.5); + let t = -0.5 * k*k / (n*r*q); + let alpha = v.ln(); + if alpha < t - rho { break; } + if alpha > t + rho { + continue; + } + + // Step 5.3: Final acceptance/rejection test. + let x1 = (y + 1) as f64; + let f1 = (m + 1) as f64; + let z = ((n as i64) + 1 - m) as f64; + let w = ((n as i64) - y + 1) as f64; + + fn stirling(a: f64) -> f64 { + let a2 = a * a; + (13860. - (462. - (132. - (99. - 140. / a2) / a2) / a2) / a2) / a / 166320. + } + + if alpha > x_m * (f1 / x1).ln() + + (n - (m as f64) + 0.5) * (z / w).ln() + + ((y - m) as f64) * (w * r / (x1 * q)).ln() + + stirling(f1) + stirling(z) + stirling(x1) + stirling(w) + { + continue; + } + + break; } - result = lresult as u64; + assert!(y >= 0); + result = y as u64; } - // invert the result for p < 0.5 + // Invert the result for p < 0.5. if p != self.p { self.n - result } else { From 4333a770ff1576344da40786c593b2b781225079 Mon Sep 17 00:00:00 2001 From: Vinzent Steinberg Date: Thu, 21 Feb 2019 15:03:12 +0100 Subject: [PATCH 2/7] Binomial test: Better output on error --- src/distributions/binomial.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/distributions/binomial.rs b/src/distributions/binomial.rs index 793b5b2d869..70beb7bea07 100644 --- a/src/distributions/binomial.rs +++ b/src/distributions/binomial.rs @@ -278,12 +278,14 @@ mod test { for i in results.iter_mut() { *i = binomial.sample(rng) as f64; } let mean = results.iter().sum::() / results.len() as f64; - assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0); + assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0, + "mean: {}, expected_mean: {}", mean, expected_mean); let variance = results.iter().map(|x| (x - mean) * (x - mean)).sum::() / results.len() as f64; - assert!((variance - expected_variance).abs() < expected_variance / 10.0); + assert!((variance - expected_variance).abs() < expected_variance / 10.0, + "variance: {}, expected_variance: {}", variance, expected_variance); } #[test] From cd822a0bc742da50d833009da53a1a917c90ba0d Mon Sep 17 00:00:00 2001 From: Vinzent Steinberg Date: Thu, 21 Feb 2019 15:49:55 +0100 Subject: [PATCH 3/7] Binomial: Simplify code --- src/distributions/binomial.rs | 65 +++++++++++++++-------------------- 1 file changed, 27 insertions(+), 38 deletions(-) diff --git a/src/distributions/binomial.rs b/src/distributions/binomial.rs index 70beb7bea07..193d515691a 100644 --- a/src/distributions/binomial.rs +++ b/src/distributions/binomial.rs @@ -65,6 +65,7 @@ impl Distribution for Binomial { }; let result; + let q = 1. - p; // For small n * min(p, 1 - p), the BINV algorithm based on the inverse // transformation of the binomial distribution is efficient. Otherwise, @@ -81,7 +82,6 @@ impl Distribution for Binomial { if (self.n as f64) * p < BINV_THRESHOLD && self.n <= (::std::i32::MAX as u64) { // Use the BINV algorithm. - let q = 1. - p; let s = p / q; let a = ((self.n + 1) as f64) * s; let mut r = q.powi(self.n as i32); @@ -102,12 +102,10 @@ impl Distribution for Binomial { // Step 0: Calculate constants as functions of `n` and `p`. let n = self.n as f64; - let r = p; - let q = 1. - r; - let f_m = n*r + r; + let f_m = n*p + p; let m = f_m as i64; // radius of triangle region, since height=1 also area of region - let p1 = (2.195 * (n*r*q).sqrt() - 4.6 * q).floor() + 0.5; + let p1 = (2.195 * (n*p*q).sqrt() - 4.6 * q).floor() + 0.5; // tip of triangle let x_m = (m as f64) + 0.5; // left edge of triangle @@ -122,7 +120,7 @@ impl Distribution for Binomial { a * (1. + 0.5 * a) } - let lambda_l = lambda((f_m - x_l) / (f_m - x_l * r)); + let lambda_l = lambda((f_m - x_l) / (f_m - x_l * p)); let lambda_r = lambda((x_r - f_m) / (x_r * q)); // p1 + area of left tail let p3 = p2 + c / lambda_l; @@ -132,7 +130,7 @@ impl Distribution for Binomial { // return value let mut y: i64; - 'outer: loop { + loop { // Step 1: Generate `u` for selecting the region. If region 1 is // selected, generate a triangularly distributed variate. let u = rng.gen_range(0., p4); @@ -142,51 +140,42 @@ impl Distribution for Binomial { break; } - loop { - // Step 2: Region 2, parallelograms. Check if region 2 is used. If - // so, generate `y`. - if !(u > p2) { - let x = x_l + (u - p1) / c; - v = v * c + 1.0 - (x - x_m).abs() / p1; - if v > 1. { - continue 'outer; - } else { - y = x as i64; - break; - } + if !(u > p2) { + // Step 2: Region 2, parallelograms. Check if region 2 is + // used. If so, generate `y`. + let x = x_l + (u - p1) / c; + v = v * c + 1.0 - (x - x_m).abs() / p1; + if v > 1. { + continue; + } else { + y = x as i64; } - + } else if !(u > p3) { // Step 3: Region 3, left exponential tail. - if !(u > p3) { - y = (x_l + v.ln() / lambda_l) as i64; - if y < 0 { - continue 'outer; - } else { - v *= (u - p2) * lambda_l; - break; - } + y = (x_l + v.ln() / lambda_l) as i64; + if y < 0 { + continue; + } else { + v *= (u - p2) * lambda_l; } - + } else { // Step 4: Region 4, right exponential tail. y = (x_r - v.ln() / lambda_r) as i64; if y > 0 && (y as u64) > self.n { - continue 'outer; + continue; } else { v *= (u - p3) * lambda_r; } - - break; } - // Step 5: Acceptance/rejection comparison. // Step 5.0: Test for appropriate method of evaluating f(y). let k = (y - m).abs(); - if !(k > SQUEEZE_THRESHOLD && (k as f64) < 0.5 * n*r*q - 1.) { + if !(k > SQUEEZE_THRESHOLD && (k as f64) < 0.5 * n*p*q - 1.) { // Step 5.1: Evaluate f(y) via the recursive relationship. Start the // search from the mode. - let s = r / q; + let s = p / q; let a = s * (n + 1.); let mut f = 1.0; if m < y { @@ -218,8 +207,8 @@ impl Distribution for Binomial { // Step 5.2: Squeezing. Check the value of ln(v) againts upper and // lower bound of ln(f(y)). let k = k as f64; - let rho = (k / (n*r*q)) * ((k * (k / 3. + 0.625) + 1./6.) / (n*r*q) + 0.5); - let t = -0.5 * k*k / (n*r*q); + let rho = (k / (n*p*q)) * ((k * (k / 3. + 0.625) + 1./6.) / (n*p*q) + 0.5); + let t = -0.5 * k*k / (n*p*q); let alpha = v.ln(); if alpha < t - rho { break; @@ -241,7 +230,7 @@ impl Distribution for Binomial { if alpha > x_m * (f1 / x1).ln() + (n - (m as f64) + 0.5) * (z / w).ln() - + ((y - m) as f64) * (w * r / (x1 * q)).ln() + + ((y - m) as f64) * (w * p / (x1 * q)).ln() + stirling(f1) + stirling(z) + stirling(x1) + stirling(w) { continue; From f9cae577f588aa42c52e547c9c3ca036d76df2e5 Mon Sep 17 00:00:00 2001 From: Vinzent Steinberg Date: Thu, 21 Feb 2019 15:53:05 +0100 Subject: [PATCH 4/7] Binomial: Eliminate some common subexpressions --- src/distributions/binomial.rs | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/distributions/binomial.rs b/src/distributions/binomial.rs index 193d515691a..348cf07fe90 100644 --- a/src/distributions/binomial.rs +++ b/src/distributions/binomial.rs @@ -102,10 +102,12 @@ impl Distribution for Binomial { // Step 0: Calculate constants as functions of `n` and `p`. let n = self.n as f64; - let f_m = n*p + p; + let np = n * p; + let npq = np * q; + let f_m = np + p; let m = f_m as i64; // radius of triangle region, since height=1 also area of region - let p1 = (2.195 * (n*p*q).sqrt() - 4.6 * q).floor() + 0.5; + let p1 = (2.195 * npq.sqrt() - 4.6 * q).floor() + 0.5; // tip of triangle let x_m = (m as f64) + 0.5; // left edge of triangle @@ -172,7 +174,7 @@ impl Distribution for Binomial { // Step 5.0: Test for appropriate method of evaluating f(y). let k = (y - m).abs(); - if !(k > SQUEEZE_THRESHOLD && (k as f64) < 0.5 * n*p*q - 1.) { + if !(k > SQUEEZE_THRESHOLD && (k as f64) < 0.5 * npq - 1.) { // Step 5.1: Evaluate f(y) via the recursive relationship. Start the // search from the mode. let s = p / q; @@ -207,8 +209,8 @@ impl Distribution for Binomial { // Step 5.2: Squeezing. Check the value of ln(v) againts upper and // lower bound of ln(f(y)). let k = k as f64; - let rho = (k / (n*p*q)) * ((k * (k / 3. + 0.625) + 1./6.) / (n*p*q) + 0.5); - let t = -0.5 * k*k / (n*p*q); + let rho = (k / npq) * ((k * (k / 3. + 0.625) + 1./6.) / npq + 0.5); + let t = -0.5 * k*k / npq; let alpha = v.ln(); if alpha < t - rho { break; From f6d740712fe1fc79c8d49cf17030a28ebee2aece Mon Sep 17 00:00:00 2001 From: Vinzent Steinberg Date: Thu, 21 Feb 2019 16:08:15 +0100 Subject: [PATCH 5/7] Binomial: Use correct sign for stirling correction This was noted by the GSL implementors. The new signs were confirmed by one of the authors publishing the original algorithm. --- src/distributions/binomial.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/distributions/binomial.rs b/src/distributions/binomial.rs index 348cf07fe90..794fd640493 100644 --- a/src/distributions/binomial.rs +++ b/src/distributions/binomial.rs @@ -233,7 +233,7 @@ impl Distribution for Binomial { if alpha > x_m * (f1 / x1).ln() + (n - (m as f64) + 0.5) * (z / w).ln() + ((y - m) as f64) * (w * p / (x1 * q)).ln() - + stirling(f1) + stirling(z) + stirling(x1) + stirling(w) + + stirling(f1) + stirling(z) - stirling(x1) - stirling(w) { continue; } From 6e8202355da9744a8699fc8a03c1c3029c00a2b0 Mon Sep 17 00:00:00 2001 From: Vinzent Steinberg Date: Fri, 29 Mar 2019 14:27:26 +0100 Subject: [PATCH 6/7] Binomial: Comment on difference to original algorithm --- src/distributions/binomial.rs | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/src/distributions/binomial.rs b/src/distributions/binomial.rs index 794fd640493..50da9c414f4 100644 --- a/src/distributions/binomial.rs +++ b/src/distributions/binomial.rs @@ -46,6 +46,13 @@ impl Binomial { } } +/// Convert a `f64` to an `i64`, panicing on overflow. +// In the future (Rust 1.34), this might be replaced with `TryFrom`. +fn f64_to_i64(x: f64) -> i64 { + assert!(x < (::std::i64::MAX as f64)); + x as i64 +} + impl Distribution for Binomial { fn sample(&self, rng: &mut R) -> u64 { // Handle these values directly. @@ -105,7 +112,7 @@ impl Distribution for Binomial { let np = n * p; let npq = np * q; let f_m = np + p; - let m = f_m as i64; + let m = f64_to_i64(f_m); // radius of triangle region, since height=1 also area of region let p1 = (2.195 * npq.sqrt() - 4.6 * q).floor() + 0.5; // tip of triangle @@ -138,7 +145,7 @@ impl Distribution for Binomial { let u = rng.gen_range(0., p4); let mut v = rng.gen_range(0., 1.); if !(u > p1) { - y = (x_m - p1 * v + u) as i64; + y = f64_to_i64(x_m - p1 * v + u); break; } @@ -150,11 +157,11 @@ impl Distribution for Binomial { if v > 1. { continue; } else { - y = x as i64; + y = f64_to_i64(x); } } else if !(u > p3) { // Step 3: Region 3, left exponential tail. - y = (x_l + v.ln() / lambda_l) as i64; + y = f64_to_i64(x_l + v.ln() / lambda_l); if y < 0 { continue; } else { @@ -162,7 +169,7 @@ impl Distribution for Binomial { } } else { // Step 4: Region 4, right exponential tail. - y = (x_r - v.ln() / lambda_r) as i64; + y = f64_to_i64(x_r - v.ln() / lambda_r); if y > 0 && (y as u64) > self.n { continue; } else { @@ -222,8 +229,8 @@ impl Distribution for Binomial { // Step 5.3: Final acceptance/rejection test. let x1 = (y + 1) as f64; let f1 = (m + 1) as f64; - let z = ((n as i64) + 1 - m) as f64; - let w = ((n as i64) - y + 1) as f64; + let z = (f64_to_i64(n) + 1 - m) as f64; + let w = (f64_to_i64(n) - y + 1) as f64; fn stirling(a: f64) -> f64 { let a2 = a * a; @@ -233,6 +240,11 @@ impl Distribution for Binomial { if alpha > x_m * (f1 / x1).ln() + (n - (m as f64) + 0.5) * (z / w).ln() + ((y - m) as f64) * (w * p / (x1 * q)).ln() + // We use the signs from the GSL implementation, which are + // different than the ones in the reference. According to + // the GSL authors, the new signs were verified to be + // correct by one of the original designers of the + // algorithm. + stirling(f1) + stirling(z) - stirling(x1) - stirling(w) { continue; From f8149abc2cbd66dbb10f9b7f8c762a638f9f858e Mon Sep 17 00:00:00 2001 From: Vinzent Steinberg Date: Fri, 29 Mar 2019 15:50:08 +0100 Subject: [PATCH 7/7] Binomial: Move distribution initialization out of loop --- src/distributions/binomial.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/distributions/binomial.rs b/src/distributions/binomial.rs index 50da9c414f4..e832c91c5ca 100644 --- a/src/distributions/binomial.rs +++ b/src/distributions/binomial.rs @@ -10,7 +10,7 @@ //! The binomial distribution. use Rng; -use distributions::Distribution; +use distributions::{Distribution, Uniform}; /// The binomial distribution `Binomial(n, p)`. /// @@ -139,11 +139,14 @@ impl Distribution for Binomial { // return value let mut y: i64; + let gen_u = Uniform::new(0., p4); + let gen_v = Uniform::new(0., 1.); + loop { // Step 1: Generate `u` for selecting the region. If region 1 is // selected, generate a triangularly distributed variate. - let u = rng.gen_range(0., p4); - let mut v = rng.gen_range(0., 1.); + let u = gen_u.sample(rng); + let mut v = gen_v.sample(rng); if !(u > p1) { y = f64_to_i64(x_m - p1 * v + u); break;