Skip to main content

rustyqlib/core/montecarlo/
sampling.rs

1//! Stratified sampling and Latin hypercube designs — seeded, exact
2//! stratification of the unit interval / hypercube.
3
4use crate::core::utils::inv_norm_cdf;
5
6use super::rng::splitmix64;
7
8/// Tiny counter-based uniform generator on top of SplitMix64.
9struct 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
26/// `n` stratified uniforms: exactly one jittered draw per stratum
27/// `[i/n, (i+1)/n)`. Cuts the variance of smooth 1-D integrands from
28/// O(1/n) to O(1/n^3). Deterministic per seed.
29pub 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
34/// `n` stratified standard normals (stratified uniforms through the
35/// inverse normal CDF).
36pub 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
43/// An `n x dims` Latin hypercube: every dimension is exactly stratified
44/// (one point per stratum), with independent random pairings across
45/// dimensions. Deterministic per seed.
46pub fn latin_hypercube(n: usize, dims: usize, seed: u64) -> Vec<Vec<f64>> {
47    let mut rng = Counter::new(seed);
48    // one stratified, independently shuffled column per dimension
49    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        // Fisher-Yates
54        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        // the stratified sample mean of N(0,1) is far tighter than the
82        // ~1/sqrt(n) of iid draws
83        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        // different dimensions are paired differently (not comonotone)
101        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}