feat: add Gaussian Process sampler with Expected Improvement

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"]`.
This commit is contained in:
Manuel Raimann
2026-02-11 21:27:40 +01:00
parent b3bf4ccb71
commit 947701e83f
5 changed files with 1112 additions and 0 deletions
+1
View File
@@ -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 = []
+6
View File
@@ -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;
+847
View File
@@ -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<f64> = 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<f64> = Study::with_sampler(Direction::Minimize, GpSampler::new());
///
/// // With seed for reproducibility
/// let study: Study<f64> = 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<f64> = Study::with_sampler(Direction::Minimize, sampler);
/// ```
pub struct GpSampler {
state: Mutex<GpState>,
}
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<usize>,
n_candidates: Option<usize>,
noise_variance: Option<f64>,
seed: Option<u64>,
}
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<ParamValue>,
/// 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<f64, nalgebra::Dyn>,
/// α = (K + σ²I)^{-1} y.
alpha: nalgebra::DVector<f64>,
/// Training inputs (each row is a data point, normalized to [0, 1]).
x_train: Vec<Vec<f64>>,
/// ARD lengthscales per dimension.
lengthscales: Vec<f64>,
/// 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<DimensionInfo>,
trial_progress: HashMap<u64, TrialProgress>,
discovery_trial_id: Option<u64>,
}
impl GpState {
fn new(
n_startup: Option<usize>,
n_candidates: Option<usize>,
noise_var: Option<f64>,
seed: Option<u64>,
) -> 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<f64>],
lengthscales: &[f64],
signal_var: f64,
noise_var: f64,
) -> DMatrix<f64> {
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<f64>],
lengthscales: &[f64],
signal_var: f64,
) -> nalgebra::DVector<f64> {
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<f64>], y_train: &[f64], noise_var: f64) -> Option<GpModel> {
let n = y_train.len();
if n == 0 {
return None;
}
// Standardize y
let y_mean = y_train.iter().sum::<f64>() / n as f64;
let y_var = if n > 1 {
y_train.iter().map(|&y| (y - y_mean).powi(2)).sum::<f64>() / (n - 1) as f64
} else {
1.0
};
let y_std = y_var.sqrt().max(1e-10);
let y_standardized: Vec<f64> = 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<f64> = (0..d)
.map(|j| {
let vals: Vec<f64> = x_train.iter().map(|x| x[j]).collect();
let mean_j = vals.iter().sum::<f64>() / n as f64;
let var_j = vals.iter().map(|&v| (v - mean_j).powi(2)).sum::<f64>() / 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<f64> {
let mut best_ei = f64::NEG_INFINITY;
let mut best_x = vec![0.5; n_dims];
for _ in 0..n_candidates {
let x: Vec<f64> = (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<Option<crate::parameter::ParamId>> {
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<f64>>, Vec<f64>) {
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<usize> = 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(&param_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<ParamValue> {
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<ParamValue> {
state
.dimensions
.iter()
.map(|dim| sample_random(&mut state.rng, &dim.distribution))
.collect()
}
+2
View File
@@ -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;
+256
View File
@@ -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<f64> = 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<f64> = 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<f64> = 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<f64> = 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<f64> = 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::<Vec<_>>()
};
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<f64> = 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::<Vec<_>>()
};
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<f64> = 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<f64> = 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<f64> = 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<f64> = 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
);
}