rustyqlib/core/montecarlo/
sampling.rs1use crate::core::utils::inv_norm_cdf;
5
6use super::rng::splitmix64;
7
8struct Counter {
10 state: u64,
11}
12
13impl Counter {
14 fn new(seed: u64) -> Self {
15 Self { state: splitmix64(seed) }
16 }
17 fn next_u64(&mut self) -> u64 {
18 self.state = self.state.wrapping_add(1);
19 splitmix64(self.state)
20 }
21 fn uniform(&mut self) -> f64 {
22 (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
23 }
24}
25
26pub fn stratified_uniforms(n: usize, seed: u64) -> Vec<f64> {
30 let mut rng = Counter::new(seed);
31 (0..n).map(|i| (i as f64 + rng.uniform()) / n as f64).collect()
32}
33
34pub fn stratified_normals(n: usize, seed: u64) -> Vec<f64> {
37 stratified_uniforms(n, seed)
38 .into_iter()
39 .map(|u| inv_norm_cdf(u.clamp(1e-15, 1.0 - 1e-15)))
40 .collect()
41}
42
43pub fn latin_hypercube(n: usize, dims: usize, seed: u64) -> Vec<Vec<f64>> {
47 let mut rng = Counter::new(seed);
48 let mut columns: Vec<Vec<f64>> = Vec::with_capacity(dims);
50 for _ in 0..dims {
51 let mut column: Vec<f64> =
52 (0..n).map(|i| (i as f64 + rng.uniform()) / n as f64).collect();
53 for i in (1..n).rev() {
55 let j = (rng.next_u64() % (i as u64 + 1)) as usize;
56 column.swap(i, j);
57 }
58 columns.push(column);
59 }
60 (0..n).map(|i| columns.iter().map(|c| c[i]).collect()).collect()
61}
62
63#[cfg(test)]
64mod tests {
65 use super::*;
66
67 #[test]
68 fn stratified_uniforms_hit_every_stratum_once() {
69 let n = 128;
70 let u = stratified_uniforms(n, 5);
71 let mut hits = vec![0u32; n];
72 for &x in &u {
73 hits[(x * n as f64) as usize] += 1;
74 }
75 assert!(hits.iter().all(|&c| c == 1));
76 assert_eq!(u, stratified_uniforms(n, 5), "seeded determinism");
77 }
78
79 #[test]
80 fn stratified_normals_beat_plain_sampling_on_the_mean() {
81 let z = stratified_normals(4096, 11);
84 let mean: f64 = z.iter().sum::<f64>() / z.len() as f64;
85 assert!(mean.abs() < 1e-3, "mean {mean}");
86 }
87
88 #[test]
89 fn latin_hypercube_stratifies_every_dimension() {
90 let (n, dims) = (64, 5);
91 let points = latin_hypercube(n, dims, 3);
92 assert_eq!(points.len(), n);
93 for d in 0..dims {
94 let mut hits = vec![0u32; n];
95 for p in &points {
96 hits[(p[d] * n as f64) as usize] += 1;
97 }
98 assert!(hits.iter().all(|&c| c == 1), "dimension {d}");
99 }
100 let same_order = points.windows(2).all(|w| (w[0][0] < w[1][0]) == (w[0][1] < w[1][1]));
102 assert!(!same_order, "columns appear perfectly rank-correlated");
103 }
104}