Skip to main content

pounce_cli/minima/
sampling.rs

1//! Start-point sampling for `--minima`: scrambled Sobol' or uniform box
2//! sampling when a finite box is known, else Gaussian jitter around `x0`.
3//! Mirrors `_sample` / `_make_sobol` in `python/pounce/_minima.py`.
4//!
5//! Reproducibility: the PRNG is a seeded `ChaCha8Rng` and the Sobol' source
6//! is `sobol_burley` (Owen-scrambled, matching scipy's
7//! `qmc.Sobol(scramble=True)`), both keyed off `--seed`.
8
9use pounce_common::types::Number;
10use rand::distributions::{Distribution, Standard};
11use rand_chacha::rand_core::SeedableRng;
12use rand_chacha::ChaCha8Rng;
13
14/// Seeded sampler shared across a `--minima` run.
15pub struct Sampler {
16    rng: ChaCha8Rng,
17    /// Sobol' scramble key (derived from the run seed).
18    sobol_seed: u32,
19    /// Next Sobol' point index (advances by one per drawn sample).
20    sobol_index: u32,
21    /// Whether to draw box samples from the Sobol' sequence.
22    sobol: bool,
23}
24
25impl Sampler {
26    pub fn new(seed: u64, sobol: bool) -> Self {
27        Self {
28            rng: ChaCha8Rng::seed_from_u64(seed),
29            // Fold the 64-bit seed into the 32-bit Sobol' scramble key.
30            sobol_seed: (seed ^ (seed >> 32)) as u32,
31            sobol_index: 0,
32            sobol,
33        }
34    }
35
36    /// A uniform deviate in `[0, 1)` (avoids the `gen` keyword/method so the
37    /// code stays edition-2024 clean).
38    pub fn uniform(&mut self) -> f64 {
39        Standard.sample(&mut self.rng)
40    }
41
42    /// A standard-normal deviate via Box–Muller (avoids a `rand_distr` dep).
43    pub fn standard_normal(&mut self) -> f64 {
44        // u1 in (0,1] so ln() is finite.
45        let u1: f64 = 1.0 - self.uniform();
46        let u2: f64 = self.uniform();
47        (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
48    }
49
50    /// Draw a fresh start point.
51    ///
52    /// * `has_box` true ⇒ sample uniformly (Sobol' or PRNG) in `[lo, hi]`.
53    /// * otherwise ⇒ `x0 + jitter · N(0, I)` (the unbounded fallback).
54    pub fn sample(
55        &mut self,
56        x0: &[Number],
57        lo: &[Number],
58        hi: &[Number],
59        has_box: bool,
60        jitter: f64,
61    ) -> Vec<Number> {
62        let n = x0.len();
63        if has_box {
64            let idx = self.sobol_index;
65            self.sobol_index += 1;
66            (0..n)
67                .map(|j| {
68                    let u = if self.sobol {
69                        // One Owen-scrambled Sobol' coordinate per dimension.
70                        sobol_burley::sample(idx, j as u32, self.sobol_seed) as f64
71                    } else {
72                        self.uniform()
73                    };
74                    lo[j] + (hi[j] - lo[j]) * u
75                })
76                .collect()
77        } else {
78            (0..n)
79                .map(|j| x0[j] + jitter * self.standard_normal())
80                .collect()
81        }
82    }
83
84    /// `anchor + scale · N(0, I)` — a local perturbation (basinhopping step,
85    /// flooding re-seed, tunnelling outward step). `scale` may be a scalar
86    /// or per-dimension (length-1 ⇒ isotropic).
87    pub fn perturb(&mut self, anchor: &[Number], scale: &[Number]) -> Vec<Number> {
88        anchor
89            .iter()
90            .enumerate()
91            .map(|(j, &a)| {
92                let s = if scale.len() == 1 { scale[0] } else { scale[j] };
93                a + s * self.standard_normal()
94            })
95            .collect()
96    }
97}
98
99/// Clip `x` into `[lo, hi]` when a finite box is known (no-op otherwise).
100pub fn clip(x: &mut [Number], lo: &[Number], hi: &[Number], has_box: bool) {
101    if !has_box {
102        return;
103    }
104    for j in 0..x.len() {
105        x[j] = x[j].clamp(lo[j], hi[j]);
106    }
107}