From a53ba4f472b4f67181185cf56882884ba5372ec2 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 17:59:48 +0100 Subject: [PATCH] feat: add Sobol quasi-random sampler behind `sobol` feature flag Add SobolSampler using Owen-scrambled Sobol sequences (sobol_burley crate) for better uniform coverage of the parameter space compared to random sampling. Supports all distribution types including log-scale and stepped parameters. --- Cargo.toml | 2 + src/lib.rs | 6 + src/sampler/mod.rs | 2 + src/sampler/sobol.rs | 420 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 430 insertions(+) create mode 100644 src/sampler/sobol.rs diff --git a/Cargo.toml b/Cargo.toml index 902400a..236cb7f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ optimizer-derive = { version = "0.1.0", path = "optimizer-derive", optional = tr serde = { version = "1", features = ["derive"], optional = true } serde_json = { version = "1", optional = true } tracing = { version = "0.1", optional = true } +sobol_burley = { version = "0.5", optional = true } [features] default = [] @@ -31,6 +32,7 @@ async = ["dep:tokio"] derive = ["dep:optimizer-derive"] serde = ["dep:serde", "dep:serde_json"] tracing = ["dep:tracing"] +sobol = ["dep:sobol_burley"] [dev-dependencies] tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } diff --git a/src/lib.rs b/src/lib.rs index 3ca1885..cf6136d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,6 +17,7 @@ //! - **Random Search** - Simple random sampling for baseline comparisons //! - **TPE (Tree-Parzen Estimator)** - Bayesian optimization for efficient search //! - **Grid Search** - Exhaustive search over a specified parameter grid +//! - **Sobol (QMC)** - Quasi-random sampling for better space coverage (requires `sobol` feature) //! //! Additional features include: //! @@ -183,6 +184,7 @@ //! - `async`: Enable async optimization methods (requires tokio) //! - `derive`: Enable `#[derive(Categorical)]` for enum parameters //! - `serde`: Enable `Serialize`/`Deserialize` on public types and `Study::save()`/`Study::load()` +//! - `sobol`: Enable the Sobol quasi-random sampler for better space coverage //! - `tracing`: Emit structured log events via the [`tracing`](https://docs.rs/tracing) crate at key optimization points /// Emit a `tracing::info!` event when the `tracing` feature is enabled. @@ -234,6 +236,8 @@ pub use pruner::{ pub use sampler::CompletedTrial; pub use sampler::grid::GridSearchSampler; pub use sampler::random::RandomSampler; +#[cfg(feature = "sobol")] +pub use sampler::sobol::SobolSampler; pub use sampler::tpe::TpeSampler; pub use study::Study; #[cfg(feature = "serde")] @@ -262,6 +266,8 @@ pub mod prelude { pub use crate::sampler::CompletedTrial; pub use crate::sampler::grid::GridSearchSampler; pub use crate::sampler::random::RandomSampler; + #[cfg(feature = "sobol")] + pub use crate::sampler::sobol::SobolSampler; pub use crate::sampler::tpe::TpeSampler; pub use crate::study::Study; #[cfg(feature = "serde")] diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index 6f4c07f..ce4b1e5 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -2,6 +2,8 @@ pub mod grid; pub mod random; +#[cfg(feature = "sobol")] +pub mod sobol; pub mod tpe; use std::collections::HashMap; diff --git a/src/sampler/sobol.rs b/src/sampler/sobol.rs new file mode 100644 index 0000000..3c4bb18 --- /dev/null +++ b/src/sampler/sobol.rs @@ -0,0 +1,420 @@ +//! Quasi-random sampler using Sobol low-discrepancy sequences. + +use parking_lot::Mutex; +use sobol_burley::sample; + +use crate::distribution::Distribution; +use crate::param::ParamValue; +use crate::sampler::{CompletedTrial, Sampler}; + +/// Internal state for tracking the dimension counter within a trial. +struct SobolState { + /// The `trial_id` of the current trial (used to reset dimension counter). + current_trial: u64, + /// Next Sobol dimension to use for the current trial. + next_dimension: u32, +} + +/// Quasi-random sampler using Sobol low-discrepancy sequences. +/// +/// Provides better uniform coverage of the parameter space than +/// [`RandomSampler`](super::random::RandomSampler). Useful as a baseline or +/// for the startup phase of model-based samplers. +/// +/// Unlike random sampling, Sobol sequences are deterministic and fill the +/// space more evenly, reducing the number of trials needed to adequately +/// cover the search space. +/// +/// Each trial uses a different Sobol sequence index, and each parameter +/// within a trial maps to a different Sobol dimension. Parameters must be +/// suggested in the same order across trials for consistent dimension +/// assignment. +/// +/// Sobol sequences are most effective in moderate dimensions (up to ~20). +/// For very high dimensions, the uniformity advantage diminishes. +/// +/// Requires the `sobol` feature flag. +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::sobol::SobolSampler; +/// use optimizer::{Direction, Study}; +/// +/// let study: Study = Study::with_sampler(Direction::Minimize, SobolSampler::new()); +/// ``` +pub struct SobolSampler { + seed: u32, + state: Mutex, +} + +impl SobolSampler { + /// Creates a new Sobol sampler with a default seed of 0. + #[must_use] + pub fn new() -> Self { + Self::with_seed(0) + } + + /// Creates a new Sobol sampler with the given seed. + /// + /// Different seeds produce statistically independent Sobol sequences. + /// Using the same seed will produce the same sequence of sampled values. + #[must_use] + #[allow(clippy::cast_possible_truncation)] + pub fn with_seed(seed: u64) -> Self { + Self { + seed: seed as u32, + state: Mutex::new(SobolState { + current_trial: u64::MAX, + next_dimension: 0, + }), + } + } +} + +impl Default for SobolSampler { + fn default() -> Self { + Self::new() + } +} + +impl Sampler for SobolSampler { + #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + _history: &[CompletedTrial], + ) -> ParamValue { + let mut state = self.state.lock(); + + // Reset dimension counter when a new trial starts. + if state.current_trial != trial_id { + state.current_trial = trial_id; + state.next_dimension = 0; + } + + let dimension = state.next_dimension; + state.next_dimension = dimension + 1; + + // Use trial_id as the Sobol sequence index. + let index = trial_id as u32; + + // Generate a quasi-random point in [0, 1). + let point = f64::from(sample(index, dimension, self.seed)); + + map_point_to_distribution(point, distribution) + } +} + +/// Maps a uniform [0, 1) point to a value within the given distribution. +#[allow( + clippy::cast_possible_truncation, + clippy::cast_precision_loss, + clippy::cast_sign_loss +)] +fn map_point_to_distribution(point: f64, distribution: &Distribution) -> ParamValue { + match distribution { + Distribution::Float(d) => { + let value = if d.log_scale { + let log_low = d.low.ln(); + let log_high = d.high.ln(); + (log_low + point * (log_high - log_low)).exp() + } else if let Some(step) = d.step { + let n_steps = ((d.high - d.low) / step).floor() as i64; + let k = (point * (n_steps + 1) as f64).floor() as i64; + let k = k.min(n_steps); + d.low + (k as f64) * step + } else { + d.low + point * (d.high - d.low) + }; + ParamValue::Float(value) + } + Distribution::Int(d) => { + let value = if d.log_scale { + let log_low = (d.low as f64).ln(); + let log_high = (d.high as f64).ln(); + let raw = (log_low + point * (log_high - log_low)).exp().round() as i64; + raw.clamp(d.low, d.high) + } else if let Some(step) = d.step { + let n_steps = (d.high - d.low) / step; + let k = (point * (n_steps + 1) as f64).floor() as i64; + let k = k.min(n_steps); + d.low + k * step + } else { + let range = d.high - d.low + 1; + let k = (point * range as f64).floor() as i64; + (d.low + k).min(d.high) + }; + ParamValue::Int(value) + } + Distribution::Categorical(d) => { + let index = (point * d.n_choices as f64).floor() as usize; + let index = index.min(d.n_choices - 1); + ParamValue::Categorical(index) + } + } +} + +#[cfg(test)] +#[allow( + clippy::cast_possible_truncation, + clippy::cast_precision_loss, + clippy::cast_sign_loss +)] +mod tests { + use super::*; + use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; + + #[test] + fn float_within_bounds() { + let sampler = SobolSampler::with_seed(42); + let dist = Distribution::Float(FloatDistribution { + low: -5.0, + high: 5.0, + log_scale: false, + step: None, + }); + + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); + if let ParamValue::Float(v) = value { + assert!( + (-5.0..=5.0).contains(&v), + "value {v} out of bounds at trial {i}" + ); + } else { + panic!("Expected Float value"); + } + } + } + + #[test] + fn float_log_scale_within_bounds() { + let sampler = SobolSampler::with_seed(42); + let dist = Distribution::Float(FloatDistribution { + low: 1e-5, + high: 1.0, + log_scale: true, + step: None, + }); + + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); + if let ParamValue::Float(v) = value { + assert!( + (1e-5..=1.0).contains(&v), + "value {v} out of bounds at trial {i}" + ); + } else { + panic!("Expected Float value"); + } + } + } + + #[test] + fn float_step_respects_grid() { + let sampler = SobolSampler::with_seed(42); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: Some(0.25), + }); + + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); + if let ParamValue::Float(v) = value { + assert!((0.0..=1.0).contains(&v), "value {v} out of bounds"); + let k = (v / 0.25).round() as i64; + let expected = k as f64 * 0.25; + assert!((v - expected).abs() < 1e-10, "value {v} not on step grid"); + } else { + panic!("Expected Float value"); + } + } + } + + #[test] + fn int_within_bounds() { + let sampler = SobolSampler::with_seed(42); + let dist = Distribution::Int(IntDistribution { + low: 0, + high: 10, + log_scale: false, + step: None, + }); + + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); + if let ParamValue::Int(v) = value { + assert!( + (0..=10).contains(&v), + "value {v} out of bounds at trial {i}" + ); + } else { + panic!("Expected Int value"); + } + } + } + + #[test] + fn int_log_scale_within_bounds() { + let sampler = SobolSampler::with_seed(42); + let dist = Distribution::Int(IntDistribution { + low: 1, + high: 1000, + log_scale: true, + step: None, + }); + + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); + if let ParamValue::Int(v) = value { + assert!( + (1..=1000).contains(&v), + "value {v} out of bounds at trial {i}" + ); + } else { + panic!("Expected Int value"); + } + } + } + + #[test] + fn int_step_respects_grid() { + let sampler = SobolSampler::with_seed(42); + let dist = Distribution::Int(IntDistribution { + low: 0, + high: 10, + log_scale: false, + step: Some(2), + }); + + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); + if let ParamValue::Int(v) = value { + assert!((0..=10).contains(&v), "value {v} out of bounds"); + assert!(v % 2 == 0, "value {v} not on step grid"); + } else { + panic!("Expected Int value"); + } + } + } + + #[test] + fn categorical_within_bounds() { + let sampler = SobolSampler::with_seed(42); + let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 5 }); + + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); + if let ParamValue::Categorical(idx) = value { + assert!(idx < 5, "index {idx} out of bounds at trial {i}"); + } else { + panic!("Expected Categorical value"); + } + } + } + + #[test] + fn deterministic_with_same_seed() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + let sampler1 = SobolSampler::with_seed(42); + let sampler2 = SobolSampler::with_seed(42); + + for i in 0..20 { + let v1 = sampler1.sample(&dist, i, &[]); + let v2 = sampler2.sample(&dist, i, &[]); + assert_eq!(v1, v2, "mismatch at trial {i}"); + } + } + + #[test] + fn different_seeds_produce_different_sequences() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + let sampler1 = SobolSampler::with_seed(0); + let sampler2 = SobolSampler::with_seed(12345); + + let mut any_different = false; + for i in 0..20 { + let v1 = sampler1.sample(&dist, i, &[]); + let v2 = sampler2.sample(&dist, i, &[]); + if v1 != v2 { + any_different = true; + break; + } + } + assert!( + any_different, + "different seeds should produce different sequences" + ); + } + + #[test] + fn better_coverage_than_random() { + // Sobol sequence should cover [0,1] more uniformly than random. + // We measure this by checking that the Sobol samples fill all + // 10 equal-width bins with only 20 samples. + let sampler = SobolSampler::with_seed(0); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + let n_bins = 10; + let n_samples = 20; + let mut bins = vec![0u32; n_bins]; + + for i in 0..n_samples { + let value = sampler.sample(&dist, i, &[]); + if let ParamValue::Float(v) = value { + let bin = ((v * n_bins as f64).floor() as usize).min(n_bins - 1); + bins[bin] += 1; + } + } + + let filled_bins = bins.iter().filter(|&&c| c > 0).count(); + assert!( + filled_bins >= 8, + "Expected at least 8/10 bins filled, got {filled_bins}: {bins:?}" + ); + } + + #[test] + fn multi_parameter_uses_different_dimensions() { + // When sampling multiple parameters per trial, each parameter + // should get a different Sobol dimension, producing different values. + let sampler = SobolSampler::with_seed(0); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Sample two parameters for trial 0. + let v1 = sampler.sample(&dist, 0, &[]); + let v2 = sampler.sample(&dist, 0, &[]); + + // They should differ (different Sobol dimensions for the same index). + assert_ne!( + v1, v2, + "multi-parameter samples should use different dimensions" + ); + } +}