From b36137bc963c01a9263f6a9ba852bf0c5a05ee74 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 19:05:33 +0100 Subject: [PATCH] feat: add benchmark suite with criterion and standard test functions Add micro-benchmarks for TPE, Random, and Grid samplers with varying history sizes, end-to-end optimization benchmarks on Sphere and Rosenbrock, and a convergence tracking example that outputs CSV data. Includes 6 standard test functions (Sphere, Rosenbrock, Rastrigin, Ackley, Branin, Hartmann6) with correctness tests at known optima. --- Cargo.toml | 9 +++ benches/optimization.rs | 112 +++++++++++++++++++++++++++ benches/samplers.rs | 118 ++++++++++++++++++++++++++++ benches/test_functions.rs | 84 ++++++++++++++++++++ examples/benchmark_convergence.rs | 124 ++++++++++++++++++++++++++++++ tests/test_functions.rs | 45 +++++++++++ 6 files changed, 492 insertions(+) create mode 100644 benches/optimization.rs create mode 100644 benches/samplers.rs create mode 100644 benches/test_functions.rs create mode 100644 examples/benchmark_convergence.rs create mode 100644 tests/test_functions.rs diff --git a/Cargo.toml b/Cargo.toml index 442f6ac..f30214e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,6 +40,15 @@ cma-es = ["dep:nalgebra"] 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" diff --git a/benches/optimization.rs b/benches/optimization.rs new file mode 100644 index 0000000..c4f58dc --- /dev/null +++ b/benches/optimization.rs @@ -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 { + (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 = params + .iter() + .map(|p| p.suggest(trial)) + .collect::>() + .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 = params + .iter() + .map(|p| p.suggest(trial)) + .collect::>() + .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 = params + .iter() + .map(|p| p.suggest(trial)) + .collect::>() + .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 = params + .iter() + .map(|p| p.suggest(trial)) + .collect::>() + .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); diff --git a/benches/samplers.rs b/benches/samplers.rs new file mode 100644 index 0000000..f2e9ec3 --- /dev/null +++ b/benches/samplers.rs @@ -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> { + let params: Vec = (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> = 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); diff --git a/benches/test_functions.rs b/benches/test_functions.rs new file mode 100644 index 0000000..a15c1c4 --- /dev/null +++ b/benches/test_functions.rs @@ -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::() +} + +/// 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 +} diff --git a/examples/benchmark_convergence.rs b/examples/benchmark_convergence.rs new file mode 100644 index 0000000..9f478a7 --- /dev/null +++ b/examples/benchmark_convergence.rs @@ -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::() + } +} + +fn run_convergence( + name: &str, + sampler_name: &str, + study: Study, + params: &[FloatParam], + objective: fn(&[f64]) -> f64, + n_trials: usize, +) { + let start = Instant::now(); + + study + .optimize_with_callback( + n_trials, + |trial| { + let x: Vec = params + .iter() + .map(|p| p.suggest(trial)) + .collect::>() + .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 = (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, + ); +} diff --git a/tests/test_functions.rs b/tests/test_functions.rs new file mode 100644 index 0000000..23fa678 --- /dev/null +++ b/tests/test_functions.rs @@ -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); +}