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