|
| 1 | +// Copyright 2013 The Rust Project Developers. See the COPYRIGHT |
| 2 | +// file at the top-level directory of this distribution and at |
| 3 | +// http://rust-lang.org/COPYRIGHT. |
| 4 | +// |
| 5 | +// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or |
| 6 | +// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license |
| 7 | +// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your |
| 8 | +// option. This file may not be copied, modified, or distributed |
| 9 | +// except according to those terms. |
| 10 | + |
| 11 | +//! The Gamma distribution. |
| 12 | +
|
| 13 | +use rand::Rng; |
| 14 | +use super::{IndependentSample, Sample, StandardNormal, Exp}; |
| 15 | +use num; |
| 16 | + |
| 17 | +/// The Gamma distribution `Gamma(shape, scale)` distribution. |
| 18 | +/// |
| 19 | +/// The density function of this distribution is |
| 20 | +/// |
| 21 | +/// ``` |
| 22 | +/// f(x) = x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k) |
| 23 | +/// ``` |
| 24 | +/// |
| 25 | +/// where `Γ` is the Gamma function, `k` is the shape and `θ` is the |
| 26 | +/// scale and both `k` and `θ` are strictly positive. |
| 27 | +/// |
| 28 | +/// The algorithm used is that described by Marsaglia & Tsang 2000[1], |
| 29 | +/// falling back to directly sampling from an Exponential for `shape |
| 30 | +/// == 1`, and using the boosting technique described in [1] for |
| 31 | +/// `shape < 1`. |
| 32 | +/// |
| 33 | +/// # Example |
| 34 | +/// |
| 35 | +/// ```rust |
| 36 | +/// use std::rand; |
| 37 | +/// use std::rand::distributions::{IndependentSample, Gamma}; |
| 38 | +/// |
| 39 | +/// fn main() { |
| 40 | +/// let gamma = Gamma::new(2.0, 5.0); |
| 41 | +/// let v = gamma.ind_sample(rand::task_rng()); |
| 42 | +/// println!("{} is from a Gamma(2, 5) distribution", v); |
| 43 | +/// } |
| 44 | +/// ``` |
| 45 | +/// |
| 46 | +/// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method |
| 47 | +/// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3 |
| 48 | +/// (September 2000), |
| 49 | +/// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414) |
| 50 | +pub enum Gamma { |
| 51 | + priv Large(GammaLargeShape), |
| 52 | + priv One(Exp), |
| 53 | + priv Small(GammaSmallShape) |
| 54 | +} |
| 55 | + |
| 56 | +// These two helpers could be made public, but saving the |
| 57 | +// match-on-Gamma-enum branch from using them directly (e.g. if one |
| 58 | +// knows that the shape is always > 1) doesn't appear to be much |
| 59 | +// faster. |
| 60 | + |
| 61 | +/// Gamma distribution where the shape parameter is less than 1. |
| 62 | +/// |
| 63 | +/// Note, samples from this require a compulsory floating-point `pow` |
| 64 | +/// call, which makes it significantly slower than sampling from a |
| 65 | +/// gamma distribution where the shape parameter is greater than or |
| 66 | +/// equal to 1. |
| 67 | +/// |
| 68 | +/// See `Gamma` for sampling from a Gamma distribution with general |
| 69 | +/// shape parameters. |
| 70 | +struct GammaSmallShape { |
| 71 | + inv_shape: f64, |
| 72 | + large_shape: GammaLargeShape |
| 73 | +} |
| 74 | + |
| 75 | +/// Gamma distribution where the shape parameter is larger than 1. |
| 76 | +/// |
| 77 | +/// See `Gamma` for sampling from a Gamma distribution with general |
| 78 | +/// shape parameters. |
| 79 | +struct GammaLargeShape { |
| 80 | + shape: f64, |
| 81 | + scale: f64, |
| 82 | + c: f64, |
| 83 | + d: f64 |
| 84 | +} |
| 85 | + |
| 86 | +impl Gamma { |
| 87 | + /// Construct an object representing the `Gamma(shape, scale)` |
| 88 | + /// distribution. |
| 89 | + /// |
| 90 | + /// Fails if `shape <= 0` or `scale <= 0`. |
| 91 | + pub fn new(shape: f64, scale: f64) -> Gamma { |
| 92 | + assert!(shape > 0.0, "Gamma::new called with shape <= 0"); |
| 93 | + assert!(scale > 0.0, "Gamma::new called with scale <= 0"); |
| 94 | + |
| 95 | + match shape { |
| 96 | + 1.0 => One(Exp::new(1.0 / scale)), |
| 97 | + 0.0 .. 1.0 => Small(GammaSmallShape::new_raw(shape, scale)), |
| 98 | + _ => Large(GammaLargeShape::new_raw(shape, scale)) |
| 99 | + } |
| 100 | + } |
| 101 | +} |
| 102 | + |
| 103 | +impl GammaSmallShape { |
| 104 | + fn new_raw(shape: f64, scale: f64) -> GammaSmallShape { |
| 105 | + GammaSmallShape { |
| 106 | + inv_shape: 1. / shape, |
| 107 | + large_shape: GammaLargeShape::new_raw(shape + 1.0, scale) |
| 108 | + } |
| 109 | + } |
| 110 | +} |
| 111 | + |
| 112 | +impl GammaLargeShape { |
| 113 | + fn new_raw(shape: f64, scale: f64) -> GammaLargeShape { |
| 114 | + let d = shape - 1. / 3.; |
| 115 | + GammaLargeShape { |
| 116 | + shape: shape, |
| 117 | + scale: scale, |
| 118 | + c: 1. / num::sqrt(9. * d), |
| 119 | + d: d |
| 120 | + } |
| 121 | + } |
| 122 | +} |
| 123 | + |
| 124 | +impl Sample<f64> for Gamma { |
| 125 | + fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) } |
| 126 | +} |
| 127 | +impl Sample<f64> for GammaSmallShape { |
| 128 | + fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) } |
| 129 | +} |
| 130 | +impl Sample<f64> for GammaLargeShape { |
| 131 | + fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) } |
| 132 | +} |
| 133 | + |
| 134 | +impl IndependentSample<f64> for Gamma { |
| 135 | + fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 { |
| 136 | + match *self { |
| 137 | + Small(ref g) => g.ind_sample(rng), |
| 138 | + One(ref g) => g.ind_sample(rng), |
| 139 | + Large(ref g) => g.ind_sample(rng), |
| 140 | + } |
| 141 | + } |
| 142 | +} |
| 143 | +impl IndependentSample<f64> for GammaSmallShape { |
| 144 | + fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 { |
| 145 | + // Need (0, 1) here. |
| 146 | + let mut u = rng.gen::<f64>(); |
| 147 | + while u == 0. { |
| 148 | + u = rng.gen(); |
| 149 | + } |
| 150 | + |
| 151 | + self.large_shape.ind_sample(rng) * num::pow(u, self.inv_shape) |
| 152 | + } |
| 153 | +} |
| 154 | +impl IndependentSample<f64> for GammaLargeShape { |
| 155 | + fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 { |
| 156 | + loop { |
| 157 | + let x = *rng.gen::<StandardNormal>(); |
| 158 | + let v_cbrt = 1.0 + self.c * x; |
| 159 | + if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0 |
| 160 | + continue |
| 161 | + } |
| 162 | + |
| 163 | + let v = v_cbrt * v_cbrt * v_cbrt; |
| 164 | + // Need (0, 1) here, not [0, 1). This would be faster if |
| 165 | + // we were generating an f64 in (0, 1) directly. |
| 166 | + let mut u = rng.gen::<f64>(); |
| 167 | + while u == 0.0 { |
| 168 | + u = rng.gen(); |
| 169 | + } |
| 170 | + |
| 171 | + let x_sqr = x * x; |
| 172 | + if u < 1.0 - 0.0331 * x_sqr * x_sqr || |
| 173 | + num::ln(u) < 0.5 * x_sqr + self.d * (1.0 - v + num::ln(v)) { |
| 174 | + return self.d * v * self.scale |
| 175 | + } |
| 176 | + } |
| 177 | + } |
| 178 | +} |
| 179 | + |
| 180 | +#[cfg(test)] |
| 181 | +mod bench { |
| 182 | + use super::*; |
| 183 | + use mem::size_of; |
| 184 | + use rand::distributions::IndependentSample; |
| 185 | + use rand::{StdRng, RAND_BENCH_N}; |
| 186 | + use extra::test::BenchHarness; |
| 187 | + use iter::range; |
| 188 | + use option::{Some, None}; |
| 189 | + |
| 190 | + |
| 191 | + #[bench] |
| 192 | + fn bench_gamma_large_shape(bh: &mut BenchHarness) { |
| 193 | + let gamma = Gamma::new(10., 1.0); |
| 194 | + let mut rng = StdRng::new(); |
| 195 | + |
| 196 | + do bh.iter { |
| 197 | + for _ in range(0, RAND_BENCH_N) { |
| 198 | + gamma.ind_sample(&mut rng); |
| 199 | + } |
| 200 | + } |
| 201 | + bh.bytes = size_of::<f64>() as u64 * RAND_BENCH_N; |
| 202 | + } |
| 203 | + |
| 204 | + #[bench] |
| 205 | + fn bench_gamma_small_shape(bh: &mut BenchHarness) { |
| 206 | + let gamma = Gamma::new(0.1, 1.0); |
| 207 | + let mut rng = StdRng::new(); |
| 208 | + |
| 209 | + do bh.iter { |
| 210 | + for _ in range(0, RAND_BENCH_N) { |
| 211 | + gamma.ind_sample(&mut rng); |
| 212 | + } |
| 213 | + } |
| 214 | + bh.bytes = size_of::<f64>() as u64 * RAND_BENCH_N; |
| 215 | + } |
| 216 | +} |
0 commit comments