pfx_trace/
sobol.rs

1pub const PATTERNS: u32 = 256;
2pub const DIMENSIONS: u32 = 4;
3pub const MIN_SAMPLES: u32 = 256;
4pub const MAX_SAMPLES: u32 = 8192;
5pub const TABLE_WIDTH: u32 = 1024;
6
7#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8pub enum Sampler {
9    Pcg,
10    Sobol { samples: u32 },
11}
12
13impl Default for Sampler {
14    fn default() -> Self {
15        Sampler::Sobol {
16            samples: MIN_SAMPLES,
17        }
18    }
19}
20
21impl Sampler {
22    pub fn sequence(self) -> Option<u32> {
23        if std::env::var("PFX_TRACE_SAMPLER").is_ok_and(|value| value == "pcg") {
24            return None;
25        }
26        match self {
27            Sampler::Pcg => None,
28            Sampler::Sobol { samples } => Some(sequence_size(samples)),
29        }
30    }
31}
32
33pub fn sequence_size(samples: u32) -> u32 {
34    samples
35        .max(1)
36        .checked_next_power_of_two()
37        .unwrap_or(MAX_SAMPLES)
38        .clamp(MIN_SAMPLES, MAX_SAMPLES)
39}
40
41const XORS: [[u32; 32]; 4] = [
42    [0; 32],
43    [
44        0x00000000, 0x00000001, 0x00000001, 0x00000007, 0x00000001, 0x00000013, 0x00000015,
45        0x0000007f, 0x00000001, 0x00000103, 0x00000105, 0x0000070f, 0x00000111, 0x00001333,
46        0x00001555, 0x00007fff, 0x00000001, 0x00010003, 0x00010005, 0x0007000f, 0x00010011,
47        0x00130033, 0x00150055, 0x007f00ff, 0x00010101, 0x01030303, 0x01050505, 0x070f0f0f,
48        0x01111111, 0x13333333, 0x15555555, 0x7fffffff,
49    ],
50    [
51        0x00000000, 0x00000001, 0x00000003, 0x00000001, 0x00000005, 0x0000001f, 0x0000002b,
52        0x0000003d, 0x00000011, 0x00000133, 0x00000377, 0x00000199, 0x00000445, 0x00001ccf,
53        0x00002ddb, 0x0000366d, 0x00000101, 0x00010303, 0x00030707, 0x00010909, 0x00051515,
54        0x001f3f3f, 0x002b6b6b, 0x003dbdbd, 0x00101011, 0x01303033, 0x03707077, 0x01909099,
55        0x04515145, 0x1cf3f3cf, 0x2db6b6db, 0x36dbdb6d,
56    ],
57    [
58        0x00000000, 0x00000001, 0x00000000, 0x00000003, 0x0000000d, 0x0000000c, 0x00000005,
59        0x0000004f, 0x00000014, 0x000000e7, 0x00000329, 0x0000039c, 0x00000011, 0x00001033,
60        0x00000044, 0x000030bb, 0x0000d1cd, 0x0000c2ec, 0x00005415, 0x0004fc3f, 0x00015054,
61        0x000e5c97, 0x0032e5b9, 0x0039725c, 0x00000101, 0x01000303, 0x00000404, 0x03000b0b,
62        0x0d001d1d, 0x0c002c2c, 0x05004545, 0x4f00cfcf,
63    ],
64];
65
66fn hash_hp_uint(mut i: u32) -> u32 {
67    i ^= i >> 16;
68    i = i.wrapping_mul(0x21f0aaad);
69    i ^= i >> 15;
70    i = i.wrapping_mul(0xd35a2d97);
71    i ^= i >> 15;
72    i ^ 0xe6fe3beb
73}
74
75fn hash_hp_seeded_uint(i: u32, mut seed: u32) -> u32 {
76    seed ^= seed << 19;
77    hash_hp_uint(i ^ seed)
78}
79
80fn hash_hp_float(i: u32) -> f32 {
81    hash_hp_uint(i) as f32 * (1.0 / 4_294_967_808.0)
82}
83
84pub fn generate_4d(points: &mut [[f32; 4]], seed: u32) {
85    let size = points.len();
86    if size == 0 {
87        return;
88    }
89    let mut rng = hash_hp_seeded_uint(seed, 0x44605a73);
90    let mut next = || {
91        let value = hash_hp_float(rng);
92        rng = rng.wrapping_add(1);
93        value
94    };
95    points[0] = [next(), next(), next(), next()];
96    let mut log_n = 0;
97    let mut n = 1;
98    while n < size {
99        let strata = (n * 2) as f32;
100        let mut i = 0;
101        while i < n && n + i < size {
102            for axis in 0..4 {
103                let occupied = (points[i ^ XORS[axis][log_n] as usize][axis] * strata) as u32;
104                points[n + i][axis] = ((occupied ^ 1) as f32 + next()) / strata;
105            }
106            i += 1;
107        }
108        log_n += 1;
109        n *= 2;
110    }
111}
112
113pub fn table(sequence: u32) -> Vec<[f32; 4]> {
114    let length = sequence as usize;
115    let mut points = vec![[0.0; 4]; length * PATTERNS as usize];
116    let threads = std::thread::available_parallelism().map_or(1, usize::from);
117    let per = (PATTERNS as usize).div_ceil(threads) * length;
118    std::thread::scope(|scope| {
119        for (chunk, block) in points.chunks_mut(per).enumerate() {
120            scope.spawn(move || {
121                for (j, pattern) in block.chunks_mut(length).enumerate() {
122                    generate_4d(pattern, (chunk * per / length + j) as u32);
123                }
124            });
125        }
126    });
127    points
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[test]
135    fn the_table_is_the_same_bytes_every_time() {
136        let a = table(MIN_SAMPLES);
137        let b = table(MIN_SAMPLES);
138        assert_eq!(a.len(), (MIN_SAMPLES * PATTERNS) as usize);
139        assert!(a.iter().flatten().all(|v| (0.0..1.0).contains(v)));
140        assert_eq!(a, b);
141        let mut single = vec![[0.0; 4]; MIN_SAMPLES as usize];
142        generate_4d(&mut single, 37);
143        assert_eq!(
144            &a[37 * MIN_SAMPLES as usize..38 * MIN_SAMPLES as usize],
145            &single[..]
146        );
147    }
148
149    #[test]
150    fn every_power_of_two_prefix_is_stratified_in_each_axis() {
151        let mut points = vec![[0.0; 4]; 1024];
152        generate_4d(&mut points, 5);
153        for n in [2usize, 16, 128, 1024] {
154            for axis in 0..4 {
155                let mut seen = vec![false; n];
156                for point in &points[..n] {
157                    seen[(point[axis] * n as f32) as usize] = true;
158                }
159                assert!(seen.iter().all(|s| *s), "{n} points, axis {axis}");
160            }
161        }
162    }
163
164    #[test]
165    fn the_sequence_follows_the_sample_count() {
166        assert_eq!(sequence_size(1), MIN_SAMPLES);
167        assert_eq!(sequence_size(300), 512);
168        assert_eq!(sequence_size(4096), 4096);
169        assert_eq!(sequence_size(u32::MAX), MAX_SAMPLES);
170        assert_eq!(Sampler::default().sequence(), Some(MIN_SAMPLES));
171        assert_eq!(Sampler::Pcg.sequence(), None);
172    }
173}
pfx_trace/sobol.rs

pfx_trace/
sobol.rs
