somatize_study/sampler/
sobol.rs

1//! Spread on purpose, and without a seam.
2
3use super::Sampler;
4use super::drawing::{draw, splitmix};
5use crate::{Point, Space};
6
7/// Cover the space evenly, one binary bisection per knob.
8///
9/// The same promise as [`Halton`](super::Halton) — uniform by construction for
10/// every prefix — paid for elsewhere: Halton reads knob `d` in base the `d`-th
11/// prime, so its cover thins as the primes grow, while this reads every knob in
12/// base two and separates them with a set of direction numbers each.
13///
14/// Chosen, and that is the price: a Sobol sequence built on the wrong numbers
15/// does not fail, it covers worse and nobody finds out. These are Joe and Kuo's
16/// (2008), and a test walks the first dimensions against published values,
17/// because reading a table is not a way of checking it. Past [`KNOBS`]
18/// dimensions `ask` answers `None` from the first trial rather than quietly in
19/// the middle. Its point is a function of the **seed and the index**.
20#[derive(Debug, Clone, PartialEq, Eq, Hash)]
21pub struct Sobol {
22    /// The seed, which here shifts the digits rather than drawing them: without
23    /// it a Sobol sequence is one fixed sequence and two studies of the same
24    /// space would walk it in exactly the same order.
25    pub seed: u64,
26}
27
28/// How many knobs the table reaches, and so how many this can search.
29pub const KNOBS: usize = 32;
30
31/// How many numbers each knob's direction gets, which is the width of the whole
32/// arithmetic and so how many trials there are before it comes round again.
33const BITS: usize = 32;
34
35impl Sobol {
36    /// The `trial`-th point, or `None` when there are more knobs than the table
37    /// has. It never runs out, and it never looks at what the finished trials
38    /// did.
39    pub fn ask(
40        &self,
41        space: &Space,
42        trial: usize,
43        _seen: &[(Point, Option<f64>)],
44    ) -> Option<Point> {
45        if space.is_empty() || space.len() > KNOBS {
46            return None;
47        }
48        // Gray code: consecutive trials differ in exactly one bit, so each point
49        // is the one before it with a single direction number flipped into it —
50        // which is what keeps every prefix balanced and not just the whole.
51        let step = trial as u64 ^ (trial as u64 >> 1);
52        Some(Point::of(
53            space
54                .dimensions()
55                .iter()
56                .enumerate()
57                .map(|(which, (name, dimension))| {
58                    let numbers = directions(which);
59                    let mut place = shift(self.seed, which);
60                    for (bit, number) in numbers.iter().enumerate() {
61                        if (step >> bit) & 1 == 1 {
62                            place ^= number;
63                        }
64                    }
65                    let u = place as f64 / (1u64 << BITS) as f64;
66                    (name.clone(), draw(dimension, u))
67                })
68                .collect(),
69        ))
70    }
71}
72
73/// This knob's direction numbers: the `i`-th is folded in when bit `i` of the
74/// step is set. The first `s` come from the table; the rest are the recurrence
75/// the primitive polynomial stands for.
76fn directions(which: usize) -> [u32; BITS] {
77    let (a, m) = DIRECTIONS[which];
78    let mut numbers = [0u32; BITS];
79    // The first knob has no polynomial and no recurrence: its numbers are a
80    // single one walking rightwards, which is plain bisection.
81    if m.is_empty() {
82        for (i, number) in numbers.iter_mut().enumerate() {
83            *number = 1 << (BITS - 1 - i);
84        }
85        return numbers;
86    }
87    let s = m.len();
88    for i in 0..s {
89        numbers[i] = m[i] << (BITS - 1 - i);
90    }
91    for i in s..BITS {
92        numbers[i] = numbers[i - s] ^ (numbers[i - s] >> s);
93        for k in 1..s {
94            numbers[i] ^= ((a >> (s - 1 - k)) & 1) * numbers[i - k];
95        }
96    }
97    numbers
98}
99
100/// What this knob's numbers are shifted by. A digital shift and not a
101/// rearrangement: xoring a constant into every point moves the whole set without
102/// disturbing which cell each one is in, so the cover survives the seed intact.
103fn shift(seed: u64, which: usize) -> u32 {
104    let mut state = seed ^ (which as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15);
105    splitmix(&mut state) as u32
106}
107
108/// Joe and Kuo's direction numbers, `(the polynomial's middle coefficients, the
109/// numbers it starts from)`, one row per knob. *Constructing Sobol sequences
110/// with better two-dimensional projections*, SIAM J. Sci. Comput. 30, 2635–2654
111/// (2008) — the `new-joe-kuo-6` table.
112const DIRECTIONS: [(u32, &[u32]); KNOBS] = [
113    (0, &[]),
114    (0, &[1]),
115    (1, &[1, 3]),
116    (1, &[1, 3, 1]),
117    (2, &[1, 1, 1]),
118    (1, &[1, 1, 3, 3]),
119    (4, &[1, 3, 5, 13]),
120    (2, &[1, 1, 5, 5, 17]),
121    (4, &[1, 1, 5, 5, 5]),
122    (7, &[1, 1, 7, 11, 19]),
123    (11, &[1, 1, 5, 1, 1]),
124    (13, &[1, 1, 1, 3, 11]),
125    (14, &[1, 3, 5, 5, 31]),
126    (1, &[1, 3, 3, 9, 7, 49]),
127    (13, &[1, 1, 1, 15, 21, 21]),
128    (16, &[1, 3, 1, 13, 27, 49]),
129    (19, &[1, 1, 1, 15, 7, 5]),
130    (22, &[1, 3, 1, 15, 13, 25]),
131    (25, &[1, 1, 5, 5, 19, 61]),
132    (1, &[1, 3, 7, 11, 23, 15, 103]),
133    (4, &[1, 3, 7, 13, 13, 15, 69]),
134    (7, &[1, 1, 3, 13, 7, 35, 63]),
135    (8, &[1, 3, 5, 9, 1, 25, 53]),
136    (14, &[1, 3, 1, 13, 9, 35, 107]),
137    (19, &[1, 3, 1, 5, 27, 61, 31]),
138    (21, &[1, 1, 5, 11, 19, 41, 61]),
139    (28, &[1, 3, 5, 3, 3, 13, 69]),
140    (31, &[1, 1, 7, 13, 1, 19, 1]),
141    (32, &[1, 3, 7, 5, 13, 19, 59]),
142    (37, &[1, 1, 3, 9, 25, 29, 41]),
143    (41, &[1, 3, 5, 13, 23, 1, 55]),
144    (42, &[1, 3, 7, 3, 13, 59, 17]),
145];
146
147impl From<Sobol> for Sampler {
148    fn from(how: Sobol) -> Self {
149        Self::Sobol(how)
150    }
151}
somatize_study/sampler/sobol.rs

somatize_study/sampler/
sobol.rs
