src/librand/distributions/gamma.rs

   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 and derived distributions.
  12
  13 use self::GammaRepr::*;
  14 use self::ChiSquaredRepr::*;
  15
  16 use core::num::Float;
  17
  18 use {Rng, Open01};
  19 use super::normal::StandardNormal;
  20 use super::{IndependentSample, Sample, Exp};
  21
  22 /// The Gamma distribution `Gamma(shape, scale)` distribution.
  23 ///
  24 /// The density function of this distribution is
  25 ///
  26 /// ```text
  27 /// f(x) =  x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k)
  28 /// ```
  29 ///
  30 /// where `Γ` is the Gamma function, `k` is the shape and `θ` is the
  31 /// scale and both `k` and `θ` are strictly positive.
  32 ///
  33 /// The algorithm used is that described by Marsaglia & Tsang 2000[1],
  34 /// falling back to directly sampling from an Exponential for `shape
  35 /// == 1`, and using the boosting technique described in [1] for
  36 /// `shape < 1`.
  37 ///
  38 /// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method
  39 /// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3
  40 /// (September 2000),
  41 /// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414)
  42 pub struct Gamma {
  43     repr: GammaRepr,
  44 }
  45
  46 enum GammaRepr {
  47     Large(GammaLargeShape),
  48     One(Exp),
  49     Small(GammaSmallShape)
  50 }
  51
  52 // These two helpers could be made public, but saving the
  53 // match-on-Gamma-enum branch from using them directly (e.g. if one
  54 // knows that the shape is always > 1) doesn't appear to be much
  55 // faster.
  56
  57 /// Gamma distribution where the shape parameter is less than 1.
  58 ///
  59 /// Note, samples from this require a compulsory floating-point `pow`
  60 /// call, which makes it significantly slower than sampling from a
  61 /// gamma distribution where the shape parameter is greater than or
  62 /// equal to 1.
  63 ///
  64 /// See `Gamma` for sampling from a Gamma distribution with general
  65 /// shape parameters.
  66 struct GammaSmallShape {
  67     inv_shape: f64,
  68     large_shape: GammaLargeShape
  69 }
  70
  71 /// Gamma distribution where the shape parameter is larger than 1.
  72 ///
  73 /// See `Gamma` for sampling from a Gamma distribution with general
  74 /// shape parameters.
  75 struct GammaLargeShape {
  76     scale: f64,
  77     c: f64,
  78     d: f64
  79 }
  80
  81 impl Gamma {
  82     /// Construct an object representing the `Gamma(shape, scale)`
  83     /// distribution.
  84     ///
  85     /// Panics if `shape <= 0` or `scale <= 0`.
  86     pub fn new(shape: f64, scale: f64) -> Gamma {
  87         assert!(shape > 0.0, "Gamma::new called with shape <= 0");
  88         assert!(scale > 0.0, "Gamma::new called with scale <= 0");
  89
  90         let repr = match shape {
  91             1.0         => One(Exp::new(1.0 / scale)),
  92             0.0 ... 1.0 => Small(GammaSmallShape::new_raw(shape, scale)),
  93             _           => Large(GammaLargeShape::new_raw(shape, scale))
  94         };
  95         Gamma { repr: repr }
  96     }
  97 }
  98
  99 impl GammaSmallShape {
 100     fn new_raw(shape: f64, scale: f64) -> GammaSmallShape {
 101         GammaSmallShape {
 102             inv_shape: 1. / shape,
 103             large_shape: GammaLargeShape::new_raw(shape + 1.0, scale)
 104         }
 105     }
 106 }
 107
 108 impl GammaLargeShape {
 109     fn new_raw(shape: f64, scale: f64) -> GammaLargeShape {
 110         let d = shape - 1. / 3.;
 111         GammaLargeShape {
 112             scale: scale,
 113             c: 1. / (9. * d).sqrt(),
 114             d: d
 115         }
 116     }
 117 }
 118
 119 impl Sample<f64> for Gamma {
 120     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
 121 }
 122 impl Sample<f64> for GammaSmallShape {
 123     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
 124 }
 125 impl Sample<f64> for GammaLargeShape {
 126     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
 127 }
 128
 129 impl IndependentSample<f64> for Gamma {
 130     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 131         match self.repr {
 132             Small(ref g) => g.ind_sample(rng),
 133             One(ref g) => g.ind_sample(rng),
 134             Large(ref g) => g.ind_sample(rng),
 135         }
 136     }
 137 }
 138 impl IndependentSample<f64> for GammaSmallShape {
 139     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 140         let Open01(u) = rng.gen::<Open01<f64>>();
 141
 142         self.large_shape.ind_sample(rng) * u.powf(self.inv_shape)
 143     }
 144 }
 145 impl IndependentSample<f64> for GammaLargeShape {
 146     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 147         loop {
 148             let StandardNormal(x) = rng.gen::<StandardNormal>();
 149             let v_cbrt = 1.0 + self.c * x;
 150             if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0
 151                 continue
 152             }
 153
 154             let v = v_cbrt * v_cbrt * v_cbrt;
 155             let Open01(u) = rng.gen::<Open01<f64>>();
 156
 157             let x_sqr = x * x;
 158             if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
 159                 u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
 160                 return self.d * v * self.scale
 161             }
 162         }
 163     }
 164 }
 165
 166 /// The chi-squared distribution `χ²(k)`, where `k` is the degrees of
 167 /// freedom.
 168 ///
 169 /// For `k > 0` integral, this distribution is the sum of the squares
 170 /// of `k` independent standard normal random variables. For other
 171 /// `k`, this uses the equivalent characterisation `χ²(k) = Gamma(k/2,
 172 /// 2)`.
 173 pub struct ChiSquared {
 174     repr: ChiSquaredRepr,
 175 }
 176
 177 enum ChiSquaredRepr {
 178     // k == 1, Gamma(alpha, ..) is particularly slow for alpha < 1,
 179     // e.g. when alpha = 1/2 as it would be for this case, so special-
 180     // casing and using the definition of N(0,1)^2 is faster.
 181     DoFExactlyOne,
 182     DoFAnythingElse(Gamma),
 183 }
 184
 185 impl ChiSquared {
 186     /// Create a new chi-squared distribution with degrees-of-freedom
 187     /// `k`. Panics if `k < 0`.
 188     pub fn new(k: f64) -> ChiSquared {
 189         let repr = if k == 1.0 {
 190             DoFExactlyOne
 191         } else {
 192             assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
 193             DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
 194         };
 195         ChiSquared { repr: repr }
 196     }
 197 }
 198 impl Sample<f64> for ChiSquared {
 199     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
 200 }
 201 impl IndependentSample<f64> for ChiSquared {
 202     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 203         match self.repr {
 204             DoFExactlyOne => {
 205                 // k == 1 => N(0,1)^2
 206                 let StandardNormal(norm) = rng.gen::<StandardNormal>();
 207                 norm * norm
 208             }
 209             DoFAnythingElse(ref g) => g.ind_sample(rng)
 210         }
 211     }
 212 }
 213
 214 /// The Fisher F distribution `F(m, n)`.
 215 ///
 216 /// This distribution is equivalent to the ratio of two normalised
 217 /// chi-squared distributions, that is, `F(m,n) = (χ²(m)/m) /
 218 /// (χ²(n)/n)`.
 219 pub struct FisherF {
 220     numer: ChiSquared,
 221     denom: ChiSquared,
 222     // denom_dof / numer_dof so that this can just be a straight
 223     // multiplication, rather than a division.
 224     dof_ratio: f64,
 225 }
 226
 227 impl FisherF {
 228     /// Create a new `FisherF` distribution, with the given
 229     /// parameter. Panics if either `m` or `n` are not positive.
 230     pub fn new(m: f64, n: f64) -> FisherF {
 231         assert!(m > 0.0, "FisherF::new called with `m < 0`");
 232         assert!(n > 0.0, "FisherF::new called with `n < 0`");
 233
 234         FisherF {
 235             numer: ChiSquared::new(m),
 236             denom: ChiSquared::new(n),
 237             dof_ratio: n / m
 238         }
 239     }
 240 }
 241 impl Sample<f64> for FisherF {
 242     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
 243 }
 244 impl IndependentSample<f64> for FisherF {
 245     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 246         self.numer.ind_sample(rng) / self.denom.ind_sample(rng) * self.dof_ratio
 247     }
 248 }
 249
 250 /// The Student t distribution, `t(nu)`, where `nu` is the degrees of
 251 /// freedom.
 252 pub struct StudentT {
 253     chi: ChiSquared,
 254     dof: f64
 255 }
 256
 257 impl StudentT {
 258     /// Create a new Student t distribution with `n` degrees of
 259     /// freedom. Panics if `n <= 0`.
 260     pub fn new(n: f64) -> StudentT {
 261         assert!(n > 0.0, "StudentT::new called with `n <= 0`");
 262         StudentT {
 263             chi: ChiSquared::new(n),
 264             dof: n
 265         }
 266     }
 267 }
 268 impl Sample<f64> for StudentT {
 269     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
 270 }
 271 impl IndependentSample<f64> for StudentT {
 272     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 273         let StandardNormal(norm) = rng.gen::<StandardNormal>();
 274         norm * (self.dof / self.chi.ind_sample(rng)).sqrt()
 275     }
 276 }
 277
 278 #[cfg(test)]
 279 mod test {
 280     use std::prelude::v1::*;
 281
 282     use distributions::{Sample, IndependentSample};
 283     use super::{ChiSquared, StudentT, FisherF};
 284
 285     #[test]
 286     fn test_chi_squared_one() {
 287         let mut chi = ChiSquared::new(1.0);
 288         let mut rng = ::test::rng();
 289         for _ in 0..1000 {
 290             chi.sample(&mut rng);
 291             chi.ind_sample(&mut rng);
 292         }
 293     }
 294     #[test]
 295     fn test_chi_squared_small() {
 296         let mut chi = ChiSquared::new(0.5);
 297         let mut rng = ::test::rng();
 298         for _ in 0..1000 {
 299             chi.sample(&mut rng);
 300             chi.ind_sample(&mut rng);
 301         }
 302     }
 303     #[test]
 304     fn test_chi_squared_large() {
 305         let mut chi = ChiSquared::new(30.0);
 306         let mut rng = ::test::rng();
 307         for _ in 0..1000 {
 308             chi.sample(&mut rng);
 309             chi.ind_sample(&mut rng);
 310         }
 311     }
 312     #[test]
 313     #[should_panic]
 314     fn test_chi_squared_invalid_dof() {
 315         ChiSquared::new(-1.0);
 316     }
 317
 318     #[test]
 319     fn test_f() {
 320         let mut f = FisherF::new(2.0, 32.0);
 321         let mut rng = ::test::rng();
 322         for _ in 0..1000 {
 323             f.sample(&mut rng);
 324             f.ind_sample(&mut rng);
 325         }
 326     }
 327
 328     #[test]
 329     fn test_t() {
 330         let mut t = StudentT::new(11.0);
 331         let mut rng = ::test::rng();
 332         for _ in 0..1000 {
 333             t.sample(&mut rng);
 334             t.ind_sample(&mut rng);
 335         }
 336     }
 337 }
 338
 339 #[cfg(test)]
 340 mod bench {
 341     extern crate test;
 342     use std::prelude::v1::*;
 343     use self::test::Bencher;
 344     use std::mem::size_of;
 345     use distributions::IndependentSample;
 346     use super::Gamma;
 347
 348
 349     #[bench]
 350     fn bench_gamma_large_shape(b: &mut Bencher) {
 351         let gamma = Gamma::new(10., 1.0);
 352         let mut rng = ::test::weak_rng();
 353
 354         b.iter(|| {
 355             for _ in 0..::RAND_BENCH_N {
 356                 gamma.ind_sample(&mut rng);
 357             }
 358         });
 359         b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
 360     }
 361
 362     #[bench]
 363     fn bench_gamma_small_shape(b: &mut Bencher) {
 364         let gamma = Gamma::new(0.1, 1.0);
 365         let mut rng = ::test::weak_rng();
 366
 367         b.iter(|| {
 368             for _ in 0..::RAND_BENCH_N {
 369                 gamma.ind_sample(&mut rng);
 370             }
 371         });
 372         b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
 373     }
 374 }