rand_distr/
utils.rs

1// Copyright 2018 Developers of the Rand project.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9//! Math helper functions
10
11use crate::ziggurat_tables;
12use rand::distributions::hidden_export::IntoFloat;
13use rand::Rng;
14use num_traits::Float;
15
16/// Calculates ln(gamma(x)) (natural logarithm of the gamma
17/// function) using the Lanczos approximation.
18///
19/// The approximation expresses the gamma function as:
20/// `gamma(z+1) = sqrt(2*pi)*(z+g+0.5)^(z+0.5)*exp(-z-g-0.5)*Ag(z)`
21/// `g` is an arbitrary constant; we use the approximation with `g=5`.
22///
23/// Noting that `gamma(z+1) = z*gamma(z)` and applying `ln` to both sides:
24/// `ln(gamma(z)) = (z+0.5)*ln(z+g+0.5)-(z+g+0.5) + ln(sqrt(2*pi)*Ag(z)/z)`
25///
26/// `Ag(z)` is an infinite series with coefficients that can be calculated
27/// ahead of time - we use just the first 6 terms, which is good enough
28/// for most purposes.
29pub(crate) fn log_gamma<F: Float>(x: F) -> F {
30    // precalculated 6 coefficients for the first 6 terms of the series
31    let coefficients: [F; 6] = [
32        F::from(76.18009172947146).unwrap(),
33        F::from(-86.50532032941677).unwrap(),
34        F::from(24.01409824083091).unwrap(),
35        F::from(-1.231739572450155).unwrap(),
36        F::from(0.1208650973866179e-2).unwrap(),
37        F::from(-0.5395239384953e-5).unwrap(),
38    ];
39
40    // (x+0.5)*ln(x+g+0.5)-(x+g+0.5)
41    let tmp = x + F::from(5.5).unwrap();
42    let log = (x + F::from(0.5).unwrap()) * tmp.ln() - tmp;
43
44    // the first few terms of the series for Ag(x)
45    let mut a = F::from(1.000000000190015).unwrap();
46    let mut denom = x;
47    for &coeff in &coefficients {
48        denom = denom + F::one();
49        a = a + (coeff / denom);
50    }
51
52    // get everything together
53    // a is Ag(x)
54    // 2.5066... is sqrt(2pi)
55    log + (F::from(2.5066282746310005).unwrap() * a / x).ln()
56}
57
58/// Sample a random number using the Ziggurat method (specifically the
59/// ZIGNOR variant from Doornik 2005). Most of the arguments are
60/// directly from the paper:
61///
62/// * `rng`: source of randomness
63/// * `symmetric`: whether this is a symmetric distribution, or one-sided with P(x < 0) = 0.
64/// * `X`: the $x_i$ abscissae.
65/// * `F`: precomputed values of the PDF at the $x_i$, (i.e. $f(x_i)$)
66/// * `F_DIFF`: precomputed values of $f(x_i) - f(x_{i+1})$
67/// * `pdf`: the probability density function
68/// * `zero_case`: manual sampling from the tail when we chose the
69///    bottom box (i.e. i == 0)
70
71// the perf improvement (25-50%) is definitely worth the extra code
72// size from force-inlining.
73#[inline(always)]
74pub(crate) fn ziggurat<R: Rng + ?Sized, P, Z>(
75    rng: &mut R,
76    symmetric: bool,
77    x_tab: ziggurat_tables::ZigTable,
78    f_tab: ziggurat_tables::ZigTable,
79    mut pdf: P,
80    mut zero_case: Z
81) -> f64
82where
83    P: FnMut(f64) -> f64,
84    Z: FnMut(&mut R, f64) -> f64,
85{
86    loop {
87        // As an optimisation we re-implement the conversion to a f64.
88        // From the remaining 12 most significant bits we use 8 to construct `i`.
89        // This saves us generating a whole extra random number, while the added
90        // precision of using 64 bits for f64 does not buy us much.
91        let bits = rng.next_u64();
92        let i = bits as usize & 0xff;
93
94        let u = if symmetric {
95            // Convert to a value in the range [2,4) and subtract to get [-1,1)
96            // We can't convert to an open range directly, that would require
97            // subtracting `3.0 - EPSILON`, which is not representable.
98            // It is possible with an extra step, but an open range does not
99            // seem necessary for the ziggurat algorithm anyway.
100            (bits >> 12).into_float_with_exponent(1) - 3.0
101        } else {
102            // Convert to a value in the range [1,2) and subtract to get (0,1)
103            (bits >> 12).into_float_with_exponent(0) - (1.0 - core::f64::EPSILON / 2.0)
104        };
105        let x = u * x_tab[i];
106
107        let test_x = if symmetric { x.abs() } else { x };
108
109        // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u < x_tab[i+1]/x_tab[i])
110        if test_x < x_tab[i + 1] {
111            return x;
112        }
113        if i == 0 {
114            return zero_case(rng, u);
115        }
116        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
117        if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen::<f64>() < pdf(x) {
118            return x;
119        }
120    }
121}