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