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