From 947701e83fa04215ada254a2dba8261c4e3b4d2f Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 21:27:40 +0100 Subject: [PATCH] feat: add Gaussian Process sampler with Expected Improvement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement a classical Bayesian optimization sampler using a GP surrogate with Matérn 5/2 kernel and Expected Improvement acquisition function. Feature-gated behind `gp = ["dep:nalgebra"]`. --- Cargo.toml | 1 + src/lib.rs | 6 + src/sampler/gp.rs | 847 +++++++++++++++++++++++++++++++++++++++++++++ src/sampler/mod.rs | 2 + tests/gp_tests.rs | 256 ++++++++++++++ 5 files changed, 1112 insertions(+) create mode 100644 src/sampler/gp.rs create mode 100644 tests/gp_tests.rs diff --git a/Cargo.toml b/Cargo.toml index ff91328..f755d86 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ serde = ["dep:serde", "dep:serde_json"] tracing = ["dep:tracing"] sobol = ["dep:sobol_burley"] cma-es = ["dep:nalgebra"] +gp = ["dep:nalgebra"] visualization = [] fanova = [] diff --git a/src/lib.rs b/src/lib.rs index 27e9ac8..3943f37 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,6 +19,7 @@ //! - **Grid Search** - Exhaustive search over a specified parameter grid //! - **Sobol (QMC)** - Quasi-random sampling for better space coverage (requires `sobol` feature) //! - **CMA-ES** - Covariance Matrix Adaptation Evolution Strategy for continuous optimization (requires `cma-es` feature) +//! - **GP** - Gaussian Process Bayesian optimization with Expected Improvement (requires `gp` feature) //! - **BOHB** - Bayesian Optimization + `HyperBand` for budget-aware TPE sampling //! - **NSGA-II** - Non-dominated Sorting Genetic Algorithm II for multi-objective optimization //! - **MOTPE** - Multi-Objective Tree-Parzen Estimator for Bayesian multi-objective optimization @@ -190,6 +191,7 @@ //! - `serde`: Enable `Serialize`/`Deserialize` on public types and `Study::save()`/`Study::load()` //! - `sobol`: Enable the Sobol quasi-random sampler for better space coverage //! - `cma-es`: Enable the CMA-ES sampler for continuous optimization +//! - `gp`: Enable the Gaussian Process sampler for Bayesian optimization //! - `visualization`: Generate self-contained HTML reports with interactive Plotly.js charts //! - `tracing`: Emit structured log events via the [`tracing`](https://docs.rs/tracing) crate at key optimization points @@ -253,6 +255,8 @@ pub use sampler::CompletedTrial; pub use sampler::bohb::BohbSampler; #[cfg(feature = "cma-es")] pub use sampler::cma_es::CmaEsSampler; +#[cfg(feature = "gp")] +pub use sampler::gp::GpSampler; pub use sampler::grid::GridSearchSampler; pub use sampler::motpe::MotpeSampler; pub use sampler::nsga2::Nsga2Sampler; @@ -293,6 +297,8 @@ pub mod prelude { pub use crate::sampler::bohb::BohbSampler; #[cfg(feature = "cma-es")] pub use crate::sampler::cma_es::CmaEsSampler; + #[cfg(feature = "gp")] + pub use crate::sampler::gp::GpSampler; pub use crate::sampler::grid::GridSearchSampler; pub use crate::sampler::motpe::MotpeSampler; pub use crate::sampler::nsga2::Nsga2Sampler; diff --git a/src/sampler/gp.rs b/src/sampler/gp.rs new file mode 100644 index 0000000..211d85e --- /dev/null +++ b/src/sampler/gp.rs @@ -0,0 +1,847 @@ +//! Gaussian Process (GP) sampler with Expected Improvement acquisition. +//! +//! A classical Bayesian optimization sampler that uses a Gaussian Process +//! surrogate model with a Matérn 5/2 kernel and Expected Improvement (EI) +//! acquisition function. Best suited for small, expensive evaluations in +//! low-dimensional continuous spaces (d ≤ 20). +//! +//! Categorical parameters are sampled uniformly at random (not part of +//! the GP model). If all parameters are categorical, the sampler falls +//! back to pure random sampling. +//! +//! Requires the `gp` feature flag. +//! +//! # Examples +//! +//! ``` +//! use optimizer::sampler::gp::GpSampler; +//! use optimizer::{Direction, Study}; +//! +//! let sampler = GpSampler::with_seed(42); +//! let study: Study = Study::with_sampler(Direction::Minimize, sampler); +//! ``` + +use std::collections::HashMap; + +use nalgebra::DMatrix; +use parking_lot::Mutex; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; + +use crate::distribution::Distribution; +use crate::param::ParamValue; +use crate::sampler::{CompletedTrial, Sampler}; + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +/// Gaussian Process sampler for Bayesian optimization. +/// +/// Uses a GP surrogate with Matérn 5/2 kernel and Expected Improvement +/// acquisition to guide sampling toward promising regions of the search +/// space. Best suited for continuous (float/int) parameters in low +/// dimensions (up to ~20). +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::gp::GpSampler; +/// use optimizer::{Direction, Study}; +/// +/// // Default configuration +/// let study: Study = Study::with_sampler(Direction::Minimize, GpSampler::new()); +/// +/// // With seed for reproducibility +/// let study: Study = Study::with_sampler(Direction::Minimize, GpSampler::with_seed(42)); +/// +/// // Custom configuration via builder +/// let sampler = GpSampler::builder() +/// .n_startup_trials(15) +/// .n_candidates(2000) +/// .noise_variance(1e-4) +/// .seed(42) +/// .build(); +/// let study: Study = Study::with_sampler(Direction::Minimize, sampler); +/// ``` +pub struct GpSampler { + state: Mutex, +} + +impl GpSampler { + /// Creates a new GP sampler with a random seed. + #[must_use] + pub fn new() -> Self { + Self { + state: Mutex::new(GpState::new(None, None, None, None)), + } + } + + /// Creates a new GP sampler with a fixed seed for reproducibility. + #[must_use] + pub fn with_seed(seed: u64) -> Self { + Self { + state: Mutex::new(GpState::new(None, None, None, Some(seed))), + } + } + + /// Creates a builder for configuring a `GpSampler`. + #[must_use] + pub fn builder() -> GpSamplerBuilder { + GpSamplerBuilder::new() + } +} + +impl Default for GpSampler { + fn default() -> Self { + Self::new() + } +} + +/// Builder for configuring a [`GpSampler`]. +/// +/// All options have sensible defaults: +/// - `n_startup_trials`: 10 +/// - `n_candidates`: 1000 +/// - `noise_variance`: 1e-6 +/// - `seed`: random +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::gp::GpSamplerBuilder; +/// +/// let sampler = GpSamplerBuilder::new() +/// .n_startup_trials(15) +/// .n_candidates(2000) +/// .noise_variance(1e-4) +/// .seed(42) +/// .build(); +/// ``` +#[derive(Debug, Clone, Default)] +pub struct GpSamplerBuilder { + n_startup_trials: Option, + n_candidates: Option, + noise_variance: Option, + seed: Option, +} + +impl GpSamplerBuilder { + /// Creates a new builder with default settings. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Sets the number of random trials before GP-guided sampling begins. + /// + /// Default: 10. + #[must_use] + pub fn n_startup_trials(mut self, n: usize) -> Self { + self.n_startup_trials = Some(n); + self + } + + /// Sets the number of random candidate points for acquisition optimization. + /// + /// More candidates improve the quality of the acquisition maximum + /// at the cost of more GP predictions per trial. + /// + /// Default: 1000. + #[must_use] + pub fn n_candidates(mut self, n: usize) -> Self { + self.n_candidates = Some(n); + self + } + + /// Sets the observation noise variance added to the kernel diagonal. + /// + /// Controls the assumed noise level. Larger values make the GP smoother. + /// + /// Default: 1e-6 (near-noiseless). + #[must_use] + pub fn noise_variance(mut self, v: f64) -> Self { + self.noise_variance = Some(v); + self + } + + /// Sets the random seed for reproducibility. + #[must_use] + pub fn seed(mut self, seed: u64) -> Self { + self.seed = Some(seed); + self + } + + /// Builds the configured [`GpSampler`]. + #[must_use] + pub fn build(self) -> GpSampler { + GpSampler { + state: Mutex::new(GpState::new( + self.n_startup_trials, + self.n_candidates, + self.noise_variance, + self.seed, + )), + } + } +} + +// --------------------------------------------------------------------------- +// Internal types +// --------------------------------------------------------------------------- + +/// Default number of random startup trials before GP kicks in. +const DEFAULT_N_STARTUP: usize = 10; +/// Default number of candidate points for EI optimization. +const DEFAULT_N_CANDIDATES: usize = 1000; +/// Default observation noise variance. +const DEFAULT_NOISE_VAR: f64 = 1e-6; + +/// Describes how a parameter dimension maps into the GP internal vector. +#[derive(Clone, Debug)] +struct DimensionInfo { + distribution: Distribution, + is_continuous: bool, + bounds: Option<(f64, f64)>, +} + +/// Tracks per-trial sampling progress. +#[derive(Clone, Debug)] +struct TrialProgress { + /// The candidate values for each dimension. + values: Vec, + /// Next dimension to return. + next_dim: usize, +} + +/// Phase of the GP state machine. +enum GpPhase { + /// Discovering the search space (first trial). + Discovery, + /// Steady-state sampling. + Active, +} + +/// A fitted GP model ready for predictions. +struct GpModel { + /// Cholesky factor L of K + σ²I. + cholesky: nalgebra::linalg::Cholesky, + /// α = (K + σ²I)^{-1} y. + alpha: nalgebra::DVector, + /// Training inputs (each row is a data point, normalized to [0, 1]). + x_train: Vec>, + /// ARD lengthscales per dimension. + lengthscales: Vec, + /// Signal variance. + signal_var: f64, + /// Mean of original y values (for un-standardization, unused but kept for diagnostics). + _y_mean: f64, + /// Std dev of original y values (unused but kept for diagnostics). + _y_std: f64, + /// Best observed (standardized) y. + f_best: f64, +} + +/// Top-level mutable state behind the `Mutex`. +struct GpState { + rng: StdRng, + n_startup_trials: usize, + n_candidates: usize, + noise_variance: f64, + phase: GpPhase, + dimensions: Vec, + trial_progress: HashMap, + discovery_trial_id: Option, +} + +impl GpState { + fn new( + n_startup: Option, + n_candidates: Option, + noise_var: Option, + seed: Option, + ) -> Self { + let rng = seed.map_or_else(rand::make_rng, StdRng::seed_from_u64); + Self { + rng, + n_startup_trials: n_startup.unwrap_or(DEFAULT_N_STARTUP), + n_candidates: n_candidates.unwrap_or(DEFAULT_N_CANDIDATES), + noise_variance: noise_var.unwrap_or(DEFAULT_NOISE_VAR), + phase: GpPhase::Discovery, + dimensions: Vec::new(), + trial_progress: HashMap::new(), + discovery_trial_id: None, + } + } +} + +// --------------------------------------------------------------------------- +// Matérn 5/2 kernel +// --------------------------------------------------------------------------- + +/// Matérn 5/2 kernel with ARD lengthscales. +/// +/// `k(x1, x2) = σ² (1 + √5 r + 5/3 r²) exp(-√5 r)` +/// where `r = sqrt(Σ ((x1_i - x2_i) / l_i)²)` +fn matern52(x1: &[f64], x2: &[f64], lengthscales: &[f64], signal_var: f64) -> f64 { + let mut r_sq = 0.0; + for i in 0..x1.len() { + let diff = (x1[i] - x2[i]) / lengthscales[i]; + r_sq += diff * diff; + } + let r = r_sq.sqrt(); + let sqrt5_r = SQRT_5 * r; + signal_var * (1.0 + sqrt5_r + 5.0 / 3.0 * r_sq) * (-sqrt5_r).exp() +} + +/// Build the kernel matrix `K + σ²I`. +fn kernel_matrix( + x: &[Vec], + lengthscales: &[f64], + signal_var: f64, + noise_var: f64, +) -> DMatrix { + let n = x.len(); + DMatrix::from_fn(n, n, |i, j| { + let k = matern52(&x[i], &x[j], lengthscales, signal_var); + if i == j { k + noise_var } else { k } + }) +} + +/// Compute the kernel vector k(x*, X) for a test point. +fn kernel_vector( + x_star: &[f64], + x_train: &[Vec], + lengthscales: &[f64], + signal_var: f64, +) -> nalgebra::DVector { + nalgebra::DVector::from_fn(x_train.len(), |i, _| { + matern52(x_star, &x_train[i], lengthscales, signal_var) + }) +} + +/// Precomputed √5 constant. +const SQRT_5: f64 = 2.236_213_562_373_095; + +// --------------------------------------------------------------------------- +// GP fitting and prediction +// --------------------------------------------------------------------------- + +/// Fit a GP model to the training data. +/// +/// Returns `None` if fitting fails (e.g. Cholesky decomposition failure). +#[allow(clippy::cast_precision_loss)] +fn fit_gp(x_train: &[Vec], y_train: &[f64], noise_var: f64) -> Option { + let n = y_train.len(); + if n == 0 { + return None; + } + + // Standardize y + let y_mean = y_train.iter().sum::() / n as f64; + let y_var = if n > 1 { + y_train.iter().map(|&y| (y - y_mean).powi(2)).sum::() / (n - 1) as f64 + } else { + 1.0 + }; + let y_std = y_var.sqrt().max(1e-10); + let y_standardized: Vec = y_train.iter().map(|&y| (y - y_mean) / y_std).collect(); + + let f_best = y_standardized.iter().copied().fold(f64::INFINITY, f64::min); + + // ARD lengthscales: per-dimension std dev of training X, clamped + let d = if x_train.is_empty() { + 0 + } else { + x_train[0].len() + }; + let lengthscales: Vec = (0..d) + .map(|j| { + let vals: Vec = x_train.iter().map(|x| x[j]).collect(); + let mean_j = vals.iter().sum::() / n as f64; + let var_j = vals.iter().map(|&v| (v - mean_j).powi(2)).sum::() / n as f64; + var_j.sqrt().max(0.01) + }) + .collect(); + + // Signal variance = 1.0 (data is standardized) + let signal_var = 1.0; + + let k = kernel_matrix(x_train, &lengthscales, signal_var, noise_var); + let cholesky = nalgebra::linalg::Cholesky::new(k)?; + + // α = (K + σ²I)^{-1} y + let y_vec = nalgebra::DVector::from_column_slice(&y_standardized); + let alpha = cholesky.solve(&y_vec); + + Some(GpModel { + cholesky, + alpha, + x_train: x_train.to_vec(), + lengthscales, + signal_var, + _y_mean: y_mean, + _y_std: y_std, + f_best, + }) +} + +/// Predict mean and standard deviation at a test point. +fn predict(model: &GpModel, x: &[f64]) -> (f64, f64) { + let k_star = kernel_vector(x, &model.x_train, &model.lengthscales, model.signal_var); + + // Mean: k*^T α + let mean = k_star.dot(&model.alpha); + + // Variance: k(x*, x*) - k*^T (K + σ²I)^{-1} k* + let k_self = model.signal_var; + let v = model.cholesky.solve(&k_star); + let var = (k_self - k_star.dot(&v)).max(0.0); + + (mean, var.sqrt()) +} + +// --------------------------------------------------------------------------- +// Normal distribution helpers (Abramowitz-Stegun approximation) +// --------------------------------------------------------------------------- + +/// Standard normal PDF. +fn norm_pdf(x: f64) -> f64 { + const INV_SQRT_2PI: f64 = 0.398_942_280_401_432_7; + INV_SQRT_2PI * (-0.5 * x * x).exp() +} + +/// Standard normal CDF (Abramowitz-Stegun rational approximation). +fn norm_cdf(x: f64) -> f64 { + // Hart approximation (higher precision than basic A&S) + if x < -8.0 { + return 0.0; + } + if x > 8.0 { + return 1.0; + } + + let abs_x = x.abs(); + let t = 1.0 / (1.0 + 0.231_641_9 * abs_x); + let t2 = t * t; + let t3 = t2 * t; + let t4 = t3 * t; + let t5 = t4 * t; + + let poly = 0.319_381_530 * t - 0.356_563_782 * t2 + 1.781_477_937 * t3 - 1.821_255_978 * t4 + + 1.330_274_429 * t5; + let pdf = norm_pdf(abs_x); + let cdf = 1.0 - pdf * poly; + + if x >= 0.0 { cdf } else { 1.0 - cdf } +} + +// --------------------------------------------------------------------------- +// Expected Improvement +// --------------------------------------------------------------------------- + +/// Compute Expected Improvement at a point. +/// +/// `EI(x) = (f_best - mean) Φ(z) + std φ(z)` +/// where `z = (f_best - mean) / std` +fn expected_improvement(mean: f64, std: f64, f_best: f64) -> f64 { + if std < 1e-12 { + return (f_best - mean).max(0.0); + } + let z = (f_best - mean) / std; + let improvement = (f_best - mean) * norm_cdf(z) + std * norm_pdf(z); + improvement.max(0.0) +} + +// --------------------------------------------------------------------------- +// Acquisition optimization +// --------------------------------------------------------------------------- + +/// Find the point in [0, 1]^d that maximizes EI via multi-start random search. +fn optimize_acquisition( + model: &GpModel, + n_dims: usize, + n_candidates: usize, + rng: &mut StdRng, +) -> Vec { + let mut best_ei = f64::NEG_INFINITY; + let mut best_x = vec![0.5; n_dims]; + + for _ in 0..n_candidates { + let x: Vec = (0..n_dims).map(|_| rng.random_range(0.0..=1.0)).collect(); + let (mean, std) = predict(model, &x); + let ei = expected_improvement(mean, std, model.f_best); + if ei > best_ei { + best_ei = ei; + best_x = x; + } + } + + best_x +} + +// --------------------------------------------------------------------------- +// Data preprocessing helpers +// --------------------------------------------------------------------------- + +/// Compute internal-space bounds for a distribution. +#[allow(clippy::cast_precision_loss)] +fn internal_bounds(distribution: &Distribution) -> Option<(f64, f64)> { + match distribution { + Distribution::Float(d) => { + if d.log_scale { + Some((d.low.ln(), d.high.ln())) + } else { + Some((d.low, d.high)) + } + } + Distribution::Int(d) => { + if d.log_scale { + Some(((d.low as f64).ln(), (d.high as f64).ln())) + } else { + Some((d.low as f64, d.high as f64)) + } + } + Distribution::Categorical(_) => None, + } +} + +/// Convert a value from internal space to a `ParamValue` in original space. +#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] +fn from_internal(value: f64, distribution: &Distribution) -> ParamValue { + match distribution { + Distribution::Float(d) => { + let v = if d.log_scale { value.exp() } else { value }; + let v = if let Some(step) = d.step { + let k = ((v - d.low) / step).round(); + d.low + k * step + } else { + v + }; + ParamValue::Float(v.clamp(d.low, d.high)) + } + Distribution::Int(d) => { + let v = if d.log_scale { value.exp() } else { value }; + let v = if let Some(step) = d.step { + let k = ((v - d.low as f64) / step as f64).round() as i64; + d.low + k * step + } else { + v.round() as i64 + }; + ParamValue::Int(v.clamp(d.low, d.high)) + } + Distribution::Categorical(_) => { + unreachable!("from_internal should not be called for categorical distributions") + } + } +} + +/// Convert an internal-space value to normalized [0, 1] using bounds. +fn to_normalized(value: f64, lo: f64, hi: f64) -> f64 { + if (hi - lo).abs() < 1e-15 { + 0.5 + } else { + (value - lo) / (hi - lo) + } +} + +/// Convert a normalized [0, 1] value back to internal space. +fn from_normalized(value: f64, lo: f64, hi: f64) -> f64 { + lo + value * (hi - lo) +} + +/// Convert a `ParamValue` to its internal-space representation. +#[allow(clippy::cast_precision_loss)] +fn to_internal(value: &ParamValue, distribution: &Distribution) -> f64 { + match (value, distribution) { + (ParamValue::Float(v), Distribution::Float(d)) => { + if d.log_scale { + v.ln() + } else { + *v + } + } + (ParamValue::Int(v), Distribution::Int(d)) => { + if d.log_scale { + (*v as f64).ln() + } else { + *v as f64 + } + } + _ => 0.0, + } +} + +/// Sample a random value for any distribution. +#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] +fn sample_random(rng: &mut StdRng, 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(); + rng.random_range(log_low..=log_high).exp() + } else if let Some(step) = d.step { + let n_steps = ((d.high - d.low) / step).floor() as i64; + let k = rng.random_range(0..=n_steps); + d.low + (k as f64) * step + } else { + rng.random_range(d.low..=d.high) + }; + 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 = rng.random_range(log_low..=log_high).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 = rng.random_range(0..=n_steps); + d.low + k * step + } else { + rng.random_range(d.low..=d.high) + }; + ParamValue::Int(value) + } + Distribution::Categorical(d) => ParamValue::Categorical(rng.random_range(0..d.n_choices)), + } +} + +// --------------------------------------------------------------------------- +// Extract training data from history +// --------------------------------------------------------------------------- + +/// Maximum number of training points to use for the GP. +/// Caps computational cost at O(`MAX_TRAIN_POINTS`^3) per trial. +const MAX_TRAIN_POINTS: usize = 100; + +/// Establish a deterministic mapping from dimension index to `ParamId` +/// using the first trial in history. +/// +/// Matches dimensions to params by distribution equality, consuming +/// matched params to correctly handle duplicate distributions. +fn establish_param_mapping( + trial: &CompletedTrial, + dimensions: &[DimensionInfo], +) -> Vec> { + use crate::parameter::ParamId; + + let mut available: Vec<(ParamId, &Distribution)> = + trial.distributions.iter().map(|(id, d)| (*id, d)).collect(); + // Sort for deterministic matching order + available.sort_by_key(|(id, _)| *id); + + let mut mapping = Vec::with_capacity(dimensions.len()); + for dim in dimensions { + let pos = available.iter().position(|(_, d)| **d == dim.distribution); + if let Some(pos) = pos { + mapping.push(Some(available.remove(pos).0)); + } else { + mapping.push(None); + } + } + mapping +} + +/// Build normalized training data from completed trials. +/// +/// Returns `(x_train, y_train)` where x values are normalized to [0, 1] +/// per dimension using the bounds from `dimensions`. Only continuous +/// dimensions are included. Uses at most [`MAX_TRAIN_POINTS`] most recent +/// trials. +#[allow(clippy::cast_precision_loss)] +fn build_training_data( + history: &[CompletedTrial], + dimensions: &[DimensionInfo], +) -> (Vec>, Vec) { + if history.is_empty() { + return (Vec::new(), Vec::new()); + } + + // Use only the most recent trials to cap GP fitting cost + let start = history.len().saturating_sub(MAX_TRAIN_POINTS); + let recent = &history[start..]; + + // Establish dimension → ParamId mapping from the first trial + let param_mapping = establish_param_mapping(&recent[0], dimensions); + + let continuous_indices: Vec = dimensions + .iter() + .enumerate() + .filter(|(_, d)| d.is_continuous) + .map(|(i, _)| i) + .collect(); + + let mut x_train = Vec::with_capacity(recent.len()); + let mut y_train = Vec::with_capacity(recent.len()); + + for trial in recent { + let mut x_row = Vec::with_capacity(continuous_indices.len()); + let mut valid = true; + + for &dim_idx in &continuous_indices { + let dim_info = &dimensions[dim_idx]; + if let Some(param_id) = param_mapping[dim_idx] { + if let Some(param_val) = trial.params.get(¶m_id) { + let internal = to_internal(param_val, &dim_info.distribution); + let (lo, hi) = dim_info.bounds.unwrap_or((0.0, 1.0)); + x_row.push(to_normalized(internal, lo, hi)); + } else { + valid = false; + break; + } + } else { + valid = false; + break; + } + } + + if valid && x_row.len() == continuous_indices.len() { + x_train.push(x_row); + y_train.push(trial.value); + } + } + + (x_train, y_train) +} + +// --------------------------------------------------------------------------- +// Sampler trait implementation +// --------------------------------------------------------------------------- + +impl Sampler for GpSampler { + #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)] + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + history: &[CompletedTrial], + ) -> ParamValue { + let mut state = self.state.lock(); + + match &state.phase { + GpPhase::Discovery => sample_discovery(&mut state, distribution, trial_id), + GpPhase::Active => sample_active(&mut state, distribution, trial_id, history), + } + } +} + +/// Handle sampling during the discovery phase. +fn sample_discovery(state: &mut GpState, distribution: &Distribution, trial_id: u64) -> ParamValue { + // A new trial_id means discovery is done + if let Some(prev_id) = state.discovery_trial_id + && trial_id != prev_id + { + finalize_discovery(state); + return sample_active(state, distribution, trial_id, &[]); + } + + state.discovery_trial_id = Some(trial_id); + + let is_continuous = !matches!(distribution, Distribution::Categorical(_)); + let bounds = internal_bounds(distribution); + state.dimensions.push(DimensionInfo { + distribution: distribution.clone(), + is_continuous, + bounds, + }); + + sample_random(&mut state.rng, distribution) +} + +/// Finalize discovery and transition to the active phase. +fn finalize_discovery(state: &mut GpState) { + state.phase = GpPhase::Active; + state.trial_progress.clear(); +} + +/// Handle sampling during the active phase. +fn sample_active( + state: &mut GpState, + distribution: &Distribution, + trial_id: u64, + history: &[CompletedTrial], +) -> ParamValue { + // If this trial already has progress, return the next pre-computed value + if let Some(progress) = state.trial_progress.get_mut(&trial_id) { + let dim_idx = progress.next_dim; + progress.next_dim += 1; + if dim_idx < progress.values.len() { + return progress.values[dim_idx].clone(); + } + // Extra dimension not seen during discovery + return sample_random(&mut state.rng, distribution); + } + + // New trial: compute all dimension values at once + let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count(); + let use_gp = n_continuous > 0 && history.len() >= state.n_startup_trials; + + let values = if use_gp { + compute_gp_candidate(state, history) + } else { + compute_random_candidate(state) + }; + + let first_value = values + .first() + .cloned() + .unwrap_or_else(|| sample_random(&mut state.rng, distribution)); + + state.trial_progress.insert( + trial_id, + TrialProgress { + values, + next_dim: 1, + }, + ); + + first_value +} + +/// Compute a candidate using the GP model. +fn compute_gp_candidate(state: &mut GpState, history: &[CompletedTrial]) -> Vec { + let (x_train, y_train) = build_training_data(history, &state.dimensions); + + // Try to fit GP; fall back to random if it fails + let model = fit_gp(&x_train, &y_train, state.noise_variance); + + let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count(); + + let normalized_candidate = if let Some(ref model) = model { + optimize_acquisition(model, n_continuous, state.n_candidates, &mut state.rng) + } else { + // GP fitting failed; use random + (0..n_continuous) + .map(|_| state.rng.random_range(0.0..=1.0)) + .collect() + }; + + // Convert normalized candidate back to parameter values + let mut values = Vec::with_capacity(state.dimensions.len()); + let mut ci = 0; // continuous dimension index + + for dim in &state.dimensions { + if dim.is_continuous { + let (lo, hi) = dim.bounds.unwrap_or((0.0, 1.0)); + let internal_val = from_normalized(normalized_candidate[ci], lo, hi); + values.push(from_internal(internal_val, &dim.distribution)); + ci += 1; + } else { + values.push(sample_random(&mut state.rng, &dim.distribution)); + } + } + + values +} + +/// Compute a random candidate for all dimensions. +fn compute_random_candidate(state: &mut GpState) -> Vec { + state + .dimensions + .iter() + .map(|dim| sample_random(&mut state.rng, &dim.distribution)) + .collect() +} diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index 7f9c78c..a29e474 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -3,6 +3,8 @@ pub mod bohb; #[cfg(feature = "cma-es")] pub mod cma_es; +#[cfg(feature = "gp")] +pub mod gp; pub mod grid; pub mod motpe; pub mod nsga2; diff --git a/tests/gp_tests.rs b/tests/gp_tests.rs new file mode 100644 index 0000000..b2dc23f --- /dev/null +++ b/tests/gp_tests.rs @@ -0,0 +1,256 @@ +#![cfg(feature = "gp")] + +use optimizer::prelude::*; +use optimizer::sampler::gp::GpSampler; + +#[test] +fn sphere_function() { + let sampler = GpSampler::with_seed(42); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + let x = FloatParam::new(-5.0, 5.0).name("x"); + let y = FloatParam::new(-5.0, 5.0).name("y"); + + study + .optimize(80, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, Error>(xv * xv + yv * yv) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert!( + best.value < 2.0, + "sphere best value should be < 2.0, got {}", + best.value + ); +} + +#[test] +fn rosenbrock_function() { + let sampler = GpSampler::builder().n_startup_trials(15).seed(42).build(); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + let x = FloatParam::new(-5.0, 5.0).name("x"); + let y = FloatParam::new(-5.0, 5.0).name("y"); + + study + .optimize(100, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + let val = (1.0 - xv).powi(2) + 100.0 * (yv - xv * xv).powi(2); + Ok::<_, Error>(val) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert!( + best.value < 100.0, + "rosenbrock best value should be < 100.0, got {}", + best.value + ); +} + +#[test] +fn bounds_respected() { + let sampler = GpSampler::with_seed(123); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + let x = FloatParam::new(-2.0, 3.0).name("x"); + let y = FloatParam::new(0.0, 10.0).name("y"); + + study + .optimize(100, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, Error>(xv + yv) + }) + .unwrap(); + + for trial in study.trials() { + let xv: f64 = trial.get(&x).unwrap(); + let yv: f64 = trial.get(&y).unwrap(); + assert!((-2.0..=3.0).contains(&xv), "x = {xv} out of bounds [-2, 3]"); + assert!((0.0..=10.0).contains(&yv), "y = {yv} out of bounds [0, 10]"); + } +} + +#[test] +fn mixed_params_float_and_categorical() { + let sampler = GpSampler::with_seed(42); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + let x = FloatParam::new(-5.0, 5.0).name("x"); + let cat = CategoricalParam::new(vec!["a", "b", "c"]).name("cat"); + + study + .optimize(50, |trial| { + let xv = x.suggest(trial)?; + let cv = cat.suggest(trial)?; + let penalty = match cv { + "a" => 0.0, + "b" => 1.0, + _ => 2.0, + }; + Ok::<_, Error>(xv * xv + penalty) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert!( + best.value < 10.0, + "best value should be < 10.0, got {}", + best.value + ); +} + +#[test] +fn seeded_reproducibility() { + let x = FloatParam::new(-5.0, 5.0).name("x"); + let y = FloatParam::new(-5.0, 5.0).name("y"); + + let run = |seed: u64| { + let sampler = GpSampler::with_seed(seed); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + study + .optimize(50, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, Error>(xv * xv + yv * yv) + }) + .unwrap(); + study.trials().iter().map(|t| t.value).collect::>() + }; + + let results1 = run(42); + let results2 = run(42); + assert_eq!(results1, results2, "same seed should produce same results"); +} + +#[test] +fn different_seeds_different_results() { + let x = FloatParam::new(-5.0, 5.0).name("x"); + let y = FloatParam::new(-5.0, 5.0).name("y"); + + let run = |seed: u64| { + let sampler = GpSampler::with_seed(seed); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + study + .optimize(20, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, Error>(xv * xv + yv * yv) + }) + .unwrap(); + study.trials().iter().map(|t| t.value).collect::>() + }; + + let results1 = run(42); + let results2 = run(99); + assert_ne!( + results1, results2, + "different seeds should produce different results" + ); +} + +#[test] +fn single_dimension() { + let sampler = GpSampler::with_seed(42); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + let x = FloatParam::new(-10.0, 10.0).name("x"); + + study + .optimize(100, |trial| { + let xv = x.suggest(trial)?; + Ok::<_, Error>((xv - 3.0).powi(2)) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert!( + best.value < 1.0, + "1-D optimization should converge, got {}", + best.value + ); +} + +#[test] +fn integer_params() { + let sampler = GpSampler::with_seed(42); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + let n = IntParam::new(1, 20).name("n"); + + study + .optimize(100, |trial| { + let nv = n.suggest(trial)?; + Ok::<_, Error>(((nv - 10) * (nv - 10)) as f64) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + let best_n: i64 = best.get(&n).unwrap(); + assert!( + (1..=20).contains(&best_n), + "integer value {best_n} out of bounds" + ); + assert!( + best.value < 10.0, + "integer optimization should converge, got {}", + best.value + ); +} + +#[test] +fn log_scale_params() { + let sampler = GpSampler::with_seed(42); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + let lr = FloatParam::new(1e-5, 1.0).log_scale().name("lr"); + + study + .optimize(100, |trial| { + let lrv = lr.suggest(trial)?; + Ok::<_, Error>((lrv.ln() - 0.01_f64.ln()).powi(2)) + }) + .unwrap(); + + for trial in study.trials() { + let lrv: f64 = trial.get(&lr).unwrap(); + assert!( + (1e-5..=1.0).contains(&lrv), + "log-scale value {lrv} out of bounds" + ); + } +} + +#[test] +fn builder_configuration() { + let sampler = GpSampler::builder() + .n_startup_trials(5) + .n_candidates(500) + .noise_variance(1e-4) + .seed(42) + .build(); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + let x = FloatParam::new(-5.0, 5.0).name("x"); + let y = FloatParam::new(-5.0, 5.0).name("y"); + + study + .optimize(100, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, Error>(xv * xv + yv * yv) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert!( + best.value < 5.0, + "custom config optimization should work, got {}", + best.value + ); +}