conspire/math/random/
mod.rs

1#[cfg(test)]
2mod test;
3
4use std::{
5    cell::Cell,
6    f64::consts::TAU,
7    time::{SystemTime, UNIX_EPOCH},
8};
9
10thread_local! {
11    static STATE: Cell<u64> = const { Cell::new(0) };
12}
13
14fn seed() -> u64 {
15    let now = SystemTime::now()
16        .duration_since(UNIX_EPOCH)
17        .unwrap_or_default();
18    let t = now.as_nanos() as u64;
19    let x = 0u8;
20    let addr = (&x as *const u8 as usize) as u64;
21    let mut s = t ^ addr.wrapping_mul(0x9E3779B97F4A7C15);
22    if s == 0 {
23        s = 1;
24    }
25    s
26}
27
28fn next_u64() -> u64 {
29    STATE.with(|st| {
30        let mut s = st.get();
31        if s == 0 {
32            s = seed();
33        }
34        s ^= s >> 12;
35        s ^= s << 25;
36        s ^= s >> 27;
37        st.set(s);
38        s.wrapping_mul(0x2545F4914F6CDD1D)
39    })
40}
41
42fn get_random() -> u8 {
43    (next_u64() >> 56) as u8
44}
45
46/// Returns a uniformly random `u8` in `0..=max`.
47pub fn random_u8(max: u8) -> u8 {
48    if max == u8::MAX {
49        return get_random();
50    }
51    let bound = (max as u16) + 1;
52    let threshold = (256u16 / bound) * bound;
53    loop {
54        let v = get_random() as u16;
55        if v < threshold {
56            return (v % bound) as u8;
57        }
58    }
59}
60
61/// Returns a uniformly random `u64`.
62pub fn random_u64() -> u64 {
63    next_u64()
64}
65
66/// Returns a uniformly random `f64` in `[0, 1)`.
67pub fn random_uniform() -> f64 {
68    let x = next_u64() >> 11;
69    (x as f64) * (1.0 / ((1u64 << 53) as f64))
70}
71
72thread_local! {
73    static NORMAL_SPARE: Cell<Option<f64>> = const { Cell::new(None) };
74}
75
76/// Returns a random sample from the standard normal distribution.
77pub fn random_normal_standard() -> f64 {
78    NORMAL_SPARE.with(|spare| {
79        if let Some(z) = spare.take() {
80            return z;
81        }
82        let mut u1 = random_uniform();
83        while u1 <= 0.0 {
84            u1 = random_uniform();
85        }
86        let u2 = random_uniform();
87        let r = (-2.0 * u1.ln()).sqrt();
88        let (s, c) = (TAU * u2).sin_cos();
89        let z0 = r * c;
90        let z1 = r * s;
91        spare.set(Some(z1));
92        z0
93    })
94}
95
96/// Returns a random sample from a given normal distribution.
97pub fn random_normal(mean: f64, std: f64) -> f64 {
98    mean + std * random_normal_standard()
99}
100
101// fn random_exp1() -> f64 {
102//     let mut u = random_uniform();
103//     while u <= 0.0 {
104//         u = random_uniform();
105//     }
106//     -u.ln()
107// }
108
109// fn random_gamma_k3_scale1() -> f64 {
110//     random_exp1() + random_exp1() + random_exp1()
111// }
112
113// pub fn random_x2_normal(mean: f64, std: f64) -> f64 {
114//     let m = mean / std;
115//     let z_star = if m >= -1.0 { m + 1.0 } else { 0.0 };
116//     let h_min = 0.5 * (z_star - m).powi(2) - z_star;
117//     loop {
118//         let z = random_gamma_k3_scale1();
119//         let h = 0.5 * (z - m).powi(2) - z;
120//         let acceptance_probability = (-(h - h_min)).exp();
121//         if random_uniform() < acceptance_probability {
122//             return std * z;
123//         }
124//     }
125// }
126
127use crate::math::special::erf;
128
129use std::f64::consts::{PI, SQRT_2};
130
131fn x2_normal_primitive(lambda: f64, mean: f64, std: f64) -> f64 {
132    let t = (lambda - mean) / (std * SQRT_2);
133    std * (PI / 2.0).sqrt() * (mean * mean + std * std) * erf(t)
134        - std * std * (lambda + mean) * (-t * t).exp()
135}
136
137fn x2_normal_norm(mean: f64, std: f64) -> f64 {
138    let at_infinity = std * (PI / 2.0).sqrt() * (mean * mean + std * std);
139    let at_zero = x2_normal_primitive(0.0, mean, std);
140    at_infinity - at_zero
141}
142
143fn x2_normal_cdf(lambda: f64, mean: f64, std: f64, norm: f64) -> f64 {
144    if lambda <= 0.0 {
145        return 0.0;
146    }
147    let at_zero = x2_normal_primitive(0.0, mean, std);
148    (x2_normal_primitive(lambda, mean, std) - at_zero) / norm
149}
150
151fn x2_normal_pdf(lambda: f64, mean: f64, std: f64, norm: f64) -> f64 {
152    if lambda <= 0.0 {
153        0.0
154    } else {
155        lambda * lambda * (-(lambda - mean).powi(2) / (2.0 * std * std)).exp() / norm
156    }
157}
158
159/// Returns a random sample from the normal distribution rectified and reweighted by `x^2`.
160pub fn random_x2_normal(mean: f64, std: f64) -> f64 {
161    let norm = x2_normal_norm(mean, std);
162    let u = random_uniform();
163
164    let mut lo = 0.0;
165    let mut hi = mean + 8.0 * std;
166    if hi <= 0.0 {
167        hi = 1.0;
168    }
169    while x2_normal_cdf(hi, mean, std, norm) < u {
170        hi *= 2.0;
171    }
172
173    let mut x = mean.max(1e-12);
174
175    for _ in 0..50 {
176        let fx = x2_normal_cdf(x, mean, std, norm) - u;
177        let dfx = x2_normal_pdf(x, mean, std, norm);
178
179        let mut x_new = if dfx > 0.0 {
180            x - fx / dfx
181        } else {
182            0.5 * (lo + hi)
183        };
184
185        if !x_new.is_finite() || x_new <= lo || x_new >= hi {
186            x_new = 0.5 * (lo + hi);
187        }
188
189        let f_new = x2_normal_cdf(x_new, mean, std, norm);
190
191        if f_new < u {
192            lo = x_new;
193        } else {
194            hi = x_new;
195        }
196
197        x = x_new;
198
199        if (hi - lo) <= 1e-14 * (1.0 + x.abs()) {
200            break;
201        }
202    }
203
204    x
205}
conspire/math/random/mod.rs

conspire/math/random/
mod.rs
