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}