statrs/distribution/
ziggurat.rs

1use super::ziggurat_tables;
2use rand::distributions::Open01;
3use rand::Rng;
4
5pub fn sample_std_normal<R: Rng + ?Sized>(rng: &mut R) -> f64 {
6    #[inline]
7    fn pdf(x: f64) -> f64 {
8        (-x * x / 2.0).exp()
9    }
10
11    #[inline]
12    fn zero_case<R: Rng + ?Sized>(rng: &mut R, u: f64) -> f64 {
13        let mut x = 1.0f64;
14        let mut y = 0.0f64;
15        while -2.0 * y < x * x {
16            let x_: f64 = rng.sample(Open01);
17            let y_: f64 = rng.sample(Open01);
18
19            x = x_.ln() / ziggurat_tables::ZIG_NORM_R;
20            y = y_.ln();
21        }
22        if u < 0.0 {
23            x - ziggurat_tables::ZIG_NORM_R
24        } else {
25            ziggurat_tables::ZIG_NORM_R - x
26        }
27    }
28
29    ziggurat(
30        rng,
31        true,
32        &ziggurat_tables::ZIG_NORM_X,
33        &ziggurat_tables::ZIG_NORM_F,
34        pdf,
35        zero_case,
36    )
37}
38
39pub fn sample_exp_1<R: Rng + ?Sized>(rng: &mut R) -> f64 {
40    #[inline]
41    fn pdf(x: f64) -> f64 {
42        (-x).exp()
43    }
44
45    #[inline]
46    fn zero_case<R: Rng + ?Sized>(rng: &mut R, _u: f64) -> f64 {
47        ziggurat_tables::ZIG_EXP_R - rng.gen::<f64>().ln()
48    }
49
50    ziggurat(
51        rng,
52        false,
53        &ziggurat_tables::ZIG_EXP_X,
54        &ziggurat_tables::ZIG_EXP_F,
55        pdf,
56        zero_case,
57    )
58}
59
60// Ziggurat method for sampling a random number based on the ZIGNOR
61// variant from Doornik 2005. Code borrowed from
62// https://github.com/rust-lang-nursery/rand/blob/master/src/distributions/mod.
63// rs#L223
64#[inline(always)]
65fn ziggurat<R: Rng + ?Sized, P, Z>(
66    rng: &mut R,
67    symmetric: bool,
68    x_tab: ziggurat_tables::ZigTable,
69    f_tab: ziggurat_tables::ZigTable,
70    mut pdf: P,
71    mut zero_case: Z,
72) -> f64
73where
74    P: FnMut(f64) -> f64,
75    Z: FnMut(&mut R, f64) -> f64,
76{
77    const SCALE: f64 = (1u64 << 53) as f64;
78    loop {
79        let bits: u64 = rng.gen();
80        let i = (bits & 0xff) as usize;
81        let f = (bits >> 11) as f64 / SCALE;
82
83        // u is either U(-1, 1) or U(0, 1) depending on if this is a
84        // symmetric distribution or not.
85        let u = if symmetric { 2.0 * f - 1.0 } else { f };
86        let x = u * x_tab[i];
87
88        let test_x = if symmetric { x.abs() } else { x };
89
90        // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u <
91        // x_tab[i+1]/x_tab[i])
92        if test_x < x_tab[i + 1] {
93            return x;
94        }
95        if i == 0 {
96            return zero_case(rng, u);
97        }
98        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
99        if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen::<f64>() < pdf(x) {
100            return x;
101        }
102    }
103}