diff --git a/src/libstd/rand/distributions/gamma.rs b/src/libstd/rand/distributions/gamma.rs new file mode 100644 index 0000000000000..735b083df75f9 --- /dev/null +++ b/src/libstd/rand/distributions/gamma.rs @@ -0,0 +1,216 @@ +// Copyright 2013 The Rust Project Developers. See the COPYRIGHT +// file at the top-level directory of this distribution and at +// http://rust-lang.org/COPYRIGHT. +// +// Licensed under the Apache License, Version 2.0 or the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + +//! The Gamma distribution. + +use rand::Rng; +use super::{IndependentSample, Sample, StandardNormal, Exp}; +use num; + +/// The Gamma distribution `Gamma(shape, scale)` distribution. +/// +/// The density function of this distribution is +/// +/// ``` +/// f(x) = x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k) +/// ``` +/// +/// where `Γ` is the Gamma function, `k` is the shape and `θ` is the +/// scale and both `k` and `θ` are strictly positive. +/// +/// The algorithm used is that described by Marsaglia & Tsang 2000[1], +/// falling back to directly sampling from an Exponential for `shape +/// == 1`, and using the boosting technique described in [1] for +/// `shape < 1`. +/// +/// # Example +/// +/// ```rust +/// use std::rand; +/// use std::rand::distributions::{IndependentSample, Gamma}; +/// +/// fn main() { +/// let gamma = Gamma::new(2.0, 5.0); +/// let v = gamma.ind_sample(rand::task_rng()); +/// println!("{} is from a Gamma(2, 5) distribution", v); +/// } +/// ``` +/// +/// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method +/// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3 +/// (September 2000), +/// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414) +pub enum Gamma { + priv Large(GammaLargeShape), + priv One(Exp), + priv Small(GammaSmallShape) +} + +// These two helpers could be made public, but saving the +// match-on-Gamma-enum branch from using them directly (e.g. if one +// knows that the shape is always > 1) doesn't appear to be much +// faster. + +/// Gamma distribution where the shape parameter is less than 1. +/// +/// Note, samples from this require a compulsory floating-point `pow` +/// call, which makes it significantly slower than sampling from a +/// gamma distribution where the shape parameter is greater than or +/// equal to 1. +/// +/// See `Gamma` for sampling from a Gamma distribution with general +/// shape parameters. +struct GammaSmallShape { + inv_shape: f64, + large_shape: GammaLargeShape +} + +/// Gamma distribution where the shape parameter is larger than 1. +/// +/// See `Gamma` for sampling from a Gamma distribution with general +/// shape parameters. +struct GammaLargeShape { + shape: f64, + scale: f64, + c: f64, + d: f64 +} + +impl Gamma { + /// Construct an object representing the `Gamma(shape, scale)` + /// distribution. + /// + /// Fails if `shape <= 0` or `scale <= 0`. + pub fn new(shape: f64, scale: f64) -> Gamma { + assert!(shape > 0.0, "Gamma::new called with shape <= 0"); + assert!(scale > 0.0, "Gamma::new called with scale <= 0"); + + match shape { + 1.0 => One(Exp::new(1.0 / scale)), + 0.0 .. 1.0 => Small(GammaSmallShape::new_raw(shape, scale)), + _ => Large(GammaLargeShape::new_raw(shape, scale)) + } + } +} + +impl GammaSmallShape { + fn new_raw(shape: f64, scale: f64) -> GammaSmallShape { + GammaSmallShape { + inv_shape: 1. / shape, + large_shape: GammaLargeShape::new_raw(shape + 1.0, scale) + } + } +} + +impl GammaLargeShape { + fn new_raw(shape: f64, scale: f64) -> GammaLargeShape { + let d = shape - 1. / 3.; + GammaLargeShape { + shape: shape, + scale: scale, + c: 1. / num::sqrt(9. * d), + d: d + } + } +} + +impl Sample for Gamma { + fn sample(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) } +} +impl Sample for GammaSmallShape { + fn sample(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) } +} +impl Sample for GammaLargeShape { + fn sample(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) } +} + +impl IndependentSample for Gamma { + fn ind_sample(&self, rng: &mut R) -> f64 { + match *self { + Small(ref g) => g.ind_sample(rng), + One(ref g) => g.ind_sample(rng), + Large(ref g) => g.ind_sample(rng), + } + } +} +impl IndependentSample for GammaSmallShape { + fn ind_sample(&self, rng: &mut R) -> f64 { + // Need (0, 1) here. + let mut u = rng.gen::(); + while u == 0. { + u = rng.gen(); + } + + self.large_shape.ind_sample(rng) * num::pow(u, self.inv_shape) + } +} +impl IndependentSample for GammaLargeShape { + fn ind_sample(&self, rng: &mut R) -> f64 { + loop { + let x = *rng.gen::(); + let v_cbrt = 1.0 + self.c * x; + if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0 + continue + } + + let v = v_cbrt * v_cbrt * v_cbrt; + // Need (0, 1) here, not [0, 1). This would be faster if + // we were generating an f64 in (0, 1) directly. + let mut u = rng.gen::(); + while u == 0.0 { + u = rng.gen(); + } + + let x_sqr = x * x; + if u < 1.0 - 0.0331 * x_sqr * x_sqr || + num::ln(u) < 0.5 * x_sqr + self.d * (1.0 - v + num::ln(v)) { + return self.d * v * self.scale + } + } + } +} + +#[cfg(test)] +mod bench { + use super::*; + use mem::size_of; + use rand::distributions::IndependentSample; + use rand::{StdRng, RAND_BENCH_N}; + use extra::test::BenchHarness; + use iter::range; + use option::{Some, None}; + + + #[bench] + fn bench_gamma_large_shape(bh: &mut BenchHarness) { + let gamma = Gamma::new(10., 1.0); + let mut rng = StdRng::new(); + + do bh.iter { + for _ in range(0, RAND_BENCH_N) { + gamma.ind_sample(&mut rng); + } + } + bh.bytes = size_of::() as u64 * RAND_BENCH_N; + } + + #[bench] + fn bench_gamma_small_shape(bh: &mut BenchHarness) { + let gamma = Gamma::new(0.1, 1.0); + let mut rng = StdRng::new(); + + do bh.iter { + for _ in range(0, RAND_BENCH_N) { + gamma.ind_sample(&mut rng); + } + } + bh.bytes = size_of::() as u64 * RAND_BENCH_N; + } +} diff --git a/src/libstd/rand/distributions.rs b/src/libstd/rand/distributions/mod.rs similarity index 98% rename from src/libstd/rand/distributions.rs rename to src/libstd/rand/distributions/mod.rs index a6fa1a2bd875f..c832438647017 100644 --- a/src/libstd/rand/distributions.rs +++ b/src/libstd/rand/distributions/mod.rs @@ -27,8 +27,10 @@ use rand::{Rng,Rand}; use clone::Clone; pub use self::range::Range; +pub use self::gamma::Gamma; pub mod range; +pub mod gamma; /// Types that can be used to create a random instance of `Support`. pub trait Sample { @@ -554,25 +556,23 @@ mod tests { #[cfg(test)] mod bench { use extra::test::BenchHarness; - use rand::*; + use rand::{XorShiftRng, RAND_BENCH_N}; use super::*; use iter::range; use option::{Some, None}; use mem::size_of; - static N: u64 = 100; - #[bench] fn rand_normal(bh: &mut BenchHarness) { let mut rng = XorShiftRng::new(); let mut normal = Normal::new(-2.71828, 3.14159); do bh.iter { - for _ in range(0, N) { + for _ in range(0, RAND_BENCH_N) { normal.sample(&mut rng); } } - bh.bytes = size_of::() as u64 * N; + bh.bytes = size_of::() as u64 * RAND_BENCH_N; } #[bench] fn rand_exp(bh: &mut BenchHarness) { @@ -580,10 +580,10 @@ mod bench { let mut exp = Exp::new(2.71828 * 3.14159); do bh.iter { - for _ in range(0, N) { + for _ in range(0, RAND_BENCH_N) { exp.sample(&mut rng); } } - bh.bytes = size_of::() as u64 * N; + bh.bytes = size_of::() as u64 * RAND_BENCH_N; } } diff --git a/src/libstd/rand/range.rs b/src/libstd/rand/distributions/range.rs similarity index 100% rename from src/libstd/rand/range.rs rename to src/libstd/rand/distributions/range.rs diff --git a/src/libstd/rand/ziggurat_tables.rs b/src/libstd/rand/distributions/ziggurat_tables.rs similarity index 100% rename from src/libstd/rand/ziggurat_tables.rs rename to src/libstd/rand/distributions/ziggurat_tables.rs diff --git a/src/libstd/rand/mod.rs b/src/libstd/rand/mod.rs index a6efdfc66db2e..e320898175180 100644 --- a/src/libstd/rand/mod.rs +++ b/src/libstd/rand/mod.rs @@ -833,58 +833,58 @@ mod test { } } +static RAND_BENCH_N: u64 = 100; + #[cfg(test)] mod bench { use extra::test::BenchHarness; - use rand::*; + use rand::{XorShiftRng, StdRng, IsaacRng, Isaac64Rng, Rng, RAND_BENCH_N}; use mem::size_of; use iter::range; use option::{Some, None}; - static N: u64 = 100; - #[bench] fn rand_xorshift(bh: &mut BenchHarness) { let mut rng = XorShiftRng::new(); do bh.iter { - for _ in range(0, N) { + for _ in range(0, RAND_BENCH_N) { rng.gen::(); } } - bh.bytes = size_of::() as u64 * N; + bh.bytes = size_of::() as u64 * RAND_BENCH_N; } #[bench] fn rand_isaac(bh: &mut BenchHarness) { let mut rng = IsaacRng::new(); do bh.iter { - for _ in range(0, N) { + for _ in range(0, RAND_BENCH_N) { rng.gen::(); } } - bh.bytes = size_of::() as u64 * N; + bh.bytes = size_of::() as u64 * RAND_BENCH_N; } #[bench] fn rand_isaac64(bh: &mut BenchHarness) { let mut rng = Isaac64Rng::new(); do bh.iter { - for _ in range(0, N) { + for _ in range(0, RAND_BENCH_N) { rng.gen::(); } } - bh.bytes = size_of::() as u64 * N; + bh.bytes = size_of::() as u64 * RAND_BENCH_N; } #[bench] fn rand_std(bh: &mut BenchHarness) { let mut rng = StdRng::new(); do bh.iter { - for _ in range(0, N) { + for _ in range(0, RAND_BENCH_N) { rng.gen::(); } } - bh.bytes = size_of::() as u64 * N; + bh.bytes = size_of::() as u64 * RAND_BENCH_N; } #[bench]