Skip to content

Commit c15038d

Browse files
committed
auto merge of #10223 : huonw/rust/gamma, r=cmr
Implements the [Gamma distribution](https://en.wikipedia.org/wiki/Gamma_distribution), using the algorithm described by Marsaglia & Tsang 2000[1]. I added tests checking that the mean and variance of this implementation is as expected for a range of values of the parameters in huonw/random-tests@5d87c00 (they pass locally, but obviously won't even build on Travis until this is merged). Also, moves `std::rand::distributions` to a subfolder, and performs a minor clean-up of the benchmarking (makes the number of iterations shared by the whole `std::rand` subtree). [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).
2 parents 894c1f6 + 3c25baa commit c15038d

File tree

5 files changed

+234
-18
lines changed

5 files changed

+234
-18
lines changed
+216
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,216 @@
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+
}

src/libstd/rand/distributions.rs renamed to src/libstd/rand/distributions/mod.rs

+7-7
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,10 @@ use rand::{Rng,Rand};
2727
use clone::Clone;
2828

2929
pub use self::range::Range;
30+
pub use self::gamma::Gamma;
3031

3132
pub mod range;
33+
pub mod gamma;
3234

3335
/// Types that can be used to create a random instance of `Support`.
3436
pub trait Sample<Support> {
@@ -554,36 +556,34 @@ mod tests {
554556
#[cfg(test)]
555557
mod bench {
556558
use extra::test::BenchHarness;
557-
use rand::*;
559+
use rand::{XorShiftRng, RAND_BENCH_N};
558560
use super::*;
559561
use iter::range;
560562
use option::{Some, None};
561563
use mem::size_of;
562564

563-
static N: u64 = 100;
564-
565565
#[bench]
566566
fn rand_normal(bh: &mut BenchHarness) {
567567
let mut rng = XorShiftRng::new();
568568
let mut normal = Normal::new(-2.71828, 3.14159);
569569

570570
do bh.iter {
571-
for _ in range(0, N) {
571+
for _ in range(0, RAND_BENCH_N) {
572572
normal.sample(&mut rng);
573573
}
574574
}
575-
bh.bytes = size_of::<f64>() as u64 * N;
575+
bh.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
576576
}
577577
#[bench]
578578
fn rand_exp(bh: &mut BenchHarness) {
579579
let mut rng = XorShiftRng::new();
580580
let mut exp = Exp::new(2.71828 * 3.14159);
581581

582582
do bh.iter {
583-
for _ in range(0, N) {
583+
for _ in range(0, RAND_BENCH_N) {
584584
exp.sample(&mut rng);
585585
}
586586
}
587-
bh.bytes = size_of::<f64>() as u64 * N;
587+
bh.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
588588
}
589589
}
File renamed without changes.

src/libstd/rand/mod.rs

+11-11
Original file line numberDiff line numberDiff line change
@@ -833,58 +833,58 @@ mod test {
833833
}
834834
}
835835

836+
static RAND_BENCH_N: u64 = 100;
837+
836838
#[cfg(test)]
837839
mod bench {
838840
use extra::test::BenchHarness;
839-
use rand::*;
841+
use rand::{XorShiftRng, StdRng, IsaacRng, Isaac64Rng, Rng, RAND_BENCH_N};
840842
use mem::size_of;
841843
use iter::range;
842844
use option::{Some, None};
843845

844-
static N: u64 = 100;
845-
846846
#[bench]
847847
fn rand_xorshift(bh: &mut BenchHarness) {
848848
let mut rng = XorShiftRng::new();
849849
do bh.iter {
850-
for _ in range(0, N) {
850+
for _ in range(0, RAND_BENCH_N) {
851851
rng.gen::<uint>();
852852
}
853853
}
854-
bh.bytes = size_of::<uint>() as u64 * N;
854+
bh.bytes = size_of::<uint>() as u64 * RAND_BENCH_N;
855855
}
856856

857857
#[bench]
858858
fn rand_isaac(bh: &mut BenchHarness) {
859859
let mut rng = IsaacRng::new();
860860
do bh.iter {
861-
for _ in range(0, N) {
861+
for _ in range(0, RAND_BENCH_N) {
862862
rng.gen::<uint>();
863863
}
864864
}
865-
bh.bytes = size_of::<uint>() as u64 * N;
865+
bh.bytes = size_of::<uint>() as u64 * RAND_BENCH_N;
866866
}
867867

868868
#[bench]
869869
fn rand_isaac64(bh: &mut BenchHarness) {
870870
let mut rng = Isaac64Rng::new();
871871
do bh.iter {
872-
for _ in range(0, N) {
872+
for _ in range(0, RAND_BENCH_N) {
873873
rng.gen::<uint>();
874874
}
875875
}
876-
bh.bytes = size_of::<uint>() as u64 * N;
876+
bh.bytes = size_of::<uint>() as u64 * RAND_BENCH_N;
877877
}
878878

879879
#[bench]
880880
fn rand_std(bh: &mut BenchHarness) {
881881
let mut rng = StdRng::new();
882882
do bh.iter {
883-
for _ in range(0, N) {
883+
for _ in range(0, RAND_BENCH_N) {
884884
rng.gen::<uint>();
885885
}
886886
}
887-
bh.bytes = size_of::<uint>() as u64 * N;
887+
bh.bytes = size_of::<uint>() as u64 * RAND_BENCH_N;
888888
}
889889

890890
#[bench]

0 commit comments

Comments
 (0)