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 FloatMath;
  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 {
 121         self.ind_sample(rng)
 122     }
 123 }
 124 impl Sample<f64> for GammaSmallShape {
 125     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 {
 126         self.ind_sample(rng)
 127     }
 128 }
 129 impl Sample<f64> for GammaLargeShape {
 130     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 {
 131         self.ind_sample(rng)
 132     }
 133 }
 134
 135 impl IndependentSample<f64> for Gamma {
 136     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 137         match self.repr {
 138             Small(ref g) => g.ind_sample(rng),
 139             One(ref g) => g.ind_sample(rng),
 140             Large(ref g) => g.ind_sample(rng),
 141         }
 142     }
 143 }
 144 impl IndependentSample<f64> for GammaSmallShape {
 145     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 146         let Open01(u) = rng.gen::<Open01<f64>>();
 147
 148         self.large_shape.ind_sample(rng) * u.powf(self.inv_shape)
 149     }
 150 }
 151 impl IndependentSample<f64> for GammaLargeShape {
 152     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 153         loop {
 154             let StandardNormal(x) = rng.gen::<StandardNormal>();
 155             let v_cbrt = 1.0 + self.c * x;
 156             if v_cbrt <= 0.0 {
 157                 // a^3 <= 0 iff a <= 0
 158                 continue;
 159             }
 160
 161             let v = v_cbrt * v_cbrt * v_cbrt;
 162             let Open01(u) = rng.gen::<Open01<f64>>();
 163
 164             let x_sqr = x * x;
 165             if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
 166                u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
 167                 return self.d * v * self.scale;
 168             }
 169         }
 170     }
 171 }
 172
 173 /// The chi-squared distribution `χ²(k)`, where `k` is the degrees of
 174 /// freedom.
 175 ///
 176 /// For `k > 0` integral, this distribution is the sum of the squares
 177 /// of `k` independent standard normal random variables. For other
 178 /// `k`, this uses the equivalent characterization `χ²(k) = Gamma(k/2,
 179 /// 2)`.
 180 pub struct ChiSquared {
 181     repr: ChiSquaredRepr,
 182 }
 183
 184 enum ChiSquaredRepr {
 185     // k == 1, Gamma(alpha, ..) is particularly slow for alpha < 1,
 186     // e.g. when alpha = 1/2 as it would be for this case, so special-
 187     // casing and using the definition of N(0,1)^2 is faster.
 188     DoFExactlyOne,
 189     DoFAnythingElse(Gamma),
 190 }
 191
 192 impl ChiSquared {
 193     /// Create a new chi-squared distribution with degrees-of-freedom
 194     /// `k`. Panics if `k < 0`.
 195     pub fn new(k: f64) -> ChiSquared {
 196         let repr = if k == 1.0 {
 197             DoFExactlyOne
 198         } else {
 199             assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
 200             DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
 201         };
 202         ChiSquared { repr: repr }
 203     }
 204 }
 205 impl Sample<f64> for ChiSquared {
 206     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 {
 207         self.ind_sample(rng)
 208     }
 209 }
 210 impl IndependentSample<f64> for ChiSquared {
 211     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 212         match self.repr {
 213             DoFExactlyOne => {
 214                 // k == 1 => N(0,1)^2
 215                 let StandardNormal(norm) = rng.gen::<StandardNormal>();
 216                 norm * norm
 217             }
 218             DoFAnythingElse(ref g) => g.ind_sample(rng),
 219         }
 220     }
 221 }
 222
 223 /// The Fisher F distribution `F(m, n)`.
 224 ///
 225 /// This distribution is equivalent to the ratio of two normalised
 226 /// chi-squared distributions, that is, `F(m,n) = (χ²(m)/m) /
 227 /// (χ²(n)/n)`.
 228 pub struct FisherF {
 229     numer: ChiSquared,
 230     denom: ChiSquared,
 231     // denom_dof / numer_dof so that this can just be a straight
 232     // multiplication, rather than a division.
 233     dof_ratio: f64,
 234 }
 235
 236 impl FisherF {
 237     /// Create a new `FisherF` distribution, with the given
 238     /// parameter. Panics if either `m` or `n` are not positive.
 239     pub fn new(m: f64, n: f64) -> FisherF {
 240         assert!(m > 0.0, "FisherF::new called with `m < 0`");
 241         assert!(n > 0.0, "FisherF::new called with `n < 0`");
 242
 243         FisherF {
 244             numer: ChiSquared::new(m),
 245             denom: ChiSquared::new(n),
 246             dof_ratio: n / m,
 247         }
 248     }
 249 }
 250 impl Sample<f64> for FisherF {
 251     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 {
 252         self.ind_sample(rng)
 253     }
 254 }
 255 impl IndependentSample<f64> for FisherF {
 256     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 257         self.numer.ind_sample(rng) / self.denom.ind_sample(rng) * self.dof_ratio
 258     }
 259 }
 260
 261 /// The Student t distribution, `t(nu)`, where `nu` is the degrees of
 262 /// freedom.
 263 pub struct StudentT {
 264     chi: ChiSquared,
 265     dof: f64,
 266 }
 267
 268 impl StudentT {
 269     /// Create a new Student t distribution with `n` degrees of
 270     /// freedom. Panics if `n <= 0`.
 271     pub fn new(n: f64) -> StudentT {
 272         assert!(n > 0.0, "StudentT::new called with `n <= 0`");
 273         StudentT {
 274             chi: ChiSquared::new(n),
 275             dof: n,
 276         }
 277     }
 278 }
 279 impl Sample<f64> for StudentT {
 280     fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 {
 281         self.ind_sample(rng)
 282     }
 283 }
 284 impl IndependentSample<f64> for StudentT {
 285     fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
 286         let StandardNormal(norm) = rng.gen::<StandardNormal>();
 287         norm * (self.dof / self.chi.ind_sample(rng)).sqrt()
 288     }
 289 }
 290
 291 #[cfg(test)]
 292 mod tests {
 293     use distributions::{Sample, IndependentSample};
 294     use super::{ChiSquared, StudentT, FisherF};
 295
 296     #[test]
 297     fn test_chi_squared_one() {
 298         let mut chi = ChiSquared::new(1.0);
 299         let mut rng = ::test::rng();
 300         for _ in 0..1000 {
 301             chi.sample(&mut rng);
 302             chi.ind_sample(&mut rng);
 303         }
 304     }
 305     #[test]
 306     fn test_chi_squared_small() {
 307         let mut chi = ChiSquared::new(0.5);
 308         let mut rng = ::test::rng();
 309         for _ in 0..1000 {
 310             chi.sample(&mut rng);
 311             chi.ind_sample(&mut rng);
 312         }
 313     }
 314     #[test]
 315     fn test_chi_squared_large() {
 316         let mut chi = ChiSquared::new(30.0);
 317         let mut rng = ::test::rng();
 318         for _ in 0..1000 {
 319             chi.sample(&mut rng);
 320             chi.ind_sample(&mut rng);
 321         }
 322     }
 323     #[test]
 324     #[should_panic]
 325     fn test_chi_squared_invalid_dof() {
 326         ChiSquared::new(-1.0);
 327     }
 328
 329     #[test]
 330     fn test_f() {
 331         let mut f = FisherF::new(2.0, 32.0);
 332         let mut rng = ::test::rng();
 333         for _ in 0..1000 {
 334             f.sample(&mut rng);
 335             f.ind_sample(&mut rng);
 336         }
 337     }
 338
 339     #[test]
 340     fn test_t() {
 341         let mut t = StudentT::new(11.0);
 342         let mut rng = ::test::rng();
 343         for _ in 0..1000 {
 344             t.sample(&mut rng);
 345             t.ind_sample(&mut rng);
 346         }
 347     }
 348 }
 349
 350 #[cfg(test)]
 351 mod bench {
 352     extern crate test;
 353     use self::test::Bencher;
 354     use std::mem::size_of;
 355     use distributions::IndependentSample;
 356     use super::Gamma;
 357
 358
 359     #[bench]
 360     fn bench_gamma_large_shape(b: &mut Bencher) {
 361         let gamma = Gamma::new(10., 1.0);
 362         let mut rng = ::test::weak_rng();
 363
 364         b.iter(|| {
 365             for _ in 0..::RAND_BENCH_N {
 366                 gamma.ind_sample(&mut rng);
 367             }
 368         });
 369         b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
 370     }
 371
 372     #[bench]
 373     fn bench_gamma_small_shape(b: &mut Bencher) {
 374         let gamma = Gamma::new(0.1, 1.0);
 375         let mut rng = ::test::weak_rng();
 376
 377         b.iter(|| {
 378             for _ in 0..::RAND_BENCH_N {
 379                 gamma.ind_sample(&mut rng);
 380             }
 381         });
 382         b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
 383     }
 384 }