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
46pub fn random_u8(max: u8) -> u8 {
47    if max == u8::MAX {
48        return get_random();
49    }
50    let bound = (max as u16) + 1;
51    let threshold = (256u16 / bound) * bound;
52    loop {
53        let v = get_random() as u16;
54        if v < threshold {
55            return (v % bound) as u8;
56        }
57    }
58}
59
60pub fn random_u64() -> u64 {
61    next_u64()
62}
63
64pub fn random_uniform() -> f64 {
65    let x = next_u64() >> 11;
66    (x as f64) * (1.0 / ((1u64 << 53) as f64))
67}
68
69thread_local! {
70    static NORMAL_SPARE: Cell<Option<f64>> = const { Cell::new(None) };
71}
72
73pub fn random_normal_standard() -> f64 {
74    NORMAL_SPARE.with(|spare| {
75        if let Some(z) = spare.take() {
76            return z;
77        }
78        let mut u1 = random_uniform();
79        while u1 <= 0.0 {
80            u1 = random_uniform();
81        }
82        let u2 = random_uniform();
83        let r = (-2.0 * u1.ln()).sqrt();
84        let (s, c) = (TAU * u2).sin_cos();
85        let z0 = r * c;
86        let z1 = r * s;
87        spare.set(Some(z1));
88        z0
89    })
90}
91
92pub fn random_normal(mean: f64, std: f64) -> f64 {
93    mean + std * random_normal_standard()
94}
95
96// fn random_exp1() -> f64 {
97//     let mut u = random_uniform();
98//     while u <= 0.0 {
99//         u = random_uniform();
100//     }
101//     -u.ln()
102// }
103
104// fn random_gamma_k3_scale1() -> f64 {
105//     random_exp1() + random_exp1() + random_exp1()
106// }
107
108// pub fn random_x2_normal(mean: f64, std: f64) -> f64 {
109//     let m = mean / std;
110//     let z_star = if m >= -1.0 { m + 1.0 } else { 0.0 };
111//     let h_min = 0.5 * (z_star - m).powi(2) - z_star;
112//     loop {
113//         let z = random_gamma_k3_scale1();
114//         let h = 0.5 * (z - m).powi(2) - z;
115//         let acceptance_probability = (-(h - h_min)).exp();
116//         if random_uniform() < acceptance_probability {
117//             return std * z;
118//         }
119//     }
120// }
121
122use crate::math::special::erf;
123
124use std::f64::consts::{PI, SQRT_2};
125
126fn x2_normal_primitive(lambda: f64, mean: f64, std: f64) -> f64 {
127    let t = (lambda - mean) / (std * SQRT_2);
128    std * (PI / 2.0).sqrt() * (mean * mean + std * std) * erf(t)
129        - std * std * (lambda + mean) * (-t * t).exp()
130}
131
132fn x2_normal_norm(mean: f64, std: f64) -> f64 {
133    let at_infinity = std * (PI / 2.0).sqrt() * (mean * mean + std * std);
134    let at_zero = x2_normal_primitive(0.0, mean, std);
135    at_infinity - at_zero
136}
137
138fn x2_normal_cdf(lambda: f64, mean: f64, std: f64, norm: f64) -> f64 {
139    if lambda <= 0.0 {
140        return 0.0;
141    }
142    let at_zero = x2_normal_primitive(0.0, mean, std);
143    (x2_normal_primitive(lambda, mean, std) - at_zero) / norm
144}
145
146fn x2_normal_pdf(lambda: f64, mean: f64, std: f64, norm: f64) -> f64 {
147    if lambda <= 0.0 {
148        0.0
149    } else {
150        lambda * lambda * (-(lambda - mean).powi(2) / (2.0 * std * std)).exp() / norm
151    }
152}
153
154pub fn random_x2_normal(mean: f64, std: f64) -> f64 {
155    let norm = x2_normal_norm(mean, std);
156    let u = random_uniform();
157
158    let mut lo = 0.0;
159    let mut hi = mean + 8.0 * std;
160    if hi <= 0.0 {
161        hi = 1.0;
162    }
163    while x2_normal_cdf(hi, mean, std, norm) < u {
164        hi *= 2.0;
165    }
166
167    let mut x = mean.max(1e-12);
168
169    for _ in 0..50 {
170        let fx = x2_normal_cdf(x, mean, std, norm) - u;
171        let dfx = x2_normal_pdf(x, mean, std, norm);
172
173        let mut x_new = if dfx > 0.0 {
174            x - fx / dfx
175        } else {
176            0.5 * (lo + hi)
177        };
178
179        if !x_new.is_finite() || x_new <= lo || x_new >= hi {
180            x_new = 0.5 * (lo + hi);
181        }
182
183        let f_new = x2_normal_cdf(x_new, mean, std, norm);
184
185        if f_new < u {
186            lo = x_new;
187        } else {
188            hi = x_new;
189        }
190
191        x = x_new;
192
193        if (hi - lo) <= 1e-14 * (1.0 + x.abs()) {
194            break;
195        }
196    }
197
198    x
199}
conspire/math/random/mod.rs

conspire/math/random/
mod.rs
