| //! The hypergeometric distribution. |
| |
| use crate::Distribution; |
| use rand::Rng; |
| use rand::distributions::uniform::Uniform; |
| use core::fmt; |
| |
| #[derive(Clone, Copy, Debug)] |
| #[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))] |
| enum SamplingMethod { |
| InverseTransform{ initial_p: f64, initial_x: i64 }, |
| RejectionAcceptance{ |
| m: f64, |
| a: f64, |
| lambda_l: f64, |
| lambda_r: f64, |
| x_l: f64, |
| x_r: f64, |
| p1: f64, |
| p2: f64, |
| p3: f64 |
| }, |
| } |
| |
| /// The hypergeometric distribution `Hypergeometric(N, K, n)`. |
| /// |
| /// This is the distribution of successes in samples of size `n` drawn without |
| /// replacement from a population of size `N` containing `K` success states. |
| /// It has the density function: |
| /// `f(k) = binomial(K, k) * binomial(N-K, n-k) / binomial(N, n)`, |
| /// where `binomial(a, b) = a! / (b! * (a - b)!)`. |
| /// |
| /// The [binomial distribution](crate::Binomial) is the analogous distribution |
| /// for sampling with replacement. It is a good approximation when the population |
| /// size is much larger than the sample size. |
| /// |
| /// # Example |
| /// |
| /// ``` |
| /// use rand_distr::{Distribution, Hypergeometric}; |
| /// |
| /// let hypergeo = Hypergeometric::new(60, 24, 7).unwrap(); |
| /// let v = hypergeo.sample(&mut rand::thread_rng()); |
| /// println!("{} is from a hypergeometric distribution", v); |
| /// ``` |
| #[derive(Copy, Clone, Debug)] |
| #[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))] |
| pub struct Hypergeometric { |
| n1: u64, |
| n2: u64, |
| k: u64, |
| offset_x: i64, |
| sign_x: i64, |
| sampling_method: SamplingMethod, |
| } |
| |
| /// Error type returned from `Hypergeometric::new`. |
| #[derive(Clone, Copy, Debug, PartialEq, Eq)] |
| pub enum Error { |
| /// `total_population_size` is too large, causing floating point underflow. |
| PopulationTooLarge, |
| /// `population_with_feature > total_population_size`. |
| ProbabilityTooLarge, |
| /// `sample_size > total_population_size`. |
| SampleSizeTooLarge, |
| } |
| |
| impl fmt::Display for Error { |
| fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| f.write_str(match self { |
| Error::PopulationTooLarge => "total_population_size is too large causing underflow in geometric distribution", |
| Error::ProbabilityTooLarge => "population_with_feature > total_population_size in geometric distribution", |
| Error::SampleSizeTooLarge => "sample_size > total_population_size in geometric distribution", |
| }) |
| } |
| } |
| |
| #[cfg(feature = "std")] |
| #[cfg_attr(doc_cfg, doc(cfg(feature = "std")))] |
| impl std::error::Error for Error {} |
| |
| // evaluate fact(numerator.0)*fact(numerator.1) / fact(denominator.0)*fact(denominator.1) |
| fn fraction_of_products_of_factorials(numerator: (u64, u64), denominator: (u64, u64)) -> f64 { |
| let min_top = u64::min(numerator.0, numerator.1); |
| let min_bottom = u64::min(denominator.0, denominator.1); |
| // the factorial of this will cancel out: |
| let min_all = u64::min(min_top, min_bottom); |
| |
| let max_top = u64::max(numerator.0, numerator.1); |
| let max_bottom = u64::max(denominator.0, denominator.1); |
| let max_all = u64::max(max_top, max_bottom); |
| |
| let mut result = 1.0; |
| for i in (min_all + 1)..=max_all { |
| if i <= min_top { |
| result *= i as f64; |
| } |
| |
| if i <= min_bottom { |
| result /= i as f64; |
| } |
| |
| if i <= max_top { |
| result *= i as f64; |
| } |
| |
| if i <= max_bottom { |
| result /= i as f64; |
| } |
| } |
| |
| result |
| } |
| |
| fn ln_of_factorial(v: f64) -> f64 { |
| // the paper calls for ln(v!), but also wants to pass in fractions, |
| // so we need to use Stirling's approximation to fill in the gaps: |
| v * v.ln() - v |
| } |
| |
| impl Hypergeometric { |
| /// Constructs a new `Hypergeometric` with the shape parameters |
| /// `N = total_population_size`, |
| /// `K = population_with_feature`, |
| /// `n = sample_size`. |
| #[allow(clippy::many_single_char_names)] // Same names as in the reference. |
| pub fn new(total_population_size: u64, population_with_feature: u64, sample_size: u64) -> Result<Self, Error> { |
| if population_with_feature > total_population_size { |
| return Err(Error::ProbabilityTooLarge); |
| } |
| |
| if sample_size > total_population_size { |
| return Err(Error::SampleSizeTooLarge); |
| } |
| |
| // set-up constants as function of original parameters |
| let n = total_population_size; |
| let (mut sign_x, mut offset_x) = (1, 0); |
| let (n1, n2) = { |
| // switch around success and failure states if necessary to ensure n1 <= n2 |
| let population_without_feature = n - population_with_feature; |
| if population_with_feature > population_without_feature { |
| sign_x = -1; |
| offset_x = sample_size as i64; |
| (population_without_feature, population_with_feature) |
| } else { |
| (population_with_feature, population_without_feature) |
| } |
| }; |
| // when sampling more than half the total population, take the smaller |
| // group as sampled instead (we can then return n1-x instead). |
| // |
| // Note: the boundary condition given in the paper is `sample_size < n / 2`; |
| // we're deviating here, because when n is even, it doesn't matter whether |
| // we switch here or not, but when n is odd `n/2 < n - n/2`, so switching |
| // when `k == n/2`, we'd actually be taking the _larger_ group as sampled. |
| let k = if sample_size <= n / 2 { |
| sample_size |
| } else { |
| offset_x += n1 as i64 * sign_x; |
| sign_x *= -1; |
| n - sample_size |
| }; |
| |
| // Algorithm H2PE has bounded runtime only if `M - max(0, k-n2) >= 10`, |
| // where `M` is the mode of the distribution. |
| // Use algorithm HIN for the remaining parameter space. |
| // |
| // Voratas Kachitvichyanukul and Bruce W. Schmeiser. 1985. Computer |
| // generation of hypergeometric random variates. |
| // J. Statist. Comput. Simul. Vol.22 (August 1985), 127-145 |
| // https://www.researchgate.net/publication/233212638 |
| const HIN_THRESHOLD: f64 = 10.0; |
| let m = ((k + 1) as f64 * (n1 + 1) as f64 / (n + 2) as f64).floor(); |
| let sampling_method = if m - f64::max(0.0, k as f64 - n2 as f64) < HIN_THRESHOLD { |
| let (initial_p, initial_x) = if k < n2 { |
| (fraction_of_products_of_factorials((n2, n - k), (n, n2 - k)), 0) |
| } else { |
| (fraction_of_products_of_factorials((n1, k), (n, k - n2)), (k - n2) as i64) |
| }; |
| |
| if initial_p <= 0.0 || !initial_p.is_finite() { |
| return Err(Error::PopulationTooLarge); |
| } |
| |
| SamplingMethod::InverseTransform { initial_p, initial_x } |
| } else { |
| let a = ln_of_factorial(m) + |
| ln_of_factorial(n1 as f64 - m) + |
| ln_of_factorial(k as f64 - m) + |
| ln_of_factorial((n2 - k) as f64 + m); |
| |
| let numerator = (n - k) as f64 * k as f64 * n1 as f64 * n2 as f64; |
| let denominator = (n - 1) as f64 * n as f64 * n as f64; |
| let d = 1.5 * (numerator / denominator).sqrt() + 0.5; |
| |
| let x_l = m - d + 0.5; |
| let x_r = m + d + 0.5; |
| |
| let k_l = f64::exp(a - |
| ln_of_factorial(x_l) - |
| ln_of_factorial(n1 as f64 - x_l) - |
| ln_of_factorial(k as f64 - x_l) - |
| ln_of_factorial((n2 - k) as f64 + x_l)); |
| let k_r = f64::exp(a - |
| ln_of_factorial(x_r - 1.0) - |
| ln_of_factorial(n1 as f64 - x_r + 1.0) - |
| ln_of_factorial(k as f64 - x_r + 1.0) - |
| ln_of_factorial((n2 - k) as f64 + x_r - 1.0)); |
| |
| let numerator = x_l * ((n2 - k) as f64 + x_l); |
| let denominator = (n1 as f64 - x_l + 1.0) * (k as f64 - x_l + 1.0); |
| let lambda_l = -((numerator / denominator).ln()); |
| |
| let numerator = (n1 as f64 - x_r + 1.0) * (k as f64 - x_r + 1.0); |
| let denominator = x_r * ((n2 - k) as f64 + x_r); |
| let lambda_r = -((numerator / denominator).ln()); |
| |
| // the paper literally gives `p2 + kL/lambdaL` where it (probably) |
| // should have been `p2 <- p1 + kL/lambdaL`; another print error?! |
| let p1 = 2.0 * d; |
| let p2 = p1 + k_l / lambda_l; |
| let p3 = p2 + k_r / lambda_r; |
| |
| SamplingMethod::RejectionAcceptance { |
| m, a, lambda_l, lambda_r, x_l, x_r, p1, p2, p3 |
| } |
| }; |
| |
| Ok(Hypergeometric { n1, n2, k, offset_x, sign_x, sampling_method }) |
| } |
| } |
| |
| impl Distribution<u64> for Hypergeometric { |
| #[allow(clippy::many_single_char_names)] // Same names as in the reference. |
| fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 { |
| use SamplingMethod::*; |
| |
| let Hypergeometric { n1, n2, k, sign_x, offset_x, sampling_method } = *self; |
| let x = match sampling_method { |
| InverseTransform { initial_p: mut p, initial_x: mut x } => { |
| let mut u = rng.gen::<f64>(); |
| while u > p && x < k as i64 { // the paper erroneously uses `until n < p`, which doesn't make any sense |
| u -= p; |
| p *= ((n1 as i64 - x as i64) * (k as i64 - x as i64)) as f64; |
| p /= ((x as i64 + 1) * (n2 as i64 - k as i64 + 1 + x as i64)) as f64; |
| x += 1; |
| } |
| x |
| }, |
| RejectionAcceptance { m, a, lambda_l, lambda_r, x_l, x_r, p1, p2, p3 } => { |
| let distr_region_select = Uniform::new(0.0, p3); |
| loop { |
| let (y, v) = loop { |
| let u = distr_region_select.sample(rng); |
| let v = rng.gen::<f64>(); // for the accept/reject decision |
| |
| if u <= p1 { |
| // Region 1, central bell |
| let y = (x_l + u).floor(); |
| break (y, v); |
| } else if u <= p2 { |
| // Region 2, left exponential tail |
| let y = (x_l + v.ln() / lambda_l).floor(); |
| if y as i64 >= i64::max(0, k as i64 - n2 as i64) { |
| let v = v * (u - p1) * lambda_l; |
| break (y, v); |
| } |
| } else { |
| // Region 3, right exponential tail |
| let y = (x_r - v.ln() / lambda_r).floor(); |
| if y as u64 <= u64::min(n1, k) { |
| let v = v * (u - p2) * lambda_r; |
| break (y, v); |
| } |
| } |
| }; |
| |
| // Step 4: Acceptance/Rejection Comparison |
| if m < 100.0 || y <= 50.0 { |
| // Step 4.1: evaluate f(y) via recursive relationship |
| let mut f = 1.0; |
| if m < y { |
| for i in (m as u64 + 1)..=(y as u64) { |
| f *= (n1 - i + 1) as f64 * (k - i + 1) as f64; |
| f /= i as f64 * (n2 - k + i) as f64; |
| } |
| } else { |
| for i in (y as u64 + 1)..=(m as u64) { |
| f *= i as f64 * (n2 - k + i) as f64; |
| f /= (n1 - i) as f64 * (k - i) as f64; |
| } |
| } |
| |
| if v <= f { break y as i64; } |
| } else { |
| // Step 4.2: Squeezing |
| let y1 = y + 1.0; |
| let ym = y - m; |
| let yn = n1 as f64 - y + 1.0; |
| let yk = k as f64 - y + 1.0; |
| let nk = n2 as f64 - k as f64 + y1; |
| let r = -ym / y1; |
| let s = ym / yn; |
| let t = ym / yk; |
| let e = -ym / nk; |
| let g = yn * yk / (y1 * nk) - 1.0; |
| let dg = if g < 0.0 { |
| 1.0 + g |
| } else { |
| 1.0 |
| }; |
| let gu = g * (1.0 + g * (-0.5 + g / 3.0)); |
| let gl = gu - g.powi(4) / (4.0 * dg); |
| let xm = m + 0.5; |
| let xn = n1 as f64 - m + 0.5; |
| let xk = k as f64 - m + 0.5; |
| let nm = n2 as f64 - k as f64 + xm; |
| let ub = xm * r * (1.0 + r * (-0.5 + r / 3.0)) + |
| xn * s * (1.0 + s * (-0.5 + s / 3.0)) + |
| xk * t * (1.0 + t * (-0.5 + t / 3.0)) + |
| nm * e * (1.0 + e * (-0.5 + e / 3.0)) + |
| y * gu - m * gl + 0.0034; |
| let av = v.ln(); |
| if av > ub { continue; } |
| let dr = if r < 0.0 { |
| xm * r.powi(4) / (1.0 + r) |
| } else { |
| xm * r.powi(4) |
| }; |
| let ds = if s < 0.0 { |
| xn * s.powi(4) / (1.0 + s) |
| } else { |
| xn * s.powi(4) |
| }; |
| let dt = if t < 0.0 { |
| xk * t.powi(4) / (1.0 + t) |
| } else { |
| xk * t.powi(4) |
| }; |
| let de = if e < 0.0 { |
| nm * e.powi(4) / (1.0 + e) |
| } else { |
| nm * e.powi(4) |
| }; |
| |
| if av < ub - 0.25*(dr + ds + dt + de) + (y + m)*(gl - gu) - 0.0078 { |
| break y as i64; |
| } |
| |
| // Step 4.3: Final Acceptance/Rejection Test |
| let av_critical = a - |
| ln_of_factorial(y) - |
| ln_of_factorial(n1 as f64 - y) - |
| ln_of_factorial(k as f64 - y) - |
| ln_of_factorial((n2 - k) as f64 + y); |
| if v.ln() <= av_critical { |
| break y as i64; |
| } |
| } |
| } |
| } |
| }; |
| |
| (offset_x + sign_x * x) as u64 |
| } |
| } |
| |
| #[cfg(test)] |
| mod test { |
| use super::*; |
| |
| #[test] |
| fn test_hypergeometric_invalid_params() { |
| assert!(Hypergeometric::new(100, 101, 5).is_err()); |
| assert!(Hypergeometric::new(100, 10, 101).is_err()); |
| assert!(Hypergeometric::new(100, 101, 101).is_err()); |
| assert!(Hypergeometric::new(100, 10, 5).is_ok()); |
| } |
| |
| fn test_hypergeometric_mean_and_variance<R: Rng>(n: u64, k: u64, s: u64, rng: &mut R) |
| { |
| let distr = Hypergeometric::new(n, k, s).unwrap(); |
| |
| let expected_mean = s as f64 * k as f64 / n as f64; |
| let expected_variance = { |
| let numerator = (s * k * (n - k) * (n - s)) as f64; |
| let denominator = (n * n * (n - 1)) as f64; |
| numerator / denominator |
| }; |
| |
| let mut results = [0.0; 1000]; |
| for i in results.iter_mut() { |
| *i = distr.sample(rng) as f64; |
| } |
| |
| let mean = results.iter().sum::<f64>() / results.len() as f64; |
| assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0); |
| |
| let variance = |
| results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>() / results.len() as f64; |
| assert!((variance - expected_variance).abs() < expected_variance / 10.0); |
| } |
| |
| #[test] |
| fn test_hypergeometric() { |
| let mut rng = crate::test::rng(737); |
| |
| // exercise algorithm HIN: |
| test_hypergeometric_mean_and_variance(500, 400, 30, &mut rng); |
| test_hypergeometric_mean_and_variance(250, 200, 230, &mut rng); |
| test_hypergeometric_mean_and_variance(100, 20, 6, &mut rng); |
| test_hypergeometric_mean_and_variance(50, 10, 47, &mut rng); |
| |
| // exercise algorithm H2PE |
| test_hypergeometric_mean_and_variance(5000, 2500, 500, &mut rng); |
| test_hypergeometric_mean_and_variance(10100, 10000, 1000, &mut rng); |
| test_hypergeometric_mean_and_variance(100100, 100, 10000, &mut rng); |
| } |
| } |
