diff --git a/src/distribution/binomial.rs b/src/distribution/binomial.rs index c9794f0a..a07603b4 100644 --- a/src/distribution/binomial.rs +++ b/src/distribution/binomial.rs @@ -115,11 +115,218 @@ impl core::fmt::Display for Binomial { #[cfg(feature = "rand")] #[cfg_attr(docsrs, doc(cfg(feature = "rand")))] impl ::rand::distr::Distribution for Binomial { + /// Samples in `O(1)` expected time (independent of `n`) via the BINV + /// inversion algorithm for `n * min(p, 1 - p) < 10` and the BTPE + /// triangle-parallelogram-exponential rejection algorithm otherwise. + /// + /// Kachitvichyanukul, V. and Schmeiser, B. W. (1988). Binomial random + /// variate generation. Communications of the ACM 31(2), 216-222. + /// fn sample(&self, rng: &mut R) -> u64 { - (0..self.n).fold(0, |acc, _| { - let n: f64 = ::rand::RngExt::random(rng); - if n < self.p { acc + 1 } else { acc } - }) + sample_unchecked(rng, self.n, self.p) + } +} + +// The BINV/BTPE implementation below is adapted from the `rand_distr` crate +// (https://github.com/rust-random/rand_distr, Copyright 2018 Developers of the +// Rand project, Copyright 2016-2017 The Rust Project Developers; MIT or +// Apache-2.0), which implements Kachitvichyanukul & Schmeiser (1988) with the +// corrected Stirling-series signs from GSL. + +/// Samples from a binomial distribution with the given `n` and `p`, without +/// validating the parameters. +#[cfg(feature = "rand")] +pub fn sample_unchecked(rng: &mut R, n: u64, p: f64) -> u64 { + if p <= 0.0 || n == 0 { + return 0; + } + if p >= 1.0 { + return n; + } + + // the distribution is symmetric under p -> 1 - p + let flipped = p > 0.5; + let p = if flipped { 1.0 - p } else { p }; + let q = 1.0 - p; + + if q == 1.0 { + // p is below ~2^-54 (and not flipped, since p <= 0.5): the + // distribution is indistinguishable from Poisson(n * p) to O(p) + return super::poisson::sample_unchecked(rng, n as f64 * p) as u64; + } + + // Threshold for preferring BINV; the paper suggests 10. + let sample = if n as f64 * p < 10.0 { + binv(rng, n, p, q) + } else { + btpe(rng, n, p, q) + }; + if flipped { n - sample } else { sample } +} + +/// BINV: sequential inversion from x = 0. Expected iterations ~ n * p. +#[cfg(feature = "rand")] +fn binv(rng: &mut R, n: u64, p: f64, q: f64) -> u64 { + // BINV can get numerically stuck accumulating `u -= r`; a result beyond + // 110 is > 31 standard deviations out for n * p < 10, so restart instead + // (same guard value as GSL and rand_distr). + const BINV_MAX_X: u64 = 110; + + let s = p / q; + let a = (n as f64 + 1.0) * s; + // q^n, via ln_1p to keep accuracy for small p + let r0 = ((-p).ln_1p() * n as f64).exp(); + + 'restart: loop { + let mut r = r0; + let mut u: f64 = ::rand::RngExt::random(rng); + let mut x = 0u64; + while u > r { + u -= r; + x += 1; + if x > BINV_MAX_X { + continue 'restart; + } + r *= a / (x as f64) - s; + } + return x; + } +} + +/// BTPE: rejection from a triangle + parallelogram + two exponential tails +/// envelope around the scaled pmf. Requires `p <= 0.5` and `n * p >= 10`. +#[cfg(feature = "rand")] +#[allow(clippy::many_single_char_names)] // same names as the reference paper +fn btpe(rng: &mut R, n: u64, p: f64, q: f64) -> u64 { + use core::cmp::Ordering; + + // Below this |y - m| the pmf ratio f(y)/f(m) is evaluated directly. + const SQUEEZE_THRESHOLD: u64 = 20; + + // Step 0: constants as functions of n and p. + let n_f = n as f64; + let np = n_f * p; + let npq = np * q; + let f_m = np + p; + let m = f_m as u64; // mode + // radius (and, height being 1, area) of the triangle region + let p1 = (2.195 * npq.sqrt() - 4.6 * q).floor() + 0.5; + let x_m = m as f64 + 0.5; // tip of the triangle + let x_l = x_m - p1; // left edge of the triangle + let x_r = x_m + p1; // right edge of the triangle + let c = 0.134 + 20.5 / (15.3 + m as f64); + // p1 + area of the parallelogram region + let p2 = p1 * (1.0 + 2.0 * c); + + let lambda = |a: f64| a * (1.0 + 0.5 * a); + let lambda_l = lambda((f_m - x_l) / (f_m - x_l * p)); + let lambda_r = lambda((x_r - f_m) / (x_r * q)); + let p3 = p2 + c / lambda_l; + let p4 = p3 + c / lambda_r; + + loop { + // Step 1: select the region via u; v decides acceptance within it. + let u: f64 = ::rand::RngExt::random::(rng) * p4; + let mut v: f64 = ::rand::RngExt::random(rng); + + let y: u64; + if u <= p1 { + // triangle: accept immediately + return (x_m - p1 * v + u) as u64; + } else if u <= p2 { + // Step 2: parallelogram + let x = x_l + (u - p1) / c; + v = v * c + 1.0 - (x - x_m).abs() / p1; + if v > 1.0 { + continue; + } + y = x as u64; + } else if u <= p3 { + // Step 3: left exponential tail (v == 0 gives -inf and retries) + let y_tmp = x_l + v.ln() / lambda_l; + if y_tmp < 0.0 { + continue; + } + y = y_tmp as u64; + v *= (u - p2) * lambda_l; + } else { + // Step 4: right exponential tail (the `as` cast saturates) + let y_tmp = x_r - v.ln() / lambda_r; + if y_tmp > n_f { + continue; + } + y = y_tmp as u64; + v *= (u - p3) * lambda_r; + } + + // Step 5: acceptance/rejection comparison of v against f(y)/f(m). + let k = y.abs_diff(m); + if k <= SQUEEZE_THRESHOLD || (k as f64) >= 0.5 * npq - 1.0 { + // Step 5.1: evaluate the ratio via the recurrence from the mode. + let s = p / q; + let a = s * (n_f + 1.0); + let mut f = 1.0; + match m.cmp(&y) { + Ordering::Less => { + for i in (m + 1)..=y { + f *= a / (i as f64) - s; + } + } + Ordering::Greater => { + for i in (y + 1)..=m { + f /= a / (i as f64) - s; + } + } + Ordering::Equal => {} + } + if v <= f { + return y; + } + continue; + } + + // Step 5.2: squeeze ln(v) between quadratic bounds on ln(f(y)/f(m)). + let kf = k as f64; + let rho = (kf / npq) * ((kf * (kf / 3.0 + 0.625) + 1.0 / 6.0) / npq + 0.5); + let t = -0.5 * kf * kf / npq; + let alpha = v.ln(); + if alpha < t - rho { + return y; + } + if alpha > t + rho { + continue; + } + + // Step 5.3: final comparison against ln(f(y)/f(m)) via Stirling series. + let x1 = (y + 1) as f64; + let f1 = (m + 1) as f64; + let z = ((n - m) + 1) as f64; + let w = ((n - y) + 1) as f64; + + // 13860/166320 = 1/12, 462/166320 = 1/360, ...: the ln k! tail + let stirling = |a: f64| { + let a2 = a * a; + (13860.0 - (462.0 - (132.0 - (99.0 - 140.0 / a2) / a2) / a2) / a2) / a / 166320.0 + }; + + let y_sub_m = if y > m { + (y - m) as f64 + } else { + -((m - y) as f64) + }; + // Sign convention on the Stirling terms follows GSL (verified correct + // by one of the algorithm's original designers), not the paper. + if alpha + <= x_m * (f1 / x1).ln() + + (((n - m) as f64) + 0.5) * (z / w).ln() + + y_sub_m * (w * p / (x1 * q)).ln() + + stirling(f1) + + stirling(z) + - stirling(x1) + - stirling(w) + { + return y; + } } } @@ -567,4 +774,131 @@ mod tests { density_util::check_discrete_distribution(&create_ok(0.3, 5), 5); density_util::check_discrete_distribution(&create_ok(0.7, 10), 10); } + /// Chi-square goodness-of-fit of the BINV/BTPE sampler against the exact + /// pmf, over parameter sets covering every code path: BINV (np < 10), BTPE + /// (np >= 10), and both flipped (p > 0.5) variants. Seeds are fixed, so + /// this is deterministic; the 6-sigma acceptance threshold means a failure + /// indicates a real sampler defect, not chance. + /// + /// Requires `std` as well as `rand`: the histogram and the pooled cells are + /// sized from `n` and `p` at run time, and the crate is `no_std` without + /// `alloc`. The sampler itself is exercised without `std` by + /// [`test_sample_extreme_parameters_moments`], which needs no collections. + #[cfg(all(feature = "rand", feature = "std"))] + #[test] + fn test_sample_chi_square_goodness_of_fit() { + use ::rand::SeedableRng; + use ::rand::distr::Distribution as _; + use ::rand::rngs::StdRng; + + const SAMPLES: usize = 100_000; + for &(n, p) in &[ + (20u64, 0.3f64), // BINV + (100, 0.4), // BTPE + (1000, 0.02), // BTPE, skewed + (100, 0.93), // BINV, flipped + (2000, 0.995), // BTPE, flipped + ] { + let dist = create_ok(p, n); + let mut rng = StdRng::seed_from_u64(0x5EED + n); + + let mean = n as f64 * p; + let sd = (mean * (1.0 - p)).sqrt(); + let lo = (mean - 6.0 * sd).floor().max(0.0) as u64; + let hi = ((mean + 6.0 * sd).ceil() as u64).min(n); + + let mut counts = vec![0u64; (hi - lo + 1) as usize]; + let mut outside = 0u64; + for _ in 0..SAMPLES { + let x: u64 = dist.sample(&mut rng); + match counts.get_mut((x.max(lo) - lo) as usize) { + Some(c) if x >= lo => *c += 1, + _ => outside += 1, + } + } + // beyond 6 sigma the expected count over 1e5 draws is << 1 + assert!(outside <= 3, "n={n} p={p}: {outside} samples beyond 6 sigma"); + + // pool adjacent cells until each has expected count >= 5 + let mut cells: Vec<(f64, u64)> = Vec::new(); + let mut acc = (0.0_f64, 0u64); + for k in lo..=hi { + acc.0 += SAMPLES as f64 * dist.pmf(k); + acc.1 += counts[(k - lo) as usize]; + if acc.0 >= 5.0 { + cells.push(acc); + acc = (0.0, 0); + } + } + if let Some(last) = cells.last_mut() { + last.0 += acc.0; + last.1 += acc.1; + } + + let chi2: f64 = cells + .iter() + .map(|&(e, o)| (o as f64 - e) * (o as f64 - e) / e) + .sum(); + let df = (cells.len() - 1) as f64; + let threshold = df + 6.0 * (2.0 * df).sqrt(); + assert!( + chi2 < threshold, + "n={n} p={p}: chi2 = {chi2:.1} over {} cells (threshold {threshold:.1})", + cells.len() + ); + } + } + + /// Moment sanity for parameter ranges the chi-square test can't cover: + /// n far too large for the old O(n) Bernoulli sampler, and the + /// Poisson-limit path (q rounds to 1.0, i.e. p below ~2^-54). + #[cfg(feature = "rand")] + #[test] + fn test_sample_extreme_parameters_moments() { + use ::rand::SeedableRng; + use ::rand::distr::Distribution as _; + use ::rand::rngs::StdRng; + + // n = 1e9: a single draw used to cost a billion Bernoulli trials + let dist = create_ok(0.4, 1_000_000_000); + let mut rng = StdRng::seed_from_u64(99); + const SAMPLES: usize = 20_000; + let mean = 4.0e8_f64; + let sd = (1.0e9_f64 * 0.4 * 0.6).sqrt(); + let mut sum = 0.0; + for _ in 0..SAMPLES { + let x: u64 = dist.sample(&mut rng); + assert!((x as f64 - mean).abs() < 8.0 * sd, "sample {x} implausibly far out"); + sum += x as f64; + } + let observed_mean = sum / SAMPLES as f64; + // sample mean of 2e4 draws is within ~6 sd / sqrt(SAMPLES) of the mean + prec::assert_abs_diff_eq!(observed_mean, mean, epsilon = 6.0 * sd / (SAMPLES as f64).sqrt()); + + // Poisson-limit path: p < 2^-54 so that 1 - p rounds to 1.0 + // n * p = 0.01 + let dist = create_ok(1e-18, 10_000_000_000_000_000); + let mut rng = StdRng::seed_from_u64(7); + let mut total = 0u64; + for _ in 0..200_000 { + total += >::sample(&dist, &mut rng); + } + // total ~ Poisson(200_000 * 0.01 = 2000); 6 sigma is +-268 + assert!((total as f64 - 2000.0).abs() < 268.0, "Poisson-limit path total {total}"); + } + + /// Degenerate parameters short-circuit. + #[cfg(feature = "rand")] + #[test] + fn test_sample_degenerate() { + use ::rand::SeedableRng; + use ::rand::rngs::StdRng; + + let mut rng = StdRng::seed_from_u64(3); + for _ in 0..10 { + assert_eq!(>::sample(&create_ok(0.0, 50), &mut rng), 0); + assert_eq!(>::sample(&create_ok(1.0, 50), &mut rng), 50); + assert_eq!(>::sample(&create_ok(0.5, 0), &mut rng), 0); + } + } }