Compare commits
14 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 147a64708d | |||
| b36137bc96 | |||
| f3941a4b3e | |||
| 0cae3b4227 | |||
| 09480a973b | |||
| aa9734587a | |||
| a1dd6d393c | |||
| a53ba4f472 | |||
| 9b4a8321b5 | |||
| db0314a1d1 | |||
| df5fcedecf | |||
| 5061df9876 | |||
| 5fe0a75f78 | |||
| dcf3403261 |
+20
-1
@@ -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"
|
||||
|
||||
@@ -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), ¶ms, |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), ¶ms, |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);
|
||||
@@ -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 ¶ms {
|
||||
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);
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)),
|
||||
¶ms,
|
||||
functions::sphere,
|
||||
n_trials,
|
||||
);
|
||||
run_convergence(
|
||||
"sphere_5d",
|
||||
"tpe",
|
||||
Study::minimize(TpeSampler::builder().seed(1).build().unwrap()),
|
||||
¶ms,
|
||||
functions::sphere,
|
||||
n_trials,
|
||||
);
|
||||
|
||||
// Rosenbrock: Random vs TPE
|
||||
run_convergence(
|
||||
"rosenbrock_5d",
|
||||
"random",
|
||||
Study::minimize(RandomSampler::with_seed(2)),
|
||||
¶ms,
|
||||
functions::rosenbrock,
|
||||
n_trials,
|
||||
);
|
||||
run_convergence(
|
||||
"rosenbrock_5d",
|
||||
"tpe",
|
||||
Study::minimize(TpeSampler::builder().seed(2).build().unwrap()),
|
||||
¶ms,
|
||||
functions::rosenbrock,
|
||||
n_trials,
|
||||
);
|
||||
|
||||
// Rastrigin: Random vs TPE
|
||||
run_convergence(
|
||||
"rastrigin_5d",
|
||||
"random",
|
||||
Study::minimize(RandomSampler::with_seed(3)),
|
||||
¶ms,
|
||||
functions::rastrigin,
|
||||
n_trials,
|
||||
);
|
||||
run_convergence(
|
||||
"rastrigin_5d",
|
||||
"tpe",
|
||||
Study::minimize(TpeSampler::builder().seed(3).build().unwrap()),
|
||||
¶ms,
|
||||
functions::rastrigin,
|
||||
n_trials,
|
||||
);
|
||||
}
|
||||
@@ -13,4 +13,3 @@ proc-macro = true
|
||||
[dependencies]
|
||||
syn = { version = "2", features = ["full"] }
|
||||
quote = "1"
|
||||
proc-macro2 = "1"
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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> {
|
||||
|
||||
@@ -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
@@ -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 ¶m_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(¶m_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
@@ -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,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,
|
||||
|
||||
@@ -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}"
|
||||
);
|
||||
}
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
@@ -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]);
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
Reference in New Issue
Block a user