14 Commits

Author SHA1 Message Date
Manuel Raimann 147a64708d chore: release v0.7.0 2026-02-11 19:05:52 +01:00
Manuel Raimann b36137bc96 feat: add benchmark suite with criterion and standard test functions
Add micro-benchmarks for TPE, Random, and Grid samplers with varying
history sizes, end-to-end optimization benchmarks on Sphere and
Rosenbrock, and a convergence tracking example that outputs CSV data.
Includes 6 standard test functions (Sphere, Rosenbrock, Rastrigin,
Ackley, Branin, Hartmann6) with correctness tests at known optima.
2026-02-11 19:05:33 +01:00
Manuel Raimann f3941a4b3e chore: remove unused proc-macro2 dependency from optimizer-derive 2026-02-11 18:57:30 +01:00
Manuel Raimann 0cae3b4227 feat: add parameter importance via Spearman rank correlation
Compute per-parameter importance scores by measuring the absolute
Spearman rank correlation between each parameter's values and the
objective across completed trials. Scores are normalized to sum to 1.0
and returned sorted by descending importance.
2026-02-11 18:56:19 +01:00
Manuel Raimann 09480a973b feat: add constraint handling with feasibility-aware trial ranking
Add constraint support so that optimization problems with constraints
(e.g., "model size < 100MB") prefer feasible solutions. Convention:
constraint value <= 0.0 means feasible.

- Add constraint_values field to Trial with set_constraints/getter
- Add constraints field to CompletedTrial with is_feasible() method
- Propagate constraints through complete_trial/prune_trial
- Make best_trial() and top_trials() constraint-aware: feasible trials
  rank above infeasible, infeasible ranked by total violation
2026-02-11 18:49:08 +01:00
Manuel Raimann aa9734587a feat: add BOHB sampler for budget-aware Bayesian optimization
BOHB combines TPE's model-guided sampling with Hyperband's budget-aware
evaluation by conditioning the TPE model on trials at specific budget
levels, producing better-calibrated parameter proposals.
2026-02-11 18:42:09 +01:00
Manuel Raimann a1dd6d393c feat: add CMA-ES sampler behind cma-es feature flag
Implement CMA-ES (Covariance Matrix Adaptation Evolution Strategy) as a
new sampler for continuous optimization. Uses nalgebra for matrix
operations and implements the full Hansen 2016 algorithm: mean update,
evolution paths, covariance matrix adaptation, and cumulative step-size
adaptation.

Key features:
- Builder API with configurable sigma0, population_size, and seed
- Three-phase state machine: discovery, active sampling, generation update
- Rejection sampling with bounds clipping fallback
- Log-scale and step parameter support via internal-space mapping
- Categorical parameters handled via random sampling fallback
- Numerical robustness: eigenvalue clamping, symmetry enforcement
2026-02-11 18:31:12 +01:00
Manuel Raimann a53ba4f472 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.
2026-02-11 17:59:48 +01:00
Manuel Raimann 9b4a8321b5 feat: implement IntoIterator for &Study and add iter() method
Enables idiomatic `for trial in &study` iteration over completed trials.
2026-02-11 17:54:09 +01:00
Manuel Raimann db0314a1d1 feat: add optimize_with_retries() for automatic retry of failed trials 2026-02-11 17:51:36 +01:00
Manuel Raimann df5fcedecf feat: add tracing integration behind tracing feature flag
Add optional structured logging via the tracing crate:
- info-level spans on all optimize* methods (direction, n_trials/duration)
- info events for trial completed, new best value, trial pruned
- debug events for trial failed and parameter sampled
- Zero overhead when the feature is disabled
2026-02-11 17:47:13 +01:00
Manuel Raimann 5061df9876 feat: add optimize_with_checkpoint() and atomic save writes
Adds a convenience method that combines optimize_with_callback + save
to automatically checkpoint every N trials. Also makes save() use
atomic writes (write to temp file, then rename) to prevent corrupt
files on crash.
2026-02-11 17:37:57 +01:00
Manuel Raimann 5fe0a75f78 feat: add serde serialization support behind serde feature flag
Add Serialize/Deserialize derives to all public data types (ParamValue,
Distribution, Direction, TrialState, ParamId, AttrValue, CompletedTrial)
gated behind a `serde` feature flag. Introduce StudySnapshot struct and
Study::save()/Study::load() for persisting and restoring study state as
human-readable JSON.
2026-02-11 17:33:48 +01:00
Manuel Raimann dcf3403261 feat: add Study::summary() and Display impl
Add a human-readable summary method and Display trait for Study<V>
where V: Display. The summary shows optimization direction, trial
counts (with complete/pruned breakdown), best value, and best
parameters with their labels.

Also derive PartialOrd + Ord on ParamId for deterministic parameter
ordering in summary output.
2026-02-11 17:27:43 +01:00
24 changed files with 4638 additions and 29 deletions
+20 -1
View File
@@ -3,7 +3,7 @@ members = ["optimizer-derive"]
[package]
name = "optimizer"
version = "0.6.0"
version = "0.7.0"
edition = "2024"
rust-version = "1.88"
license = "MIT"
@@ -21,15 +21,34 @@ thiserror = "2"
parking_lot = "0.12"
tokio = { version = "1", features = ["sync", "rt-multi-thread"], optional = true }
optimizer-derive = { version = "0.1.0", path = "optimizer-derive", optional = true }
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 }
nalgebra = { version = "0.33", optional = true }
[features]
default = []
async = ["dep:tokio"]
derive = ["dep:optimizer-derive"]
serde = ["dep:serde", "dep:serde_json"]
tracing = ["dep:tracing"]
sobol = ["dep:sobol_burley"]
cma-es = ["dep:nalgebra"]
[dev-dependencies]
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
optimizer-derive = { version = "0.1.0", path = "optimizer-derive" }
serde_json = "1"
criterion = { version = "0.8", features = ["html_reports"] }
[[bench]]
name = "samplers"
harness = false
[[bench]]
name = "optimization"
harness = false
[[example]]
name = "async_api_optimization"
+112
View File
@@ -0,0 +1,112 @@
#[allow(dead_code)]
mod test_functions;
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use optimizer::parameter::Parameter;
use optimizer::sampler::random::RandomSampler;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{FloatParam, Study};
fn make_params(dims: usize) -> Vec<FloatParam> {
(0..dims)
.map(|i| FloatParam::new(-5.0, 5.0).name(format!("x{i}")))
.collect()
}
fn bench_tpe_sphere(c: &mut Criterion) {
let mut group = c.benchmark_group("tpe_sphere");
group.sample_size(10);
for dims in [2, 10, 50] {
let params = make_params(dims);
group.bench_with_input(BenchmarkId::new("dims", dims), &params, |b, params| {
b.iter(|| {
let study = Study::minimize(TpeSampler::builder().seed(42).build().unwrap());
study
.optimize(100, |trial| {
let x: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()
.unwrap();
Ok::<_, optimizer::Error>(test_functions::sphere(&x))
})
.unwrap();
});
});
}
group.finish();
}
fn bench_tpe_rosenbrock(c: &mut Criterion) {
let mut group = c.benchmark_group("tpe_rosenbrock");
group.sample_size(10);
for dims in [2, 10] {
let params = make_params(dims);
group.bench_with_input(BenchmarkId::new("dims", dims), &params, |b, params| {
b.iter(|| {
let study = Study::minimize(TpeSampler::builder().seed(42).build().unwrap());
study
.optimize(100, |trial| {
let x: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()
.unwrap();
Ok::<_, optimizer::Error>(test_functions::rosenbrock(&x))
})
.unwrap();
});
});
}
group.finish();
}
fn bench_random_vs_tpe(c: &mut Criterion) {
let mut group = c.benchmark_group("random_vs_tpe");
group.sample_size(10);
let params = make_params(5);
group.bench_function("random_5d", |b| {
b.iter(|| {
let study = Study::minimize(RandomSampler::with_seed(42));
study
.optimize(100, |trial| {
let x: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()
.unwrap();
Ok::<_, optimizer::Error>(test_functions::sphere(&x))
})
.unwrap();
});
});
group.bench_function("tpe_5d", |b| {
b.iter(|| {
let study = Study::minimize(TpeSampler::builder().seed(42).build().unwrap());
study
.optimize(100, |trial| {
let x: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()
.unwrap();
Ok::<_, optimizer::Error>(test_functions::sphere(&x))
})
.unwrap();
});
});
group.finish();
}
criterion_group!(
benches,
bench_tpe_sphere,
bench_tpe_rosenbrock,
bench_random_vs_tpe
);
criterion_main!(benches);
+118
View File
@@ -0,0 +1,118 @@
use std::collections::HashMap;
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use optimizer::parameter::Parameter;
use optimizer::sampler::Sampler;
use optimizer::sampler::grid::GridSearchSampler;
use optimizer::sampler::random::RandomSampler;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{CompletedTrial, FloatParam};
/// Build a synthetic history of `n` completed trials over `dims` float parameters.
fn build_history(n: usize, dims: usize) -> Vec<CompletedTrial<f64>> {
let params: Vec<FloatParam> = (0..dims)
.map(|i| FloatParam::new(-5.0, 5.0).name(format!("x{i}")))
.collect();
let mut history = Vec::with_capacity(n);
let sampler = RandomSampler::with_seed(42);
for trial_id in 0..n {
let id = trial_id as u64;
let mut param_values = HashMap::new();
let mut distributions = HashMap::new();
let mut param_labels = HashMap::new();
for p in &params {
let dist = p.distribution();
let val = sampler.sample(&dist, id, &history);
param_values.insert(p.id(), val);
distributions.insert(p.id(), dist);
param_labels.insert(p.id(), p.label());
}
// Use sphere function value as objective
let value: f64 = param_values
.values()
.map(|v| {
let optimizer::ParamValue::Float(f) = v else {
unreachable!()
};
f * f
})
.sum();
history.push(CompletedTrial::new(
id,
param_values,
distributions,
param_labels,
value,
));
}
history
}
fn bench_tpe_sample(c: &mut Criterion) {
let mut group = c.benchmark_group("tpe_sample");
let dist = FloatParam::new(-5.0, 5.0).distribution();
let tpe = TpeSampler::builder().seed(42).build().unwrap();
for history_size in [10, 100, 1000] {
let history = build_history(history_size, 2);
group.bench_with_input(
BenchmarkId::new("history", history_size),
&history,
|b, history| {
b.iter(|| tpe.sample(&dist, history.len() as u64, history));
},
);
}
group.finish();
}
fn bench_random_sample(c: &mut Criterion) {
let mut group = c.benchmark_group("random_sample");
let dist = FloatParam::new(-5.0, 5.0).distribution();
let sampler = RandomSampler::with_seed(42);
for history_size in [10, 100, 1000] {
let history = build_history(history_size, 2);
group.bench_with_input(
BenchmarkId::new("history", history_size),
&history,
|b, history| {
b.iter(|| sampler.sample(&dist, history.len() as u64, history));
},
);
}
group.finish();
}
fn bench_grid_sample(c: &mut Criterion) {
let mut group = c.benchmark_group("grid_sample");
let dist = FloatParam::new(-5.0, 5.0).distribution();
let history: Vec<CompletedTrial<f64>> = Vec::new();
for grid_points in [5, 10, 50] {
group.bench_with_input(
BenchmarkId::new("points", grid_points),
&grid_points,
|b, _| {
b.iter(|| {
// Fresh sampler each iteration since grid tracks used points
let sampler = GridSearchSampler::builder()
.n_points_per_param(grid_points)
.build();
sampler.sample(&dist, 0, &history)
});
},
);
}
group.finish();
}
criterion_group!(
benches,
bench_tpe_sample,
bench_random_sample,
bench_grid_sample
);
criterion_main!(benches);
+84
View File
@@ -0,0 +1,84 @@
//! Standard optimization test functions for benchmarking.
/// Sphere function: unimodal, convex. Global minimum f(0,...,0) = 0.
pub fn sphere(x: &[f64]) -> f64 {
x.iter().map(|xi| xi * xi).sum()
}
/// Rosenbrock function: narrow valley. Global minimum f(1,...,1) = 0.
pub fn rosenbrock(x: &[f64]) -> f64 {
x.windows(2)
.map(|w| 100.0 * (w[1] - w[0] * w[0]).powi(2) + (1.0 - w[0]).powi(2))
.sum()
}
/// Rastrigin function: highly multimodal. Global minimum f(0,...,0) = 0.
pub fn rastrigin(x: &[f64]) -> f64 {
let n = x.len() as f64;
10.0 * n
+ x.iter()
.map(|xi| xi * xi - 10.0 * (2.0 * std::f64::consts::PI * xi).cos())
.sum::<f64>()
}
/// Ackley function: nearly flat with a deep well. Global minimum f(0,...,0) = 0.
pub fn ackley(x: &[f64]) -> f64 {
let n = x.len() as f64;
let sum_sq: f64 = x.iter().map(|xi| xi * xi).sum();
let sum_cos: f64 = x
.iter()
.map(|xi| (2.0 * std::f64::consts::PI * xi).cos())
.sum();
-20.0 * (-0.2 * (sum_sq / n).sqrt()).exp() - (sum_cos / n).exp() + 20.0 + std::f64::consts::E
}
/// Branin function (2D only). Three global minima with f* ≈ 0.397887.
///
/// # Panics
///
/// Panics if `x` does not have exactly 2 elements.
pub fn branin(x: &[f64]) -> f64 {
assert!(x.len() == 2, "Branin requires exactly 2 dimensions");
let (x1, x2) = (x[0], x[1]);
let pi = std::f64::consts::PI;
let a = 1.0;
let b = 5.1 / (4.0 * pi * pi);
let c = 5.0 / pi;
let r = 6.0;
let s = 10.0;
let t = 1.0 / (8.0 * pi);
a * (x2 - b * x1 * x1 + c * x1 - r).powi(2) + s * (1.0 - t) * x1.cos() + s
}
/// Hartmann 6D function. Global minimum f* ≈ -3.3224.
///
/// # Panics
///
/// Panics if `x` does not have exactly 6 elements.
pub fn hartmann6(x: &[f64]) -> f64 {
assert!(x.len() == 6, "Hartmann6 requires exactly 6 dimensions");
let alpha = [1.0, 1.2, 3.0, 3.2];
let a_matrix = [
[10.0, 3.0, 17.0, 3.5, 1.7, 8.0],
[0.05, 10.0, 17.0, 0.1, 8.0, 14.0],
[3.0, 3.5, 1.7, 10.0, 17.0, 8.0],
[17.0, 8.0, 0.05, 10.0, 0.1, 14.0],
];
let p_matrix = [
[0.1312, 0.1696, 0.5569, 0.0124, 0.8283, 0.5886],
[0.2329, 0.4135, 0.8307, 0.3736, 0.1004, 0.9991],
[0.2348, 0.1451, 0.3522, 0.2883, 0.3047, 0.6650],
[0.4047, 0.8828, 0.8732, 0.5743, 0.1091, 0.0381],
];
let mut result = 0.0;
for i in 0..4 {
let mut inner = 0.0;
for (j, xj) in x.iter().enumerate() {
inner += a_matrix[i][j] * (xj - p_matrix[i][j]).powi(2);
}
result -= alpha[i] * (-inner).exp();
}
result
}
+124
View File
@@ -0,0 +1,124 @@
use std::ops::ControlFlow;
use std::time::Instant;
use optimizer::parameter::Parameter;
use optimizer::sampler::random::RandomSampler;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{FloatParam, Study};
/// Standard optimization test functions.
mod functions {
pub fn sphere(x: &[f64]) -> f64 {
x.iter().map(|xi| xi * xi).sum()
}
pub fn rosenbrock(x: &[f64]) -> f64 {
x.windows(2)
.map(|w| 100.0 * (w[1] - w[0] * w[0]).powi(2) + (1.0 - w[0]).powi(2))
.sum()
}
pub fn rastrigin(x: &[f64]) -> f64 {
let n = x.len() as f64;
10.0 * n
+ x.iter()
.map(|xi| xi * xi - 10.0 * (2.0 * std::f64::consts::PI * xi).cos())
.sum::<f64>()
}
}
fn run_convergence(
name: &str,
sampler_name: &str,
study: Study<f64>,
params: &[FloatParam],
objective: fn(&[f64]) -> f64,
n_trials: usize,
) {
let start = Instant::now();
study
.optimize_with_callback(
n_trials,
|trial| {
let x: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()
.unwrap();
Ok::<_, optimizer::Error>(objective(&x))
},
|study, _trial| {
let elapsed = start.elapsed().as_millis();
let best = study.best_value().unwrap();
let n = study.n_trials();
println!("{n},{best},{elapsed},{sampler_name},{name}");
ControlFlow::Continue(())
},
)
.unwrap();
}
fn main() {
println!("trial,best_value,elapsed_ms,sampler,function");
let dims = 5;
let params: Vec<FloatParam> = (0..dims)
.map(|i| FloatParam::new(-5.0, 5.0).name(format!("x{i}")))
.collect();
let n_trials = 200;
// Sphere: Random vs TPE
run_convergence(
"sphere_5d",
"random",
Study::minimize(RandomSampler::with_seed(1)),
&params,
functions::sphere,
n_trials,
);
run_convergence(
"sphere_5d",
"tpe",
Study::minimize(TpeSampler::builder().seed(1).build().unwrap()),
&params,
functions::sphere,
n_trials,
);
// Rosenbrock: Random vs TPE
run_convergence(
"rosenbrock_5d",
"random",
Study::minimize(RandomSampler::with_seed(2)),
&params,
functions::rosenbrock,
n_trials,
);
run_convergence(
"rosenbrock_5d",
"tpe",
Study::minimize(TpeSampler::builder().seed(2).build().unwrap()),
&params,
functions::rosenbrock,
n_trials,
);
// Rastrigin: Random vs TPE
run_convergence(
"rastrigin_5d",
"random",
Study::minimize(RandomSampler::with_seed(3)),
&params,
functions::rastrigin,
n_trials,
);
run_convergence(
"rastrigin_5d",
"tpe",
Study::minimize(TpeSampler::builder().seed(3).build().unwrap()),
&params,
functions::rastrigin,
n_trials,
);
}
-1
View File
@@ -13,4 +13,3 @@ proc-macro = true
[dependencies]
syn = { version = "2", features = ["full"] }
quote = "1"
proc-macro2 = "1"
+4
View File
@@ -2,6 +2,7 @@
/// Distribution for floating-point parameters.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct FloatDistribution {
/// Lower bound (inclusive).
pub low: f64,
@@ -15,6 +16,7 @@ pub struct FloatDistribution {
/// Distribution for integer parameters.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct IntDistribution {
/// Lower bound (inclusive).
pub low: i64,
@@ -28,6 +30,7 @@ pub struct IntDistribution {
/// Distribution for categorical parameters.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct CategoricalDistribution {
/// Number of choices available.
pub n_choices: usize,
@@ -35,6 +38,7 @@ pub struct CategoricalDistribution {
/// Enum wrapping all parameter distribution types.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum Distribution {
/// A floating-point distribution.
Float(FloatDistribution),
+93
View File
@@ -0,0 +1,93 @@
//! Parameter importance via Spearman rank correlation.
/// Assign average ranks to a slice of `f64` values (handles ties).
#[allow(clippy::cast_precision_loss, clippy::float_cmp)]
pub(crate) fn rank(values: &[f64]) -> Vec<f64> {
let n = values.len();
let mut indexed: Vec<(usize, f64)> = values.iter().copied().enumerate().collect();
indexed.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
let mut ranks = vec![0.0; n];
let mut i = 0;
while i < n {
// Find the run of tied values.
let mut j = i + 1;
while j < n && indexed[j].1 == indexed[i].1 {
j += 1;
}
// Average rank for the tie group (1-based ranks).
let avg = (i + 1..=j).sum::<usize>() as f64 / (j - i) as f64;
for item in &indexed[i..j] {
ranks[item.0] = avg;
}
i = j;
}
ranks
}
/// Pearson correlation coefficient on two equal-length slices.
#[allow(clippy::cast_precision_loss)]
fn pearson(x: &[f64], y: &[f64]) -> f64 {
let n = x.len() as f64;
let mean_x = x.iter().sum::<f64>() / n;
let mean_y = y.iter().sum::<f64>() / n;
let mut cov = 0.0;
let mut var_x = 0.0;
let mut var_y = 0.0;
for (xi, yi) in x.iter().zip(y.iter()) {
let dx = xi - mean_x;
let dy = yi - mean_y;
cov += dx * dy;
var_x += dx * dx;
var_y += dy * dy;
}
let denom = (var_x * var_y).sqrt();
if denom == 0.0 { 0.0 } else { cov / denom }
}
/// Spearman rank correlation (Pearson on ranks).
pub(crate) fn spearman(x: &[f64], y: &[f64]) -> f64 {
pearson(&rank(x), &rank(y))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rank_no_ties() {
let ranks = rank(&[30.0, 10.0, 20.0]);
assert_eq!(ranks, vec![3.0, 1.0, 2.0]);
}
#[test]
fn rank_with_ties() {
let ranks = rank(&[10.0, 20.0, 20.0, 30.0]);
assert_eq!(ranks, vec![1.0, 2.5, 2.5, 4.0]);
}
#[test]
fn perfect_positive_correlation() {
let x = vec![1.0, 2.0, 3.0, 4.0, 5.0];
let y = vec![2.0, 4.0, 6.0, 8.0, 10.0];
let r = spearman(&x, &y);
assert!((r - 1.0).abs() < 1e-10);
}
#[test]
fn perfect_negative_correlation() {
let x = vec![1.0, 2.0, 3.0, 4.0, 5.0];
let y = vec![10.0, 8.0, 6.0, 4.0, 2.0];
let r = spearman(&x, &y);
assert!((r + 1.0).abs() < 1e-10);
}
#[test]
fn zero_variance_returns_zero() {
let x = vec![5.0, 5.0, 5.0];
let y = vec![1.0, 2.0, 3.0];
assert!(spearman(&x, &y).abs() < f64::EPSILON);
}
}
+46
View File
@@ -17,6 +17,9 @@
//! - **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)
//! - **CMA-ES** - Covariance Matrix Adaptation Evolution Strategy for continuous optimization (requires `cma-es` feature)
//! - **BOHB** - Bayesian Optimization + `HyperBand` for budget-aware TPE sampling
//!
//! Additional features include:
//!
@@ -182,9 +185,38 @@
//!
//! - `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
//! - `cma-es`: Enable the CMA-ES sampler for continuous optimization
//! - `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.
/// No-op otherwise.
#[cfg(feature = "tracing")]
macro_rules! trace_info {
($($arg:tt)*) => { tracing::info!($($arg)*) };
}
#[cfg(not(feature = "tracing"))]
macro_rules! trace_info {
($($arg:tt)*) => {};
}
/// Emit a `tracing::debug!` event when the `tracing` feature is enabled.
/// No-op otherwise.
#[cfg(feature = "tracing")]
macro_rules! trace_debug {
($($arg:tt)*) => { tracing::debug!($($arg)*) };
}
#[cfg(not(feature = "tracing"))]
macro_rules! trace_debug {
($($arg:tt)*) => {};
}
mod distribution;
mod error;
mod importance;
mod kde;
mod param;
pub mod parameter;
@@ -206,10 +238,17 @@ pub use pruner::{
SuccessiveHalvingPruner, ThresholdPruner, WilcoxonPruner,
};
pub use sampler::CompletedTrial;
pub use sampler::bohb::BohbSampler;
#[cfg(feature = "cma-es")]
pub use sampler::cma_es::CmaEsSampler;
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")]
pub use study::StudySnapshot;
pub use trial::{AttrValue, Trial};
pub use types::{Direction, TrialState};
@@ -232,10 +271,17 @@ pub mod prelude {
SuccessiveHalvingPruner, ThresholdPruner,
};
pub use crate::sampler::CompletedTrial;
pub use crate::sampler::bohb::BohbSampler;
#[cfg(feature = "cma-es")]
pub use crate::sampler::cma_es::CmaEsSampler;
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")]
pub use crate::study::StudySnapshot;
pub use crate::trial::{AttrValue, Trial};
pub use crate::types::Direction;
}
+1
View File
@@ -6,6 +6,7 @@
/// For categorical parameters, the `Categorical` variant stores
/// the index into the choices array.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum ParamValue {
/// A floating-point parameter value.
Float(f64),
+2 -1
View File
@@ -37,7 +37,8 @@ static NEXT_PARAM_ID: AtomicU64 = AtomicU64::new(0);
///
/// Each parameter is assigned a unique `ParamId` at creation time. Cloning a parameter
/// copies its `ParamId`, so clones refer to the same logical parameter.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ParamId(u64);
impl ParamId {
+691
View File
@@ -0,0 +1,691 @@
//! BOHB (Bayesian Optimization + `HyperBand`) sampler.
//!
//! BOHB combines TPE's model-guided sampling with Hyperband's budget-aware
//! evaluation. Instead of building one global TPE model, BOHB conditions
//! its TPE model on trials evaluated at a specific budget level, giving
//! better-calibrated proposals for each rung of the Hyperband schedule.
//!
//! # How it works
//!
//! 1. Compute all Hyperband rung steps (budget levels) from the config.
//! 2. On each `sample()` call, scan the history's `intermediate_values`
//! to find the **largest budget level** with enough observations
//! (`>= min_points_in_model`).
//! 3. Build a filtered history where each trial's `value` is replaced
//! with its intermediate value at that budget level.
//! 4. Delegate to an internal [`TpeSampler`] for the actual sampling.
//! 5. Fall back to random sampling if no budget level has enough data.
//!
//! # Examples
//!
//! ```
//! use optimizer::sampler::bohb::BohbSampler;
//! use optimizer::{Direction, Study};
//!
//! let bohb = BohbSampler::new();
//! let pruner = bohb.matching_pruner(Direction::Minimize);
//! let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner);
//! ```
//!
//! Using the builder for custom configuration:
//!
//! ```
//! use optimizer::sampler::bohb::BohbSampler;
//!
//! let bohb = BohbSampler::builder()
//! .min_resource(1)
//! .max_resource(81)
//! .reduction_factor(3)
//! .min_points_in_model(10)
//! .seed(42)
//! .build()
//! .unwrap();
//! ```
use crate::distribution::Distribution;
use crate::error::Result;
use crate::param::ParamValue;
use crate::pruner::HyperbandPruner;
use crate::sampler::tpe::TpeSampler;
use crate::sampler::{CompletedTrial, Sampler};
use crate::types::Direction;
/// A BOHB sampler that combines TPE with Hyperband budget awareness.
///
/// BOHB filters trial history by budget level before delegating to TPE,
/// so the surrogate model is conditioned on trials evaluated at the same
/// resource level. This produces better-calibrated parameter proposals
/// than using a single global model across all budgets.
///
/// Use [`BohbSampler::matching_pruner`] to create a [`HyperbandPruner`]
/// with matching parameters.
pub struct BohbSampler {
min_resource: u64,
max_resource: u64,
reduction_factor: u64,
min_points_in_model: usize,
tpe: TpeSampler,
}
impl BohbSampler {
/// Creates a new BOHB sampler with default settings.
///
/// Defaults:
/// - `min_resource`: 1
/// - `max_resource`: 81
/// - `reduction_factor`: 3
/// - `min_points_in_model`: 10
/// - TPE: default settings
#[must_use]
pub fn new() -> Self {
Self {
min_resource: 1,
max_resource: 81,
reduction_factor: 3,
min_points_in_model: 10,
tpe: TpeSampler::new(),
}
}
/// Creates a builder for configuring a BOHB sampler.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::bohb::BohbSampler;
///
/// let sampler = BohbSampler::builder()
/// .min_resource(1)
/// .max_resource(27)
/// .reduction_factor(3)
/// .min_points_in_model(5)
/// .seed(42)
/// .build()
/// .unwrap();
/// ```
#[must_use]
pub fn builder() -> BohbSamplerBuilder {
BohbSamplerBuilder::new()
}
/// Creates a [`HyperbandPruner`] with matching Hyperband parameters.
///
/// This ensures the pruner's budget schedule is consistent with the
/// budget levels used by BOHB for model conditioning.
#[must_use]
pub fn matching_pruner(&self, direction: Direction) -> HyperbandPruner {
HyperbandPruner::new()
.min_resource(self.min_resource)
.max_resource(self.max_resource)
.reduction_factor(self.reduction_factor)
.direction(direction)
}
/// Compute all unique budget levels (rung steps) across all Hyperband brackets.
///
/// Returns sorted ascending.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
fn all_budget_levels(&self) -> Vec<u64> {
let eta = self.reduction_factor as f64;
let ratio = self.max_resource as f64 / self.min_resource as f64;
let s_max = (ratio.ln() / eta.ln()).floor() as u64;
let mut levels = Vec::new();
for bracket in 0..=s_max {
let exponent = s_max.saturating_sub(bracket);
let min_resource_bracket =
(self.max_resource as f64 / eta.powi(exponent as i32)).ceil() as u64;
let mut rung: u32 = 0;
while let Some(power) = self.reduction_factor.checked_pow(rung) {
let step = min_resource_bracket.saturating_mul(power);
if step > self.max_resource {
break;
}
levels.push(step);
rung += 1;
}
}
levels.sort_unstable();
levels.dedup();
levels
}
/// Build a filtered history for a specific budget level.
///
/// For each trial that has an intermediate value at the given budget step,
/// creates a new `CompletedTrial` with `value` replaced by the intermediate
/// value at that step.
fn filter_history_for_budget(history: &[CompletedTrial], budget: u64) -> Vec<CompletedTrial> {
history
.iter()
.filter_map(|trial| {
trial
.intermediate_values
.iter()
.find(|(step, _)| *step == budget)
.map(|(_, iv)| CompletedTrial {
id: trial.id,
params: trial.params.clone(),
distributions: trial.distributions.clone(),
param_labels: trial.param_labels.clone(),
value: *iv,
intermediate_values: trial.intermediate_values.clone(),
state: trial.state,
user_attrs: trial.user_attrs.clone(),
constraints: trial.constraints.clone(),
})
})
.collect()
}
}
impl Default for BohbSampler {
fn default() -> Self {
Self::new()
}
}
impl Sampler for BohbSampler {
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[CompletedTrial],
) -> ParamValue {
// Find the largest budget level with enough observations
let levels = self.all_budget_levels();
for &budget in levels.iter().rev() {
let count = history
.iter()
.filter(|t| {
t.intermediate_values
.iter()
.any(|(step, _)| *step == budget)
})
.count();
if count >= self.min_points_in_model {
let filtered = Self::filter_history_for_budget(history, budget);
return self.tpe.sample(distribution, trial_id, &filtered);
}
}
// Not enough data at any budget level: delegate to TPE with empty history
// which triggers its uniform-random startup behavior.
self.tpe.sample(distribution, trial_id, &[])
}
}
/// Builder for configuring a [`BohbSampler`].
///
/// # Examples
///
/// ```
/// use optimizer::sampler::bohb::BohbSamplerBuilder;
///
/// let sampler = BohbSamplerBuilder::new()
/// .min_resource(1)
/// .max_resource(81)
/// .reduction_factor(3)
/// .gamma(0.15)
/// .seed(42)
/// .build()
/// .unwrap();
/// ```
pub struct BohbSamplerBuilder {
min_resource: u64,
max_resource: u64,
reduction_factor: u64,
min_points_in_model: usize,
tpe_builder: crate::sampler::tpe::TpeSamplerBuilder,
}
impl BohbSamplerBuilder {
/// Creates a new builder with default settings.
#[must_use]
pub fn new() -> Self {
Self {
min_resource: 1,
max_resource: 81,
reduction_factor: 3,
min_points_in_model: 10,
tpe_builder: crate::sampler::tpe::TpeSamplerBuilder::new(),
}
}
/// Sets the minimum resource (budget) per trial.
///
/// # Panics
///
/// Panics if `r` is 0.
#[must_use]
pub fn min_resource(mut self, r: u64) -> Self {
assert!(r > 0, "min_resource must be > 0, got {r}");
self.min_resource = r;
self
}
/// Sets the maximum resource (budget) per trial.
///
/// # Panics
///
/// Panics if `r` is 0.
#[must_use]
pub fn max_resource(mut self, r: u64) -> Self {
assert!(r > 0, "max_resource must be > 0, got {r}");
self.max_resource = r;
self
}
/// Sets the reduction factor (eta).
///
/// # Panics
///
/// Panics if `eta` is less than 2.
#[must_use]
pub fn reduction_factor(mut self, eta: u64) -> Self {
assert!(eta >= 2, "reduction_factor must be >= 2, got {eta}");
self.reduction_factor = eta;
self
}
/// Sets the minimum number of observations at a budget level before
/// BOHB uses TPE instead of random sampling.
#[must_use]
pub fn min_points_in_model(mut self, n: usize) -> Self {
self.min_points_in_model = n;
self
}
/// Sets a fixed gamma value for the internal TPE sampler.
#[must_use]
pub fn gamma(mut self, gamma: f64) -> Self {
self.tpe_builder = self.tpe_builder.gamma(gamma);
self
}
/// Sets a custom gamma strategy for the internal TPE sampler.
#[must_use]
pub fn gamma_strategy<G: crate::sampler::tpe::GammaStrategy + 'static>(
mut self,
strategy: G,
) -> Self {
self.tpe_builder = self.tpe_builder.gamma_strategy(strategy);
self
}
/// Sets the number of EI candidates for the internal TPE sampler.
#[must_use]
pub fn n_ei_candidates(mut self, n: usize) -> Self {
self.tpe_builder = self.tpe_builder.n_ei_candidates(n);
self
}
/// Sets a fixed KDE bandwidth for the internal TPE sampler.
#[must_use]
pub fn kde_bandwidth(mut self, bandwidth: f64) -> Self {
self.tpe_builder = self.tpe_builder.kde_bandwidth(bandwidth);
self
}
/// Sets a seed for reproducible sampling.
#[must_use]
pub fn seed(mut self, seed: u64) -> Self {
self.tpe_builder = self.tpe_builder.seed(seed);
self
}
/// Builds the configured [`BohbSampler`].
///
/// # Errors
///
/// Returns an error if the TPE configuration is invalid (e.g. gamma
/// not in (0, 1) or bandwidth not positive).
pub fn build(self) -> Result<BohbSampler> {
let tpe = self.tpe_builder.build()?;
Ok(BohbSampler {
min_resource: self.min_resource,
max_resource: self.max_resource,
reduction_factor: self.reduction_factor,
min_points_in_model: self.min_points_in_model,
tpe,
})
}
}
impl Default for BohbSamplerBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
#[allow(clippy::cast_precision_loss)]
mod tests {
use std::collections::HashMap;
use super::*;
use crate::distribution::{FloatDistribution, IntDistribution};
use crate::parameter::ParamId;
use crate::types::TrialState;
fn make_trial_with_intermediates(
id: u64,
value: f64,
params: Vec<(ParamId, ParamValue, Distribution)>,
intermediate_values: Vec<(u64, f64)>,
) -> CompletedTrial {
let mut param_map = HashMap::new();
let mut dist_map = HashMap::new();
for (param_id, pv, dist) in params {
param_map.insert(param_id, pv);
dist_map.insert(param_id, dist);
}
CompletedTrial {
id,
params: param_map,
distributions: dist_map,
param_labels: HashMap::new(),
value,
intermediate_values,
state: TrialState::Complete,
user_attrs: HashMap::new(),
constraints: Vec::new(),
}
}
#[test]
fn budget_levels_default() {
let bohb = BohbSampler::new();
let levels = bohb.all_budget_levels();
// With min=1, max=81, eta=3:
// bracket 0: 1, 3, 9, 27, 81
// bracket 1: 3, 9, 27, 81
// bracket 2: 9, 27, 81
// bracket 3: 27, 81
// bracket 4: 81
// Unique sorted: [1, 3, 9, 27, 81]
assert_eq!(levels, vec![1, 3, 9, 27, 81]);
}
#[test]
fn budget_levels_eta2() {
let bohb = BohbSampler::builder()
.min_resource(1)
.max_resource(16)
.reduction_factor(2)
.build()
.unwrap();
let levels = bohb.all_budget_levels();
// s_max = floor(ln(16)/ln(2)) = 4
// bracket 0: 1, 2, 4, 8, 16
// bracket 1: 2, 4, 8, 16
// bracket 2: 4, 8, 16
// bracket 3: 8, 16
// bracket 4: 16
// Unique sorted: [1, 2, 4, 8, 16]
assert_eq!(levels, vec![1, 2, 4, 8, 16]);
}
#[test]
fn filter_history_selects_correct_budget() {
let x_id = ParamId::new();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let history = vec![
make_trial_with_intermediates(
0,
0.5,
vec![(x_id, ParamValue::Float(0.3), dist.clone())],
vec![(1, 0.9), (3, 0.7), (9, 0.5)],
),
make_trial_with_intermediates(
1,
0.4,
vec![(x_id, ParamValue::Float(0.6), dist.clone())],
vec![(1, 0.8), (3, 0.4)],
),
make_trial_with_intermediates(
2,
0.3,
vec![(x_id, ParamValue::Float(0.1), dist.clone())],
vec![(1, 0.7)],
),
];
// Budget 3: trials 0 and 1 have intermediate values at step 3
let filtered = BohbSampler::filter_history_for_budget(&history, 3);
assert_eq!(filtered.len(), 2);
assert!((filtered[0].value - 0.7).abs() < f64::EPSILON);
assert!((filtered[1].value - 0.4).abs() < f64::EPSILON);
// Budget 9: only trial 0
let filtered = BohbSampler::filter_history_for_budget(&history, 9);
assert_eq!(filtered.len(), 1);
assert!((filtered[0].value - 0.5).abs() < f64::EPSILON);
// Budget 27: nobody
let filtered = BohbSampler::filter_history_for_budget(&history, 27);
assert!(filtered.is_empty());
}
#[test]
fn matching_pruner_has_same_params() {
let bohb = BohbSampler::builder()
.min_resource(2)
.max_resource(64)
.reduction_factor(4)
.build()
.unwrap();
let pruner = bohb.matching_pruner(Direction::Minimize);
// We can't directly inspect HyperbandPruner fields, but we can
// verify it was created without panicking with the same params.
// The pruner's rung steps should match BOHB's budget levels.
// Just verify it doesn't panic.
drop(pruner);
}
#[test]
fn fallback_to_random_when_insufficient_data() {
let bohb = BohbSampler::builder()
.min_points_in_model(10)
.seed(42)
.build()
.unwrap();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
// Only 3 trials with intermediate values (< min_points_in_model=10)
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..3)
.map(|i| {
make_trial_with_intermediates(
i,
i as f64,
vec![(x_id, ParamValue::Float(i as f64 / 3.0), dist.clone())],
vec![(1, i as f64)],
)
})
.collect();
// Should not panic, should sample within bounds
for trial_id in 0..20 {
let val = bohb.sample(&dist, trial_id, &history);
if let ParamValue::Float(v) = val {
assert!((0.0..=1.0).contains(&v));
} else {
panic!("Expected Float");
}
}
}
#[test]
fn uses_budget_level_when_enough_data() {
let bohb = BohbSampler::builder()
.min_points_in_model(5)
.seed(42)
.build()
.unwrap();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 10.0,
log_scale: false,
step: None,
});
// Create 20 trials with intermediate values at budget 1.
// Good trials have x near 2.0, bad trials have x far from 2.0.
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..20)
.map(|i| {
let x = i as f64 / 2.0;
let iv_at_1 = (x - 2.0).powi(2);
make_trial_with_intermediates(
i,
iv_at_1, // final value same as intermediate for simplicity
vec![(x_id, ParamValue::Float(x), dist.clone())],
vec![(1, iv_at_1)],
)
})
.collect();
// Should use TPE on filtered history at budget 1
let val = bohb.sample(&dist, 100, &history);
if let ParamValue::Float(v) = val {
assert!((0.0..=10.0).contains(&v), "Value {v} out of bounds");
} else {
panic!("Expected Float");
}
}
#[test]
fn prefers_largest_budget_level() {
let bohb = BohbSampler::builder()
.min_resource(1)
.max_resource(9)
.reduction_factor(3)
.min_points_in_model(3)
.seed(42)
.build()
.unwrap();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 10.0,
log_scale: false,
step: None,
});
// Budget levels: [1, 3, 9]
assert_eq!(bohb.all_budget_levels(), vec![1, 3, 9]);
// Create 5 trials with intermediates at budget 1 and 3
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..5)
.map(|i| {
let x = i as f64;
make_trial_with_intermediates(
i,
x,
vec![(x_id, ParamValue::Float(x), dist.clone())],
vec![(1, x * 2.0), (3, x)],
)
})
.collect();
// Budget 3 has 5 observations (>= 3), budget 9 has 0.
// BOHB should pick budget 3 (largest with enough data).
// The filtered history at budget 3 has values [0, 1, 2, 3, 4].
let filtered_3 = BohbSampler::filter_history_for_budget(&history, 3);
assert_eq!(filtered_3.len(), 5);
let filtered_9 = BohbSampler::filter_history_for_budget(&history, 9);
assert_eq!(filtered_9.len(), 0);
// Should sample successfully
let val = bohb.sample(&dist, 100, &history);
assert!(matches!(val, ParamValue::Float(_)));
}
#[test]
fn builder_validates_tpe_params() {
// Invalid gamma
let result = BohbSampler::builder().gamma(1.5).build();
assert!(result.is_err());
// Invalid bandwidth
let result = BohbSampler::builder().kde_bandwidth(-1.0).build();
assert!(result.is_err());
}
#[test]
#[should_panic(expected = "min_resource must be > 0")]
fn builder_rejects_zero_min_resource() {
let _ = BohbSampler::builder().min_resource(0);
}
#[test]
#[should_panic(expected = "max_resource must be > 0")]
fn builder_rejects_zero_max_resource() {
let _ = BohbSampler::builder().max_resource(0);
}
#[test]
#[should_panic(expected = "reduction_factor must be >= 2")]
fn builder_rejects_small_reduction_factor() {
let _ = BohbSampler::builder().reduction_factor(1);
}
#[test]
fn int_distribution_works() {
let bohb = BohbSampler::builder()
.min_points_in_model(3)
.seed(42)
.build()
.unwrap();
let dist = Distribution::Int(IntDistribution {
low: 0,
high: 100,
log_scale: false,
step: None,
});
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..10)
.map(|i| {
make_trial_with_intermediates(
i,
i as f64,
vec![(x_id, ParamValue::Int(i.cast_signed() * 10), dist.clone())],
vec![(1, i as f64)],
)
})
.collect();
let val = bohb.sample(&dist, 100, &history);
if let ParamValue::Int(v) = val {
assert!((0..=100).contains(&v));
} else {
panic!("Expected Int");
}
}
}
File diff suppressed because it is too large Load Diff
+19
View File
@@ -1,7 +1,12 @@
//! Sampler trait and implementations for parameter sampling.
pub mod bohb;
#[cfg(feature = "cma-es")]
pub mod cma_es;
pub mod grid;
pub mod random;
#[cfg(feature = "sobol")]
pub mod sobol;
pub mod tpe;
use std::collections::HashMap;
@@ -18,6 +23,7 @@ use crate::types::TrialState;
/// parameter values, their distributions, and the objective value returned
/// by the objective function.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct CompletedTrial<V = f64> {
/// The unique identifier for this trial.
pub id: u64,
@@ -35,6 +41,9 @@ pub struct CompletedTrial<V = f64> {
pub state: TrialState,
/// User-defined attributes stored during the trial.
pub user_attrs: HashMap<String, AttrValue>,
/// Constraint values for this trial (<=0.0 means feasible).
#[cfg_attr(feature = "serde", serde(default))]
pub constraints: Vec<f64>,
}
impl<V> CompletedTrial<V> {
@@ -55,6 +64,7 @@ impl<V> CompletedTrial<V> {
intermediate_values: Vec::new(),
state: TrialState::Complete,
user_attrs: HashMap::new(),
constraints: Vec::new(),
}
}
@@ -77,6 +87,7 @@ impl<V> CompletedTrial<V> {
intermediate_values,
state: TrialState::Complete,
user_attrs,
constraints: Vec::new(),
}
}
@@ -121,6 +132,14 @@ impl<V> CompletedTrial<V> {
})
}
/// Returns `true` if all constraints are satisfied (values <= 0.0).
///
/// A trial with no constraints is considered feasible.
#[must_use]
pub fn is_feasible(&self) -> bool {
self.constraints.iter().all(|&c| c <= 0.0)
}
/// Gets a user attribute by key.
#[must_use]
pub fn user_attr(&self, key: &str) -> Option<&AttrValue> {
+420
View File
@@ -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<f64> = Study::with_sampler(Direction::Minimize, SobolSampler::new());
/// ```
pub struct SobolSampler {
seed: u32,
state: Mutex<SobolState>,
}
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"
);
}
}
+615 -25
View File
@@ -1,6 +1,7 @@
//! Study implementation for managing optimization trials.
use core::any::Any;
use core::fmt;
#[cfg(feature = "async")]
use core::future::Future;
use core::ops::ControlFlow;
@@ -339,6 +340,33 @@ where
self.enqueued_params.lock().push_back(params);
}
/// Returns the trial ID of the current best trial from the given slice.
#[cfg(feature = "tracing")]
fn best_id(&self, trials: &[CompletedTrial<V>]) -> Option<u64> {
let direction = self.direction;
trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.max_by(|a, b| Self::compare_trials(a, b, direction))
.map(|t| t.id)
}
/// Creates a new trial with pre-set parameter values.
///
/// The trial gets a new unique ID but reuses the given parameters. When
/// `suggest_param` is called on the resulting trial, fixed values are
/// returned instead of sampling.
fn create_trial_with_params(&self, params: HashMap<ParamId, ParamValue>) -> Trial {
let id = self.next_trial_id();
let mut trial = if let Some(factory) = &self.trial_factory {
factory(id)
} else {
Trial::new(id)
};
trial.set_fixed_params(params);
trial
}
/// Returns the number of enqueued parameter configurations.
#[must_use]
pub fn n_enqueued(&self) -> usize {
@@ -426,6 +454,7 @@ where
trial.user_attrs().clone(),
);
completed.state = TrialState::Complete;
completed.constraints = trial.constraint_values().to_vec();
self.completed_trials.write().push(completed);
}
@@ -535,6 +564,7 @@ where
trial.user_attrs().clone(),
);
completed.state = TrialState::Pruned;
completed.constraints = trial.constraint_values().to_vec();
self.completed_trials.write().push(completed);
}
@@ -601,12 +631,46 @@ where
.count()
}
/// Compares two completed trials using constraint-aware ranking.
///
/// 1. Feasible trials always rank above infeasible trials.
/// 2. Among feasible trials, rank by objective value (respecting direction).
/// 3. Among infeasible trials, rank by total constraint violation (lower is better).
fn compare_trials(
a: &CompletedTrial<V>,
b: &CompletedTrial<V>,
direction: Direction,
) -> core::cmp::Ordering {
match (a.is_feasible(), b.is_feasible()) {
(true, false) => core::cmp::Ordering::Greater,
(false, true) => core::cmp::Ordering::Less,
(false, false) => {
let va: f64 = a.constraints.iter().map(|c| c.max(0.0)).sum();
let vb: f64 = b.constraints.iter().map(|c| c.max(0.0)).sum();
vb.partial_cmp(&va).unwrap_or(core::cmp::Ordering::Equal)
}
(true, true) => {
let ordering = a.value.partial_cmp(&b.value);
match direction {
Direction::Minimize => {
ordering.map_or(core::cmp::Ordering::Equal, core::cmp::Ordering::reverse)
}
Direction::Maximize => ordering.unwrap_or(core::cmp::Ordering::Equal),
}
}
}
}
/// Returns the trial with the best objective value.
///
/// The "best" trial depends on the optimization direction:
/// - `Direction::Minimize`: Returns the trial with the lowest objective value.
/// - `Direction::Maximize`: Returns the trial with the highest objective value.
///
/// When constraints are present, feasible trials always rank above infeasible
/// trials. Among infeasible trials, those with lower total constraint violation
/// are preferred.
///
/// # Errors
///
/// Returns `Error::NoCompletedTrials` if no trials have been completed.
@@ -640,25 +704,12 @@ where
V: Clone,
{
let trials = self.completed_trials.read();
let direction = self.direction;
let best = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.max_by(|a, b| {
// For Minimize, we want the smallest value to be "max" in ordering
// For Maximize, we want the largest value to be "max" in ordering
let ordering = a.value.partial_cmp(&b.value);
match self.direction {
Direction::Minimize => {
// Reverse ordering: smaller values are "greater" for max_by
ordering.map_or(core::cmp::Ordering::Equal, core::cmp::Ordering::reverse)
}
Direction::Maximize => {
// Normal ordering: larger values are "greater" for max_by
ordering.unwrap_or(core::cmp::Ordering::Equal)
}
}
})
.max_by(|a, b| Self::compare_trials(a, b, direction))
.ok_or(crate::Error::NoCompletedTrials)?;
Ok(best.clone())
@@ -717,21 +768,14 @@ where
V: Clone,
{
let trials = self.completed_trials.read();
let direction = self.direction;
let mut completed: Vec<_> = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.cloned()
.collect();
completed.sort_by(|a, b| match self.direction {
Direction::Minimize => a
.value
.partial_cmp(&b.value)
.unwrap_or(core::cmp::Ordering::Equal),
Direction::Maximize => b
.value
.partial_cmp(&a.value)
.unwrap_or(core::cmp::Ordering::Equal),
});
// Sort best-first: reverse the compare_trials ordering (which is designed for max_by)
completed.sort_by(|a, b| Self::compare_trials(b, a, direction));
completed.truncate(n);
completed
}
@@ -788,18 +832,43 @@ where
E: ToString + 'static,
V: Default,
{
#[cfg(feature = "tracing")]
let _span =
tracing::info_span!("optimize", n_trials, direction = ?self.direction).entered();
for _ in 0..n_trials {
let mut trial = self.create_trial();
match objective(&mut trial) {
Ok(value) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
self.complete_trial(trial, value);
#[cfg(feature = "tracing")]
{
tracing::info!(trial_id, "trial completed");
let trials = self.completed_trials.read();
if trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.count()
== 1
|| trials.last().map(|t| t.id) == self.best_id(&trials)
{
tracing::info!(trial_id, "new best value found");
}
}
}
Err(e) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
if is_trial_pruned(&e) {
self.prune_trial(trial);
trace_info!(trial_id, "trial pruned");
} else {
self.fail_trial(trial, e.to_string());
trace_debug!(trial_id, "trial failed");
}
}
}
@@ -881,17 +950,25 @@ where
Fut: Future<Output = core::result::Result<(Trial, V), E>>,
E: ToString,
{
#[cfg(feature = "tracing")]
let _span =
tracing::info_span!("optimize_async", n_trials, direction = ?self.direction).entered();
for _ in 0..n_trials {
let trial = self.create_trial();
#[cfg(feature = "tracing")]
let trial_id = trial.id();
match objective(trial).await {
Ok((trial, value)) => {
self.complete_trial(trial, value);
trace_info!(trial_id, "trial completed");
}
Err(e) => {
// For async, we don't have the trial back on error
// We'll just count this as a failed trial without recording it
let _ = e.to_string();
trace_debug!(trial_id, "trial failed");
}
}
}
@@ -978,6 +1055,9 @@ where
{
use tokio::sync::Semaphore;
#[cfg(feature = "tracing")]
let _span = tracing::info_span!("optimize_parallel", n_trials, concurrency, direction = ?self.direction).entered();
let semaphore = Arc::new(Semaphore::new(concurrency));
let objective = Arc::new(objective);
@@ -1008,7 +1088,10 @@ where
.map_err(|e| crate::Error::TaskError(e.to_string()))?
{
Ok((trial, value)) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
self.complete_trial(trial, value);
trace_info!(trial_id, "trial completed");
}
Err(e) => {
let _ = e.to_string();
@@ -1098,13 +1181,34 @@ where
C: FnMut(&Study<V>, &CompletedTrial<V>) -> ControlFlow<()>,
E: ToString + 'static,
{
#[cfg(feature = "tracing")]
let _span =
tracing::info_span!("optimize", n_trials, direction = ?self.direction).entered();
for _ in 0..n_trials {
let mut trial = self.create_trial();
match objective(&mut trial) {
Ok(value) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
self.complete_trial(trial, value);
#[cfg(feature = "tracing")]
{
tracing::info!(trial_id, "trial completed");
let trials = self.completed_trials.read();
if trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.count()
== 1
|| trials.last().map(|t| t.id) == self.best_id(&trials)
{
tracing::info!(trial_id, "new best value found");
}
}
// Get the just-completed trial for the callback
let trials = self.completed_trials.read();
let Some(completed) = trials.last() else {
@@ -1124,10 +1228,14 @@ where
}
}
Err(e) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
if is_trial_pruned(&e) {
self.prune_trial(trial);
trace_info!(trial_id, "trial pruned");
} else {
self.fail_trial(trial, e.to_string());
trace_debug!(trial_id, "trial failed");
}
}
}
@@ -1191,19 +1299,29 @@ where
E: ToString + 'static,
V: Default,
{
#[cfg(feature = "tracing")]
let _span = tracing::info_span!("optimize", duration_secs = duration.as_secs(), direction = ?self.direction).entered();
let deadline = Instant::now() + duration;
while Instant::now() < deadline {
let mut trial = self.create_trial();
match objective(&mut trial) {
Ok(value) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
self.complete_trial(trial, value);
trace_info!(trial_id, "trial completed");
}
Err(e) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
if is_trial_pruned(&e) {
self.prune_trial(trial);
trace_info!(trial_id, "trial pruned");
} else {
self.fail_trial(trial, e.to_string());
trace_debug!(trial_id, "trial failed");
}
}
}
@@ -1286,14 +1404,34 @@ where
C: FnMut(&Study<V>, &CompletedTrial<V>) -> ControlFlow<()>,
E: ToString + 'static,
{
#[cfg(feature = "tracing")]
let _span = tracing::info_span!("optimize", duration_secs = duration.as_secs(), direction = ?self.direction).entered();
let deadline = Instant::now() + duration;
while Instant::now() < deadline {
let mut trial = self.create_trial();
match objective(&mut trial) {
Ok(value) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
self.complete_trial(trial, value);
#[cfg(feature = "tracing")]
{
tracing::info!(trial_id, "trial completed");
let trials = self.completed_trials.read();
if trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.count()
== 1
|| trials.last().map(|t| t.id) == self.best_id(&trials)
{
tracing::info!(trial_id, "new best value found");
}
}
let trials = self.completed_trials.read();
let Some(completed) = trials.last() else {
return Err(crate::Error::Internal(
@@ -1309,10 +1447,14 @@ where
}
}
Err(e) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
if is_trial_pruned(&e) {
self.prune_trial(trial);
trace_info!(trial_id, "trial pruned");
} else {
self.fail_trial(trial, e.to_string());
trace_debug!(trial_id, "trial failed");
}
}
}
@@ -1355,16 +1497,23 @@ where
Fut: Future<Output = core::result::Result<(Trial, V), E>>,
E: ToString,
{
#[cfg(feature = "tracing")]
let _span = tracing::info_span!("optimize_until_async", duration_secs = duration.as_secs(), direction = ?self.direction).entered();
let deadline = Instant::now() + duration;
while Instant::now() < deadline {
let trial = self.create_trial();
#[cfg(feature = "tracing")]
let trial_id = trial.id();
match objective(trial).await {
Ok((trial, value)) => {
self.complete_trial(trial, value);
trace_info!(trial_id, "trial completed");
}
Err(e) => {
let _ = e.to_string();
trace_debug!(trial_id, "trial failed");
}
}
}
@@ -1414,6 +1563,9 @@ where
{
use tokio::sync::Semaphore;
#[cfg(feature = "tracing")]
let _span = tracing::info_span!("optimize_until_parallel", duration_secs = duration.as_secs(), concurrency, direction = ?self.direction).entered();
let deadline = Instant::now() + duration;
let semaphore = Arc::new(Semaphore::new(concurrency));
let objective = Arc::new(objective);
@@ -1444,7 +1596,10 @@ where
.map_err(|e| crate::Error::TaskError(e.to_string()))?
{
Ok((trial, value)) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
self.complete_trial(trial, value);
trace_info!(trial_id, "trial completed");
}
Err(e) => {
let _ = e.to_string();
@@ -1463,6 +1618,323 @@ where
Ok(())
}
/// Runs optimization with automatic retry for failed trials.
///
/// If the objective function returns an error, the same parameter
/// configuration is retried up to `max_retries` times. Only after all
/// retries are exhausted is the trial recorded as permanently failed.
///
/// `n_trials` counts unique parameter configurations, not total
/// evaluations. A trial retried 3 times still counts as 1 toward the
/// `n_trials` limit.
///
/// # Arguments
///
/// * `n_trials` - The number of unique configurations to evaluate.
/// * `max_retries` - Maximum retry attempts per failed trial.
/// * `objective` - A closure that takes a mutable reference to a `Trial`
/// and returns the objective value or an error.
///
/// # Errors
///
/// Returns `Error::NoCompletedTrials` if no trials completed successfully.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// let call_count = std::cell::Cell::new(0u32);
/// study
/// .optimize_with_retries(5, 2, |trial| {
/// let x = x_param.suggest(trial)?;
/// call_count.set(call_count.get() + 1);
/// // Fail once every other call to exercise retry
/// if call_count.get() % 2 == 0 {
/// Err::<f64, _>(optimizer::Error::Internal("transient"))
/// } else {
/// Ok(x * x)
/// }
/// })
/// .unwrap();
///
/// assert_eq!(study.n_trials(), 5);
/// ```
pub fn optimize_with_retries<F, E>(
&self,
n_trials: usize,
max_retries: usize,
mut objective: F,
) -> crate::Result<()>
where
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
E: ToString + 'static,
V: Default,
{
#[cfg(feature = "tracing")]
let _span = tracing::info_span!("optimize_with_retries", n_trials, max_retries, direction = ?self.direction).entered();
for _ in 0..n_trials {
let mut trial = self.create_trial();
let mut retries = 0;
loop {
match objective(&mut trial) {
Ok(value) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
self.complete_trial(trial, value);
trace_info!(trial_id, "trial completed");
break;
}
Err(_) if retries < max_retries => {
retries += 1;
// Create a new trial with the same parameters
trial = self.create_trial_with_params(trial.params().clone());
}
Err(e) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
if is_trial_pruned(&e) {
self.prune_trial(trial);
trace_info!(trial_id, "trial pruned");
} else {
self.fail_trial(trial, e.to_string());
trace_debug!(trial_id, "trial permanently failed");
}
break;
}
}
}
}
// Return error if no trials completed successfully
let has_complete = self
.completed_trials
.read()
.iter()
.any(|t| t.state == TrialState::Complete);
if !has_complete {
return Err(crate::Error::NoCompletedTrials);
}
Ok(())
}
}
impl<V> Study<V>
where
V: PartialOrd + Clone + fmt::Display,
{
/// Returns a human-readable summary of the study.
///
/// The summary includes:
/// - Optimization direction and total trial count
/// - Breakdown by state (complete, pruned) when applicable
/// - Best trial value and parameters (if any completed trials exist)
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x = FloatParam::new(0.0, 10.0).name("x");
///
/// let mut trial = study.create_trial();
/// let _ = x.suggest(&mut trial).unwrap();
/// study.complete_trial(trial, 0.42);
///
/// let summary = study.summary();
/// assert!(summary.contains("Minimize"));
/// assert!(summary.contains("0.42"));
/// ```
#[must_use]
pub fn summary(&self) -> String {
use fmt::Write;
let trials = self.completed_trials.read();
let n_complete = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.count();
let n_pruned = trials
.iter()
.filter(|t| t.state == TrialState::Pruned)
.count();
let direction_str = match self.direction {
Direction::Minimize => "Minimize",
Direction::Maximize => "Maximize",
};
let mut s = format!("Study: {direction_str} | {n} trials", n = trials.len());
if n_pruned > 0 {
let _ = write!(s, " ({n_complete} complete, {n_pruned} pruned)");
}
drop(trials);
if let Ok(best) = self.best_trial() {
let _ = write!(s, "\nBest value: {} (trial #{})", best.value, best.id);
if !best.params.is_empty() {
s.push_str("\nBest parameters:");
let mut params: Vec<_> = best.params.iter().collect();
params.sort_by_key(|(id, _)| *id);
for (id, value) in params {
let label = best.param_labels.get(id).map_or("?", String::as_str);
let _ = write!(s, "\n {label} = {value}");
}
}
}
s
}
}
impl<V> Study<V>
where
V: PartialOrd + Clone,
{
/// Returns an iterator over all completed trials.
///
/// This clones the internal trial list, so it is suitable for
/// analysis and iteration but not for hot paths.
pub fn iter(&self) -> std::vec::IntoIter<CompletedTrial<V>> {
self.trials().into_iter()
}
}
impl<V> Study<V>
where
V: PartialOrd + Clone + Into<f64>,
{
/// Computes parameter importance scores using Spearman rank correlation.
///
/// For each parameter, the absolute Spearman correlation between its values
/// and the objective values is computed across all completed trials. Scores
/// are normalized so they sum to 1.0 and sorted in descending order.
///
/// Parameters that appear in fewer than 2 trials are omitted.
/// Returns an empty `Vec` if the study has fewer than 2 completed trials.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x = FloatParam::new(0.0, 10.0).name("x");
///
/// study
/// .optimize(20, |trial| {
/// let xv = x.suggest(trial)?;
/// Ok::<_, optimizer::Error>(xv * xv)
/// })
/// .unwrap();
///
/// let importance = study.param_importance();
/// assert_eq!(importance.len(), 1);
/// assert_eq!(importance[0].0, "x");
/// ```
#[must_use]
#[allow(clippy::cast_precision_loss)]
pub fn param_importance(&self) -> Vec<(String, f64)> {
use std::collections::BTreeSet;
use crate::importance::spearman;
use crate::param::ParamValue;
use crate::types::TrialState;
let trials = self.completed_trials.read();
let complete: Vec<_> = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.collect();
if complete.len() < 2 {
return Vec::new();
}
// Collect all parameter IDs across trials.
let all_param_ids: BTreeSet<_> = complete.iter().flat_map(|t| t.params.keys()).collect();
let mut scores: Vec<(String, f64)> = Vec::new();
for &param_id in &all_param_ids {
// Collect (param_value_f64, objective_f64) for trials that have this param.
let mut param_vals = Vec::new();
let mut obj_vals = Vec::new();
for trial in &complete {
if let Some(pv) = trial.params.get(param_id) {
let f = match *pv {
ParamValue::Float(v) => v,
ParamValue::Int(v) => v as f64,
ParamValue::Categorical(v) => v as f64,
};
param_vals.push(f);
obj_vals.push(trial.value.clone().into());
}
}
if param_vals.len() < 2 {
continue;
}
let corr = spearman(&param_vals, &obj_vals).abs();
// Determine label: use param_labels if available, else "param_{id}".
let label = complete
.iter()
.find_map(|t| t.param_labels.get(param_id))
.map_or_else(|| param_id.to_string(), Clone::clone);
scores.push((label, corr));
}
// Normalize so scores sum to 1.0.
let sum: f64 = scores.iter().map(|(_, s)| *s).sum();
if sum > 0.0 {
for entry in &mut scores {
entry.1 /= sum;
}
}
// Sort descending by score.
scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(core::cmp::Ordering::Equal));
scores
}
}
impl<V> IntoIterator for &Study<V>
where
V: PartialOrd + Clone,
{
type Item = CompletedTrial<V>;
type IntoIter = std::vec::IntoIter<CompletedTrial<V>>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<V> fmt::Display for Study<V>
where
V: PartialOrd + Clone + fmt::Display,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.summary())
}
}
// Specialized implementation for Study<f64> that provides deprecated `_with_sampler` aliases.
@@ -1569,6 +2041,124 @@ impl Study<f64> {
}
}
/// A serializable snapshot of a study's state.
///
/// Since [`Study`] contains non-serializable fields (samplers, atomics, etc.),
/// this struct captures the essential state needed to save and restore a study.
///
/// # Schema versioning
///
/// The `version` field enables future schema evolution without breaking existing files.
/// The current version is `1`.
///
/// # Sampler state
///
/// Sampler state is **not** included in the snapshot. After loading, the study
/// uses a default `RandomSampler`. Call [`Study::set_sampler`] to restore
/// the desired sampler configuration.
#[cfg(feature = "serde")]
#[derive(serde::Serialize, serde::Deserialize)]
pub struct StudySnapshot<V> {
/// Schema version for forward compatibility.
pub version: u32,
/// The optimization direction.
pub direction: Direction,
/// All completed (and pruned) trials.
pub trials: Vec<CompletedTrial<V>>,
/// The next trial ID to assign.
pub next_trial_id: u64,
/// Optional metadata (creation timestamp, sampler description, etc.).
pub metadata: HashMap<String, String>,
}
#[cfg(feature = "serde")]
impl<V: PartialOrd + Clone + serde::Serialize> Study<V> {
/// Saves the study state to a JSON file.
///
/// # Errors
///
/// Returns an I/O error if the file cannot be created or written.
pub fn save(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()> {
let path = path.as_ref();
let snapshot = StudySnapshot {
version: 1,
direction: self.direction,
trials: self.trials(),
next_trial_id: self.next_trial_id.load(Ordering::Relaxed),
metadata: HashMap::new(),
};
// Atomic write: write to a temp file in the same directory, then rename.
// This prevents corrupt files if the process crashes mid-write.
let parent = path.parent().unwrap_or(std::path::Path::new("."));
let tmp_path = parent.join(format!(
".{}.tmp",
path.file_name().unwrap_or_default().to_string_lossy()
));
let file = std::fs::File::create(&tmp_path)?;
serde_json::to_writer_pretty(file, &snapshot).map_err(std::io::Error::other)?;
std::fs::rename(&tmp_path, path)
}
}
#[cfg(feature = "serde")]
impl<V: PartialOrd + Clone + Default + serde::Serialize> Study<V> {
/// Runs optimization with automatic checkpointing every `interval` trials.
///
/// This is convenience sugar over [`optimize_with_callback`](Self::optimize_with_callback)
/// combined with [`save`](Self::save). The checkpoint is written atomically so
/// a crash mid-write will never leave a corrupt file.
///
/// # Errors
///
/// Returns an error if the optimization itself fails (see
/// [`optimize`](Self::optimize) for details). Checkpoint I/O errors are
/// silently ignored (best-effort).
pub fn optimize_with_checkpoint<F, E>(
&self,
n_trials: usize,
checkpoint_interval: usize,
checkpoint_path: impl AsRef<std::path::Path>,
objective: F,
) -> crate::Result<()>
where
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
E: ToString + 'static,
{
let path = checkpoint_path.as_ref().to_owned();
self.optimize_with_callback(n_trials, objective, |study, _trial| {
if study.n_trials().is_multiple_of(checkpoint_interval) {
let _ = study.save(&path);
}
ControlFlow::Continue(())
})
}
}
#[cfg(feature = "serde")]
impl<V: PartialOrd + Clone + serde::de::DeserializeOwned + 'static> Study<V> {
/// Loads a study from a JSON file.
///
/// The loaded study uses a `RandomSampler` by default. Call
/// [`set_sampler()`](Self::set_sampler) to restore the original sampler
/// configuration.
///
/// # Errors
///
/// Returns an I/O error if the file cannot be read or parsed.
pub fn load(path: impl AsRef<std::path::Path>) -> std::io::Result<Self> {
let file = std::fs::File::open(path)?;
let snapshot: StudySnapshot<V> = serde_json::from_reader(file)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
let study = Study::new(snapshot.direction);
*study.completed_trials.write() = snapshot.trials;
study
.next_trial_id
.store(snapshot.next_trial_id, Ordering::Relaxed);
Ok(study)
}
}
/// Returns `true` if the error represents a pruned trial.
///
/// Checks via `Any` downcasting whether `e` is `Error::TrialPruned` or
+32 -1
View File
@@ -15,6 +15,7 @@ use crate::types::TrialState;
/// A user attribute value that can be stored on a trial.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum AttrValue {
/// A floating-point attribute.
Float(f64),
@@ -88,6 +89,8 @@ pub struct Trial {
user_attrs: HashMap<String, AttrValue>,
/// Pre-filled parameter values from enqueue (used instead of sampling).
fixed_params: HashMap<ParamId, ParamValue>,
/// Constraint values for this trial (<=0.0 means feasible).
constraint_values: Vec<f64>,
}
impl core::fmt::Debug for Trial {
@@ -104,6 +107,7 @@ impl core::fmt::Debug for Trial {
.field("has_pruner", &self.pruner.is_some())
.field("user_attrs", &self.user_attrs)
.field("fixed_params", &self.fixed_params)
.field("constraint_values", &self.constraint_values)
.finish()
}
}
@@ -143,6 +147,7 @@ impl Trial {
pruner: None,
user_attrs: HashMap::new(),
fixed_params: HashMap::new(),
constraint_values: Vec::new(),
}
}
@@ -174,6 +179,7 @@ impl Trial {
pruner: Some(pruner),
user_attrs: HashMap::new(),
fixed_params: HashMap::new(),
constraint_values: Vec::new(),
}
}
@@ -261,7 +267,11 @@ impl Trial {
return false;
};
let history_guard = history.read();
pruner.should_prune(self.id, step, &self.intermediate_values, &history_guard)
let prune = pruner.should_prune(self.id, step, &self.intermediate_values, &history_guard);
if prune {
trace_info!(trial_id = self.id, step, "pruner recommends stopping");
}
prune
}
/// Returns all intermediate values reported so far.
@@ -287,6 +297,20 @@ impl Trial {
&self.user_attrs
}
/// Sets constraint values for this trial.
///
/// Each value represents a constraint; a value <= 0.0 means the constraint
/// is satisfied (feasible). A value > 0.0 means the constraint is violated.
pub fn set_constraints(&mut self, values: Vec<f64>) {
self.constraint_values = values;
}
/// Returns the constraint values for this trial.
#[must_use]
pub fn constraint_values(&self) -> &[f64] {
&self.constraint_values
}
/// Sets the trial state to Complete.
pub(crate) fn set_complete(&mut self) {
self.state = TrialState::Complete;
@@ -366,6 +390,13 @@ impl Trial {
let result = param.cast_param_value(&value)?;
trace_debug!(
trial_id = self.id,
param = %param.label(),
value = %value,
"parameter sampled"
);
// Store distribution, value, and label
self.distributions.insert(param_id, distribution);
self.params.insert(param_id, value);
+2
View File
@@ -2,6 +2,7 @@
/// The direction of optimization.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum Direction {
/// Minimize the objective value.
Minimize,
@@ -11,6 +12,7 @@ pub enum Direction {
/// The state of a trial in its lifecycle.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum TrialState {
/// The trial is currently running.
Running,
+133
View File
@@ -0,0 +1,133 @@
//! Integration tests for the BOHB sampler.
#![allow(
clippy::cast_sign_loss,
clippy::cast_precision_loss,
clippy::cast_possible_truncation
)]
use optimizer::parameter::{FloatParam, Parameter};
use optimizer::sampler::bohb::BohbSampler;
use optimizer::{Direction, Error, Study, TrialPruned};
#[test]
fn bohb_converges_on_quadratic() {
let bohb = BohbSampler::builder()
.min_resource(1)
.max_resource(9)
.reduction_factor(3)
.min_points_in_model(5)
.seed(42)
.build()
.unwrap();
let pruner = bohb.matching_pruner(Direction::Minimize);
let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner);
let x_param = FloatParam::new(-10.0, 10.0);
study
.optimize(60, |trial| {
let x = x_param.suggest(trial)?;
// Report intermediate values at budget steps 1, 3, 9
let obj = (x - 3.0).powi(2);
// Simulate budget-based evaluation with noise decreasing at higher budgets
trial.report(1, obj + 5.0);
trial.report(3, obj + 1.0);
trial.report(9, obj);
Ok::<_, Error>(obj)
})
.expect("optimization should succeed");
let best = study.best_trial().expect("should have trials");
assert!(
best.value < 10.0,
"BOHB should find a reasonable solution, got {}",
best.value
);
}
#[test]
fn bohb_with_pruning() {
let bohb = BohbSampler::builder()
.min_resource(1)
.max_resource(27)
.reduction_factor(3)
.min_points_in_model(3)
.seed(123)
.build()
.unwrap();
let pruner = bohb.matching_pruner(Direction::Minimize);
let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner);
let x_param = FloatParam::new(-5.0, 5.0);
study
.optimize(40, |trial| {
let x = x_param.suggest(trial)?;
let obj = x * x;
// Report at each rung step and check for pruning
for &step in &[1u64, 3, 9, 27] {
let noisy_obj = obj + 10.0 / step as f64;
trial.report(step, noisy_obj);
if trial.should_prune() {
return Err(TrialPruned.into());
}
}
Ok::<_, Error>(obj)
})
.expect("optimization should succeed");
// Verify we have completed trials
let best = study.best_trial().expect("should have at least one trial");
assert!(
best.value < 25.0,
"best value {} should be reasonable",
best.value
);
}
#[test]
fn bohb_uses_budget_conditioned_history() {
// Verify that BOHB conditions on budget level by testing that samples
// are influenced by intermediate values, not just final values.
let bohb = BohbSampler::builder()
.min_resource(1)
.max_resource(9)
.reduction_factor(3)
.min_points_in_model(3)
.seed(42)
.build()
.unwrap();
let pruner = bohb.matching_pruner(Direction::Minimize);
let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner);
let x_param = FloatParam::new(0.0, 10.0);
study
.optimize(30, |trial| {
let x = x_param.suggest(trial)?;
// Intermediate values that guide optimization toward x=2
trial.report(1, (x - 2.0).powi(2) + 1.0);
trial.report(3, (x - 2.0).powi(2) + 0.5);
trial.report(9, (x - 2.0).powi(2));
Ok::<_, Error>((x - 2.0).powi(2))
})
.expect("optimization should succeed");
let best = study.best_trial().unwrap();
let best_x: f64 = best.get(&x_param).unwrap();
// Should find x reasonably close to 2.0
assert!(
(best_x - 2.0).abs() < 5.0,
"BOHB should explore near x=2, got x={best_x}"
);
}
+259
View File
@@ -0,0 +1,259 @@
#![cfg(feature = "cma-es")]
use optimizer::prelude::*;
use optimizer::sampler::cma_es::CmaEsSampler;
#[test]
fn sphere_function() {
let sampler = CmaEsSampler::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(200, |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 < 1.0,
"sphere best value should be < 1.0, got {}",
best.value
);
}
#[test]
fn rosenbrock_function() {
let sampler = CmaEsSampler::builder().population_size(20).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(300, |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();
// Rosenbrock minimum is 0 at (1, 1); we just check reasonable convergence
assert!(
best.value < 50.0,
"rosenbrock best value should be < 50.0, got {}",
best.value
);
}
#[test]
fn bounds_respected() {
let sampler = CmaEsSampler::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 = CmaEsSampler::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();
// Should find a reasonable value
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 = CmaEsSampler::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 = CmaEsSampler::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 = CmaEsSampler::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 = CmaEsSampler::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)?;
// Minimum at n = 10
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 = CmaEsSampler::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)?;
// Minimum at lr = 0.01
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 custom_population_size_and_sigma() {
let sampler = CmaEsSampler::builder()
.sigma0(1.0)
.population_size(10)
.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
);
}
+166
View File
@@ -0,0 +1,166 @@
use optimizer::parameter::{CategoricalParam, FloatParam, IntParam, Parameter};
use optimizer::{Direction, Study};
#[test]
fn known_perfect_correlation() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(0.0, 100.0).name("x");
// Objective = x, so perfect correlation.
for _ in 0..30 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 1);
assert_eq!(importance[0].0, "x");
assert!(
(importance[0].1 - 1.0).abs() < 1e-10,
"single param should be 1.0"
);
}
#[test]
fn no_effect_parameter() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(0.0, 100.0).name("x");
let noise = FloatParam::new(0.0, 100.0).name("noise");
// Objective depends only on x; noise is unused in objective.
for _ in 0..50 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
let _nv = noise.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 2);
// x should have much higher importance than noise.
let x_score = importance.iter().find(|(l, _)| l == "x").unwrap().1;
let noise_score = importance.iter().find(|(l, _)| l == "noise").unwrap().1;
assert!(
x_score > noise_score,
"x ({x_score}) should outrank noise ({noise_score})"
);
// x should dominate
assert!(x_score > 0.7, "x importance {x_score} should be dominant");
}
#[test]
fn multiple_parameters_varying_importance() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(0.0, 10.0).name("x");
let y = FloatParam::new(0.0, 10.0).name("y");
// Objective = 10*x + 0.01*y → x should be far more important.
for _ in 0..50 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
let yv = y.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(10.0 * xv + 0.01 * yv));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 2);
assert_eq!(importance[0].0, "x", "x should rank first");
assert!(importance[0].1 > importance[1].1);
}
#[test]
fn fewer_than_two_trials_returns_empty() {
let study: Study<f64> = Study::new(Direction::Minimize);
// 0 trials
assert!(study.param_importance().is_empty());
// 1 trial
let x = FloatParam::new(0.0, 1.0).name("x");
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv));
assert!(study.param_importance().is_empty());
}
#[test]
fn int_parameter() {
let study: Study<f64> = Study::new(Direction::Minimize);
let n = IntParam::new(1, 100).name("n");
for _ in 0..30 {
let mut trial = study.ask();
let nv = n.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(nv as f64));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 1);
assert_eq!(importance[0].0, "n");
assert!(importance[0].1 > 0.9);
}
#[test]
fn categorical_parameter() {
let study: Study<f64> = Study::new(Direction::Minimize);
let cat = CategoricalParam::new(vec!["a", "b", "c"]).name("cat");
let x = FloatParam::new(0.0, 100.0).name("x");
// Objective depends only on x; categorical is random noise.
for _ in 0..50 {
let mut trial = study.ask();
let _c = cat.suggest(&mut trial).unwrap();
let xv = x.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 2);
let x_score = importance.iter().find(|(l, _)| l == "x").unwrap().1;
assert!(x_score > 0.5, "x should dominate over categorical noise");
}
#[test]
fn normalization_sums_to_one() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(0.0, 10.0).name("x");
let y = FloatParam::new(0.0, 10.0).name("y");
let z = FloatParam::new(0.0, 10.0).name("z");
for _ in 0..50 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
let yv = y.suggest(&mut trial).unwrap();
let zv = z.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv + 0.5 * yv + 0.1 * zv));
}
let importance = study.param_importance();
let sum: f64 = importance.iter().map(|(_, s)| *s).sum();
assert!(
(sum - 1.0).abs() < 1e-10,
"scores should sum to 1.0, got {sum}"
);
}
#[test]
fn label_when_unnamed_uses_debug() {
let study: Study<f64> = Study::new(Direction::Minimize);
// No .name() call → label defaults to Debug representation.
let x = FloatParam::new(0.0, 10.0);
for _ in 0..10 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 1);
assert!(
importance[0].0.starts_with("FloatParam"),
"expected Debug label, got {:?}",
importance[0].0
);
}
+353
View File
@@ -1892,3 +1892,356 @@ fn test_enqueue_counted_in_n_trials() {
// All 5 trials count, including the 2 enqueued ones
assert_eq!(study.n_trials(), 5);
}
// =============================================================================
// Test: Study summary and Display
// =============================================================================
#[test]
fn test_summary_with_completed_trials() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(1));
let x = FloatParam::new(0.0, 10.0).name("x");
study
.optimize(5, |trial| {
let val = x.suggest(trial)?;
Ok::<_, Error>(val * val)
})
.unwrap();
let summary = study.summary();
assert!(summary.contains("Minimize"));
assert!(summary.contains("5 trials"));
assert!(summary.contains("Best value:"));
assert!(summary.contains("x = "));
}
#[test]
fn test_summary_no_completed_trials() {
let study: Study<f64> = Study::new(Direction::Maximize);
let summary = study.summary();
assert!(summary.contains("Maximize"));
assert!(summary.contains("0 trials"));
assert!(!summary.contains("Best value:"));
}
#[test]
fn test_summary_with_pruned_trials() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(1));
let x = FloatParam::new(0.0, 10.0).name("x");
// Manually create some complete and pruned trials
for _ in 0..3 {
let mut trial = study.create_trial();
let val = x.suggest(&mut trial).unwrap();
study.complete_trial(trial, val);
}
for _ in 0..2 {
let mut trial = study.create_trial();
let _ = x.suggest(&mut trial).unwrap();
study.prune_trial(trial);
}
let summary = study.summary();
// Should show breakdown when there are pruned trials
if study.n_pruned_trials() > 0 {
assert!(summary.contains("complete"));
assert!(summary.contains("pruned"));
}
}
#[test]
fn test_display_matches_summary() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(1));
let x = FloatParam::new(0.0, 10.0).name("x");
study
.optimize(3, |trial| {
let val = x.suggest(trial)?;
Ok::<_, Error>(val)
})
.unwrap();
assert_eq!(format!("{study}"), study.summary());
}
// =============================================================================
// Tests: optimize_with_retries
// =============================================================================
#[test]
fn test_retries_successful_trials_not_retried() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x_param = FloatParam::new(0.0, 10.0);
let call_count = std::cell::Cell::new(0u32);
study
.optimize_with_retries(5, 3, |trial| {
let x = x_param.suggest(trial)?;
call_count.set(call_count.get() + 1);
Ok::<_, Error>(x * x)
})
.unwrap();
// All trials succeed on first try — exactly 5 calls
assert_eq!(call_count.get(), 5);
assert_eq!(study.n_trials(), 5);
}
#[test]
fn test_retries_failed_trials_retried_up_to_max() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x_param = FloatParam::new(0.0, 10.0);
let call_count = std::cell::Cell::new(0u32);
let result = study.optimize_with_retries(1, 3, |trial| {
let _ = x_param.suggest(trial).unwrap();
call_count.set(call_count.get() + 1);
Err::<f64, _>("always fails")
});
// 1 initial attempt + 3 retries = 4 total calls
assert_eq!(call_count.get(), 4);
// No trials completed
assert!(matches!(result, Err(Error::NoCompletedTrials)));
}
#[test]
fn test_retries_permanently_failed_after_exhaustion() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x_param = FloatParam::new(0.0, 10.0);
let result = study.optimize_with_retries(3, 2, |trial| {
let _ = x_param.suggest(trial).unwrap();
Err::<f64, _>("transient error")
});
assert!(
matches!(result, Err(Error::NoCompletedTrials)),
"all trials should permanently fail"
);
assert_eq!(
study.n_trials(),
0,
"no completed trials should be recorded"
);
}
#[test]
fn test_retries_uses_same_parameters() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x_param = FloatParam::new(0.0, 10.0);
let seen_values = std::cell::RefCell::new(Vec::new());
let call_count = std::cell::Cell::new(0u32);
study
.optimize_with_retries(1, 2, |trial| {
let x = x_param.suggest(trial).map_err(|e| e.to_string())?;
seen_values.borrow_mut().push(x);
call_count.set(call_count.get() + 1);
// Fail first two attempts, succeed on third
if call_count.get() < 3 {
Err::<f64, _>("transient".to_string())
} else {
Ok(x * x)
}
})
.unwrap();
let values = seen_values.borrow();
assert_eq!(values.len(), 3, "should be called 3 times (1 + 2 retries)");
// All three calls should have gotten the same parameter value
assert_eq!(values[0], values[1]);
assert_eq!(values[1], values[2]);
}
#[test]
fn test_retries_n_trials_counts_unique_configs() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x_param = FloatParam::new(0.0, 10.0);
let call_count = std::cell::Cell::new(0u32);
study
.optimize_with_retries(3, 2, |trial| {
let x = x_param.suggest(trial).map_err(|e| e.to_string())?;
call_count.set(call_count.get() + 1);
// Fail first attempt of each config, succeed on retry
if call_count.get() % 2 == 1 {
Err::<f64, _>("transient".to_string())
} else {
Ok(x * x)
}
})
.unwrap();
// 3 unique configs, each needing 2 calls = 6 total calls
assert_eq!(call_count.get(), 6);
// But only 3 completed trials
assert_eq!(study.n_trials(), 3);
}
#[test]
fn test_retries_with_zero_max_retries_same_as_optimize() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x_param = FloatParam::new(0.0, 10.0);
let call_count = std::cell::Cell::new(0u32);
study
.optimize_with_retries(5, 0, |trial| {
let x = x_param.suggest(trial)?;
call_count.set(call_count.get() + 1);
Ok::<_, Error>(x * x)
})
.unwrap();
assert_eq!(call_count.get(), 5);
assert_eq!(study.n_trials(), 5);
}
// =============================================================================
// Tests: IntoIterator for &Study
// =============================================================================
#[test]
fn test_into_iterator_iterates_all_trials() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x_param = FloatParam::new(0.0, 10.0);
for _ in 0..5 {
let mut trial = study.create_trial();
let x = x_param.suggest(&mut trial).unwrap();
study.complete_trial(trial, x * x);
}
let mut count = 0;
for trial in &study {
assert_eq!(trial.state, optimizer::TrialState::Complete);
count += 1;
}
assert_eq!(count, 5);
}
#[test]
fn test_into_iterator_empty_study() {
let study: Study<f64> = Study::new(Direction::Minimize);
let count = (&study).into_iter().count();
assert_eq!(count, 0);
}
#[test]
fn test_into_iterator_preserves_insertion_order() {
let study: Study<f64> = Study::new(Direction::Minimize);
for i in 0..3 {
let trial = study.create_trial();
study.complete_trial(trial, f64::from(i));
}
let ids: Vec<u64> = (&study).into_iter().map(|t| t.id).collect();
assert_eq!(ids, vec![0, 1, 2]);
}
// =============================================================================
// Tests: Constraint handling
// =============================================================================
#[test]
fn test_is_feasible_all_satisfied() {
let study: Study<f64> = Study::new(Direction::Minimize);
let mut trial = study.create_trial();
trial.set_constraints(vec![-1.0, 0.0, -0.5]);
study.complete_trial(trial, 1.0);
let completed = study.best_trial().unwrap();
assert!(completed.is_feasible());
}
#[test]
fn test_is_feasible_one_violated() {
let study: Study<f64> = Study::new(Direction::Minimize);
let mut trial = study.create_trial();
trial.set_constraints(vec![-1.0, 0.5, -0.5]);
study.complete_trial(trial, 1.0);
let completed = study.best_trial().unwrap();
assert!(!completed.is_feasible());
}
#[test]
fn test_is_feasible_empty_constraints() {
let study: Study<f64> = Study::new(Direction::Minimize);
let trial = study.create_trial();
study.complete_trial(trial, 1.0);
let completed = study.best_trial().unwrap();
assert!(completed.is_feasible());
}
#[test]
fn test_best_trial_prefers_feasible() {
let study: Study<f64> = Study::new(Direction::Minimize);
// Infeasible trial with better objective
let mut trial1 = study.create_trial();
trial1.set_constraints(vec![1.0]);
study.complete_trial(trial1, 0.1);
// Feasible trial with worse objective
let mut trial2 = study.create_trial();
trial2.set_constraints(vec![-1.0]);
study.complete_trial(trial2, 100.0);
let best = study.best_trial().unwrap();
assert_eq!(best.id, 1); // feasible trial wins
assert_eq!(best.value, 100.0);
}
#[test]
fn test_best_trial_feasible_by_objective() {
let study: Study<f64> = Study::new(Direction::Minimize);
// Feasible, worse objective
let mut trial1 = study.create_trial();
trial1.set_constraints(vec![-1.0]);
study.complete_trial(trial1, 10.0);
// Feasible, better objective
let mut trial2 = study.create_trial();
trial2.set_constraints(vec![-0.5]);
study.complete_trial(trial2, 2.0);
let best = study.best_trial().unwrap();
assert_eq!(best.id, 1); // lower objective wins among feasible
assert_eq!(best.value, 2.0);
}
#[test]
fn test_top_trials_ranks_feasible_above_infeasible() {
let study: Study<f64> = Study::new(Direction::Minimize);
// Infeasible, low violation
let mut t0 = study.create_trial();
t0.set_constraints(vec![0.5]);
study.complete_trial(t0, 1.0);
// Feasible, worst objective among feasible
let mut t1 = study.create_trial();
t1.set_constraints(vec![-1.0]);
study.complete_trial(t1, 50.0);
// Feasible, best objective among feasible
let mut t2 = study.create_trial();
t2.set_constraints(vec![-0.1]);
study.complete_trial(t2, 5.0);
// Infeasible, high violation
let mut t3 = study.create_trial();
t3.set_constraints(vec![3.0]);
study.complete_trial(t3, 0.5);
let top = study.top_trials(4);
let ids: Vec<u64> = top.iter().map(|t| t.id).collect();
// Feasible sorted by objective first (5.0, 50.0), then infeasible by violation (0.5, 3.0)
assert_eq!(ids, vec![2, 1, 0, 3]);
}
+295
View File
@@ -0,0 +1,295 @@
#![cfg(feature = "serde")]
use std::collections::HashMap;
use optimizer::parameter::{FloatParam, IntParam, Parameter};
use optimizer::sampler::CompletedTrial;
use optimizer::{Direction, ParamValue, Study, StudySnapshot, TrialState};
#[test]
fn round_trip_save_load() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(-10.0, 10.0).name("x");
let n = IntParam::new(1, 100).name("n");
study
.optimize(5, |trial| {
let x_val = x.suggest(trial)?;
let n_val = n.suggest(trial)?;
Ok::<_, optimizer::Error>(x_val * x_val + n_val as f64)
})
.unwrap();
let dir = tempdir();
let path = dir.join("study.json");
study.save(&path).unwrap();
let loaded: Study<f64> = Study::load(&path).unwrap();
assert_eq!(loaded.direction(), study.direction());
assert_eq!(loaded.n_trials(), study.n_trials());
let orig_trials = study.trials();
let loaded_trials = loaded.trials();
for (orig, loaded) in orig_trials.iter().zip(loaded_trials.iter()) {
assert_eq!(orig.id, loaded.id);
assert!((orig.value - loaded.value).abs() < 1e-10);
assert_eq!(orig.state, loaded.state);
assert_eq!(orig.params.len(), loaded.params.len());
assert_eq!(orig.distributions, loaded.distributions);
assert_eq!(orig.param_labels, loaded.param_labels);
}
// Clean up
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn json_output_is_human_readable() {
let study: Study<f64> = Study::new(Direction::Maximize);
let x = FloatParam::new(0.0, 1.0).name("x");
study
.optimize(2, |trial| {
let v = x.suggest(trial)?;
Ok::<_, optimizer::Error>(v)
})
.unwrap();
let dir = tempdir();
let path = dir.join("study.json");
study.save(&path).unwrap();
let contents = std::fs::read_to_string(&path).unwrap();
// Verify it's pretty-printed JSON with recognizable fields
assert!(contents.contains("\"version\""));
assert!(contents.contains("\"direction\""));
assert!(contents.contains("\"trials\""));
assert!(contents.contains("\"next_trial_id\""));
assert!(contents.contains("\"Maximize\""));
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn round_trip_empty_study() {
let study: Study<f64> = Study::new(Direction::Minimize);
let dir = tempdir();
let path = dir.join("empty.json");
study.save(&path).unwrap();
let loaded: Study<f64> = Study::load(&path).unwrap();
assert_eq!(loaded.direction(), Direction::Minimize);
assert_eq!(loaded.n_trials(), 0);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn snapshot_version_field_is_present() {
let study: Study<f64> = Study::new(Direction::Minimize);
let dir = tempdir();
let path = dir.join("version.json");
study.save(&path).unwrap();
let contents = std::fs::read_to_string(&path).unwrap();
let snapshot: StudySnapshot<f64> = serde_json::from_str(&contents).unwrap();
assert_eq!(snapshot.version, 1);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn completed_trial_serde_round_trip() {
let trial = CompletedTrial::new(42, HashMap::new(), HashMap::new(), HashMap::new(), 2.78);
let json = serde_json::to_string(&trial).unwrap();
let deserialized: CompletedTrial<f64> = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.id, 42);
assert_eq!(deserialized.value, 2.78);
assert_eq!(deserialized.state, TrialState::Complete);
}
#[test]
fn param_value_serde_round_trip() {
let values = vec![
ParamValue::Float(1.23),
ParamValue::Int(42),
ParamValue::Categorical(2),
];
for val in &values {
let json = serde_json::to_string(val).unwrap();
let deserialized: ParamValue = serde_json::from_str(&json).unwrap();
assert_eq!(&deserialized, val);
}
}
#[test]
fn direction_serde_round_trip() {
let min_json = serde_json::to_string(&Direction::Minimize).unwrap();
let max_json = serde_json::to_string(&Direction::Maximize).unwrap();
assert_eq!(
serde_json::from_str::<Direction>(&min_json).unwrap(),
Direction::Minimize
);
assert_eq!(
serde_json::from_str::<Direction>(&max_json).unwrap(),
Direction::Maximize
);
}
#[test]
fn round_trip_preserves_trial_id_counter() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(10, |trial| {
let v = x.suggest(trial)?;
Ok::<_, optimizer::Error>(v)
})
.unwrap();
let dir = tempdir();
let path = dir.join("counter.json");
study.save(&path).unwrap();
let loaded: Study<f64> = Study::load(&path).unwrap();
// Creating a new trial should use an ID >= 10
let trial = loaded.create_trial();
assert!(trial.id() >= 10);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn checkpoint_file_created_at_interval() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(-10.0, 10.0).name("x");
let dir = tempdir();
let checkpoint = dir.join("checkpoint.json");
study
.optimize_with_checkpoint(10, 3, &checkpoint, |trial| {
let v = x.suggest(trial)?;
Ok::<_, optimizer::Error>(v * v)
})
.unwrap();
// Checkpoint should exist (written at trials 3, 6, 9)
assert!(checkpoint.exists(), "checkpoint file was not created");
// Load it and verify it's valid
let loaded: Study<f64> = Study::load(&checkpoint).unwrap();
// Last checkpoint was at trial 9, so it should have 9 trials
assert_eq!(loaded.n_trials(), 9);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn checkpoint_overwrites_previous() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(0.0, 1.0);
let dir = tempdir();
let checkpoint = dir.join("checkpoint.json");
study
.optimize_with_checkpoint(6, 3, &checkpoint, |trial| {
let v = x.suggest(trial)?;
Ok::<_, optimizer::Error>(v)
})
.unwrap();
// The checkpoint at trial 6 should overwrite the one from trial 3
let loaded: Study<f64> = Study::load(&checkpoint).unwrap();
assert_eq!(loaded.n_trials(), 6);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn resume_from_checkpoint_continues_trial_ids() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(-5.0, 5.0).name("x");
let dir = tempdir();
let checkpoint = dir.join("resume.json");
// Run 10 trials with checkpointing
study
.optimize_with_checkpoint(10, 5, &checkpoint, |trial| {
let v = x.suggest(trial)?;
Ok::<_, optimizer::Error>(v * v)
})
.unwrap();
// Load and continue
let loaded: Study<f64> = Study::load(&checkpoint).unwrap();
assert_eq!(loaded.n_trials(), 10);
let remaining = 15 - loaded.n_trials();
loaded
.optimize(remaining, |trial| {
let v = x.suggest(trial)?;
Ok::<_, optimizer::Error>(v * v)
})
.unwrap();
assert_eq!(loaded.n_trials(), 15);
// Verify no duplicate trial IDs
let trials = loaded.trials();
let mut ids: Vec<u64> = trials.iter().map(|t| t.id).collect();
ids.sort_unstable();
ids.dedup();
assert_eq!(ids.len(), 15, "duplicate trial IDs found");
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn atomic_write_no_temp_file_left_behind() {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(0.0, 1.0);
let dir = tempdir();
let checkpoint = dir.join("atomic.json");
study
.optimize_with_checkpoint(3, 3, &checkpoint, |trial| {
let v = x.suggest(trial)?;
Ok::<_, optimizer::Error>(v)
})
.unwrap();
// The temp file should have been renamed, not left behind
let tmp_path = dir.join(".atomic.json.tmp");
assert!(!tmp_path.exists(), "temp file was not cleaned up");
assert!(checkpoint.exists(), "checkpoint file was not created");
std::fs::remove_dir_all(&dir).ok();
}
/// Helper to create a unique temporary directory.
fn tempdir() -> std::path::PathBuf {
use std::sync::atomic::{AtomicU64, Ordering};
static COUNTER: AtomicU64 = AtomicU64::new(0);
let id = COUNTER.fetch_add(1, Ordering::Relaxed);
let dir =
std::env::temp_dir().join(format!("optimizer_serde_test_{}_{id}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
dir
}
+45
View File
@@ -0,0 +1,45 @@
#[path = "../benches/test_functions.rs"]
mod test_functions;
use test_functions::*;
const TOL: f64 = 1e-10;
#[test]
fn sphere_at_optimum() {
assert!(sphere(&[0.0, 0.0]).abs() < TOL);
assert!(sphere(&[0.0; 10]).abs() < TOL);
}
#[test]
fn rosenbrock_at_optimum() {
assert!(rosenbrock(&[1.0, 1.0]).abs() < TOL);
assert!(rosenbrock(&[1.0; 5]).abs() < TOL);
}
#[test]
fn rastrigin_at_optimum() {
assert!(rastrigin(&[0.0, 0.0]).abs() < TOL);
assert!(rastrigin(&[0.0; 10]).abs() < TOL);
}
#[test]
fn ackley_at_optimum() {
assert!(ackley(&[0.0, 0.0]).abs() < 1e-8);
assert!(ackley(&[0.0; 10]).abs() < 1e-8);
}
#[test]
fn branin_at_optimum() {
let target = 0.397_887_357_729_738_1;
let val = branin(&[std::f64::consts::PI, 2.275]);
assert!((val - target).abs() < 1e-3);
}
#[test]
fn hartmann6_at_optimum() {
let target = -3.3224;
let x_opt = [0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573];
let val = hartmann6(&x_opt);
assert!((val - target).abs() < 0.01);
}