Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Manuel Raimann
|
762f1b3cde | ||
|
Manuel Raimann
|
b3534b3a3b | ||
|
Manuel Raimann
|
9944e69920 | ||
|
Manuel Raimann
|
e8e446fb0e |
@@ -61,21 +61,6 @@ jobs:
|
||||
cargo test --verbose --features "${{ matrix.features }}"
|
||||
fi
|
||||
|
||||
examples:
|
||||
name: Examples
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- name: Install Rust
|
||||
run: |
|
||||
rustup override set stable
|
||||
rustup update stable
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
- name: Run sync example (ml_hyperparameter_tuning)
|
||||
run: cargo run --example ml_hyperparameter_tuning
|
||||
- name: Run async example (async_api_optimization)
|
||||
run: cargo run --example async_api_optimization --features async
|
||||
|
||||
docs:
|
||||
name: Docs
|
||||
runs-on: ubuntu-latest
|
||||
@@ -258,7 +243,6 @@ jobs:
|
||||
- fmt
|
||||
- clippy
|
||||
- test
|
||||
- examples
|
||||
- docs
|
||||
- feature-check
|
||||
- coverage
|
||||
|
||||
+3
-42
@@ -1,9 +1,6 @@
|
||||
[workspace]
|
||||
members = ["optimizer-derive"]
|
||||
|
||||
[package]
|
||||
name = "optimizer"
|
||||
version = "0.7.2"
|
||||
version = "0.3.1"
|
||||
edition = "2024"
|
||||
rust-version = "1.88"
|
||||
license = "MIT"
|
||||
@@ -16,50 +13,14 @@ categories = ["algorithm", "science", "data-structures"]
|
||||
readme = "README.md"
|
||||
|
||||
[dependencies]
|
||||
rand = "0.10"
|
||||
rand = "0.9"
|
||||
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.29", optional = true }
|
||||
sobol_burley = { version = "0.5", optional = true }
|
||||
nalgebra = { version = "0.34", 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"
|
||||
path = "examples/async_api_optimization.rs"
|
||||
required-features = ["async"]
|
||||
|
||||
[[example]]
|
||||
name = "ml_hyperparameter_tuning"
|
||||
path = "examples/ml_hyperparameter_tuning.rs"
|
||||
|
||||
[[example]]
|
||||
name = "parameter_api"
|
||||
path = "examples/parameter_api.rs"
|
||||
required-features = ["derive"]
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
|
||||
|
||||
@@ -13,39 +13,28 @@ A Rust library for black-box optimization with multiple sampling strategies.
|
||||
- **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
|
||||
- Float, integer, categorical, boolean, and enum parameter types
|
||||
- Float, integer, and categorical parameter types
|
||||
- Log-scale and stepped parameter sampling
|
||||
- Sync and async optimization with parallel trial evaluation
|
||||
- `#[derive(Categorical)]` for enum parameters
|
||||
|
||||
## Quick Start
|
||||
|
||||
```rust
|
||||
use optimizer::parameter::{FloatParam, Parameter};
|
||||
use optimizer::sampler::tpe::TpeSampler;
|
||||
use optimizer::{Direction, Study};
|
||||
use optimizer::sampler::tpe::TpeSampler;
|
||||
|
||||
// Create a study with TPE sampler
|
||||
let sampler = TpeSampler::builder().seed(42).build().unwrap();
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
// Define parameter search space
|
||||
let x_param = FloatParam::new(-10.0, 10.0);
|
||||
|
||||
// Optimize x^2 for 20 trials
|
||||
study
|
||||
.optimize_with_sampler(20, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
Ok::<_, optimizer::Error>(x * x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// Get the best result
|
||||
let best = study.best_trial().unwrap();
|
||||
println!("Best value: {}", best.value);
|
||||
for (id, label) in &best.param_labels {
|
||||
println!(" {}: {:?}", label, best.params[id]);
|
||||
}
|
||||
println!("Best value: {} at x={:?}", best.value, best.params);
|
||||
```
|
||||
|
||||
## Samplers
|
||||
@@ -79,56 +68,6 @@ let sampler = TpeSampler::builder()
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
```
|
||||
|
||||
#### Gamma Strategies
|
||||
|
||||
The gamma parameter controls what fraction of trials are considered "good" when building the TPE model. Instead of a fixed value, you can use adaptive strategies:
|
||||
|
||||
| Strategy | Description | Formula |
|
||||
|----------|-------------|---------|
|
||||
| `FixedGamma` | Constant value (default: 0.25) | `γ = constant` |
|
||||
| `LinearGamma` | Linear interpolation over trials | `γ = γ_min + (γ_max - γ_min) * min(n/n_max, 1)` |
|
||||
| `SqrtGamma` | Optuna-style inverse sqrt scaling | `γ = min(γ_max, factor/√n / n)` |
|
||||
| `HyperoptGamma` | Hyperopt-style adaptive | `γ = min(γ_max, (base + 1) / n)` |
|
||||
|
||||
```rust
|
||||
use optimizer::sampler::tpe::{TpeSampler, SqrtGamma, LinearGamma};
|
||||
|
||||
// Optuna-style gamma that decreases with more trials
|
||||
let sampler = TpeSampler::builder()
|
||||
.gamma_strategy(SqrtGamma::default())
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
// Linear interpolation from 0.1 to 0.3 over 100 trials
|
||||
let sampler = TpeSampler::builder()
|
||||
.gamma_strategy(LinearGamma::new(0.1, 0.3, 100).unwrap())
|
||||
.build()
|
||||
.unwrap();
|
||||
```
|
||||
|
||||
You can also implement custom strategies:
|
||||
|
||||
```rust
|
||||
use optimizer::sampler::tpe::{TpeSampler, GammaStrategy};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct MyGamma { base: f64 }
|
||||
|
||||
impl GammaStrategy for MyGamma {
|
||||
fn gamma(&self, n_trials: usize) -> f64 {
|
||||
(self.base + 0.01 * n_trials as f64).min(0.5)
|
||||
}
|
||||
fn clone_box(&self) -> Box<dyn GammaStrategy> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
}
|
||||
|
||||
let sampler = TpeSampler::builder()
|
||||
.gamma_strategy(MyGamma { base: 0.1 })
|
||||
.build()
|
||||
.unwrap();
|
||||
```
|
||||
|
||||
### Grid Search
|
||||
|
||||
```rust
|
||||
@@ -145,7 +84,6 @@ let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
## Feature Flags
|
||||
|
||||
- `async` - Enable async optimization methods (requires tokio)
|
||||
- `derive` - Enable `#[derive(Categorical)]` for enum parameters
|
||||
|
||||
## Documentation
|
||||
|
||||
|
||||
@@ -1,112 +0,0 @@
|
||||
#[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);
|
||||
@@ -1,118 +0,0 @@
|
||||
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);
|
||||
@@ -1,84 +0,0 @@
|
||||
//! 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
|
||||
}
|
||||
@@ -1,11 +1,7 @@
|
||||
[advisories]
|
||||
version = 2
|
||||
db-path = "~/.cargo/advisory-db"
|
||||
ignore = [
|
||||
# paste is unmaintained but it's a transitive dep via simba -> nalgebra
|
||||
# with no fix available upstream yet
|
||||
"RUSTSEC-2024-0436",
|
||||
]
|
||||
ignore = []
|
||||
|
||||
[licenses]
|
||||
version = 2
|
||||
@@ -13,6 +9,7 @@ allow = [
|
||||
"MIT",
|
||||
"Apache-2.0",
|
||||
"Apache-2.0 WITH LLVM-exception",
|
||||
"BSD-2-Clause",
|
||||
"Unicode-3.0",
|
||||
]
|
||||
confidence-threshold = 0.8
|
||||
|
||||
@@ -1,368 +0,0 @@
|
||||
//! Async API Parameter Optimization Example
|
||||
//!
|
||||
//! This example shows how to use async/parallel optimization to tune
|
||||
//! configuration parameters for a web service. Each evaluation simulates
|
||||
//! an async operation (like deploying and load-testing a service).
|
||||
//!
|
||||
//! # Key Concepts Demonstrated
|
||||
//!
|
||||
//! - Async optimization with `optimize_parallel`
|
||||
//! - Running multiple trials concurrently for faster optimization
|
||||
//! - Boolean and categorical parameter types
|
||||
//! - Measuring speedup from parallelism
|
||||
//!
|
||||
//! # When to Use Async Optimization
|
||||
//!
|
||||
//! Use async/parallel optimization when your objective function involves:
|
||||
//! - Network requests (API calls, database queries)
|
||||
//! - File I/O operations
|
||||
//! - External service calls
|
||||
//! - Any operation where you're waiting for I/O rather than computing
|
||||
//!
|
||||
//! With parallelism, you can evaluate multiple configurations simultaneously,
|
||||
//! significantly reducing total optimization time.
|
||||
//!
|
||||
//! Run with: `cargo run --example async_api_optimization --features async`
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use optimizer::prelude::*;
|
||||
|
||||
// ============================================================================
|
||||
// Configuration: Service parameters we want to tune
|
||||
// ============================================================================
|
||||
|
||||
/// Configuration for a web service.
|
||||
///
|
||||
/// In a real application, these parameters would control:
|
||||
/// - Memory allocation (cache sizes)
|
||||
/// - Connection management (pool sizes, timeouts)
|
||||
/// - Request handling (batching, compression)
|
||||
/// - Protocol options (HTTP version, load balancing)
|
||||
struct ServiceConfig {
|
||||
cache_size_mb: i64,
|
||||
connection_pool_size: i64,
|
||||
request_timeout_ms: i64,
|
||||
retry_count: i64,
|
||||
batch_size: i64,
|
||||
compression_level: i64,
|
||||
use_http2: bool,
|
||||
load_balancing: String,
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Objective Function: Evaluate a service configuration
|
||||
// ============================================================================
|
||||
|
||||
/// Simulates deploying and load-testing a service configuration.
|
||||
///
|
||||
/// In a real scenario, this function might:
|
||||
/// 1. Deploy the configuration to a staging environment
|
||||
/// 2. Run load tests against the service
|
||||
/// 3. Collect metrics (latency, throughput, error rate)
|
||||
/// 4. Return a composite score
|
||||
///
|
||||
/// The async sleep simulates the I/O time of these operations.
|
||||
/// This is where parallel execution helps - while one trial is waiting
|
||||
/// for I/O, other trials can run.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn evaluate_service(config: &ServiceConfig) -> f64 {
|
||||
// Simulate async I/O (deployment, load testing, metric collection)
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
|
||||
// Calculate a score based on how close we are to optimal values
|
||||
// Lower score = better configuration
|
||||
let mut score = 0.0;
|
||||
|
||||
// Cache size: too small = cache misses, too large = wasted memory
|
||||
// Optimal around 512MB
|
||||
let cache_optimal = 512.0;
|
||||
score += ((config.cache_size_mb as f64 - cache_optimal) / 256.0).powi(2);
|
||||
|
||||
// Connection pool: too small = contention, too large = resource waste
|
||||
// Optimal around 100
|
||||
let pool_optimal = 100.0;
|
||||
score += ((config.connection_pool_size as f64 - pool_optimal) / 50.0).powi(2);
|
||||
|
||||
// Timeout: too short = false failures, too long = slow recovery
|
||||
// Optimal around 5000ms
|
||||
let timeout_optimal = 5000.0;
|
||||
score += ((config.request_timeout_ms as f64 - timeout_optimal) / 2000.0).powi(2);
|
||||
|
||||
// Retries: too few = fragile, too many = amplifies failures
|
||||
// Optimal around 3
|
||||
let retry_optimal = 3.0;
|
||||
score += ((config.retry_count as f64 - retry_optimal) / 2.0).powi(2);
|
||||
|
||||
// Batch size: trade-off between latency and throughput
|
||||
// Optimal around 64
|
||||
let batch_optimal = 64.0;
|
||||
score += ((config.batch_size as f64 - batch_optimal) / 32.0).powi(2);
|
||||
|
||||
// Compression level: trade-off between CPU and bandwidth
|
||||
// Optimal around 6
|
||||
let compression_optimal = 6.0;
|
||||
score += ((config.compression_level as f64 - compression_optimal) / 3.0).powi(2);
|
||||
|
||||
// HTTP/2 is generally better for our use case
|
||||
if !config.use_http2 {
|
||||
score += 0.5;
|
||||
}
|
||||
|
||||
// Load balancing strategy affects performance
|
||||
score += match config.load_balancing.as_str() {
|
||||
"round_robin" => 0.0, // Best for our use case
|
||||
"least_connections" => 0.1, // Good alternative
|
||||
"ip_hash" => 0.2, // OK for session affinity
|
||||
"random" => 0.3, // Not ideal
|
||||
_ => 1.0,
|
||||
};
|
||||
|
||||
// Add noise to simulate real-world variability
|
||||
let noise = (config.cache_size_mb as f64 * 0.1).sin() * 0.05;
|
||||
|
||||
score + noise
|
||||
}
|
||||
|
||||
/// The async objective function for each trial.
|
||||
///
|
||||
/// For async optimization, the objective function must:
|
||||
/// 1. Take ownership of the Trial (not a mutable reference)
|
||||
/// 2. Return a Future
|
||||
/// 3. Return both the Trial and the result value as a tuple
|
||||
///
|
||||
/// This ownership pattern allows the trial to be used across await points.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn objective(
|
||||
mut trial: Trial,
|
||||
cache_size_mb_param: &IntParam,
|
||||
connection_pool_size_param: &IntParam,
|
||||
request_timeout_ms_param: &IntParam,
|
||||
retry_count_param: &IntParam,
|
||||
batch_size_param: &IntParam,
|
||||
compression_level_param: &IntParam,
|
||||
use_http2_param: &BoolParam,
|
||||
load_balancing_param: &CategoricalParam<&str>,
|
||||
) -> optimizer::Result<(Trial, f64)> {
|
||||
// Sample configuration parameters using parameter definitions
|
||||
let cache_size_mb = cache_size_mb_param.suggest(&mut trial)?;
|
||||
let connection_pool_size = connection_pool_size_param.suggest(&mut trial)?;
|
||||
let request_timeout_ms = request_timeout_ms_param.suggest(&mut trial)?;
|
||||
let retry_count = retry_count_param.suggest(&mut trial)?;
|
||||
let batch_size = batch_size_param.suggest(&mut trial)?;
|
||||
let compression_level = compression_level_param.suggest(&mut trial)?;
|
||||
let use_http2 = use_http2_param.suggest(&mut trial)?;
|
||||
let load_balancing = load_balancing_param.suggest(&mut trial)?;
|
||||
|
||||
// Build configuration
|
||||
let config = ServiceConfig {
|
||||
cache_size_mb,
|
||||
connection_pool_size,
|
||||
request_timeout_ms,
|
||||
retry_count,
|
||||
batch_size,
|
||||
compression_level,
|
||||
use_http2,
|
||||
load_balancing: load_balancing.to_string(),
|
||||
};
|
||||
|
||||
// Evaluate (this is the async part)
|
||||
let score = evaluate_service(&config).await;
|
||||
|
||||
// Return both the trial and the score
|
||||
Ok((trial, score))
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Helper Functions
|
||||
// ============================================================================
|
||||
|
||||
/// Prints the results of the optimization.
|
||||
fn print_results(study: &Study<f64>, elapsed: Duration, n_trials: usize) {
|
||||
println!("\n{}", "=".repeat(60));
|
||||
println!("\nOptimization completed!");
|
||||
println!("Total trials: {}", study.n_trials());
|
||||
println!("Time elapsed: {elapsed:.2?}");
|
||||
|
||||
// Calculate speedup from parallelism
|
||||
// Each trial takes ~50ms, so sequential would take n_trials * 50ms
|
||||
let sequential_time = n_trials as f64 * 0.050;
|
||||
let actual_time = elapsed.as_secs_f64();
|
||||
println!(
|
||||
"Effective parallelism: {:.1}x speedup",
|
||||
sequential_time / actual_time
|
||||
);
|
||||
}
|
||||
|
||||
/// Prints the best configuration found.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn print_best_config(
|
||||
study: &Study<f64>,
|
||||
cache_size_mb_param: &IntParam,
|
||||
connection_pool_size_param: &IntParam,
|
||||
request_timeout_ms_param: &IntParam,
|
||||
retry_count_param: &IntParam,
|
||||
batch_size_param: &IntParam,
|
||||
compression_level_param: &IntParam,
|
||||
use_http2_param: &BoolParam,
|
||||
load_balancing_param: &CategoricalParam<&str>,
|
||||
) -> optimizer::Result<()> {
|
||||
let best = study.best_trial()?;
|
||||
|
||||
println!("\nBest configuration found:");
|
||||
println!(" Score: {:.6}", best.value);
|
||||
println!("\n Parameters:");
|
||||
println!(
|
||||
" cache_size_mb: {}",
|
||||
best.get(cache_size_mb_param).unwrap()
|
||||
);
|
||||
println!(
|
||||
" connection_pool_size: {}",
|
||||
best.get(connection_pool_size_param).unwrap()
|
||||
);
|
||||
println!(
|
||||
" request_timeout_ms: {}",
|
||||
best.get(request_timeout_ms_param).unwrap()
|
||||
);
|
||||
println!(" retry_count: {}", best.get(retry_count_param).unwrap());
|
||||
println!(" batch_size: {}", best.get(batch_size_param).unwrap());
|
||||
println!(
|
||||
" compression_level: {}",
|
||||
best.get(compression_level_param).unwrap()
|
||||
);
|
||||
println!(" use_http2: {}", best.get(use_http2_param).unwrap());
|
||||
println!(
|
||||
" load_balancing: {}",
|
||||
best.get(load_balancing_param).unwrap()
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Prints the top N trials.
|
||||
fn print_top_trials(study: &Study<f64>, n: usize) {
|
||||
println!("\nTop {n} trials:");
|
||||
|
||||
let mut trials = study.trials();
|
||||
trials.sort_by(|a, b| a.value.partial_cmp(&b.value).unwrap());
|
||||
|
||||
for (i, trial) in trials.iter().take(n).enumerate() {
|
||||
println!(
|
||||
" {}. Trial #{}: score = {:.6}",
|
||||
i + 1,
|
||||
trial.id,
|
||||
trial.value
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Main: Set up and run the async optimization
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> optimizer::Result<()> {
|
||||
println!("=== Async API Parameter Optimization Example ===\n");
|
||||
|
||||
// Step 1: Create a TPE sampler
|
||||
let sampler = TpeSampler::builder()
|
||||
.n_startup_trials(8)
|
||||
.gamma(0.2)
|
||||
.seed(123)
|
||||
.build()
|
||||
.expect("Failed to build TPE sampler");
|
||||
|
||||
// Step 2: Create a study to minimize the score
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
// Step 3: Define parameter search spaces
|
||||
let cache_size_mb_param = IntParam::new(64, 1024).name("cache_size_mb").step(64);
|
||||
let connection_pool_size_param = IntParam::new(10, 200).name("connection_pool_size").step(10);
|
||||
let request_timeout_ms_param = IntParam::new(1000, 10000)
|
||||
.name("request_timeout_ms")
|
||||
.step(500);
|
||||
let retry_count_param = IntParam::new(0, 5).name("retry_count");
|
||||
let batch_size_param = IntParam::new(1, 256).name("batch_size").log_scale();
|
||||
let compression_level_param = IntParam::new(0, 9).name("compression_level");
|
||||
let use_http2_param = BoolParam::new().name("use_http2");
|
||||
let load_balancing_param = CategoricalParam::new(vec![
|
||||
"round_robin",
|
||||
"least_connections",
|
||||
"random",
|
||||
"ip_hash",
|
||||
])
|
||||
.name("load_balancing");
|
||||
|
||||
// Clone params for use after the closure moves them
|
||||
let cache_size_mb_p = cache_size_mb_param.clone();
|
||||
let connection_pool_size_p = connection_pool_size_param.clone();
|
||||
let request_timeout_ms_p = request_timeout_ms_param.clone();
|
||||
let retry_count_p = retry_count_param.clone();
|
||||
let batch_size_p = batch_size_param.clone();
|
||||
let compression_level_p = compression_level_param.clone();
|
||||
let use_http2_p = use_http2_param.clone();
|
||||
let load_balancing_p = load_balancing_param.clone();
|
||||
|
||||
// Step 4: Configure optimization
|
||||
let n_trials = 40;
|
||||
let concurrency = 4; // Run 4 trials in parallel
|
||||
|
||||
println!("Starting parallel optimization with {concurrency} concurrent evaluations...\n");
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
// Step 5: Run parallel async optimization
|
||||
//
|
||||
// optimize_parallel:
|
||||
// - Runs up to `concurrency` trials simultaneously
|
||||
// - Each trial calls the objective function
|
||||
// - Uses a semaphore to limit concurrent evaluations
|
||||
// - Collects results as trials complete
|
||||
//
|
||||
// The sampler gets access to trial history for informed sampling.
|
||||
study
|
||||
.optimize_parallel(n_trials, concurrency, move |trial| {
|
||||
let cache_size_mb_param = cache_size_mb_param.clone();
|
||||
let connection_pool_size_param = connection_pool_size_param.clone();
|
||||
let request_timeout_ms_param = request_timeout_ms_param.clone();
|
||||
let retry_count_param = retry_count_param.clone();
|
||||
let batch_size_param = batch_size_param.clone();
|
||||
let compression_level_param = compression_level_param.clone();
|
||||
let use_http2_param = use_http2_param.clone();
|
||||
let load_balancing_param = load_balancing_param.clone();
|
||||
async move {
|
||||
objective(
|
||||
trial,
|
||||
&cache_size_mb_param,
|
||||
&connection_pool_size_param,
|
||||
&request_timeout_ms_param,
|
||||
&retry_count_param,
|
||||
&batch_size_param,
|
||||
&compression_level_param,
|
||||
&use_http2_param,
|
||||
&load_balancing_param,
|
||||
)
|
||||
.await
|
||||
}
|
||||
})
|
||||
.await?;
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
// Step 5: Print results
|
||||
print_results(&study, elapsed, n_trials);
|
||||
print_best_config(
|
||||
&study,
|
||||
&cache_size_mb_p,
|
||||
&connection_pool_size_p,
|
||||
&request_timeout_ms_p,
|
||||
&retry_count_p,
|
||||
&batch_size_p,
|
||||
&compression_level_p,
|
||||
&use_http2_p,
|
||||
&load_balancing_p,
|
||||
)?;
|
||||
print_top_trials(&study, 5);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,124 +0,0 @@
|
||||
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,
|
||||
);
|
||||
}
|
||||
@@ -1,275 +0,0 @@
|
||||
//! Machine Learning Hyperparameter Tuning Example
|
||||
//!
|
||||
//! This example shows how to use the optimizer library to find the best
|
||||
//! hyperparameters for a machine learning model. We simulate a gradient
|
||||
//! boosting model (like XGBoost or LightGBM) and search for optimal settings.
|
||||
//!
|
||||
//! # Key Concepts Demonstrated
|
||||
//!
|
||||
//! - Creating a Study with a TPE (Tree-Parzen Estimator) sampler
|
||||
//! - Defining an objective function that the optimizer will minimize
|
||||
//! - Using different parameter types: floats, integers, log-scale, stepped
|
||||
//! - Using callbacks to monitor progress and implement early stopping
|
||||
//!
|
||||
//! # How It Works
|
||||
//!
|
||||
//! 1. Create a `Study` - this manages the optimization process
|
||||
//! 2. Define an objective function that takes a `Trial` and returns a score
|
||||
//! 3. Inside the objective, use `trial.suggest_*()` to sample parameters
|
||||
//! 4. The optimizer runs many trials, learning which parameter regions work best
|
||||
//! 5. After optimization, retrieve the best parameters found
|
||||
//!
|
||||
//! Run with: `cargo run --example ml_hyperparameter_tuning`
|
||||
|
||||
use std::ops::ControlFlow;
|
||||
|
||||
use optimizer::prelude::*;
|
||||
|
||||
// ============================================================================
|
||||
// Configuration: Hyperparameters we want to tune
|
||||
// ============================================================================
|
||||
|
||||
/// Holds all the hyperparameters for our model.
|
||||
///
|
||||
/// In a real application, you would pass these to your ML framework
|
||||
/// (e.g., XGBoost, LightGBM, scikit-learn).
|
||||
struct ModelConfig {
|
||||
learning_rate: f64,
|
||||
max_depth: i64,
|
||||
n_estimators: i64,
|
||||
subsample: f64,
|
||||
colsample_bytree: f64,
|
||||
min_child_weight: i64,
|
||||
reg_alpha: f64,
|
||||
reg_lambda: f64,
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Objective Function: What we want to optimize
|
||||
// ============================================================================
|
||||
|
||||
/// Simulates training a model and returns the validation loss.
|
||||
///
|
||||
/// In a real scenario, this function would:
|
||||
/// 1. Create a model with the given hyperparameters
|
||||
/// 2. Train it on your training data
|
||||
/// 3. Evaluate it on validation data
|
||||
/// 4. Return the validation metric (e.g., RMSE, log loss, accuracy)
|
||||
///
|
||||
/// The optimizer will try to MINIMIZE this value (we set Direction::Minimize).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn evaluate_model(config: &ModelConfig) -> f64 {
|
||||
// Simulated optimal hyperparameters:
|
||||
// learning_rate ~ 0.05, max_depth ~ 6, n_estimators ~ 200
|
||||
// subsample ~ 0.8, colsample_bytree ~ 0.8, min_child_weight ~ 3
|
||||
// reg_alpha ~ 0.1, reg_lambda ~ 1.0
|
||||
|
||||
let mut loss = 0.15; // Base loss
|
||||
|
||||
// Each term penalizes deviation from the optimal value
|
||||
loss += (config.learning_rate - 0.05).powi(2) * 100.0;
|
||||
loss += ((config.max_depth - 6) as f64).powi(2) * 0.01;
|
||||
loss += ((config.n_estimators - 200) as f64).powi(2) * 0.00001;
|
||||
loss += (config.subsample - 0.8).powi(2) * 10.0;
|
||||
loss += (config.colsample_bytree - 0.8).powi(2) * 10.0;
|
||||
loss += ((config.min_child_weight - 3) as f64).powi(2) * 0.05;
|
||||
loss += (config.reg_alpha - 0.1).powi(2) * 5.0;
|
||||
loss += (config.reg_lambda - 1.0).powi(2) * 2.0;
|
||||
|
||||
// Add some noise to simulate real-world variability
|
||||
let noise = (config.learning_rate * 1000.0).sin() * 0.01;
|
||||
|
||||
loss + noise
|
||||
}
|
||||
|
||||
/// The objective function that the optimizer calls for each trial.
|
||||
///
|
||||
/// This function:
|
||||
/// 1. Uses parameter definitions passed as arguments
|
||||
/// 2. Builds a model configuration from the suggested values
|
||||
/// 3. Evaluates the model and returns the loss
|
||||
///
|
||||
/// The optimizer learns from the results to suggest better parameters
|
||||
/// in future trials.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn objective(
|
||||
trial: &mut Trial,
|
||||
learning_rate_param: &FloatParam,
|
||||
max_depth_param: &IntParam,
|
||||
n_estimators_param: &IntParam,
|
||||
subsample_param: &FloatParam,
|
||||
colsample_bytree_param: &FloatParam,
|
||||
min_child_weight_param: &IntParam,
|
||||
reg_alpha_param: &FloatParam,
|
||||
reg_lambda_param: &FloatParam,
|
||||
) -> optimizer::Result<f64> {
|
||||
let learning_rate = learning_rate_param.suggest(trial)?;
|
||||
let max_depth = max_depth_param.suggest(trial)?;
|
||||
let n_estimators = n_estimators_param.suggest(trial)?;
|
||||
let subsample = subsample_param.suggest(trial)?;
|
||||
let colsample_bytree = colsample_bytree_param.suggest(trial)?;
|
||||
let min_child_weight = min_child_weight_param.suggest(trial)?;
|
||||
let reg_alpha = reg_alpha_param.suggest(trial)?;
|
||||
let reg_lambda = reg_lambda_param.suggest(trial)?;
|
||||
|
||||
// Build configuration and evaluate
|
||||
let config = ModelConfig {
|
||||
learning_rate,
|
||||
max_depth,
|
||||
n_estimators,
|
||||
subsample,
|
||||
colsample_bytree,
|
||||
min_child_weight,
|
||||
reg_alpha,
|
||||
reg_lambda,
|
||||
};
|
||||
|
||||
let loss = evaluate_model(&config);
|
||||
|
||||
Ok(loss)
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Callback Function: Monitor progress and implement early stopping
|
||||
// ============================================================================
|
||||
|
||||
/// Called after each successful trial completes.
|
||||
///
|
||||
/// Use callbacks to:
|
||||
/// - Log progress to console or file
|
||||
/// - Save checkpoints
|
||||
/// - Implement early stopping when a good solution is found
|
||||
/// - Track metrics over time
|
||||
///
|
||||
/// Return `ControlFlow::Continue(())` to keep optimizing.
|
||||
/// Return `ControlFlow::Break(())` to stop early.
|
||||
fn on_trial_complete(study: &Study<f64>, trial: &CompletedTrial<f64>) -> ControlFlow<()> {
|
||||
// Print trial number and objective value
|
||||
print!("{:>5} ", study.n_trials());
|
||||
for value in trial.params.values() {
|
||||
print!("{value:>12} ");
|
||||
}
|
||||
println!("{:>12.6}", trial.value);
|
||||
|
||||
// Early stopping: if we find an excellent solution, stop early
|
||||
if trial.value < 0.16 {
|
||||
println!("\nEarly stopping: found excellent solution!");
|
||||
return ControlFlow::Break(());
|
||||
}
|
||||
|
||||
ControlFlow::Continue(())
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Main: Set up and run the optimization
|
||||
// ============================================================================
|
||||
|
||||
fn main() -> optimizer::Result<()> {
|
||||
println!("=== ML Hyperparameter Tuning Example ===\n");
|
||||
|
||||
// Step 1: Create a sampler
|
||||
//
|
||||
// TPE (Tree-Parzen Estimator) is a Bayesian optimization algorithm.
|
||||
// It learns from previous trials to suggest better parameters.
|
||||
// - n_startup_trials: Number of random trials before TPE kicks in
|
||||
// - gamma: What fraction of trials are considered "good" (lower = more selective)
|
||||
// - seed: For reproducibility
|
||||
let sampler = TpeSampler::builder()
|
||||
.n_startup_trials(10)
|
||||
.gamma(0.25)
|
||||
.seed(42)
|
||||
.build()
|
||||
.expect("Failed to build TPE sampler");
|
||||
|
||||
// Step 2: Create a study
|
||||
//
|
||||
// The study manages the optimization process. We want to MINIMIZE
|
||||
// the loss (lower is better). Use Direction::Maximize for metrics
|
||||
// where higher is better (like accuracy).
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
// Print header
|
||||
println!("Starting hyperparameter optimization...\n");
|
||||
println!(
|
||||
"{:>5} {:>12} (parameters...) {:>12}",
|
||||
"Trial", "Params", "Loss"
|
||||
);
|
||||
println!("{}", "-".repeat(60));
|
||||
|
||||
// Step 3: Define parameter search spaces
|
||||
let learning_rate_param = FloatParam::new(0.001, 0.3)
|
||||
.name("learning_rate")
|
||||
.log_scale();
|
||||
let max_depth_param = IntParam::new(3, 12).name("max_depth");
|
||||
let n_estimators_param = IntParam::new(50, 500).name("n_estimators").step(50);
|
||||
let subsample_param = FloatParam::new(0.5, 1.0).name("subsample");
|
||||
let colsample_bytree_param = FloatParam::new(0.5, 1.0).name("colsample_bytree");
|
||||
let min_child_weight_param = IntParam::new(1, 10).name("min_child_weight");
|
||||
let reg_alpha_param = FloatParam::new(1e-3, 10.0).name("reg_alpha").log_scale();
|
||||
let reg_lambda_param = FloatParam::new(1e-3, 10.0).name("reg_lambda").log_scale();
|
||||
|
||||
// Step 4: Run optimization
|
||||
//
|
||||
// optimize_with_callback runs the objective function for up to
|
||||
// n_trials iterations. After each trial, it calls the callback.
|
||||
// The sampler gets access to trial history for informed sampling.
|
||||
let n_trials = 50;
|
||||
|
||||
study.optimize_with_callback(
|
||||
n_trials,
|
||||
|trial| {
|
||||
objective(
|
||||
trial,
|
||||
&learning_rate_param,
|
||||
&max_depth_param,
|
||||
&n_estimators_param,
|
||||
&subsample_param,
|
||||
&colsample_bytree_param,
|
||||
&min_child_weight_param,
|
||||
®_alpha_param,
|
||||
®_lambda_param,
|
||||
)
|
||||
},
|
||||
on_trial_complete,
|
||||
)?;
|
||||
|
||||
// Step 4: Get the best result
|
||||
println!("\n{}", "=".repeat(110));
|
||||
println!("\nOptimization completed!");
|
||||
println!("Total trials: {}", study.n_trials());
|
||||
|
||||
let best = study.best_trial()?;
|
||||
println!("\nBest trial:");
|
||||
println!(" Loss: {:.6}", best.value);
|
||||
println!(" Parameters:");
|
||||
println!(
|
||||
" learning_rate: {:.6}",
|
||||
best.get(&learning_rate_param).unwrap()
|
||||
);
|
||||
println!(" max_depth: {}", best.get(&max_depth_param).unwrap());
|
||||
println!(
|
||||
" n_estimators: {}",
|
||||
best.get(&n_estimators_param).unwrap()
|
||||
);
|
||||
println!(" subsample: {:.6}", best.get(&subsample_param).unwrap());
|
||||
println!(
|
||||
" colsample_bytree: {:.6}",
|
||||
best.get(&colsample_bytree_param).unwrap()
|
||||
);
|
||||
println!(
|
||||
" min_child_weight: {}",
|
||||
best.get(&min_child_weight_param).unwrap()
|
||||
);
|
||||
println!(" reg_alpha: {:.6}", best.get(®_alpha_param).unwrap());
|
||||
println!(
|
||||
" reg_lambda: {:.6}",
|
||||
best.get(®_lambda_param).unwrap()
|
||||
);
|
||||
|
||||
// Step 5: Use the best parameters (in a real app)
|
||||
//
|
||||
// Now you would take best.params and use them to train your final model
|
||||
// on the full dataset.
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,56 +0,0 @@
|
||||
use optimizer::prelude::*;
|
||||
use optimizer_derive::Categorical;
|
||||
|
||||
#[derive(Clone, Debug, Categorical)]
|
||||
enum Activation {
|
||||
Relu,
|
||||
Sigmoid,
|
||||
Tanh,
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
|
||||
// Define parameters outside the objective function
|
||||
let lr_param = FloatParam::new(1e-5, 1e-1).name("lr").log_scale();
|
||||
let n_layers_param = IntParam::new(1, 5).name("n_layers");
|
||||
let units_param = IntParam::new(32, 512).name("units").step(32);
|
||||
let optimizer_param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]).name("optimizer");
|
||||
let activation_param = EnumParam::<Activation>::new().name("activation");
|
||||
let batch_size_param = IntParam::new(16, 256).name("batch_size").log_scale();
|
||||
let use_dropout_param = BoolParam::new().name("use_dropout");
|
||||
|
||||
study
|
||||
.optimize(20, |trial| {
|
||||
let lr = lr_param.suggest(trial)?;
|
||||
let n_layers = n_layers_param.suggest(trial)?;
|
||||
let units = units_param.suggest(trial)?;
|
||||
let optimizer = optimizer_param.suggest(trial)?;
|
||||
let use_dropout = use_dropout_param.suggest(trial)?;
|
||||
let activation = activation_param.suggest(trial)?;
|
||||
let batch_size = batch_size_param.suggest(trial)?;
|
||||
|
||||
// Simulate a loss function
|
||||
let loss = lr * (n_layers as f64) + (units as f64) * 0.001
|
||||
- if use_dropout { 0.1 } else { 0.0 };
|
||||
|
||||
println!(
|
||||
"Trial {}: lr={lr:.6}, layers={n_layers}, units={units}, opt={optimizer}, \
|
||||
dropout={use_dropout}, activation={activation:?}, batch={batch_size} -> loss={loss:.4}",
|
||||
trial.id()
|
||||
);
|
||||
|
||||
Ok::<_, Error>(loss)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
println!("\nBest trial: value={:.4}", best.value);
|
||||
println!(" lr: {:.6}", best.get(&lr_param).unwrap());
|
||||
println!(" n_layers: {}", best.get(&n_layers_param).unwrap());
|
||||
println!(" units: {}", best.get(&units_param).unwrap());
|
||||
println!(" optimizer: {}", best.get(&optimizer_param).unwrap());
|
||||
println!(" activation: {:?}", best.get(&activation_param).unwrap());
|
||||
println!(" batch_size: {}", best.get(&batch_size_param).unwrap());
|
||||
println!(" use_dropout: {}", best.get(&use_dropout_param).unwrap());
|
||||
}
|
||||
@@ -1,15 +0,0 @@
|
||||
[package]
|
||||
name = "optimizer-derive"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.88"
|
||||
license = "MIT"
|
||||
description = "Derive macros for the optimizer crate"
|
||||
repository = "https://github.com/raimannma/rust-optimizer"
|
||||
|
||||
[lib]
|
||||
proc-macro = true
|
||||
|
||||
[dependencies]
|
||||
syn = { version = "2", features = ["full"] }
|
||||
quote = "1"
|
||||
@@ -1,71 +0,0 @@
|
||||
use proc_macro::TokenStream;
|
||||
use quote::quote;
|
||||
use syn::{Data, DeriveInput, Fields, parse_macro_input};
|
||||
|
||||
/// Derive macro for the `Categorical` trait on fieldless enums.
|
||||
///
|
||||
/// Generates an implementation of `optimizer::Categorical` that maps
|
||||
/// enum variants to/from sequential indices.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```ignore
|
||||
/// use optimizer::Categorical;
|
||||
///
|
||||
/// #[derive(Clone, Categorical)]
|
||||
/// enum Color {
|
||||
/// Red,
|
||||
/// Green,
|
||||
/// Blue,
|
||||
/// }
|
||||
/// ```
|
||||
#[proc_macro_derive(Categorical)]
|
||||
pub fn derive_categorical(input: TokenStream) -> TokenStream {
|
||||
let input = parse_macro_input!(input as DeriveInput);
|
||||
let name = &input.ident;
|
||||
|
||||
let Data::Enum(data_enum) = &input.data else {
|
||||
return syn::Error::new_spanned(&input, "Categorical can only be derived for enums")
|
||||
.to_compile_error()
|
||||
.into();
|
||||
};
|
||||
|
||||
// Validate all variants are fieldless
|
||||
for variant in &data_enum.variants {
|
||||
if !matches!(variant.fields, Fields::Unit) {
|
||||
return syn::Error::new_spanned(
|
||||
variant,
|
||||
"Categorical can only be derived for enums with unit variants (no fields)",
|
||||
)
|
||||
.to_compile_error()
|
||||
.into();
|
||||
}
|
||||
}
|
||||
|
||||
let n_choices = data_enum.variants.len();
|
||||
let variant_names: Vec<_> = data_enum.variants.iter().map(|v| &v.ident).collect();
|
||||
let indices: Vec<usize> = (0..n_choices).collect();
|
||||
|
||||
let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
|
||||
|
||||
let expanded = quote! {
|
||||
impl #impl_generics optimizer::Categorical for #name #ty_generics #where_clause {
|
||||
const N_CHOICES: usize = #n_choices;
|
||||
|
||||
fn from_index(index: usize) -> Self {
|
||||
match index {
|
||||
#(#indices => #name::#variant_names,)*
|
||||
_ => panic!("invalid index {} for {} with {} variants", index, stringify!(#name), #n_choices),
|
||||
}
|
||||
}
|
||||
|
||||
fn to_index(&self) -> usize {
|
||||
match self {
|
||||
#(#name::#variant_names => #indices,)*
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
expanded.into()
|
||||
}
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
/// 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,
|
||||
@@ -16,7 +15,6 @@ 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,
|
||||
@@ -30,7 +28,6 @@ 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,
|
||||
@@ -38,7 +35,6 @@ 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),
|
||||
|
||||
@@ -46,36 +46,6 @@ pub enum Error {
|
||||
#[error("KDE requires at least one sample")]
|
||||
EmptySamples,
|
||||
|
||||
/// Returned when multivariate KDE samples have zero dimensions.
|
||||
#[error("multivariate KDE samples must have at least one dimension")]
|
||||
ZeroDimensions,
|
||||
|
||||
/// Returned when multivariate KDE samples have inconsistent dimensions.
|
||||
#[error(
|
||||
"dimension mismatch: expected {expected} dimensions but sample {sample_index} has {got}"
|
||||
)]
|
||||
DimensionMismatch {
|
||||
/// The expected number of dimensions.
|
||||
expected: usize,
|
||||
/// The actual number of dimensions in the sample.
|
||||
got: usize,
|
||||
/// The index of the sample with mismatched dimensions.
|
||||
sample_index: usize,
|
||||
},
|
||||
|
||||
/// Returned when bandwidth vector length doesn't match the number of dimensions.
|
||||
#[error("bandwidth dimension mismatch: expected {expected} bandwidths but got {got}")]
|
||||
BandwidthDimensionMismatch {
|
||||
/// The expected number of bandwidths.
|
||||
expected: usize,
|
||||
/// The actual number of bandwidths provided.
|
||||
got: usize,
|
||||
},
|
||||
|
||||
/// Returned when a trial is pruned (stopped early by the objective function).
|
||||
#[error("trial was pruned")]
|
||||
TrialPruned,
|
||||
|
||||
/// Returned when an internal invariant is violated.
|
||||
#[error("internal error: {0}")]
|
||||
Internal(&'static str),
|
||||
@@ -87,33 +57,3 @@ pub enum Error {
|
||||
}
|
||||
|
||||
pub type Result<T> = core::result::Result<T, Error>;
|
||||
|
||||
/// Convenience type for signalling a pruned trial from an objective function.
|
||||
///
|
||||
/// Implements `Into<Error>` so it can be used with `?` in objectives that
|
||||
/// return `Result<V, Error>`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::{Error, TrialPruned};
|
||||
///
|
||||
/// fn objective_that_prunes() -> Result<f64, Error> {
|
||||
/// // ... some computation ...
|
||||
/// Err(TrialPruned)?
|
||||
/// }
|
||||
/// ```
|
||||
#[derive(Debug)]
|
||||
pub struct TrialPruned;
|
||||
|
||||
impl core::fmt::Display for TrialPruned {
|
||||
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
|
||||
write!(f, "trial was pruned")
|
||||
}
|
||||
}
|
||||
|
||||
impl From<TrialPruned> for Error {
|
||||
fn from(_: TrialPruned) -> Self {
|
||||
Error::TrialPruned
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,93 +0,0 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,7 @@
|
||||
//! This module provides a Gaussian kernel density estimator used by the TPE
|
||||
//! sampler to model probability distributions over good and bad trial regions.
|
||||
|
||||
use rand::{Rng, RngExt};
|
||||
use rand::Rng;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
@@ -1,13 +0,0 @@
|
||||
//! Kernel Density Estimation for parameter distributions.
|
||||
//!
|
||||
//! This module provides kernel density estimators used by TPE samplers
|
||||
//! to model probability distributions over good and bad trial regions.
|
||||
//!
|
||||
//! - [`univariate`] - Univariate (single-parameter) KDE
|
||||
//! - [`multivariate`] - Multivariate (joint-parameter) KDE for capturing parameter dependencies
|
||||
|
||||
mod multivariate;
|
||||
mod univariate;
|
||||
|
||||
pub(crate) use multivariate::MultivariateKDE;
|
||||
pub(crate) use univariate::KernelDensityEstimator;
|
||||
@@ -1,921 +0,0 @@
|
||||
//! Multivariate Kernel Density Estimation for joint parameter distributions.
|
||||
//!
|
||||
//! This module provides a multivariate kernel density estimator that captures
|
||||
//! dependencies between parameters. Unlike the univariate KDE which models each
|
||||
//! parameter independently, the multivariate KDE models the joint distribution
|
||||
//! to better capture correlations between parameters.
|
||||
|
||||
use rand::{Rng, RngExt};
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// A multivariate Gaussian kernel density estimator for joint distributions.
|
||||
///
|
||||
/// `MultivariateKDE` estimates a joint probability density function from a set
|
||||
/// of multi-dimensional samples. This is used in multivariate TPE to model
|
||||
/// the joint distributions l(x) and g(x) across multiple parameters simultaneously,
|
||||
/// capturing their dependencies.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```ignore
|
||||
/// use crate::kde::MultivariateKDE;
|
||||
///
|
||||
/// // Create samples: 3 samples with 2 dimensions each
|
||||
/// let samples = vec![
|
||||
/// vec![1.0, 2.0],
|
||||
/// vec![1.5, 2.5],
|
||||
/// vec![2.0, 3.0],
|
||||
/// ];
|
||||
/// let kde = MultivariateKDE::new(samples).unwrap();
|
||||
///
|
||||
/// // Get dimensionality
|
||||
/// assert_eq!(kde.n_dims(), 2);
|
||||
/// ```
|
||||
#[allow(dead_code)] // Fields and methods will be used in subsequent stories (US-003, US-004)
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct MultivariateKDE {
|
||||
/// The sample points used to construct the KDE.
|
||||
/// Each inner Vec is one sample with `n_dims` values.
|
||||
samples: Vec<Vec<f64>>,
|
||||
/// The bandwidth (standard deviation) for each dimension.
|
||||
/// Uses a diagonal bandwidth matrix (independent bandwidths per dimension).
|
||||
bandwidths: Vec<f64>,
|
||||
/// The number of dimensions.
|
||||
n_dims: usize,
|
||||
}
|
||||
|
||||
#[allow(dead_code)] // Methods will be used in subsequent stories (US-003, US-004)
|
||||
impl MultivariateKDE {
|
||||
/// Creates a new multivariate KDE with automatic bandwidth selection using Scott's rule.
|
||||
///
|
||||
/// Scott's rule for multivariate KDE sets bandwidth per dimension as:
|
||||
/// `h_j = n^(-1/(d+4)) * sigma_j`
|
||||
///
|
||||
/// where n is the number of samples, d is the dimensionality, and `sigma_j` is the
|
||||
/// standard deviation of the j-th dimension.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::EmptySamples` if `samples` is empty.
|
||||
/// Returns `Error::DimensionMismatch` if samples have inconsistent dimensions.
|
||||
/// Returns `Error::ZeroDimensions` if samples have zero dimensions.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
pub(crate) fn new(samples: Vec<Vec<f64>>) -> Result<Self> {
|
||||
if samples.is_empty() {
|
||||
return Err(Error::EmptySamples);
|
||||
}
|
||||
|
||||
let n_dims = samples[0].len();
|
||||
if n_dims == 0 {
|
||||
return Err(Error::ZeroDimensions);
|
||||
}
|
||||
|
||||
// Check all samples have the same dimensionality
|
||||
for (i, sample) in samples.iter().enumerate() {
|
||||
if sample.len() != n_dims {
|
||||
return Err(Error::DimensionMismatch {
|
||||
expected: n_dims,
|
||||
got: sample.len(),
|
||||
sample_index: i,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let bandwidths = Self::scotts_rule_multivariate(&samples, n_dims);
|
||||
|
||||
Ok(Self {
|
||||
samples,
|
||||
bandwidths,
|
||||
n_dims,
|
||||
})
|
||||
}
|
||||
|
||||
/// Creates a new multivariate KDE with specified bandwidths.
|
||||
///
|
||||
/// Use this when you want explicit control over the smoothing parameters.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::EmptySamples` if `samples` is empty.
|
||||
/// Returns `Error::DimensionMismatch` if samples have inconsistent dimensions.
|
||||
/// Returns `Error::ZeroDimensions` if samples have zero dimensions.
|
||||
/// Returns `Error::BandwidthDimensionMismatch` if bandwidths length doesn't match dimensions.
|
||||
/// Returns `Error::InvalidBandwidth` if any bandwidth is not positive.
|
||||
pub(crate) fn with_bandwidths(samples: Vec<Vec<f64>>, bandwidths: Vec<f64>) -> Result<Self> {
|
||||
if samples.is_empty() {
|
||||
return Err(Error::EmptySamples);
|
||||
}
|
||||
|
||||
let n_dims = samples[0].len();
|
||||
if n_dims == 0 {
|
||||
return Err(Error::ZeroDimensions);
|
||||
}
|
||||
|
||||
// Check all samples have the same dimensionality
|
||||
for (i, sample) in samples.iter().enumerate() {
|
||||
if sample.len() != n_dims {
|
||||
return Err(Error::DimensionMismatch {
|
||||
expected: n_dims,
|
||||
got: sample.len(),
|
||||
sample_index: i,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Check bandwidths length matches dimensions
|
||||
if bandwidths.len() != n_dims {
|
||||
return Err(Error::BandwidthDimensionMismatch {
|
||||
expected: n_dims,
|
||||
got: bandwidths.len(),
|
||||
});
|
||||
}
|
||||
|
||||
// Check all bandwidths are positive
|
||||
for &bw in &bandwidths {
|
||||
if bw <= 0.0 {
|
||||
return Err(Error::InvalidBandwidth(bw));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
samples,
|
||||
bandwidths,
|
||||
n_dims,
|
||||
})
|
||||
}
|
||||
|
||||
/// Computes bandwidths using Scott's rule for multivariate KDE.
|
||||
///
|
||||
/// Scott's rule for d dimensions: `h_j = n^(-1/(d+4)) * sigma_j`
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
fn scotts_rule_multivariate(samples: &[Vec<f64>], n_dims: usize) -> Vec<f64> {
|
||||
let n = samples.len() as f64;
|
||||
let d = n_dims as f64;
|
||||
|
||||
// Scott's rule exponent for multivariate: -1/(d+4)
|
||||
let exponent = -1.0 / (d + 4.0);
|
||||
let scale_factor = n.powf(exponent);
|
||||
|
||||
(0..n_dims)
|
||||
.map(|dim| {
|
||||
let std_dev = Self::dimension_std_dev(samples, dim);
|
||||
// For degenerate case where all samples are identical in this dimension,
|
||||
// use a small positive bandwidth
|
||||
if std_dev < f64::EPSILON {
|
||||
1.0
|
||||
} else {
|
||||
scale_factor * std_dev
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Computes the sample standard deviation for a single dimension.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
fn dimension_std_dev(samples: &[Vec<f64>], dim: usize) -> f64 {
|
||||
let n = samples.len() as f64;
|
||||
let values: Vec<f64> = samples.iter().map(|s| s[dim]).collect();
|
||||
let mean = values.iter().sum::<f64>() / n;
|
||||
let variance = values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / n;
|
||||
variance.sqrt()
|
||||
}
|
||||
|
||||
/// Returns the number of dimensions.
|
||||
pub(crate) fn n_dims(&self) -> usize {
|
||||
self.n_dims
|
||||
}
|
||||
|
||||
/// Returns the number of samples.
|
||||
#[allow(dead_code)] // Will be used in subsequent stories
|
||||
pub(crate) fn n_samples(&self) -> usize {
|
||||
self.samples.len()
|
||||
}
|
||||
|
||||
/// Returns the bandwidths for each dimension.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn bandwidths(&self) -> &[f64] {
|
||||
&self.bandwidths
|
||||
}
|
||||
|
||||
/// Returns a reference to the samples.
|
||||
#[allow(dead_code)] // Will be used in subsequent stories
|
||||
pub(crate) fn samples(&self) -> &[Vec<f64>] {
|
||||
&self.samples
|
||||
}
|
||||
|
||||
/// Returns the log probability density at point `x`.
|
||||
///
|
||||
/// This computes the log-density for numerical stability, using the formula:
|
||||
///
|
||||
/// `log f(x) = log((1/n) * Σ_i Π_j K_hj((x_j - x_ij) / h_j))`
|
||||
///
|
||||
/// The computation uses the log-sum-exp trick for numerical stability:
|
||||
///
|
||||
/// `log(Σ exp(a_i)) = max(a) + log(Σ exp(a_i - max(a)))`
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `x.len() != self.n_dims`.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
pub(crate) fn log_pdf(&self, x: &[f64]) -> f64 {
|
||||
assert_eq!(
|
||||
x.len(),
|
||||
self.n_dims,
|
||||
"Point dimension {} doesn't match KDE dimension {}",
|
||||
x.len(),
|
||||
self.n_dims
|
||||
);
|
||||
|
||||
let n = self.samples.len() as f64;
|
||||
|
||||
// Precompute log normalization constant for each dimension
|
||||
// For a Gaussian kernel: K_h(z) = (1/(h*sqrt(2*pi))) * exp(-0.5*z^2)
|
||||
// log(K_h(z)) = -log(h) - 0.5*log(2*pi) - 0.5*z^2
|
||||
let log_2pi = (2.0 * core::f64::consts::PI).ln();
|
||||
let log_norm_per_dim: Vec<f64> = self
|
||||
.bandwidths
|
||||
.iter()
|
||||
.map(|&h| -h.ln() - 0.5 * log_2pi)
|
||||
.collect();
|
||||
|
||||
// Compute log of kernel contribution for each sample
|
||||
// log(prod_j K_hj(z_j)) = sum_j log(K_hj(z_j))
|
||||
let log_kernels: Vec<f64> = self
|
||||
.samples
|
||||
.iter()
|
||||
.map(|sample| {
|
||||
let mut log_kernel_sum = 0.0;
|
||||
for j in 0..self.n_dims {
|
||||
let z = (x[j] - sample[j]) / self.bandwidths[j];
|
||||
log_kernel_sum += log_norm_per_dim[j] - 0.5 * z * z;
|
||||
}
|
||||
log_kernel_sum
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Use log-sum-exp trick for numerical stability
|
||||
// log(sum(exp(log_kernels))) = max + log(sum(exp(log_kernels - max)))
|
||||
let max_log_kernel = log_kernels
|
||||
.iter()
|
||||
.copied()
|
||||
.fold(f64::NEG_INFINITY, f64::max);
|
||||
|
||||
// Handle case where all log_kernels are -inf (extremely unlikely point)
|
||||
if max_log_kernel.is_infinite() && max_log_kernel < 0.0 {
|
||||
return f64::NEG_INFINITY;
|
||||
}
|
||||
|
||||
let sum_exp: f64 = log_kernels
|
||||
.iter()
|
||||
.map(|&lk| (lk - max_log_kernel).exp())
|
||||
.sum();
|
||||
|
||||
// log((1/n) * sum) = -log(n) + log(sum)
|
||||
-n.ln() + max_log_kernel + sum_exp.ln()
|
||||
}
|
||||
|
||||
/// Returns the probability density at point `x`.
|
||||
///
|
||||
/// The density is computed as the average of multivariate Gaussian kernels
|
||||
/// centered at each sample point:
|
||||
///
|
||||
/// `f(x) = (1/n) * Σ_i Π_j K_hj((x_j - x_ij) / h_j)`
|
||||
///
|
||||
/// where `K_hj` is the univariate Gaussian kernel with bandwidth `h_j`.
|
||||
///
|
||||
/// This method computes in log-space for numerical stability and then
|
||||
/// exponentiates the result.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `x.len() != self.n_dims`.
|
||||
pub(crate) fn pdf(&self, x: &[f64]) -> f64 {
|
||||
self.log_pdf(x).exp()
|
||||
}
|
||||
|
||||
/// Samples a point from the estimated joint density distribution.
|
||||
///
|
||||
/// Sampling works by:
|
||||
/// 1. Uniformly selecting one of the kernel centers (samples)
|
||||
/// 2. Adding independent Gaussian noise to each dimension with that
|
||||
/// dimension's bandwidth as the standard deviation
|
||||
///
|
||||
/// This is equivalent to sampling from a mixture of multivariate Gaussians
|
||||
/// with diagonal covariance matrices, where each mixture component is
|
||||
/// centered at a sample point.
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// A `Vec<f64>` of length `n_dims` representing a sample from the KDE.
|
||||
pub(crate) fn sample<R: Rng>(&self, rng: &mut R) -> Vec<f64> {
|
||||
// Select a random sample to center the kernel on
|
||||
let idx = rng.random_range(0..self.samples.len());
|
||||
let center = &self.samples[idx];
|
||||
|
||||
// Add independent Gaussian noise to each dimension
|
||||
// Using Box-Muller transform for Gaussian sampling
|
||||
center
|
||||
.iter()
|
||||
.zip(self.bandwidths.iter())
|
||||
.map(|(¢er_j, &bandwidth_j)| {
|
||||
let u1: f64 = rng.random();
|
||||
let u2: f64 = rng.random();
|
||||
|
||||
// Box-Muller transform: generates standard normal variate
|
||||
let z = (-2.0 * u1.ln()).sqrt() * (2.0 * core::f64::consts::PI * u2).cos();
|
||||
|
||||
// Scale by bandwidth and shift by center
|
||||
center_j + z * bandwidth_j
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_new_basic() {
|
||||
let samples = vec![vec![1.0, 2.0], vec![1.5, 2.5], vec![2.0, 3.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
assert_eq!(kde.n_dims(), 2);
|
||||
assert_eq!(kde.n_samples(), 3);
|
||||
assert_eq!(kde.bandwidths().len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_new_single_sample() {
|
||||
let samples = vec![vec![1.0, 2.0, 3.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
assert_eq!(kde.n_dims(), 3);
|
||||
assert_eq!(kde.n_samples(), 1);
|
||||
// With single sample, bandwidths should default to 1.0 (degenerate case)
|
||||
for &bw in kde.bandwidths() {
|
||||
assert!((bw - 1.0).abs() < f64::EPSILON);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_new_single_dimension() {
|
||||
let samples = vec![vec![1.0], vec![2.0], vec![3.0], vec![4.0], vec![5.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
assert_eq!(kde.n_dims(), 1);
|
||||
assert_eq!(kde.n_samples(), 5);
|
||||
assert_eq!(kde.bandwidths().len(), 1);
|
||||
// Bandwidth should be positive
|
||||
assert!(kde.bandwidths()[0] > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_scotts_rule() {
|
||||
// Create samples with known statistics
|
||||
// 10 samples, 2 dimensions
|
||||
let samples: Vec<Vec<f64>> = (0..10)
|
||||
.map(|i| {
|
||||
let x = f64::from(i);
|
||||
vec![x, x * 2.0] // Second dimension has 2x variance
|
||||
})
|
||||
.collect();
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// Scott's rule: h = n^(-1/(d+4)) * sigma
|
||||
// n=10, d=2: exponent = -1/6 ≈ -0.167
|
||||
// 10^(-1/6) ≈ 0.681
|
||||
// First dim std_dev ≈ 2.87, second ≈ 5.74
|
||||
// Expected bandwidths: ~1.95 and ~3.91
|
||||
|
||||
let bw = kde.bandwidths();
|
||||
assert!(
|
||||
bw[0] > 1.0 && bw[0] < 3.0,
|
||||
"First bandwidth {} unexpected",
|
||||
bw[0]
|
||||
);
|
||||
assert!(
|
||||
bw[1] > 2.0 && bw[1] < 6.0,
|
||||
"Second bandwidth {} unexpected",
|
||||
bw[1]
|
||||
);
|
||||
// Second bandwidth should be approximately 2x the first
|
||||
assert!(
|
||||
(bw[1] / bw[0] - 2.0).abs() < 0.1,
|
||||
"Ratio {} not close to 2",
|
||||
bw[1] / bw[0]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_with_bandwidths() {
|
||||
let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
|
||||
let bandwidths = vec![0.5, 1.0];
|
||||
let kde = MultivariateKDE::with_bandwidths(samples, bandwidths).unwrap();
|
||||
|
||||
assert_eq!(kde.n_dims(), 2);
|
||||
assert!((kde.bandwidths()[0] - 0.5).abs() < f64::EPSILON);
|
||||
assert!((kde.bandwidths()[1] - 1.0).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_empty_samples() {
|
||||
let samples: Vec<Vec<f64>> = vec![];
|
||||
let result = MultivariateKDE::new(samples);
|
||||
assert!(matches!(result, Err(Error::EmptySamples)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_zero_dimensions() {
|
||||
let samples = vec![vec![], vec![]];
|
||||
let result = MultivariateKDE::new(samples);
|
||||
assert!(matches!(result, Err(Error::ZeroDimensions)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_dimension_mismatch() {
|
||||
let samples = vec![vec![1.0, 2.0], vec![3.0]]; // Second sample has wrong dimensions
|
||||
let result = MultivariateKDE::new(samples);
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::DimensionMismatch {
|
||||
expected: 2,
|
||||
got: 1,
|
||||
sample_index: 1
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_with_bandwidths_wrong_length() {
|
||||
let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
|
||||
let bandwidths = vec![0.5]; // Only 1 bandwidth for 2 dimensions
|
||||
let result = MultivariateKDE::with_bandwidths(samples, bandwidths);
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::BandwidthDimensionMismatch {
|
||||
expected: 2,
|
||||
got: 1
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_with_bandwidths_zero() {
|
||||
let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
|
||||
let bandwidths = vec![0.5, 0.0]; // Second bandwidth is zero
|
||||
let result = MultivariateKDE::with_bandwidths(samples, bandwidths);
|
||||
assert!(matches!(result, Err(Error::InvalidBandwidth(bw)) if bw == 0.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_with_bandwidths_negative() {
|
||||
let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
|
||||
let bandwidths = vec![0.5, -1.0]; // Negative bandwidth
|
||||
let result = MultivariateKDE::with_bandwidths(samples, bandwidths);
|
||||
assert!(
|
||||
matches!(result, Err(Error::InvalidBandwidth(bw)) if (bw - (-1.0)).abs() < f64::EPSILON)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_identical_samples() {
|
||||
// All samples identical - should handle degenerate case
|
||||
let samples = vec![vec![5.0, 10.0], vec![5.0, 10.0], vec![5.0, 10.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// Bandwidths should default to 1.0 for degenerate dimensions
|
||||
for &bw in kde.bandwidths() {
|
||||
assert!((bw - 1.0).abs() < f64::EPSILON);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_high_dimensional() {
|
||||
// Test with higher dimensions
|
||||
let samples: Vec<Vec<f64>> = (0..20)
|
||||
.map(|i| {
|
||||
let x = f64::from(i);
|
||||
vec![x, x * 0.5, x * 2.0, x * 0.1, x * 10.0]
|
||||
})
|
||||
.collect();
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
assert_eq!(kde.n_dims(), 5);
|
||||
assert_eq!(kde.n_samples(), 20);
|
||||
assert_eq!(kde.bandwidths().len(), 5);
|
||||
|
||||
// All bandwidths should be positive
|
||||
for &bw in kde.bandwidths() {
|
||||
assert!(bw > 0.0);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_samples_accessor() {
|
||||
let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
|
||||
let kde = MultivariateKDE::new(samples.clone()).unwrap();
|
||||
|
||||
assert_eq!(kde.samples(), &samples);
|
||||
}
|
||||
|
||||
// ==================== PDF Tests ====================
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_pdf_basic() {
|
||||
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0], vec![2.0, 2.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// Density should be positive everywhere
|
||||
assert!(kde.pdf(&[0.0, 0.0]) > 0.0);
|
||||
assert!(kde.pdf(&[1.0, 1.0]) > 0.0);
|
||||
assert!(kde.pdf(&[2.0, 2.0]) > 0.0);
|
||||
assert!(kde.pdf(&[0.5, 0.5]) > 0.0);
|
||||
|
||||
// Density should be higher near sample points
|
||||
let near_density = kde.pdf(&[1.0, 1.0]);
|
||||
let far_density = kde.pdf(&[10.0, 10.0]);
|
||||
assert!(near_density > far_density);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_pdf_with_custom_bandwidths() {
|
||||
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0]];
|
||||
let bandwidths = vec![0.5, 0.5];
|
||||
let kde = MultivariateKDE::with_bandwidths(samples, bandwidths).unwrap();
|
||||
|
||||
// Density should be positive
|
||||
assert!(kde.pdf(&[0.5, 0.5]) > 0.0);
|
||||
assert!(kde.pdf(&[0.0, 0.0]) > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_pdf_single_sample() {
|
||||
let samples = vec![vec![5.0, 10.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// Should have positive density near the sample
|
||||
assert!(kde.pdf(&[5.0, 10.0]) > 0.0);
|
||||
assert!(kde.pdf(&[4.5, 9.5]) > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_log_pdf_consistency() {
|
||||
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0], vec![2.0, 2.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// log_pdf and pdf should be consistent: exp(log_pdf(x)) == pdf(x)
|
||||
let test_points = vec![
|
||||
vec![0.0, 0.0],
|
||||
vec![1.0, 1.0],
|
||||
vec![0.5, 0.5],
|
||||
vec![3.0, 3.0],
|
||||
];
|
||||
|
||||
for point in test_points {
|
||||
let log_p = kde.log_pdf(&point);
|
||||
let p = kde.pdf(&point);
|
||||
let p_from_log = log_p.exp();
|
||||
assert!(
|
||||
(p - p_from_log).abs() < 1e-10,
|
||||
"pdf={p}, exp(log_pdf)={p_from_log}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_pdf_integrates_to_one_1d() {
|
||||
// Test with 1D case first (should match univariate KDE behavior)
|
||||
let samples = vec![vec![0.0], vec![1.0], vec![2.0], vec![3.0], vec![4.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// Numerical integration over a wide range
|
||||
let n_points = 1000;
|
||||
let low = -10.0;
|
||||
let high = 15.0;
|
||||
let dx = (high - low) / f64::from(n_points);
|
||||
|
||||
let integral: f64 = (0..n_points)
|
||||
.map(|i| {
|
||||
let x = low + (f64::from(i) + 0.5) * dx;
|
||||
kde.pdf(&[x]) * dx
|
||||
})
|
||||
.sum();
|
||||
|
||||
// Should be approximately 1.0 (within numerical error)
|
||||
assert!(
|
||||
(integral - 1.0).abs() < 0.02,
|
||||
"1D integral = {integral}, expected ~1.0"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_pdf_integrates_to_one_2d() {
|
||||
// Test with 2D case
|
||||
let samples = vec![
|
||||
vec![0.0, 0.0],
|
||||
vec![1.0, 0.0],
|
||||
vec![0.0, 1.0],
|
||||
vec![1.0, 1.0],
|
||||
vec![0.5, 0.5],
|
||||
];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// Numerical integration over a 2D grid
|
||||
let n_points = 100; // 100x100 = 10000 points
|
||||
let low = -5.0;
|
||||
let high = 6.0;
|
||||
let dx = (high - low) / f64::from(n_points);
|
||||
|
||||
let mut integral = 0.0;
|
||||
for i in 0..n_points {
|
||||
for j in 0..n_points {
|
||||
let x = low + (f64::from(i) + 0.5) * dx;
|
||||
let y = low + (f64::from(j) + 0.5) * dx;
|
||||
integral += kde.pdf(&[x, y]) * dx * dx;
|
||||
}
|
||||
}
|
||||
|
||||
// Should be approximately 1.0 (within numerical error)
|
||||
// 2D integration has more error, so use larger tolerance
|
||||
assert!(
|
||||
(integral - 1.0).abs() < 0.05,
|
||||
"2D integral = {integral}, expected ~1.0"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_pdf_symmetry() {
|
||||
// With symmetric samples, PDF should be symmetric
|
||||
let samples = vec![
|
||||
vec![1.0, 0.0],
|
||||
vec![-1.0, 0.0],
|
||||
vec![0.0, 1.0],
|
||||
vec![0.0, -1.0],
|
||||
];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// Density at symmetric points should be equal
|
||||
let d1 = kde.pdf(&[0.5, 0.0]);
|
||||
let d2 = kde.pdf(&[-0.5, 0.0]);
|
||||
assert!(
|
||||
(d1 - d2).abs() < 1e-10,
|
||||
"Symmetric points have different densities: {d1} vs {d2}"
|
||||
);
|
||||
|
||||
let d3 = kde.pdf(&[0.0, 0.5]);
|
||||
let d4 = kde.pdf(&[0.0, -0.5]);
|
||||
assert!(
|
||||
(d3 - d4).abs() < 1e-10,
|
||||
"Symmetric points have different densities: {d3} vs {d4}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_pdf_high_dimensional() {
|
||||
// Test with higher dimensions
|
||||
let samples: Vec<Vec<f64>> = (0..10)
|
||||
.map(|i| {
|
||||
let x = f64::from(i) * 0.1;
|
||||
vec![x, x, x, x, x] // 5D
|
||||
})
|
||||
.collect();
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// Density should be positive
|
||||
assert!(kde.pdf(&[0.5, 0.5, 0.5, 0.5, 0.5]) > 0.0);
|
||||
assert!(kde.pdf(&[0.0, 0.0, 0.0, 0.0, 0.0]) > 0.0);
|
||||
|
||||
// Log PDF should be finite
|
||||
let log_p = kde.log_pdf(&[0.5, 0.5, 0.5, 0.5, 0.5]);
|
||||
assert!(log_p.is_finite());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_pdf_numerical_stability() {
|
||||
// Test numerical stability with points far from samples
|
||||
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// Very far point should have very small but non-negative density
|
||||
let far_pdf = kde.pdf(&[100.0, 100.0]);
|
||||
assert!(far_pdf >= 0.0);
|
||||
assert!(far_pdf.is_finite() || far_pdf == 0.0);
|
||||
|
||||
// Log PDF should be finite (or -inf for zero density)
|
||||
let far_log_pdf = kde.log_pdf(&[100.0, 100.0]);
|
||||
assert!(far_log_pdf.is_finite() || far_log_pdf.is_infinite());
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "Point dimension")]
|
||||
fn test_multivariate_kde_pdf_wrong_dimension() {
|
||||
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
|
||||
// Should panic with wrong dimension
|
||||
let _ = kde.pdf(&[0.0]); // Only 1 value for 2D KDE
|
||||
}
|
||||
|
||||
// ==================== Sampling Tests ====================
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_sample_basic() {
|
||||
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0], vec![2.0, 2.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
let mut rng = rand::rng();
|
||||
|
||||
// Sample should have correct dimensionality
|
||||
let sample = kde.sample(&mut rng);
|
||||
assert_eq!(sample.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_sample_in_reasonable_range() {
|
||||
let samples = vec![
|
||||
vec![0.0, 0.0],
|
||||
vec![1.0, 1.0],
|
||||
vec![2.0, 2.0],
|
||||
vec![3.0, 3.0],
|
||||
vec![4.0, 4.0],
|
||||
];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
let mut rng = rand::rng();
|
||||
|
||||
// Samples should generally be in a reasonable range around the data
|
||||
for _ in 0..100 {
|
||||
let s = kde.sample(&mut rng);
|
||||
// With high probability, samples should be within a few bandwidths
|
||||
// of the data range. Use a generous range to avoid flaky tests.
|
||||
assert!(
|
||||
s[0] > -10.0 && s[0] < 15.0,
|
||||
"Sample dimension 0: {} outside expected range",
|
||||
s[0]
|
||||
);
|
||||
assert!(
|
||||
s[1] > -10.0 && s[1] < 15.0,
|
||||
"Sample dimension 1: {} outside expected range",
|
||||
s[1]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_sample_single_sample() {
|
||||
// When KDE has only one sample, all samples should be centered around it
|
||||
let samples = vec![vec![5.0, 10.0]];
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
let mut rng = rand::rng();
|
||||
|
||||
// Generate many samples and check they cluster around (5.0, 10.0)
|
||||
let n_samples = 100;
|
||||
let mut sum_x = 0.0;
|
||||
let mut sum_y = 0.0;
|
||||
for _ in 0..n_samples {
|
||||
let s = kde.sample(&mut rng);
|
||||
sum_x += s[0];
|
||||
sum_y += s[1];
|
||||
}
|
||||
|
||||
let mean_x = sum_x / f64::from(n_samples);
|
||||
let mean_y = sum_y / f64::from(n_samples);
|
||||
|
||||
// Mean should be close to the single sample point
|
||||
// With 100 samples and bandwidth=1.0, we expect mean within ~0.3 of center
|
||||
assert!(
|
||||
(mean_x - 5.0).abs() < 1.0,
|
||||
"Mean x={mean_x}, expected close to 5.0"
|
||||
);
|
||||
assert!(
|
||||
(mean_y - 10.0).abs() < 1.0,
|
||||
"Mean y={mean_y}, expected close to 10.0"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_sample_high_dimensional() {
|
||||
// Test sampling in higher dimensions
|
||||
let samples: Vec<Vec<f64>> = (0..10)
|
||||
.map(|i| {
|
||||
let x = f64::from(i) * 0.5;
|
||||
vec![x, x * 2.0, x * 0.5, x + 1.0, x - 1.0] // 5D
|
||||
})
|
||||
.collect();
|
||||
let kde = MultivariateKDE::new(samples).unwrap();
|
||||
let mut rng = rand::rng();
|
||||
|
||||
// Sample should have correct dimensionality
|
||||
for _ in 0..50 {
|
||||
let sample = kde.sample(&mut rng);
|
||||
assert_eq!(sample.len(), 5);
|
||||
// All values should be finite
|
||||
for &val in &sample {
|
||||
assert!(val.is_finite(), "Sample value is not finite: {val}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
fn test_multivariate_kde_sample_respects_bandwidth() {
|
||||
// Create samples all at origin with large custom bandwidth
|
||||
let data = vec![vec![0.0, 0.0], vec![0.0, 0.0], vec![0.0, 0.0]];
|
||||
let bandwidths = vec![0.1, 10.0]; // Small bandwidth in x, large in y
|
||||
let kde = MultivariateKDE::with_bandwidths(data, bandwidths).unwrap();
|
||||
let mut rng = rand::rng();
|
||||
|
||||
// Generate samples and check variance in each dimension
|
||||
let n_samples = 1000;
|
||||
let mut values_x: Vec<f64> = Vec::with_capacity(n_samples);
|
||||
let mut values_y: Vec<f64> = Vec::with_capacity(n_samples);
|
||||
|
||||
for _ in 0..n_samples {
|
||||
let s = kde.sample(&mut rng);
|
||||
values_x.push(s[0]);
|
||||
values_y.push(s[1]);
|
||||
}
|
||||
|
||||
// Compute sample variances
|
||||
let n = n_samples as f64;
|
||||
let mean_x: f64 = values_x.iter().sum::<f64>() / n;
|
||||
let mean_y: f64 = values_y.iter().sum::<f64>() / n;
|
||||
|
||||
let var_x: f64 = values_x.iter().map(|x| (x - mean_x).powi(2)).sum::<f64>() / n;
|
||||
let var_y: f64 = values_y.iter().map(|y| (y - mean_y).powi(2)).sum::<f64>() / n;
|
||||
|
||||
// Variance should be approximately bandwidth^2
|
||||
// x: bandwidth=0.1, expected variance ~0.01
|
||||
// y: bandwidth=10.0, expected variance ~100.0
|
||||
assert!(
|
||||
var_x < 0.05,
|
||||
"X variance {var_x} too large for bandwidth 0.1"
|
||||
);
|
||||
assert!(
|
||||
var_y > 50.0 && var_y < 200.0,
|
||||
"Y variance {var_y} unexpected for bandwidth 10.0"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_sample_distribution_shape() {
|
||||
// Create samples along a diagonal line
|
||||
let data = vec![
|
||||
vec![0.0, 0.0],
|
||||
vec![1.0, 1.0],
|
||||
vec![2.0, 2.0],
|
||||
vec![3.0, 3.0],
|
||||
vec![4.0, 4.0],
|
||||
];
|
||||
let kde = MultivariateKDE::new(data).unwrap();
|
||||
let mut rng = rand::rng();
|
||||
|
||||
// Sample many points and verify the mean is near the center
|
||||
let n_samples = 500;
|
||||
let mut sum = [0.0, 0.0];
|
||||
|
||||
for _ in 0..n_samples {
|
||||
let s = kde.sample(&mut rng);
|
||||
sum[0] += s[0];
|
||||
sum[1] += s[1];
|
||||
}
|
||||
|
||||
let mean_x = sum[0] / f64::from(n_samples);
|
||||
let mean_y = sum[1] / f64::from(n_samples);
|
||||
|
||||
// Mean should be close to (2.0, 2.0) - the center of the samples
|
||||
assert!(
|
||||
(mean_x - 2.0).abs() < 0.5,
|
||||
"Mean x={mean_x}, expected close to 2.0"
|
||||
);
|
||||
assert!(
|
||||
(mean_y - 2.0).abs() < 0.5,
|
||||
"Mean y={mean_y}, expected close to 2.0"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_kde_sample_deterministic_with_seeded_rng() {
|
||||
use rand::SeedableRng;
|
||||
|
||||
let data = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
|
||||
let kde = MultivariateKDE::new(data).unwrap();
|
||||
|
||||
// Use a seeded RNG for reproducibility
|
||||
let mut rng1 = rand::rngs::StdRng::seed_from_u64(42);
|
||||
let mut rng2 = rand::rngs::StdRng::seed_from_u64(42);
|
||||
|
||||
// Same seed should produce same samples
|
||||
let result1 = kde.sample(&mut rng1);
|
||||
let result2 = kde.sample(&mut rng2);
|
||||
|
||||
assert!(
|
||||
(result1[0] - result2[0]).abs() < f64::EPSILON,
|
||||
"Samples with same seed differ in dimension 0"
|
||||
);
|
||||
assert!(
|
||||
(result1[1] - result2[1]).abs() < f64::EPSILON,
|
||||
"Samples with same seed differ in dimension 1"
|
||||
);
|
||||
}
|
||||
}
|
||||
+28
-134
@@ -17,9 +17,6 @@
|
||||
//! - **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:
|
||||
//!
|
||||
@@ -31,26 +28,24 @@
|
||||
//! # Quick Start
|
||||
//!
|
||||
//! ```
|
||||
//! use optimizer::prelude::*;
|
||||
//! use optimizer::sampler::tpe::TpeSampler;
|
||||
//! use optimizer::{Direction, Study};
|
||||
//!
|
||||
//! // Create a study with TPE sampler
|
||||
//! let sampler = TpeSampler::builder().seed(42).build().unwrap();
|
||||
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
//!
|
||||
//! // Define parameter search space
|
||||
//! let x = FloatParam::new(-10.0, 10.0).name("x");
|
||||
//!
|
||||
//! // Optimize x^2 for 20 trials
|
||||
//! study
|
||||
//! .optimize(20, |trial| {
|
||||
//! let x_val = x.suggest(trial)?;
|
||||
//! Ok::<_, Error>(x_val * x_val)
|
||||
//! .optimize_with_sampler(20, |trial| {
|
||||
//! let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
//! Ok::<_, optimizer::Error>(x * x)
|
||||
//! })
|
||||
//! .unwrap();
|
||||
//!
|
||||
//! // Get the best result
|
||||
//! let best = study.best_trial().unwrap();
|
||||
//! println!("x = {}", best.get(&x).unwrap());
|
||||
//! println!("Best value: {} at x={:?}", best.value, best.params);
|
||||
//! ```
|
||||
//!
|
||||
//! # Creating a Study
|
||||
@@ -74,35 +69,29 @@
|
||||
//!
|
||||
//! # Suggesting Parameters
|
||||
//!
|
||||
//! Within the objective function, use parameter types to suggest values:
|
||||
//! Within the objective function, use [`Trial`] to suggest parameter values:
|
||||
//!
|
||||
//! ```
|
||||
//! use optimizer::parameter::{BoolParam, CategoricalParam, FloatParam, IntParam, Parameter};
|
||||
//! use optimizer::{Direction, Study};
|
||||
//!
|
||||
//! let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
//!
|
||||
//! // Define parameter search spaces
|
||||
//! let x_param = FloatParam::new(0.0, 1.0);
|
||||
//! let lr_param = FloatParam::new(1e-5, 1e-1).log_scale();
|
||||
//! let step_param = FloatParam::new(0.0, 1.0).step(0.1);
|
||||
//! let n_param = IntParam::new(1, 10);
|
||||
//! let batch_param = IntParam::new(16, 256).log_scale();
|
||||
//! let units_param = IntParam::new(32, 512).step(32);
|
||||
//! let flag_param = BoolParam::new();
|
||||
//! let optimizer_param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]);
|
||||
//!
|
||||
//! study
|
||||
//! .optimize(10, |trial| {
|
||||
//! let x = x_param.suggest(trial)?;
|
||||
//! let lr = lr_param.suggest(trial)?;
|
||||
//! let step = step_param.suggest(trial)?;
|
||||
//! let n = n_param.suggest(trial)?;
|
||||
//! let batch = batch_param.suggest(trial)?;
|
||||
//! let units = units_param.suggest(trial)?;
|
||||
//! let flag = flag_param.suggest(trial)?;
|
||||
//! let optimizer = optimizer_param.suggest(trial)?;
|
||||
//! // Float parameters
|
||||
//! let x = trial.suggest_float("x", 0.0, 1.0)?;
|
||||
//! let lr = trial.suggest_float_log("learning_rate", 1e-5, 1e-1)?;
|
||||
//! let step = trial.suggest_float_step("step", 0.0, 1.0, 0.1)?;
|
||||
//!
|
||||
//! // Integer parameters
|
||||
//! let n = trial.suggest_int("n_layers", 1, 10)?;
|
||||
//! let batch = trial.suggest_int_log("batch_size", 16, 256)?;
|
||||
//! let units = trial.suggest_int_step("units", 32, 512, 32)?;
|
||||
//!
|
||||
//! // Categorical parameters
|
||||
//! let optimizer = trial.suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])?;
|
||||
//!
|
||||
//! // Return objective value
|
||||
//! Ok::<_, optimizer::Error>(x * n as f64)
|
||||
//! })
|
||||
//! .unwrap();
|
||||
@@ -158,130 +147,35 @@
|
||||
//!
|
||||
//! ```ignore
|
||||
//! use optimizer::{Study, Direction};
|
||||
//! use optimizer::parameter::{FloatParam, Parameter};
|
||||
//!
|
||||
//! let x_param = FloatParam::new(0.0, 1.0);
|
||||
//!
|
||||
//! // Sequential async
|
||||
//! study.optimize_async(10, |mut trial| {
|
||||
//! let x_param = x_param.clone();
|
||||
//! async move {
|
||||
//! let x = x_param.suggest(&mut trial)?;
|
||||
//! Ok((trial, x * x))
|
||||
//! }
|
||||
//! study.optimize_async(10, |mut trial| async move {
|
||||
//! let x = trial.suggest_float("x", 0.0, 1.0)?;
|
||||
//! Ok((trial, x * x))
|
||||
//! }).await?;
|
||||
//!
|
||||
//! // Parallel with bounded concurrency
|
||||
//! study.optimize_parallel(10, 4, |mut trial| {
|
||||
//! let x_param = x_param.clone();
|
||||
//! async move {
|
||||
//! let x = x_param.suggest(&mut trial)?;
|
||||
//! Ok((trial, x * x))
|
||||
//! }
|
||||
//! study.optimize_parallel(10, 4, |mut trial| async move {
|
||||
//! let x = trial.suggest_float("x", 0.0, 1.0)?;
|
||||
//! Ok((trial, x * x))
|
||||
//! }).await?;
|
||||
//! ```
|
||||
//!
|
||||
//! # Feature Flags
|
||||
//!
|
||||
//! - `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;
|
||||
pub mod pruner;
|
||||
pub mod sampler;
|
||||
mod study;
|
||||
mod trial;
|
||||
mod types;
|
||||
|
||||
pub use error::{Error, Result, TrialPruned};
|
||||
#[cfg(feature = "derive")]
|
||||
pub use optimizer_derive::Categorical;
|
||||
pub use error::{Error, Result};
|
||||
pub use param::ParamValue;
|
||||
pub use parameter::{
|
||||
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, ParamId, Parameter,
|
||||
};
|
||||
pub use pruner::{
|
||||
HyperbandPruner, MedianPruner, NopPruner, PatientPruner, PercentilePruner, 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 trial::{SuggestableRange, Trial};
|
||||
pub use types::{Direction, TrialState};
|
||||
|
||||
/// Convenient wildcard import for the most common types.
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::prelude::*;
|
||||
/// ```
|
||||
pub mod prelude {
|
||||
#[cfg(feature = "derive")]
|
||||
pub use optimizer_derive::Categorical as DeriveCategory;
|
||||
|
||||
pub use crate::error::{Error, Result, TrialPruned};
|
||||
pub use crate::param::ParamValue;
|
||||
pub use crate::parameter::{
|
||||
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter,
|
||||
};
|
||||
pub use crate::pruner::{
|
||||
HyperbandPruner, MedianPruner, NopPruner, PatientPruner, PercentilePruner, Pruner,
|
||||
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,7 +6,6 @@
|
||||
/// 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),
|
||||
@@ -15,13 +14,3 @@ pub enum ParamValue {
|
||||
/// A categorical parameter value, stored as an index into the choices array.
|
||||
Categorical(usize),
|
||||
}
|
||||
|
||||
impl core::fmt::Display for ParamValue {
|
||||
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
|
||||
match self {
|
||||
Self::Float(v) => write!(f, "{v}"),
|
||||
Self::Int(v) => write!(f, "{v}"),
|
||||
Self::Categorical(v) => write!(f, "category({v})"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
-1039
File diff suppressed because it is too large
Load Diff
@@ -1,559 +0,0 @@
|
||||
use core::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
use crate::types::{Direction, TrialState};
|
||||
|
||||
/// Hyperband pruner that manages multiple Successive Halving brackets.
|
||||
///
|
||||
/// Hyperband addresses SHA's sensitivity to the `min_resource` choice by
|
||||
/// running multiple brackets, each with a different tradeoff between the
|
||||
/// number of configurations and the starting budget:
|
||||
///
|
||||
/// - Bracket 0: many trials, very small starting budget (aggressive pruning)
|
||||
/// - Bracket 1: fewer trials, larger starting budget (moderate pruning)
|
||||
/// - ...
|
||||
/// - Bracket `s_max`: few trials, full budget (no pruning)
|
||||
///
|
||||
/// Trials are assigned to brackets in round-robin fashion. Each bracket
|
||||
/// runs SHA with its own `min_resource` and rung schedule.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Direction;
|
||||
/// use optimizer::pruner::HyperbandPruner;
|
||||
///
|
||||
/// let pruner = HyperbandPruner::new()
|
||||
/// .min_resource(1)
|
||||
/// .max_resource(81)
|
||||
/// .reduction_factor(3)
|
||||
/// .direction(Direction::Minimize);
|
||||
/// ```
|
||||
pub struct HyperbandPruner {
|
||||
min_resource: u64,
|
||||
max_resource: u64,
|
||||
reduction_factor: u64,
|
||||
direction: Direction,
|
||||
/// Tracks which bracket each trial belongs to.
|
||||
trial_brackets: Mutex<HashMap<u64, usize>>,
|
||||
/// Counter for round-robin bracket assignment.
|
||||
next_bracket: AtomicU64,
|
||||
}
|
||||
|
||||
impl HyperbandPruner {
|
||||
/// Create a new `HyperbandPruner` with default parameters.
|
||||
///
|
||||
/// Defaults: `min_resource=1`, `max_resource=81`, `reduction_factor=3`,
|
||||
/// `direction=Minimize`.
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
min_resource: 1,
|
||||
max_resource: 81,
|
||||
reduction_factor: 3,
|
||||
direction: Direction::Minimize,
|
||||
trial_brackets: Mutex::new(HashMap::new()),
|
||||
next_bracket: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set 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
|
||||
}
|
||||
|
||||
/// Set 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
|
||||
}
|
||||
|
||||
/// Set the reduction factor (eta). At each rung, the top 1/eta trials survive.
|
||||
///
|
||||
/// # 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
|
||||
}
|
||||
|
||||
/// Set the optimization direction.
|
||||
#[must_use]
|
||||
pub fn direction(mut self, d: Direction) -> Self {
|
||||
self.direction = d;
|
||||
self
|
||||
}
|
||||
|
||||
/// Compute `s_max = floor(log(max_resource / min_resource) / log(eta))`.
|
||||
#[allow(
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
fn s_max(&self) -> u64 {
|
||||
let eta = self.reduction_factor as f64;
|
||||
let ratio = self.max_resource as f64 / self.min_resource as f64;
|
||||
(ratio.ln() / eta.ln()).floor() as u64
|
||||
}
|
||||
|
||||
/// Compute the rung steps for a given bracket `s`.
|
||||
///
|
||||
/// For bracket `s`, the starting resource is `max_resource / eta^(s_max - s)`,
|
||||
/// and rungs are spaced at powers of eta from there up to `max_resource`.
|
||||
#[allow(
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
fn rung_steps_for_bracket(&self, bracket: usize) -> Vec<u64> {
|
||||
let s_max = self.s_max();
|
||||
let eta = self.reduction_factor as f64;
|
||||
|
||||
// Starting resource for this bracket
|
||||
let exponent = s_max.saturating_sub(bracket as u64);
|
||||
let min_resource_bracket =
|
||||
(self.max_resource as f64 / eta.powi(exponent as i32)).ceil() as u64;
|
||||
|
||||
let mut steps = Vec::new();
|
||||
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;
|
||||
}
|
||||
steps.push(step);
|
||||
rung += 1;
|
||||
}
|
||||
steps
|
||||
}
|
||||
|
||||
/// Assign a trial to a bracket (round-robin) and return the bracket index.
|
||||
#[allow(clippy::cast_possible_truncation)]
|
||||
fn assign_bracket(&self, trial_id: u64) -> usize {
|
||||
let n_brackets = (self.s_max() + 1) as usize;
|
||||
let mut map = self.trial_brackets.lock().expect("lock poisoned");
|
||||
*map.entry(trial_id).or_insert_with(|| {
|
||||
let idx = self.next_bracket.fetch_add(1, Ordering::Relaxed);
|
||||
(idx as usize) % n_brackets
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for HyperbandPruner {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
impl Pruner for HyperbandPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let bracket = self.assign_bracket(trial_id);
|
||||
let rungs = self.rung_steps_for_bracket(bracket);
|
||||
|
||||
// Find the highest rung step <= current step
|
||||
let Some(&rung_step) = rungs.iter().rev().find(|&&r| r <= step) else {
|
||||
return false;
|
||||
};
|
||||
|
||||
// Never prune at the last rung (full budget)
|
||||
if rung_step >= self.max_resource {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Get the current trial's value at this rung step
|
||||
let current_value =
|
||||
if let Some(&(_, v)) = intermediate_values.iter().find(|(s, _)| *s == rung_step) {
|
||||
v
|
||||
} else if let Some(&(_, v)) = intermediate_values
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|(s, _)| *s <= rung_step)
|
||||
{
|
||||
v
|
||||
} else {
|
||||
return false;
|
||||
};
|
||||
|
||||
self.is_pruned_at_rung(current_value, rung_step, bracket, completed_trials)
|
||||
}
|
||||
}
|
||||
|
||||
impl HyperbandPruner {
|
||||
/// Determine whether a trial should be pruned at the given rung within its bracket.
|
||||
///
|
||||
/// Only compares against other trials in the same bracket.
|
||||
#[allow(
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
fn is_pruned_at_rung(
|
||||
&self,
|
||||
current_value: f64,
|
||||
rung_step: u64,
|
||||
bracket: usize,
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let eta = self.reduction_factor as usize;
|
||||
|
||||
// Collect values at this rung step from trials in the same bracket
|
||||
let map = self.trial_brackets.lock().expect("lock poisoned");
|
||||
let mut values_at_rung: Vec<f64> = completed_trials
|
||||
.iter()
|
||||
.filter(|t| t.state == TrialState::Complete || t.state == TrialState::Pruned)
|
||||
.filter(|t| map.get(&t.id).copied() == Some(bracket))
|
||||
.filter_map(|t| {
|
||||
t.intermediate_values
|
||||
.iter()
|
||||
.find(|(s, _)| *s == rung_step)
|
||||
.map(|(_, v)| *v)
|
||||
})
|
||||
.collect();
|
||||
drop(map);
|
||||
|
||||
// Need at least eta trials to make a meaningful comparison
|
||||
if values_at_rung.len() < eta {
|
||||
return false;
|
||||
}
|
||||
|
||||
values_at_rung.push(current_value);
|
||||
|
||||
values_at_rung
|
||||
.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal));
|
||||
if self.direction == Direction::Maximize {
|
||||
values_at_rung.reverse();
|
||||
}
|
||||
|
||||
let n_keep = (values_at_rung.len() as f64 / eta as f64).ceil() as usize;
|
||||
let threshold_idx = n_keep.max(1) - 1;
|
||||
let threshold = values_at_rung[threshold_idx];
|
||||
|
||||
match self.direction {
|
||||
Direction::Minimize => current_value > threshold,
|
||||
Direction::Maximize => current_value < threshold,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::parameter::ParamId;
|
||||
|
||||
CompletedTrial::with_intermediate_values(
|
||||
id,
|
||||
HashMap::<ParamId, crate::ParamValue>::new(),
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
0.0,
|
||||
values.to_vec(),
|
||||
HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
fn make_pruned_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
|
||||
let mut t = make_trial(id, values);
|
||||
t.state = TrialState::Pruned;
|
||||
t
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn s_max_default() {
|
||||
let pruner = HyperbandPruner::new();
|
||||
// s_max = floor(ln(81/1) / ln(3)) = floor(4.0) = 4
|
||||
assert_eq!(pruner.s_max(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn s_max_custom() {
|
||||
let pruner = HyperbandPruner::new()
|
||||
.min_resource(1)
|
||||
.max_resource(16)
|
||||
.reduction_factor(2);
|
||||
// s_max = floor(ln(16) / ln(2)) = floor(4.0) = 4
|
||||
assert_eq!(pruner.s_max(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bracket_count() {
|
||||
let pruner = HyperbandPruner::new();
|
||||
// s_max=4, so brackets 0..=4 → 5 brackets
|
||||
assert_eq!(pruner.s_max() + 1, 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_bracket_0_default() {
|
||||
let pruner = HyperbandPruner::new();
|
||||
// Bracket 0: min_resource_bracket = ceil(81 / 3^4) = ceil(81/81) = 1
|
||||
// Rungs: 1, 3, 9, 27, 81
|
||||
assert_eq!(pruner.rung_steps_for_bracket(0), vec![1, 3, 9, 27, 81]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_bracket_2_default() {
|
||||
let pruner = HyperbandPruner::new();
|
||||
// Bracket 2: min_resource_bracket = ceil(81 / 3^(4-2)) = ceil(81/9) = 9
|
||||
// Rungs: 9, 27, 81
|
||||
assert_eq!(pruner.rung_steps_for_bracket(2), vec![9, 27, 81]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_bracket_4_default() {
|
||||
let pruner = HyperbandPruner::new();
|
||||
// Bracket 4 (s_max): min_resource_bracket = ceil(81 / 3^0) = 81
|
||||
// Rungs: 81 only (no pruning, full budget)
|
||||
assert_eq!(pruner.rung_steps_for_bracket(4), vec![81]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_eta2() {
|
||||
let pruner = HyperbandPruner::new()
|
||||
.min_resource(1)
|
||||
.max_resource(16)
|
||||
.reduction_factor(2);
|
||||
// s_max = 4
|
||||
// Bracket 0: min=ceil(16/2^4)=1, rungs: 1,2,4,8,16
|
||||
assert_eq!(pruner.rung_steps_for_bracket(0), vec![1, 2, 4, 8, 16]);
|
||||
// Bracket 2: min=ceil(16/2^2)=4, rungs: 4,8,16
|
||||
assert_eq!(pruner.rung_steps_for_bracket(2), vec![4, 8, 16]);
|
||||
// Bracket 4: min=16, rungs: 16
|
||||
assert_eq!(pruner.rung_steps_for_bracket(4), vec![16]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_robin_bracket_assignment() {
|
||||
let pruner = HyperbandPruner::new(); // 5 brackets (0..=4)
|
||||
// Trials get assigned in round-robin: 0→0, 1→1, 2→2, 3→3, 4→4, 5→0, ...
|
||||
assert_eq!(pruner.assign_bracket(100), 0);
|
||||
assert_eq!(pruner.assign_bracket(101), 1);
|
||||
assert_eq!(pruner.assign_bracket(102), 2);
|
||||
assert_eq!(pruner.assign_bracket(103), 3);
|
||||
assert_eq!(pruner.assign_bracket(104), 4);
|
||||
assert_eq!(pruner.assign_bracket(105), 0); // wraps around
|
||||
|
||||
// Repeated calls for same trial return same bracket
|
||||
assert_eq!(pruner.assign_bracket(100), 0);
|
||||
assert_eq!(pruner.assign_bracket(103), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_before_first_rung() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
|
||||
// Assign trial 0 to bracket 0 (rungs: 1, 3, 9, 27, 81)
|
||||
pruner.assign_bracket(0);
|
||||
|
||||
// Register completed trials in bracket 0
|
||||
let mut completed = Vec::new();
|
||||
for i in 1..=9 {
|
||||
pruner.assign_bracket(i);
|
||||
completed.push(make_trial(i, &[(1, i as f64)]));
|
||||
}
|
||||
|
||||
// Trial at step 0 (before rung 1) → don't prune
|
||||
assert!(!pruner.should_prune(0, 0, &[(0, 100.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_at_max_resource() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// Put all trials in bracket 0
|
||||
let mut completed = Vec::new();
|
||||
for i in 0..9 {
|
||||
pruner.assign_bracket(i);
|
||||
completed.push(make_trial(i, &[(81, (i + 1) as f64)]));
|
||||
}
|
||||
|
||||
let trial_id = 9;
|
||||
pruner.assign_bracket(trial_id);
|
||||
// At max_resource (81), never prune
|
||||
assert!(!pruner.should_prune(trial_id, 81, &[(81, 100.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_worst_in_bracket_minimize() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// Force all trials into bracket 0 by assigning sequentially
|
||||
// With 5 brackets, trials 0,5,10,... go to bracket 0
|
||||
let bracket_0_ids: Vec<u64> = (0..5).map(|i| i * 5).collect();
|
||||
// Assign all 25 trial IDs to fill brackets
|
||||
for i in 0..25 {
|
||||
pruner.assign_bracket(i);
|
||||
}
|
||||
|
||||
// Create 9 completed trials in bracket 0 at rung step=1
|
||||
let completed: Vec<_> = bracket_0_ids
|
||||
.iter()
|
||||
.take(3)
|
||||
.enumerate()
|
||||
.map(|(idx, &id)| make_trial(id, &[(1, (idx + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// Trial 25 → bracket 0 (25 % 5 == 0)
|
||||
let test_id = 25;
|
||||
pruner.assign_bracket(test_id);
|
||||
assert_eq!(pruner.assign_bracket(test_id), 0);
|
||||
|
||||
// 3 completed + 1 current = 4. eta=3. ceil(4/3)=2. Threshold = 2.0
|
||||
// Value 2.0 → keep
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 2.0)], &completed));
|
||||
// Value 3.0 → prune
|
||||
assert!(pruner.should_prune(test_id, 1, &[(1, 3.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_worst_in_bracket_maximize() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Maximize);
|
||||
|
||||
// Assign trials so they end up in bracket 0
|
||||
for i in 0..25 {
|
||||
pruner.assign_bracket(i);
|
||||
}
|
||||
|
||||
let completed: Vec<_> = [0u64, 5, 10]
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, &id)| make_trial(id, &[(1, (idx + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
let test_id = 25;
|
||||
pruner.assign_bracket(test_id);
|
||||
|
||||
// For maximize, best = highest. Values: 1,2,3 + current
|
||||
// Value 2.0 → keep (threshold = 2.0 when sorted desc: 3,2,current,1)
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 2.0)], &completed));
|
||||
// Value 1.0 → prune
|
||||
assert!(pruner.should_prune(test_id, 1, &[(1, 0.5)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_brackets_have_different_aggressiveness() {
|
||||
let pruner = HyperbandPruner::new()
|
||||
.min_resource(1)
|
||||
.max_resource(81)
|
||||
.reduction_factor(3)
|
||||
.direction(Direction::Minimize);
|
||||
|
||||
let rungs_0 = pruner.rung_steps_for_bracket(0);
|
||||
let rungs_2 = pruner.rung_steps_for_bracket(2);
|
||||
let rungs_4 = pruner.rung_steps_for_bracket(4);
|
||||
|
||||
// Bracket 0 has the most rungs (most aggressive)
|
||||
assert!(rungs_0.len() > rungs_2.len());
|
||||
// Bracket 4 has just 1 rung (no pruning)
|
||||
assert_eq!(rungs_4.len(), 1);
|
||||
// Bracket 0 starts earliest
|
||||
assert!(rungs_0[0] < rungs_2[0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trials_in_different_brackets_independent() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// Assign trials: bracket 0 gets IDs 0,5,10,15,20
|
||||
for i in 0..25 {
|
||||
pruner.assign_bracket(i);
|
||||
}
|
||||
|
||||
// Bracket 0 trials: bad values at rung step=1
|
||||
let bracket_0_trials: Vec<_> = [0u64, 5, 10]
|
||||
.iter()
|
||||
.map(|&id| make_trial(id, &[(1, 100.0)]))
|
||||
.collect();
|
||||
|
||||
// Bracket 1 trials: good values at rung step=1
|
||||
let bracket_1_trials: Vec<_> = [1u64, 6, 11]
|
||||
.iter()
|
||||
.map(|&id| make_trial(id, &[(1, 1.0)]))
|
||||
.collect();
|
||||
|
||||
let mut all_trials = bracket_0_trials;
|
||||
all_trials.extend(bracket_1_trials);
|
||||
|
||||
// A new bracket-0 trial with value 50 should be compared against
|
||||
// bracket-0 peers (100,100,100), not bracket-1 peers (1,1,1)
|
||||
let test_id = 25; // bracket 0
|
||||
pruner.assign_bracket(test_id);
|
||||
// 3 peers at 100.0 + current at 50.0. ceil(4/3)=2. Sorted: 50,100,100,100. Threshold=100.0
|
||||
// Value 50.0 < 100.0 → keep
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 50.0)], &all_trials));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn includes_pruned_trials() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
|
||||
|
||||
for i in 0..25 {
|
||||
pruner.assign_bracket(i);
|
||||
}
|
||||
|
||||
let completed = vec![
|
||||
make_trial(0, &[(1, 1.0)]),
|
||||
make_pruned_trial(5, &[(1, 8.0)]),
|
||||
make_pruned_trial(10, &[(1, 9.0)]),
|
||||
];
|
||||
|
||||
let test_id = 25;
|
||||
pruner.assign_bracket(test_id);
|
||||
|
||||
// Values: 1.0, 8.0, 9.0 + current. eta=3.
|
||||
// Value 1.0 → keep
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 1.0)], &completed));
|
||||
// Value 5.0 → prune (sorted: 1,5,8,9 → keep ceil(4/3)=2 → threshold=5.0, 5.0 not > 5.0 → keep)
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 5.0)], &completed));
|
||||
// Value 6.0 → prune (sorted: 1,6,8,9 → threshold=6.0, 6.0 not > 6.0 → keep)
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 6.0)], &completed));
|
||||
// Value 9.5 → prune
|
||||
assert!(pruner.should_prune(test_id, 1, &[(1, 9.5)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "min_resource must be > 0")]
|
||||
fn rejects_zero_min_resource() {
|
||||
let _ = HyperbandPruner::new().min_resource(0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "max_resource must be > 0")]
|
||||
fn rejects_zero_max_resource() {
|
||||
let _ = HyperbandPruner::new().max_resource(0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "reduction_factor must be >= 2")]
|
||||
fn rejects_reduction_factor_one() {
|
||||
let _ = HyperbandPruner::new().reduction_factor(1);
|
||||
}
|
||||
}
|
||||
@@ -1,127 +0,0 @@
|
||||
use super::Pruner;
|
||||
use super::percentile::compute_percentile;
|
||||
use crate::sampler::CompletedTrial;
|
||||
use crate::types::{Direction, TrialState};
|
||||
|
||||
/// Prune trials that are performing worse than the median of completed trials
|
||||
/// at the same step.
|
||||
///
|
||||
/// This is the most commonly used pruner. It compares the current trial's
|
||||
/// intermediate value at each step with the median of all completed trials'
|
||||
/// values at that same step.
|
||||
///
|
||||
/// Equivalent to `PercentilePruner::new(50.0, direction)`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Direction;
|
||||
/// use optimizer::pruner::MedianPruner;
|
||||
///
|
||||
/// // Prune trials worse than median when minimizing, after 5 warmup steps
|
||||
/// let pruner = MedianPruner::new(Direction::Minimize)
|
||||
/// .n_warmup_steps(5)
|
||||
/// .n_min_trials(3);
|
||||
/// ```
|
||||
pub struct MedianPruner {
|
||||
/// The optimization direction.
|
||||
direction: Direction,
|
||||
/// Don't prune in the first N steps (let the trial warm up).
|
||||
n_warmup_steps: u64,
|
||||
/// Require at least N completed trials before pruning.
|
||||
n_min_trials: usize,
|
||||
}
|
||||
|
||||
impl MedianPruner {
|
||||
/// Create a new `MedianPruner` for the given optimization direction.
|
||||
///
|
||||
/// By default, `n_warmup_steps` is 0 and `n_min_trials` is 1.
|
||||
#[must_use]
|
||||
pub fn new(direction: Direction) -> Self {
|
||||
Self {
|
||||
direction,
|
||||
n_warmup_steps: 0,
|
||||
n_min_trials: 1,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the number of warmup steps. No pruning occurs before this step.
|
||||
#[must_use]
|
||||
pub fn n_warmup_steps(mut self, n: u64) -> Self {
|
||||
self.n_warmup_steps = n;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the minimum number of completed trials required before pruning.
|
||||
#[must_use]
|
||||
pub fn n_min_trials(mut self, n: usize) -> Self {
|
||||
self.n_min_trials = n;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Pruner for MedianPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
// 1. Don't prune during warmup
|
||||
if step < self.n_warmup_steps {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Get the current trial's latest value
|
||||
let Some(&(_, current_value)) = intermediate_values.last() else {
|
||||
return false;
|
||||
};
|
||||
|
||||
// 2. Collect values at this step from completed (non-pruned) trials
|
||||
let mut values_at_step: Vec<f64> = completed_trials
|
||||
.iter()
|
||||
.filter(|t| t.state == TrialState::Complete)
|
||||
.filter_map(|t| {
|
||||
t.intermediate_values
|
||||
.iter()
|
||||
.find(|(s, _)| *s == step)
|
||||
.map(|(_, v)| *v)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 3. Not enough trials
|
||||
if values_at_step.len() < self.n_min_trials {
|
||||
return false;
|
||||
}
|
||||
|
||||
// 4. Compute median (50th percentile)
|
||||
let median = compute_percentile(&mut values_at_step, 50.0);
|
||||
|
||||
// 5. Compare against median based on direction
|
||||
match self.direction {
|
||||
Direction::Minimize => current_value > median,
|
||||
Direction::Maximize => current_value < median,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn compute_median_odd() {
|
||||
assert!((compute_percentile(&mut [3.0, 1.0, 2.0], 50.0) - 2.0).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compute_median_even() {
|
||||
assert!((compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 50.0) - 2.5).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compute_median_single() {
|
||||
assert!((compute_percentile(&mut [5.0], 50.0) - 5.0).abs() < f64::EPSILON);
|
||||
}
|
||||
}
|
||||
@@ -1,74 +0,0 @@
|
||||
//! Pruner trait and implementations for trial pruning.
|
||||
//!
|
||||
//! Pruners decide whether to stop (prune) a trial early based on its
|
||||
//! intermediate values compared to other trials. This is useful for
|
||||
//! discarding unpromising trials before they complete, saving compute.
|
||||
|
||||
mod hyperband;
|
||||
mod median;
|
||||
mod nop;
|
||||
mod patient;
|
||||
pub(crate) mod percentile;
|
||||
mod successive_halving;
|
||||
mod threshold;
|
||||
mod wilcoxon;
|
||||
|
||||
pub use hyperband::HyperbandPruner;
|
||||
pub use median::MedianPruner;
|
||||
pub use nop::NopPruner;
|
||||
pub use patient::PatientPruner;
|
||||
pub use percentile::PercentilePruner;
|
||||
pub use successive_halving::SuccessiveHalvingPruner;
|
||||
pub use threshold::ThresholdPruner;
|
||||
pub use wilcoxon::WilcoxonPruner;
|
||||
|
||||
use crate::sampler::CompletedTrial;
|
||||
|
||||
/// Trait for pluggable trial pruning strategies.
|
||||
///
|
||||
/// Pruners are consulted after each intermediate value is reported to
|
||||
/// decide whether the trial should be stopped early. The trait requires
|
||||
/// `Send + Sync` to support concurrent and async optimization.
|
||||
///
|
||||
/// # Implementing a custom pruner
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::pruner::Pruner;
|
||||
/// use optimizer::sampler::CompletedTrial;
|
||||
///
|
||||
/// struct MyPruner {
|
||||
/// threshold: f64,
|
||||
/// }
|
||||
///
|
||||
/// impl Pruner for MyPruner {
|
||||
/// fn should_prune(
|
||||
/// &self,
|
||||
/// _trial_id: u64,
|
||||
/// _step: u64,
|
||||
/// intermediate_values: &[(u64, f64)],
|
||||
/// _completed_trials: &[CompletedTrial],
|
||||
/// ) -> bool {
|
||||
/// // Prune if the latest value exceeds the threshold
|
||||
/// intermediate_values
|
||||
/// .last()
|
||||
/// .is_some_and(|&(_, v)| v > self.threshold)
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub trait Pruner: Send + Sync {
|
||||
/// Decide whether to prune a trial at the given step.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `trial_id` - The current trial's ID.
|
||||
/// * `step` - The step at which the intermediate value was reported.
|
||||
/// * `intermediate_values` - All `(step, value)` pairs reported so far for this trial.
|
||||
/// * `completed_trials` - History of all completed trials (for comparison).
|
||||
fn should_prune(
|
||||
&self,
|
||||
trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool;
|
||||
}
|
||||
@@ -1,17 +0,0 @@
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
|
||||
/// A pruner that never prunes. This is the default when no pruner is configured.
|
||||
pub struct NopPruner;
|
||||
|
||||
impl Pruner for NopPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
_step: u64,
|
||||
_intermediate_values: &[(u64, f64)],
|
||||
_completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
@@ -1,160 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
|
||||
/// Wraps another pruner and adds a patience window.
|
||||
///
|
||||
/// The inner pruner must recommend pruning for `patience` consecutive
|
||||
/// steps before this pruner actually prunes the trial. This is useful
|
||||
/// to prevent premature pruning when intermediate values are noisy.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::pruner::{PatientPruner, ThresholdPruner};
|
||||
///
|
||||
/// // Only prune after the threshold pruner recommends pruning 3 times in a row
|
||||
/// let inner = ThresholdPruner::new().upper(100.0);
|
||||
/// let pruner = PatientPruner::new(inner, 3);
|
||||
/// ```
|
||||
pub struct PatientPruner {
|
||||
inner: Box<dyn Pruner>,
|
||||
patience: u64,
|
||||
/// Track consecutive prune recommendations per trial.
|
||||
consecutive_counts: Mutex<HashMap<u64, u64>>,
|
||||
}
|
||||
|
||||
impl PatientPruner {
|
||||
/// Create a new `PatientPruner` wrapping the given inner pruner.
|
||||
///
|
||||
/// The inner pruner must recommend pruning for `patience` consecutive
|
||||
/// calls before this pruner returns `true`.
|
||||
pub fn new(inner: impl Pruner + 'static, patience: u64) -> Self {
|
||||
Self {
|
||||
inner: Box::new(inner),
|
||||
patience,
|
||||
consecutive_counts: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Pruner for PatientPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let inner_says_prune =
|
||||
self.inner
|
||||
.should_prune(trial_id, step, intermediate_values, completed_trials);
|
||||
let mut counts = self.consecutive_counts.lock().expect("lock poisoned");
|
||||
let count = counts.entry(trial_id).or_insert(0);
|
||||
if inner_says_prune {
|
||||
*count += 1;
|
||||
*count >= self.patience
|
||||
} else {
|
||||
*count = 0;
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::pruner::ThresholdPruner;
|
||||
|
||||
/// A test pruner that always returns the given value.
|
||||
struct ConstPruner(bool);
|
||||
|
||||
impl Pruner for ConstPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
_step: u64,
|
||||
_intermediate_values: &[(u64, f64)],
|
||||
_completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
/// A pruner that returns values from a sequence.
|
||||
struct SequencePruner(Mutex<Vec<bool>>);
|
||||
|
||||
impl Pruner for SequencePruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
_step: u64,
|
||||
_intermediate_values: &[(u64, f64)],
|
||||
_completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
self.0.lock().expect("lock poisoned").remove(0)
|
||||
}
|
||||
}
|
||||
|
||||
fn call(pruner: &PatientPruner, trial_id: u64, step: u64) -> bool {
|
||||
pruner.should_prune(trial_id, step, &[(step, 0.0)], &[])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn patience_1_behaves_like_inner() {
|
||||
let pruner = PatientPruner::new(ConstPruner(true), 1);
|
||||
assert!(call(&pruner, 0, 0));
|
||||
assert!(call(&pruner, 0, 1));
|
||||
|
||||
let pruner = PatientPruner::new(ConstPruner(false), 1);
|
||||
assert!(!call(&pruner, 0, 0));
|
||||
assert!(!call(&pruner, 0, 1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn patience_3_requires_consecutive_recommendations() {
|
||||
let pruner = PatientPruner::new(ConstPruner(true), 3);
|
||||
assert!(!call(&pruner, 0, 0)); // count=1
|
||||
assert!(!call(&pruner, 0, 1)); // count=2
|
||||
assert!(call(&pruner, 0, 2)); // count=3 → prune
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn counter_resets_on_no_prune() {
|
||||
// Sequence: prune, prune, no-prune, prune, prune, prune
|
||||
let seq = vec![true, true, false, true, true, true];
|
||||
let pruner = PatientPruner::new(SequencePruner(Mutex::new(seq)), 3);
|
||||
|
||||
assert!(!call(&pruner, 0, 0)); // count=1
|
||||
assert!(!call(&pruner, 0, 1)); // count=2
|
||||
assert!(!call(&pruner, 0, 2)); // reset → count=0
|
||||
assert!(!call(&pruner, 0, 3)); // count=1
|
||||
assert!(!call(&pruner, 0, 4)); // count=2
|
||||
assert!(call(&pruner, 0, 5)); // count=3 → prune
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn independent_per_trial() {
|
||||
let pruner = PatientPruner::new(ConstPruner(true), 2);
|
||||
assert!(!call(&pruner, 0, 0)); // trial 0: count=1
|
||||
assert!(!call(&pruner, 1, 0)); // trial 1: count=1
|
||||
assert!(call(&pruner, 0, 1)); // trial 0: count=2 → prune
|
||||
assert!(!call(&pruner, 2, 0)); // trial 2: count=1
|
||||
assert!(call(&pruner, 1, 1)); // trial 1: count=2 → prune
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn works_with_threshold_pruner() {
|
||||
let inner = ThresholdPruner::new().upper(10.0);
|
||||
let pruner = PatientPruner::new(inner, 2);
|
||||
|
||||
// Value below threshold → inner says no
|
||||
assert!(!pruner.should_prune(0, 0, &[(0, 5.0)], &[]));
|
||||
// Value above threshold → inner says yes, count=1
|
||||
assert!(!pruner.should_prune(0, 1, &[(0, 5.0), (1, 15.0)], &[]));
|
||||
// Value above threshold again → count=2 → prune
|
||||
assert!(pruner.should_prune(0, 2, &[(0, 5.0), (1, 15.0), (2, 20.0)], &[]));
|
||||
}
|
||||
}
|
||||
@@ -1,294 +0,0 @@
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
use crate::types::{Direction, TrialState};
|
||||
|
||||
/// Prune trials that are not in the top `percentile`% of completed trials
|
||||
/// at the same training step.
|
||||
///
|
||||
/// `PercentilePruner::new(50.0, direction)` is equivalent to `MedianPruner`.
|
||||
/// `PercentilePruner::new(25.0, direction)` keeps only the top 25% of trials.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Direction;
|
||||
/// use optimizer::pruner::PercentilePruner;
|
||||
///
|
||||
/// // Keep only the top 25% of trials (aggressive pruning)
|
||||
/// let pruner = PercentilePruner::new(25.0, Direction::Minimize)
|
||||
/// .n_warmup_steps(5)
|
||||
/// .n_min_trials(3);
|
||||
/// ```
|
||||
pub struct PercentilePruner {
|
||||
/// Keep trials in the top `percentile`%. Range: (0.0, 100.0).
|
||||
percentile: f64,
|
||||
/// Don't prune in the first N steps (let the trial warm up).
|
||||
n_warmup_steps: u64,
|
||||
/// Require at least N completed trials before pruning.
|
||||
n_min_trials: usize,
|
||||
/// The optimization direction.
|
||||
direction: Direction,
|
||||
}
|
||||
|
||||
impl PercentilePruner {
|
||||
/// Create a new `PercentilePruner` for the given percentile and direction.
|
||||
///
|
||||
/// The `percentile` value must be in `(0.0, 100.0)`.
|
||||
/// A percentile of 50.0 is equivalent to median pruning.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `percentile` is not in `(0.0, 100.0)`.
|
||||
#[must_use]
|
||||
pub fn new(percentile: f64, direction: Direction) -> Self {
|
||||
assert!(
|
||||
percentile > 0.0 && percentile < 100.0,
|
||||
"percentile must be in (0.0, 100.0), got {percentile}"
|
||||
);
|
||||
Self {
|
||||
percentile,
|
||||
n_warmup_steps: 0,
|
||||
n_min_trials: 1,
|
||||
direction,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the number of warmup steps. No pruning occurs before this step.
|
||||
#[must_use]
|
||||
pub fn n_warmup_steps(mut self, n: u64) -> Self {
|
||||
self.n_warmup_steps = n;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the minimum number of completed trials required before pruning.
|
||||
#[must_use]
|
||||
pub fn n_min_trials(mut self, n: usize) -> Self {
|
||||
self.n_min_trials = n;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Pruner for PercentilePruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
// 1. Don't prune during warmup
|
||||
if step < self.n_warmup_steps {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Get the current trial's latest value
|
||||
let Some(&(_, current_value)) = intermediate_values.last() else {
|
||||
return false;
|
||||
};
|
||||
|
||||
// 2. Collect values at this step from completed (non-pruned) trials
|
||||
let mut values_at_step: Vec<f64> = completed_trials
|
||||
.iter()
|
||||
.filter(|t| t.state == TrialState::Complete)
|
||||
.filter_map(|t| {
|
||||
t.intermediate_values
|
||||
.iter()
|
||||
.find(|(s, _)| *s == step)
|
||||
.map(|(_, v)| *v)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 3. Not enough trials
|
||||
if values_at_step.len() < self.n_min_trials {
|
||||
return false;
|
||||
}
|
||||
|
||||
// 4. Compute percentile threshold
|
||||
let threshold = compute_percentile(&mut values_at_step, self.percentile);
|
||||
|
||||
// 5. Compare against threshold based on direction
|
||||
match self.direction {
|
||||
Direction::Minimize => current_value > threshold,
|
||||
Direction::Maximize => current_value < threshold,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute the given percentile of a non-empty slice. Sorts the slice in place.
|
||||
///
|
||||
/// Uses linear interpolation between the two nearest ranks.
|
||||
#[allow(
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
pub(crate) fn compute_percentile(values: &mut [f64], percentile: f64) -> f64 {
|
||||
values.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal));
|
||||
let len = values.len();
|
||||
if len == 1 {
|
||||
return values[0];
|
||||
}
|
||||
// Rank in [0, len-1] range
|
||||
let rank = percentile / 100.0 * (len - 1) as f64;
|
||||
let lower = rank.floor() as usize;
|
||||
let upper = rank.ceil() as usize;
|
||||
if lower == upper {
|
||||
values[lower]
|
||||
} else {
|
||||
let frac = rank - lower as f64;
|
||||
values[lower] * (1.0 - frac) + values[upper] * frac
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn compute_percentile_median_odd() {
|
||||
// Percentile 50 on odd-length slice = median
|
||||
let val = compute_percentile(&mut [3.0, 1.0, 2.0], 50.0);
|
||||
assert!((val - 2.0).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compute_percentile_median_even() {
|
||||
// Percentile 50 on even-length slice = median (interpolated)
|
||||
let val = compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 50.0);
|
||||
assert!((val - 2.5).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compute_percentile_25() {
|
||||
// [1.0, 2.0, 3.0, 4.0], rank = 0.25 * 3 = 0.75
|
||||
// interpolate: 1.0 * 0.25 + 2.0 * 0.75 = 1.75
|
||||
let val = compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 25.0);
|
||||
assert!((val - 1.75).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compute_percentile_75() {
|
||||
// [1.0, 2.0, 3.0, 4.0], rank = 0.75 * 3 = 2.25
|
||||
// interpolate: 3.0 * 0.75 + 4.0 * 0.25 = 3.25
|
||||
let val = compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 75.0);
|
||||
assert!((val - 3.25).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compute_percentile_single() {
|
||||
let val = compute_percentile(&mut [5.0], 50.0);
|
||||
assert!((val - 5.0).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "percentile must be in (0.0, 100.0)")]
|
||||
fn new_rejects_zero() {
|
||||
let _ = PercentilePruner::new(0.0, Direction::Minimize);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "percentile must be in (0.0, 100.0)")]
|
||||
fn new_rejects_hundred() {
|
||||
let _ = PercentilePruner::new(100.0, Direction::Minimize);
|
||||
}
|
||||
|
||||
fn make_completed_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::parameter::ParamId;
|
||||
|
||||
CompletedTrial::with_intermediate_values(
|
||||
id,
|
||||
HashMap::<ParamId, crate::ParamValue>::new(),
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
0.0,
|
||||
values.to_vec(),
|
||||
HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn percentile_50_matches_median_behavior() {
|
||||
let pruner = PercentilePruner::new(50.0, Direction::Minimize);
|
||||
let completed = vec![
|
||||
make_completed_trial(0, &[(0, 1.0), (1, 2.0)]),
|
||||
make_completed_trial(1, &[(0, 3.0), (1, 4.0)]),
|
||||
make_completed_trial(2, &[(0, 5.0), (1, 6.0)]),
|
||||
];
|
||||
// Median at step 1 is 4.0
|
||||
// Value 5.0 > 4.0 → prune
|
||||
assert!(pruner.should_prune(3, 1, &[(0, 3.0), (1, 5.0)], &completed));
|
||||
// Value 3.0 < 4.0 → keep
|
||||
assert!(!pruner.should_prune(3, 1, &[(0, 3.0), (1, 3.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn percentile_25_is_more_aggressive() {
|
||||
let pruner_25 = PercentilePruner::new(25.0, Direction::Minimize);
|
||||
let pruner_75 = PercentilePruner::new(75.0, Direction::Minimize);
|
||||
let completed = vec![
|
||||
make_completed_trial(0, &[(0, 1.0)]),
|
||||
make_completed_trial(1, &[(0, 2.0)]),
|
||||
make_completed_trial(2, &[(0, 3.0)]),
|
||||
make_completed_trial(3, &[(0, 4.0)]),
|
||||
];
|
||||
// 25th percentile at step 0: 1.75
|
||||
// 75th percentile at step 0: 3.25
|
||||
// Value 2.5: above 25th (prune), below 75th (keep)
|
||||
assert!(pruner_25.should_prune(4, 0, &[(0, 2.5)], &completed));
|
||||
assert!(!pruner_75.should_prune(4, 0, &[(0, 2.5)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warmup_prevents_pruning() {
|
||||
let pruner = PercentilePruner::new(50.0, Direction::Minimize).n_warmup_steps(5);
|
||||
let completed = vec![make_completed_trial(0, &[(0, 1.0)])];
|
||||
// Step 3 < warmup 5 → no prune even with bad value
|
||||
assert!(!pruner.should_prune(1, 3, &[(3, 100.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn n_min_trials_prevents_pruning() {
|
||||
let pruner = PercentilePruner::new(50.0, Direction::Minimize).n_min_trials(5);
|
||||
let completed = vec![
|
||||
make_completed_trial(0, &[(0, 1.0)]),
|
||||
make_completed_trial(1, &[(0, 2.0)]),
|
||||
];
|
||||
// Only 2 trials, need 5 → no prune
|
||||
assert!(!pruner.should_prune(2, 0, &[(0, 100.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maximize_direction() {
|
||||
let pruner = PercentilePruner::new(50.0, Direction::Maximize);
|
||||
let completed = vec![
|
||||
make_completed_trial(0, &[(0, 1.0)]),
|
||||
make_completed_trial(1, &[(0, 3.0)]),
|
||||
make_completed_trial(2, &[(0, 5.0)]),
|
||||
];
|
||||
// Median at step 0 is 3.0
|
||||
// Value 2.0 < 3.0 → prune (maximize wants higher)
|
||||
assert!(pruner.should_prune(3, 0, &[(0, 2.0)], &completed));
|
||||
// Value 4.0 > 3.0 → keep
|
||||
assert!(!pruner.should_prune(3, 0, &[(0, 4.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn near_boundary_percentiles() {
|
||||
let pruner_low = PercentilePruner::new(1.0, Direction::Minimize);
|
||||
let pruner_high = PercentilePruner::new(99.0, Direction::Minimize);
|
||||
let completed = vec![
|
||||
make_completed_trial(0, &[(0, 1.0)]),
|
||||
make_completed_trial(1, &[(0, 2.0)]),
|
||||
make_completed_trial(2, &[(0, 3.0)]),
|
||||
make_completed_trial(3, &[(0, 100.0)]),
|
||||
];
|
||||
// Percentile 1 is very aggressive (threshold near 1.0)
|
||||
// Value 1.5 should be pruned
|
||||
assert!(pruner_low.should_prune(4, 0, &[(0, 1.5)], &completed));
|
||||
// Percentile 99 is very lenient (threshold near 100.0)
|
||||
// Value 50.0 should not be pruned
|
||||
assert!(!pruner_high.should_prune(4, 0, &[(0, 50.0)], &completed));
|
||||
}
|
||||
}
|
||||
@@ -1,466 +0,0 @@
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
use crate::types::{Direction, TrialState};
|
||||
|
||||
/// Successive Halving pruner based on the SHA algorithm.
|
||||
///
|
||||
/// Trials are evaluated at exponentially-spaced "rungs". At each rung,
|
||||
/// only the top 1/eta fraction of trials survive to the next rung.
|
||||
///
|
||||
/// For example, with `min_resource=1`, `max_resource=81`, `reduction_factor=3`:
|
||||
/// - Rung 0: evaluate at step 1, keep top 1/3
|
||||
/// - Rung 1: evaluate at step 3, keep top 1/3
|
||||
/// - Rung 2: evaluate at step 9, keep top 1/3
|
||||
/// - Rung 3: evaluate at step 27, keep top 1/3
|
||||
/// - Rung 4: evaluate at step 81 (full budget)
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Direction;
|
||||
/// use optimizer::pruner::SuccessiveHalvingPruner;
|
||||
///
|
||||
/// let pruner = SuccessiveHalvingPruner::new()
|
||||
/// .min_resource(1)
|
||||
/// .max_resource(81)
|
||||
/// .reduction_factor(3)
|
||||
/// .direction(Direction::Minimize);
|
||||
/// ```
|
||||
pub struct SuccessiveHalvingPruner {
|
||||
min_resource: u64,
|
||||
max_resource: u64,
|
||||
reduction_factor: u64,
|
||||
min_early_stopping_rate: u64,
|
||||
direction: Direction,
|
||||
}
|
||||
|
||||
impl SuccessiveHalvingPruner {
|
||||
/// Create a new `SuccessiveHalvingPruner` with default parameters.
|
||||
///
|
||||
/// Defaults: `min_resource=1`, `max_resource=81`, `reduction_factor=3`,
|
||||
/// `min_early_stopping_rate=0`, `direction=Minimize`.
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
min_resource: 1,
|
||||
max_resource: 81,
|
||||
reduction_factor: 3,
|
||||
min_early_stopping_rate: 0,
|
||||
direction: Direction::Minimize,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set 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
|
||||
}
|
||||
|
||||
/// Set 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
|
||||
}
|
||||
|
||||
/// Set the reduction factor (eta). At each rung, the top 1/eta trials survive.
|
||||
///
|
||||
/// # 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
|
||||
}
|
||||
|
||||
/// Set the minimum early stopping rate. Skips the first N rungs.
|
||||
#[must_use]
|
||||
pub fn min_early_stopping_rate(mut self, n: u64) -> Self {
|
||||
self.min_early_stopping_rate = n;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the optimization direction.
|
||||
#[must_use]
|
||||
pub fn direction(mut self, d: Direction) -> Self {
|
||||
self.direction = d;
|
||||
self
|
||||
}
|
||||
|
||||
/// Compute the rung steps: `[min_resource * eta^(s), ...]` up to `max_resource`,
|
||||
/// skipping the first `min_early_stopping_rate` rungs.
|
||||
fn rung_steps(&self) -> Vec<u64> {
|
||||
let eta = self.reduction_factor;
|
||||
let mut steps = Vec::new();
|
||||
let mut rung: u32 = 0;
|
||||
while let Some(power) = eta.checked_pow(rung) {
|
||||
let step = self.min_resource.saturating_mul(power);
|
||||
if step > self.max_resource {
|
||||
break;
|
||||
}
|
||||
if u64::from(rung) >= self.min_early_stopping_rate {
|
||||
steps.push(step);
|
||||
}
|
||||
rung += 1;
|
||||
}
|
||||
steps
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SuccessiveHalvingPruner {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
impl Pruner for SuccessiveHalvingPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let rungs = self.rung_steps();
|
||||
|
||||
// Find the highest rung step <= current step
|
||||
let Some(&rung_step) = rungs.iter().rev().find(|&&r| r <= step) else {
|
||||
// No rung matches (before the first rung) → don't prune
|
||||
return false;
|
||||
};
|
||||
|
||||
// If this is the last rung (full budget), don't prune
|
||||
if rung_step >= self.max_resource {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Get the current trial's value at this rung step
|
||||
let Some(&(_, current_value)) = intermediate_values.iter().find(|(s, _)| *s == rung_step)
|
||||
else {
|
||||
// Trial hasn't reported a value at this exact rung step.
|
||||
// Use the latest intermediate value at or before the rung step instead.
|
||||
let Some(&(_, current_value)) = intermediate_values
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|(s, _)| *s <= rung_step)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
return self.is_pruned_at_rung(current_value, rung_step, completed_trials);
|
||||
};
|
||||
|
||||
self.is_pruned_at_rung(current_value, rung_step, completed_trials)
|
||||
}
|
||||
}
|
||||
|
||||
impl SuccessiveHalvingPruner {
|
||||
/// Determine whether a trial with `current_value` should be pruned at the given rung.
|
||||
#[allow(
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
fn is_pruned_at_rung(
|
||||
&self,
|
||||
current_value: f64,
|
||||
rung_step: u64,
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let eta = self.reduction_factor as usize;
|
||||
|
||||
// Collect values at this rung step from all trials that reached it
|
||||
let mut values_at_rung: Vec<f64> = completed_trials
|
||||
.iter()
|
||||
.filter(|t| t.state == TrialState::Complete || t.state == TrialState::Pruned)
|
||||
.filter_map(|t| {
|
||||
t.intermediate_values
|
||||
.iter()
|
||||
.find(|(s, _)| *s == rung_step)
|
||||
.map(|(_, v)| *v)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Need at least eta trials to make a meaningful comparison
|
||||
// (with fewer trials, we can't determine the top 1/eta fraction)
|
||||
if values_at_rung.len() < eta {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Include the current trial's value for ranking
|
||||
values_at_rung.push(current_value);
|
||||
|
||||
// Sort based on direction: best values first
|
||||
values_at_rung
|
||||
.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal));
|
||||
if self.direction == Direction::Maximize {
|
||||
values_at_rung.reverse();
|
||||
}
|
||||
|
||||
// Keep top 1/eta fraction
|
||||
let n_keep = (values_at_rung.len() as f64 / eta as f64).ceil() as usize;
|
||||
let threshold_idx = n_keep.max(1) - 1;
|
||||
let threshold = values_at_rung[threshold_idx];
|
||||
|
||||
// Prune if current value is worse than the threshold
|
||||
match self.direction {
|
||||
Direction::Minimize => current_value > threshold,
|
||||
Direction::Maximize => current_value < threshold,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::parameter::ParamId;
|
||||
|
||||
CompletedTrial::with_intermediate_values(
|
||||
id,
|
||||
HashMap::<ParamId, crate::ParamValue>::new(),
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
0.0,
|
||||
values.to_vec(),
|
||||
HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
fn make_pruned_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
|
||||
let mut t = make_trial(id, values);
|
||||
t.state = TrialState::Pruned;
|
||||
t
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_default() {
|
||||
let pruner = SuccessiveHalvingPruner::new();
|
||||
let rungs = pruner.rung_steps();
|
||||
// min=1, max=81, eta=3 → 1, 3, 9, 27, 81
|
||||
assert_eq!(rungs, vec![1, 3, 9, 27, 81]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_custom() {
|
||||
let pruner = SuccessiveHalvingPruner::new()
|
||||
.min_resource(2)
|
||||
.max_resource(32)
|
||||
.reduction_factor(2);
|
||||
let rungs = pruner.rung_steps();
|
||||
// 2, 4, 8, 16, 32
|
||||
assert_eq!(rungs, vec![2, 4, 8, 16, 32]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_with_early_stopping_rate() {
|
||||
let pruner = SuccessiveHalvingPruner::new().min_early_stopping_rate(2);
|
||||
let rungs = pruner.rung_steps();
|
||||
// Skip rung 0 (step=1) and rung 1 (step=3), keep rung 2+ (9, 27, 81)
|
||||
assert_eq!(rungs, vec![9, 27, 81]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_before_first_rung() {
|
||||
let pruner = SuccessiveHalvingPruner::new()
|
||||
.min_resource(10)
|
||||
.max_resource(100)
|
||||
.reduction_factor(3);
|
||||
let completed = vec![
|
||||
make_trial(0, &[(5, 1.0)]),
|
||||
make_trial(1, &[(5, 2.0)]),
|
||||
make_trial(2, &[(5, 3.0)]),
|
||||
];
|
||||
// Step 5 is before the first rung (10)
|
||||
assert!(!pruner.should_prune(3, 5, &[(5, 100.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_with_single_trial() {
|
||||
let pruner = SuccessiveHalvingPruner::new();
|
||||
let completed = vec![make_trial(0, &[(1, 5.0)]), make_trial(1, &[(1, 3.0)])];
|
||||
// Only 2 completed trials at rung + 1 current = 3 total, threshold = ceil(3/3) = 1
|
||||
// With eta=3, we need at least 3 completed trials
|
||||
assert!(!pruner.should_prune(2, 1, &[(1, 10.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_worst_trials_at_rung() {
|
||||
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// 9 completed trials at rung step=1, with values 1..=9
|
||||
let completed: Vec<_> = (0..9)
|
||||
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// With eta=3, keep top 1/3. 10 total values → ceil(10/3) = 4 kept
|
||||
// Best 4 values: 1, 2, 3, 4. Threshold = 4.0
|
||||
// Value 3.0 → keep (in top 1/3)
|
||||
assert!(!pruner.should_prune(9, 1, &[(1, 3.0)], &completed));
|
||||
// Value 5.0 → prune (not in top 1/3)
|
||||
assert!(pruner.should_prune(9, 1, &[(1, 5.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn top_fraction_survives() {
|
||||
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// 6 completed trials at step=1
|
||||
let completed: Vec<_> = (0..6)
|
||||
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// 7 total (6 + current). ceil(7/3) = 3 keep. Threshold = 3.0
|
||||
// Value 2.0 → keep
|
||||
assert!(!pruner.should_prune(6, 1, &[(1, 2.0)], &completed));
|
||||
// Value 3.0 → keep (at threshold)
|
||||
assert!(!pruner.should_prune(6, 1, &[(1, 3.0)], &completed));
|
||||
// Value 4.0 → prune
|
||||
assert!(pruner.should_prune(6, 1, &[(1, 4.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maximize_direction() {
|
||||
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Maximize);
|
||||
|
||||
let completed: Vec<_> = (0..6)
|
||||
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// 7 total. For maximize, best = highest. ceil(7/3)=3. Top 3: 6,5,4. Threshold=4.0
|
||||
// Value 5.0 → keep
|
||||
assert!(!pruner.should_prune(6, 1, &[(1, 5.0)], &completed));
|
||||
// Value 4.0 → keep (at threshold)
|
||||
assert!(!pruner.should_prune(6, 1, &[(1, 4.0)], &completed));
|
||||
// Value 3.0 → prune
|
||||
assert!(pruner.should_prune(6, 1, &[(1, 3.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reduction_factor_2() {
|
||||
let pruner = SuccessiveHalvingPruner::new()
|
||||
.reduction_factor(2)
|
||||
.min_resource(1)
|
||||
.max_resource(16)
|
||||
.direction(Direction::Minimize);
|
||||
|
||||
// Rungs: 1, 2, 4, 8, 16
|
||||
assert_eq!(pruner.rung_steps(), vec![1, 2, 4, 8, 16]);
|
||||
|
||||
// 4 completed trials at rung step=1
|
||||
let completed: Vec<_> = (0..4)
|
||||
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// With eta=2, 5 total. ceil(5/2) = 3 keep. Threshold = 3.0
|
||||
// Value 3.0 → keep
|
||||
assert!(!pruner.should_prune(4, 1, &[(1, 3.0)], &completed));
|
||||
// Value 4.0 → prune
|
||||
assert!(pruner.should_prune(4, 1, &[(1, 4.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reduction_factor_4() {
|
||||
let pruner = SuccessiveHalvingPruner::new()
|
||||
.reduction_factor(4)
|
||||
.min_resource(1)
|
||||
.max_resource(64)
|
||||
.direction(Direction::Minimize);
|
||||
|
||||
// Rungs: 1, 4, 16, 64
|
||||
assert_eq!(pruner.rung_steps(), vec![1, 4, 16, 64]);
|
||||
|
||||
// 12 completed trials at rung step=1
|
||||
let completed: Vec<_> = (0..12)
|
||||
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// With eta=4, 13 total. ceil(13/4) = 4 keep. Threshold = 4.0
|
||||
// Value 4.0 → keep
|
||||
assert!(!pruner.should_prune(12, 1, &[(1, 4.0)], &completed));
|
||||
// Value 5.0 → prune
|
||||
assert!(pruner.should_prune(12, 1, &[(1, 5.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_contiguous_steps() {
|
||||
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// Trials reporting at rung step=3 (not step=1)
|
||||
let completed: Vec<_> = (0..6)
|
||||
.map(|i| make_trial(i, &[(3, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// Current trial reports at step 5 (between rung 3 and rung 9)
|
||||
// Highest rung <= 5 is 3. Use value at rung step 3.
|
||||
// Trial has value at step 3 → use it
|
||||
assert!(!pruner.should_prune(6, 5, &[(3, 2.0)], &completed));
|
||||
assert!(pruner.should_prune(6, 5, &[(3, 5.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_at_max_resource() {
|
||||
let pruner = SuccessiveHalvingPruner::new();
|
||||
let completed: Vec<_> = (0..9)
|
||||
.map(|i| make_trial(i, &[(81, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// At the max resource rung, never prune (trial should complete)
|
||||
assert!(!pruner.should_prune(9, 81, &[(81, 100.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn includes_pruned_trials_in_comparison() {
|
||||
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// Mix of completed and pruned trials at rung step=1
|
||||
let completed = vec![
|
||||
make_trial(0, &[(1, 1.0)]),
|
||||
make_trial(1, &[(1, 2.0)]),
|
||||
make_pruned_trial(2, &[(1, 8.0)]),
|
||||
make_pruned_trial(3, &[(1, 9.0)]),
|
||||
make_pruned_trial(4, &[(1, 10.0)]),
|
||||
];
|
||||
|
||||
// 6 total. ceil(6/3) = 2 keep. Threshold = 2.0
|
||||
// Value 2.0 → keep
|
||||
assert!(!pruner.should_prune(5, 1, &[(1, 2.0)], &completed));
|
||||
// Value 3.0 → prune
|
||||
assert!(pruner.should_prune(5, 1, &[(1, 3.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "min_resource must be > 0")]
|
||||
fn rejects_zero_min_resource() {
|
||||
let _ = SuccessiveHalvingPruner::new().min_resource(0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "max_resource must be > 0")]
|
||||
fn rejects_zero_max_resource() {
|
||||
let _ = SuccessiveHalvingPruner::new().max_resource(0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "reduction_factor must be >= 2")]
|
||||
fn rejects_reduction_factor_one() {
|
||||
let _ = SuccessiveHalvingPruner::new().reduction_factor(1);
|
||||
}
|
||||
}
|
||||
@@ -1,83 +0,0 @@
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
|
||||
/// Prune trials whose intermediate values exceed fixed thresholds.
|
||||
///
|
||||
/// Useful for cutting off trials that are clearly diverging or stuck
|
||||
/// at bad values early in training.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::pruner::ThresholdPruner;
|
||||
///
|
||||
/// // Prune if the intermediate value exceeds 100.0 or falls below 0.0
|
||||
/// let pruner = ThresholdPruner::new().upper(100.0).lower(0.0);
|
||||
/// ```
|
||||
pub struct ThresholdPruner {
|
||||
/// Prune if intermediate value is greater than this. `None` = no upper bound.
|
||||
upper: Option<f64>,
|
||||
/// Prune if intermediate value is less than this. `None` = no lower bound.
|
||||
lower: Option<f64>,
|
||||
}
|
||||
|
||||
impl ThresholdPruner {
|
||||
/// Create a new `ThresholdPruner` with no thresholds set.
|
||||
///
|
||||
/// By default, no pruning occurs. Use [`upper`](Self::upper) and
|
||||
/// [`lower`](Self::lower) to set bounds.
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
upper: None,
|
||||
lower: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the upper threshold. Trials with intermediate values above this
|
||||
/// will be pruned.
|
||||
#[must_use]
|
||||
pub fn upper(mut self, threshold: f64) -> Self {
|
||||
self.upper = Some(threshold);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the lower threshold. Trials with intermediate values below this
|
||||
/// will be pruned.
|
||||
#[must_use]
|
||||
pub fn lower(mut self, threshold: f64) -> Self {
|
||||
self.lower = Some(threshold);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ThresholdPruner {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Pruner for ThresholdPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
_step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
_completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let Some(&(_, latest_value)) = intermediate_values.last() else {
|
||||
return false;
|
||||
};
|
||||
if let Some(upper) = self.upper
|
||||
&& latest_value > upper
|
||||
{
|
||||
return true;
|
||||
}
|
||||
if let Some(lower) = self.lower
|
||||
&& latest_value < lower
|
||||
{
|
||||
return true;
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
@@ -1,531 +0,0 @@
|
||||
use core::cmp::Ordering;
|
||||
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
use crate::types::{Direction, TrialState};
|
||||
|
||||
/// Prune trials using a Wilcoxon signed-rank test comparing intermediate
|
||||
/// values against the best completed trial.
|
||||
///
|
||||
/// More principled than `MedianPruner` for noisy objectives — it accounts
|
||||
/// for the paired nature of step-aligned comparisons and doesn't prune
|
||||
/// on random fluctuations.
|
||||
///
|
||||
/// The test compares intermediate values at matching steps between the
|
||||
/// current trial and the best completed trial. If the current trial is
|
||||
/// statistically significantly worse (p < threshold), it is pruned.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Direction;
|
||||
/// use optimizer::pruner::WilcoxonPruner;
|
||||
///
|
||||
/// let pruner = WilcoxonPruner::new(Direction::Minimize)
|
||||
/// .p_value_threshold(0.05)
|
||||
/// .n_warmup_steps(5)
|
||||
/// .n_min_trials(1);
|
||||
/// ```
|
||||
pub struct WilcoxonPruner {
|
||||
/// Significance level (default 0.05). Lower = more conservative.
|
||||
p_value_threshold: f64,
|
||||
/// Don't prune in the first N steps (let the trial warm up).
|
||||
n_warmup_steps: u64,
|
||||
/// Require at least N completed trials before pruning.
|
||||
n_min_trials: usize,
|
||||
/// The optimization direction.
|
||||
direction: Direction,
|
||||
}
|
||||
|
||||
impl WilcoxonPruner {
|
||||
/// Create a new `WilcoxonPruner` for the given optimization direction.
|
||||
///
|
||||
/// By default, `p_value_threshold` is 0.05, `n_warmup_steps` is 0,
|
||||
/// and `n_min_trials` is 1.
|
||||
#[must_use]
|
||||
pub fn new(direction: Direction) -> Self {
|
||||
Self {
|
||||
p_value_threshold: 0.05,
|
||||
n_warmup_steps: 0,
|
||||
n_min_trials: 1,
|
||||
direction,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the p-value threshold for significance.
|
||||
///
|
||||
/// Must be in (0.0, 1.0). Lower values are more conservative (harder to prune).
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `p` is not in the open interval (0.0, 1.0).
|
||||
#[must_use]
|
||||
pub fn p_value_threshold(mut self, p: f64) -> Self {
|
||||
assert!(
|
||||
p > 0.0 && p < 1.0,
|
||||
"p_value_threshold must be in (0.0, 1.0)"
|
||||
);
|
||||
self.p_value_threshold = p;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the number of warmup steps. No pruning occurs before this step.
|
||||
#[must_use]
|
||||
pub fn n_warmup_steps(mut self, n: u64) -> Self {
|
||||
self.n_warmup_steps = n;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the minimum number of completed trials required before pruning.
|
||||
#[must_use]
|
||||
pub fn n_min_trials(mut self, n: usize) -> Self {
|
||||
self.n_min_trials = n;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Pruner for WilcoxonPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
if step < self.n_warmup_steps {
|
||||
return false;
|
||||
}
|
||||
|
||||
let completed: Vec<&CompletedTrial> = completed_trials
|
||||
.iter()
|
||||
.filter(|t| t.state == TrialState::Complete)
|
||||
.collect();
|
||||
|
||||
if completed.len() < self.n_min_trials {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Find the best completed trial by final objective value.
|
||||
let best = match self.direction {
|
||||
Direction::Minimize => completed
|
||||
.iter()
|
||||
.min_by(|a, b| a.value.partial_cmp(&b.value).unwrap_or(Ordering::Equal)),
|
||||
Direction::Maximize => completed
|
||||
.iter()
|
||||
.max_by(|a, b| a.value.partial_cmp(&b.value).unwrap_or(Ordering::Equal)),
|
||||
};
|
||||
|
||||
let Some(best) = best else {
|
||||
return false;
|
||||
};
|
||||
|
||||
// Pair intermediate values at matching steps.
|
||||
let pairs: Vec<(f64, f64)> = intermediate_values
|
||||
.iter()
|
||||
.filter_map(|&(s, current_v)| {
|
||||
best.intermediate_values
|
||||
.iter()
|
||||
.find(|(bs, _)| *bs == s)
|
||||
.map(|&(_, best_v)| (current_v, best_v))
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Need at least 6 pairs for a meaningful test.
|
||||
if pairs.len() < 6 {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Compute signed differences: current - best.
|
||||
// For minimization: positive diff means current is worse.
|
||||
// For maximization: negative diff means current is worse.
|
||||
let differences: Vec<f64> = pairs
|
||||
.iter()
|
||||
.map(|&(current, best_v)| current - best_v)
|
||||
.collect();
|
||||
|
||||
// Run the Wilcoxon signed-rank test.
|
||||
let p_value = wilcoxon_signed_rank_test(&differences, self.direction);
|
||||
|
||||
p_value < self.p_value_threshold
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform a one-sided Wilcoxon signed-rank test.
|
||||
///
|
||||
/// Tests whether the values tend to be worse than zero (positive for
|
||||
/// minimization, negative for maximization).
|
||||
///
|
||||
/// Returns a p-value. Small p-values indicate the current trial is
|
||||
/// significantly worse.
|
||||
fn wilcoxon_signed_rank_test(differences: &[f64], direction: Direction) -> f64 {
|
||||
// 1. Remove zero differences.
|
||||
let nonzero: Vec<f64> = differences.iter().copied().filter(|d| *d != 0.0).collect();
|
||||
let n = nonzero.len();
|
||||
|
||||
if n < 6 {
|
||||
return 1.0; // Not enough data
|
||||
}
|
||||
|
||||
// 2. Rank by absolute value.
|
||||
let mut abs_ranked: Vec<(usize, f64, f64)> = nonzero
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &d)| (i, d.abs(), d))
|
||||
.collect();
|
||||
abs_ranked.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
|
||||
|
||||
// 3. Assign ranks with tie correction.
|
||||
let ranks = assign_ranks(&abs_ranked);
|
||||
|
||||
// 4. Compute W+ (sum of ranks for positive differences) and
|
||||
// W- (sum of ranks for negative differences).
|
||||
let mut w_plus = 0.0;
|
||||
let mut w_minus = 0.0;
|
||||
for (i, &(_, _, orig)) in abs_ranked.iter().enumerate() {
|
||||
if orig > 0.0 {
|
||||
w_plus += ranks[i];
|
||||
} else {
|
||||
w_minus += ranks[i];
|
||||
}
|
||||
}
|
||||
|
||||
// For one-sided test:
|
||||
// - Minimization: we want to detect positive diffs (current worse).
|
||||
// Large W+ means significantly worse. Test statistic = W-.
|
||||
// - Maximization: we want to detect negative diffs (current worse).
|
||||
// Large W- means significantly worse. Test statistic = W+.
|
||||
let w = match direction {
|
||||
Direction::Minimize => w_minus,
|
||||
Direction::Maximize => w_plus,
|
||||
};
|
||||
|
||||
// 5. Normal approximation for the p-value.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
let n_f = n as f64;
|
||||
let mean = n_f * (n_f + 1.0) / 4.0;
|
||||
let variance = n_f * (n_f + 1.0) * (2.0 * n_f + 1.0) / 24.0;
|
||||
|
||||
// Tie correction for variance.
|
||||
let tie_correction = compute_tie_correction(&ranks);
|
||||
let adjusted_variance = variance - tie_correction;
|
||||
|
||||
if adjusted_variance <= 0.0 {
|
||||
return 1.0;
|
||||
}
|
||||
|
||||
let std_dev = adjusted_variance.sqrt();
|
||||
|
||||
// Continuity correction: shift W by 0.5 towards mean.
|
||||
let continuity = if w < mean { 0.5 } else { -0.5 };
|
||||
let z = (w + continuity - mean) / std_dev;
|
||||
|
||||
// One-sided p-value (lower tail): probability that the test statistic
|
||||
// is this small or smaller under H0.
|
||||
normal_cdf(z)
|
||||
}
|
||||
|
||||
/// Assign average ranks, handling ties.
|
||||
fn assign_ranks(sorted: &[(usize, f64, f64)]) -> Vec<f64> {
|
||||
let n = sorted.len();
|
||||
let mut ranks = vec![0.0; n];
|
||||
let mut i = 0;
|
||||
|
||||
while i < n {
|
||||
let mut j = i;
|
||||
// Find all items tied with sorted[i].
|
||||
while j < n
|
||||
&& (sorted[j].1 - sorted[i].1).abs() < f64::EPSILON * sorted[i].1.max(1.0) * 100.0
|
||||
{
|
||||
j += 1;
|
||||
}
|
||||
// Average rank for the tie group. Ranks are 1-based.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
let avg_rank = (i + 1 + j) as f64 / 2.0;
|
||||
for rank in ranks.iter_mut().take(j).skip(i) {
|
||||
*rank = avg_rank;
|
||||
}
|
||||
i = j;
|
||||
}
|
||||
|
||||
ranks
|
||||
}
|
||||
|
||||
/// Compute the tie correction term for the variance.
|
||||
/// For each tie group of size t, subtract t^3 - t from the sum,
|
||||
/// then divide by 48.
|
||||
fn compute_tie_correction(ranks: &[f64]) -> f64 {
|
||||
let mut correction = 0.0;
|
||||
let mut i = 0;
|
||||
|
||||
while i < ranks.len() {
|
||||
let mut j = i;
|
||||
while j < ranks.len() && (ranks[j] - ranks[i]).abs() < f64::EPSILON {
|
||||
j += 1;
|
||||
}
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
let t = (j - i) as f64;
|
||||
if t > 1.0 {
|
||||
correction += t * t * t - t;
|
||||
}
|
||||
i = j;
|
||||
}
|
||||
|
||||
correction / 48.0
|
||||
}
|
||||
|
||||
/// Standard normal CDF using an approximation (Abramowitz & Stegun).
|
||||
fn normal_cdf(x: f64) -> f64 {
|
||||
// Use the complementary error function relationship:
|
||||
// Φ(x) = 0.5 * erfc(-x / √2)
|
||||
0.5 * erfc(-x / core::f64::consts::SQRT_2)
|
||||
}
|
||||
|
||||
/// Complementary error function approximation.
|
||||
/// Maximum error: 1.5 × 10⁻⁷ (Abramowitz & Stegun formula 7.1.26).
|
||||
fn erfc(x: f64) -> f64 {
|
||||
let t = 1.0 / (1.0 + 0.327_591_1 * x.abs());
|
||||
let poly = t
|
||||
* (0.254_829_592
|
||||
+ t * (-0.284_496_736
|
||||
+ t * (1.421_413_741 + t * (-1.453_152_027 + t * 1.061_405_429))));
|
||||
let result = poly * (-x * x).exp();
|
||||
|
||||
if x >= 0.0 { result } else { 2.0 - result }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn trial_with_values(
|
||||
id: u64,
|
||||
value: f64,
|
||||
intermediate_values: Vec<(u64, f64)>,
|
||||
) -> CompletedTrial {
|
||||
CompletedTrial::with_intermediate_values(
|
||||
id,
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
value,
|
||||
intermediate_values,
|
||||
HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_during_warmup() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize).n_warmup_steps(10);
|
||||
let completed = vec![trial_with_values(
|
||||
0,
|
||||
0.1,
|
||||
(0..20).map(|s| (s, 0.1)).collect(),
|
||||
)];
|
||||
let current: Vec<(u64, f64)> = (0..8).map(|s| (s, 100.0)).collect();
|
||||
assert!(!pruner.should_prune(1, 7, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_with_insufficient_trials() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize).n_min_trials(5);
|
||||
let completed = vec![trial_with_values(
|
||||
0,
|
||||
0.1,
|
||||
(0..20).map(|s| (s, 0.1)).collect(),
|
||||
)];
|
||||
let current: Vec<(u64, f64)> = (0..10).map(|s| (s, 100.0)).collect();
|
||||
assert!(!pruner.should_prune(1, 9, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_with_fewer_than_6_pairs() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
let completed = vec![trial_with_values(
|
||||
0,
|
||||
0.1,
|
||||
(0..5).map(|s| (s, 0.1)).collect(),
|
||||
)];
|
||||
// Only 5 matching steps
|
||||
let current: Vec<(u64, f64)> = (0..5).map(|s| (s, 100.0)).collect();
|
||||
assert!(!pruner.should_prune(1, 4, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_when_consistently_worse_minimize() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
|
||||
// Best trial has low values.
|
||||
let best_values: Vec<(u64, f64)> = (0..20).map(|s| (s, 0.1)).collect();
|
||||
let completed = vec![trial_with_values(0, 0.1, best_values)];
|
||||
|
||||
// Current trial is consistently much worse.
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 10.0)).collect();
|
||||
assert!(pruner.should_prune(1, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_when_consistently_worse_maximize() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Maximize);
|
||||
|
||||
// Best trial has high values.
|
||||
let best_values: Vec<(u64, f64)> = (0..20).map(|s| (s, 10.0)).collect();
|
||||
let completed = vec![trial_with_values(0, 10.0, best_values)];
|
||||
|
||||
// Current trial is consistently much worse.
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 0.1)).collect();
|
||||
assert!(pruner.should_prune(1, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_statistically_similar() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
|
||||
// Best trial and current trial have very similar values with noise.
|
||||
let best_values: Vec<(u64, f64)> = (0..20_u64)
|
||||
.map(|s| {
|
||||
let noise = if s.is_multiple_of(2) { 0.01 } else { -0.01 };
|
||||
(s, 1.0 + noise)
|
||||
})
|
||||
.collect();
|
||||
let completed = vec![trial_with_values(0, 1.0, best_values)];
|
||||
|
||||
// Current trial is similar — alternating above/below.
|
||||
let current: Vec<(u64, f64)> = (0..20_u64)
|
||||
.map(|s| {
|
||||
let noise = if s.is_multiple_of(2) { -0.01 } else { 0.01 };
|
||||
(s, 1.0 + noise)
|
||||
})
|
||||
.collect();
|
||||
assert!(!pruner.should_prune(1, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_best_trial_minimize() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
|
||||
// Two completed trials: trial 0 is better (lower).
|
||||
let completed = vec![
|
||||
trial_with_values(0, 0.1, (0..20).map(|s| (s, 0.1)).collect()),
|
||||
trial_with_values(1, 5.0, (0..20).map(|s| (s, 5.0)).collect()),
|
||||
];
|
||||
|
||||
// Current trial is worse than the best but similar to the second.
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 5.0)).collect();
|
||||
assert!(pruner.should_prune(2, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_best_trial_maximize() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Maximize);
|
||||
|
||||
// Two completed trials: trial 1 is better (higher).
|
||||
let completed = vec![
|
||||
trial_with_values(0, 0.1, (0..20).map(|s| (s, 0.1)).collect()),
|
||||
trial_with_values(1, 10.0, (0..20).map(|s| (s, 10.0)).collect()),
|
||||
];
|
||||
|
||||
// Current trial is worse than the best.
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 0.1)).collect();
|
||||
assert!(pruner.should_prune(2, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_pruned_trials() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
|
||||
// Only a pruned trial — no complete trials.
|
||||
let mut trial = trial_with_values(0, 0.1, (0..20).map(|s| (s, 0.1)).collect());
|
||||
trial.state = TrialState::Pruned;
|
||||
let completed = vec![trial];
|
||||
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 100.0)).collect();
|
||||
assert!(!pruner.should_prune(1, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lower_p_value_is_more_conservative() {
|
||||
let strict = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(0.001);
|
||||
let lenient = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(0.1);
|
||||
|
||||
let completed = vec![trial_with_values(
|
||||
0,
|
||||
0.1,
|
||||
(0..20).map(|s| (s, 0.1)).collect(),
|
||||
)];
|
||||
|
||||
// Moderately worse — should pass lenient but maybe not strict.
|
||||
let current: Vec<(u64, f64)> = (0..20)
|
||||
.map(|s| if s < 15 { (s, 0.2) } else { (s, 0.15) })
|
||||
.collect();
|
||||
|
||||
let lenient_prunes = lenient.should_prune(1, 19, ¤t, &completed);
|
||||
let strict_prunes = strict.should_prune(1, 19, ¤t, &completed);
|
||||
|
||||
// A stricter threshold should never prune when a lenient one doesn't.
|
||||
if !lenient_prunes {
|
||||
assert!(!strict_prunes);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "p_value_threshold must be in (0.0, 1.0)")]
|
||||
fn panics_on_zero_p_value() {
|
||||
let _ = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "p_value_threshold must be in (0.0, 1.0)")]
|
||||
fn panics_on_one_p_value() {
|
||||
let _ = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn correct_signed_rank_statistic() {
|
||||
// Known example: differences [1, 2, 3, 4, 5, 6] (all positive).
|
||||
// Ranks: 1, 2, 3, 4, 5, 6. W+ = 21, W- = 0.
|
||||
// For minimization (testing if positive = worse), W- = 0.
|
||||
// This should give a very small p-value.
|
||||
let diffs = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
|
||||
let p = wilcoxon_signed_rank_test(&diffs, Direction::Minimize);
|
||||
assert!(
|
||||
p < 0.05,
|
||||
"p-value {p} should be < 0.05 for all-positive diffs"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn symmetric_differences_not_significant() {
|
||||
// Balanced differences: half positive, half negative.
|
||||
let diffs = vec![1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0, -4.0];
|
||||
let p = wilcoxon_signed_rank_test(&diffs, Direction::Minimize);
|
||||
assert!(p > 0.05, "p-value {p} should be > 0.05 for symmetric diffs");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normal_cdf_known_values() {
|
||||
assert!((normal_cdf(0.0) - 0.5).abs() < 1e-6);
|
||||
assert!(normal_cdf(-10.0) < 1e-6);
|
||||
assert!((normal_cdf(10.0) - 1.0).abs() < 1e-6);
|
||||
assert!((normal_cdf(-1.96) - 0.025).abs() < 0.001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_intermediate_values() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
let completed = vec![trial_with_values(
|
||||
0,
|
||||
0.1,
|
||||
(0..20).map(|s| (s, 0.1)).collect(),
|
||||
)];
|
||||
assert!(!pruner.should_prune(1, 0, &[], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_completed_trials() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 1.0)).collect();
|
||||
assert!(!pruner.should_prune(1, 19, ¤t, &[]));
|
||||
}
|
||||
}
|
||||
@@ -1,691 +0,0 @@
|
||||
//! 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
+28
-151
@@ -1,21 +1,13 @@
|
||||
//! 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;
|
||||
|
||||
use crate::distribution::Distribution;
|
||||
use crate::param::ParamValue;
|
||||
use crate::parameter::{ParamId, Parameter};
|
||||
use crate::trial::AttrValue;
|
||||
use crate::types::TrialState;
|
||||
|
||||
/// A completed trial with its parameters, distributions, and objective value.
|
||||
///
|
||||
@@ -23,167 +15,30 @@ 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,
|
||||
/// The sampled parameter values, keyed by parameter id.
|
||||
pub params: HashMap<ParamId, ParamValue>,
|
||||
/// The parameter distributions used, keyed by parameter id.
|
||||
pub distributions: HashMap<ParamId, Distribution>,
|
||||
/// Human-readable labels for parameters, keyed by parameter id.
|
||||
pub param_labels: HashMap<ParamId, String>,
|
||||
/// The sampled parameter values, keyed by parameter name.
|
||||
pub params: HashMap<String, ParamValue>,
|
||||
/// The parameter distributions used, keyed by parameter name.
|
||||
pub distributions: HashMap<String, Distribution>,
|
||||
/// The objective value returned by the objective function.
|
||||
pub value: V,
|
||||
/// Intermediate objective values reported during the trial.
|
||||
pub intermediate_values: Vec<(u64, f64)>,
|
||||
/// The state of the trial (Complete, Pruned, or Failed).
|
||||
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> {
|
||||
/// Creates a new completed trial.
|
||||
pub fn new(
|
||||
id: u64,
|
||||
params: HashMap<ParamId, ParamValue>,
|
||||
distributions: HashMap<ParamId, Distribution>,
|
||||
param_labels: HashMap<ParamId, String>,
|
||||
params: HashMap<String, ParamValue>,
|
||||
distributions: HashMap<String, Distribution>,
|
||||
value: V,
|
||||
) -> Self {
|
||||
Self {
|
||||
id,
|
||||
params,
|
||||
distributions,
|
||||
param_labels,
|
||||
value,
|
||||
intermediate_values: Vec::new(),
|
||||
state: TrialState::Complete,
|
||||
user_attrs: HashMap::new(),
|
||||
constraints: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a new completed trial with intermediate values and user attributes.
|
||||
pub fn with_intermediate_values(
|
||||
id: u64,
|
||||
params: HashMap<ParamId, ParamValue>,
|
||||
distributions: HashMap<ParamId, Distribution>,
|
||||
param_labels: HashMap<ParamId, String>,
|
||||
value: V,
|
||||
intermediate_values: Vec<(u64, f64)>,
|
||||
user_attrs: HashMap<String, AttrValue>,
|
||||
) -> Self {
|
||||
Self {
|
||||
id,
|
||||
params,
|
||||
distributions,
|
||||
param_labels,
|
||||
value,
|
||||
intermediate_values,
|
||||
state: TrialState::Complete,
|
||||
user_attrs,
|
||||
constraints: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the typed value for the given parameter.
|
||||
///
|
||||
/// Looks up the parameter by its unique id and casts the stored
|
||||
/// [`ParamValue`] to the parameter's typed value.
|
||||
///
|
||||
/// Returns `None` if the parameter was not used in this trial.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if the stored value is incompatible with the parameter type
|
||||
/// (e.g., a `Float` value stored for an `IntParam`). This indicates
|
||||
/// a bug in the program, not a runtime error.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::parameter::{FloatParam, Parameter};
|
||||
/// use optimizer::{Direction, Study};
|
||||
///
|
||||
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
/// let x = FloatParam::new(-10.0, 10.0);
|
||||
///
|
||||
/// study
|
||||
/// .optimize(5, |trial| {
|
||||
/// let val = x.suggest(trial)?;
|
||||
/// Ok::<_, optimizer::Error>(val * val)
|
||||
/// })
|
||||
/// .unwrap();
|
||||
///
|
||||
/// let best = study.best_trial().unwrap();
|
||||
/// let x_val: f64 = best.get(&x).unwrap();
|
||||
/// assert!((-10.0..=10.0).contains(&x_val));
|
||||
/// ```
|
||||
pub fn get<P: Parameter>(&self, param: &P) -> Option<P::Value> {
|
||||
self.params.get(¶m.id()).map(|v| {
|
||||
param
|
||||
.cast_param_value(v)
|
||||
.expect("parameter type mismatch: stored value incompatible with parameter")
|
||||
})
|
||||
}
|
||||
|
||||
/// 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> {
|
||||
self.user_attrs.get(key)
|
||||
}
|
||||
|
||||
/// Returns all user attributes.
|
||||
#[must_use]
|
||||
pub fn user_attrs(&self) -> &HashMap<String, AttrValue> {
|
||||
&self.user_attrs
|
||||
}
|
||||
}
|
||||
|
||||
/// A pending (running) trial with its parameters and distributions, but no objective value yet.
|
||||
///
|
||||
/// This struct represents a trial that has been started and has sampled parameters,
|
||||
/// but is still running and hasn't returned an objective value. It is used with the
|
||||
/// constant liar strategy for parallel optimization.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PendingTrial {
|
||||
/// The unique identifier for this trial.
|
||||
pub id: u64,
|
||||
/// The sampled parameter values, keyed by parameter id.
|
||||
pub params: HashMap<ParamId, ParamValue>,
|
||||
/// The parameter distributions used, keyed by parameter id.
|
||||
pub distributions: HashMap<ParamId, Distribution>,
|
||||
/// Human-readable labels for parameters, keyed by parameter id.
|
||||
pub param_labels: HashMap<ParamId, String>,
|
||||
}
|
||||
|
||||
impl PendingTrial {
|
||||
/// Creates a new pending trial.
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
id: u64,
|
||||
params: HashMap<ParamId, ParamValue>,
|
||||
distributions: HashMap<ParamId, Distribution>,
|
||||
param_labels: HashMap<ParamId, String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
id,
|
||||
params,
|
||||
distributions,
|
||||
param_labels,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -193,6 +48,28 @@ impl PendingTrial {
|
||||
/// Samplers are responsible for generating parameter values based on
|
||||
/// the distribution and historical trial data. The trait requires
|
||||
/// `Send + Sync` to support concurrent and async optimization.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// Implementing a custom sampler:
|
||||
///
|
||||
/// ```ignore
|
||||
/// use optimizer::{Sampler, ParamValue, Distribution, CompletedTrial};
|
||||
///
|
||||
/// struct MySampler;
|
||||
///
|
||||
/// impl Sampler for MySampler {
|
||||
/// fn sample(
|
||||
/// &self,
|
||||
/// distribution: &Distribution,
|
||||
/// trial_id: u64,
|
||||
/// history: &[CompletedTrial],
|
||||
/// ) -> ParamValue {
|
||||
/// // Custom sampling logic here
|
||||
/// todo!()
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub trait Sampler: Send + Sync {
|
||||
/// Samples a parameter value from the given distribution.
|
||||
///
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
use parking_lot::Mutex;
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{RngExt, SeedableRng};
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
use crate::distribution::Distribution;
|
||||
use crate::param::ParamValue;
|
||||
@@ -34,7 +34,7 @@ impl RandomSampler {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
rng: Mutex::new(rand::make_rng()),
|
||||
rng: Mutex::new(StdRng::from_os_rng()),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,420 +0,0 @@
|
||||
//! 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"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -3,80 +3,17 @@
|
||||
//! TPE is a Bayesian optimization algorithm that models the objective function
|
||||
//! using two probability distributions: one for promising (good) parameter values
|
||||
//! and one for unpromising (bad) parameter values.
|
||||
//!
|
||||
//! # Gamma Strategies
|
||||
//!
|
||||
//! The gamma parameter controls what fraction of trials are considered "good".
|
||||
//! This module provides several built-in strategies via the [`GammaStrategy`] trait:
|
||||
//!
|
||||
//! - [`FixedGamma`]: Constant gamma value (default: 0.25)
|
||||
//! - [`LinearGamma`]: Linear interpolation between min and max based on trial count
|
||||
//! - [`SqrtGamma`]: Gamma decreases as 1/√n (similar to Optuna)
|
||||
//! - [`HyperoptGamma`]: Hyperopt-style adaptive gamma
|
||||
//!
|
||||
//! You can also implement your own strategy by implementing the [`GammaStrategy`] trait.
|
||||
//!
|
||||
//! # Examples
|
||||
//!
|
||||
//! Using a built-in gamma strategy:
|
||||
//!
|
||||
//! ```
|
||||
//! use optimizer::sampler::tpe::{SqrtGamma, TpeSampler};
|
||||
//!
|
||||
//! let sampler = TpeSampler::builder()
|
||||
//! .gamma_strategy(SqrtGamma::default())
|
||||
//! .build()
|
||||
//! .unwrap();
|
||||
//! ```
|
||||
//!
|
||||
//! Implementing a custom gamma strategy:
|
||||
//!
|
||||
//! ```
|
||||
//! use optimizer::sampler::tpe::{GammaStrategy, TpeSampler};
|
||||
//!
|
||||
//! #[derive(Debug, Clone)]
|
||||
//! struct MyGamma {
|
||||
//! base: f64,
|
||||
//! }
|
||||
//!
|
||||
//! impl GammaStrategy for MyGamma {
|
||||
//! fn gamma(&self, n_trials: usize) -> f64 {
|
||||
//! (self.base + 0.01 * n_trials as f64).min(0.5)
|
||||
//! }
|
||||
//!
|
||||
//! fn clone_box(&self) -> Box<dyn GammaStrategy> {
|
||||
//! Box::new(self.clone())
|
||||
//! }
|
||||
//! }
|
||||
//!
|
||||
//! let sampler = TpeSampler::builder()
|
||||
//! .gamma_strategy(MyGamma { base: 0.1 })
|
||||
//! .build()
|
||||
//! .unwrap();
|
||||
//! ```
|
||||
|
||||
use core::fmt::Debug;
|
||||
use std::sync::Arc;
|
||||
|
||||
use parking_lot::Mutex;
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{RngExt, SeedableRng};
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
use crate::distribution::Distribution;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::kde::KernelDensityEstimator;
|
||||
use crate::param::ParamValue;
|
||||
use crate::sampler::tpe::gamma::{FixedGamma, GammaStrategy};
|
||||
use crate::sampler::{CompletedTrial, Sampler};
|
||||
|
||||
// ============================================================================
|
||||
// Gamma Strategy Trait and Implementations
|
||||
// ============================================================================
|
||||
|
||||
// ============================================================================
|
||||
// TPE Sampler
|
||||
// ============================================================================
|
||||
|
||||
/// A Tree-Parzen Estimator (TPE) sampler for Bayesian optimization.
|
||||
///
|
||||
/// TPE works by splitting completed trials into two groups based on their
|
||||
@@ -88,42 +25,26 @@ use crate::sampler::{CompletedTrial, Sampler};
|
||||
/// During the startup phase (when fewer than `n_startup_trials` are completed),
|
||||
/// TPE falls back to random sampling to gather initial data.
|
||||
///
|
||||
/// # Gamma Strategies
|
||||
///
|
||||
/// The gamma quantile can be configured using different strategies via the
|
||||
/// [`GammaStrategy`] trait.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSampler;
|
||||
///
|
||||
/// // Create with default settings (FixedGamma at 0.25)
|
||||
/// // Create with default settings
|
||||
/// let sampler = TpeSampler::new();
|
||||
///
|
||||
/// // Create with custom settings using the builder
|
||||
/// let sampler = TpeSampler::builder()
|
||||
/// .gamma(0.15) // Shorthand for Fixednew(0.15)
|
||||
/// .gamma(0.15)
|
||||
/// .n_startup_trials(20)
|
||||
/// .n_ei_candidates(32)
|
||||
/// .seed(42)
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// ```
|
||||
///
|
||||
/// Using a different gamma strategy:
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::{SqrtGamma, TpeSampler};
|
||||
///
|
||||
/// let sampler = TpeSampler::builder()
|
||||
/// .gamma_strategy(SqrtGamma::default())
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// ```
|
||||
pub struct TpeSampler {
|
||||
/// Strategy for computing the gamma quantile.
|
||||
gamma_strategy: Arc<dyn GammaStrategy>,
|
||||
/// Fraction of trials to consider as "good" (gamma quantile).
|
||||
gamma: f64,
|
||||
/// Number of trials before TPE kicks in (uses random sampling before this).
|
||||
n_startup_trials: usize,
|
||||
/// Number of candidate samples to evaluate when selecting the next point.
|
||||
@@ -138,18 +59,18 @@ impl TpeSampler {
|
||||
/// Creates a new TPE sampler with default settings.
|
||||
///
|
||||
/// Default settings:
|
||||
/// - gamma strategy: [`FixedGamma`] with gamma = 0.25
|
||||
/// - gamma: 0.25 (top 25% of trials are considered "good")
|
||||
/// - `n_startup_trials`: 10 (random sampling for first 10 trials)
|
||||
/// - `n_ei_candidates`: 24 (evaluate 24 candidates per sample)
|
||||
/// - `kde_bandwidth`: None (uses Scott's rule for automatic bandwidth)
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
gamma_strategy: Arc::new(FixedGamma::default()),
|
||||
gamma: 0.25,
|
||||
n_startup_trials: 10,
|
||||
n_ei_candidates: 24,
|
||||
kde_bandwidth: None,
|
||||
rng: Mutex::new(rand::make_rng()),
|
||||
rng: Mutex::new(StdRng::from_os_rng()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -175,10 +96,6 @@ impl TpeSampler {
|
||||
|
||||
/// Creates a new TPE sampler with custom configuration.
|
||||
///
|
||||
/// This method uses a fixed gamma value. For more advanced gamma strategies,
|
||||
/// use [`TpeSampler::with_strategy`] or the builder pattern with
|
||||
/// [`TpeSamplerBuilder::gamma_strategy`].
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `gamma` - Fraction of trials to consider "good" (0.0 to 1.0).
|
||||
@@ -198,51 +115,9 @@ impl TpeSampler {
|
||||
kde_bandwidth: Option<f64>,
|
||||
seed: Option<u64>,
|
||||
) -> Result<Self> {
|
||||
let gamma_strategy = FixedGamma::new(gamma)?;
|
||||
Self::with_strategy(
|
||||
gamma_strategy,
|
||||
n_startup_trials,
|
||||
n_ei_candidates,
|
||||
kde_bandwidth,
|
||||
seed,
|
||||
)
|
||||
}
|
||||
|
||||
/// Creates a new TPE sampler with a custom gamma strategy.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `gamma_strategy` - The strategy for computing the gamma quantile.
|
||||
/// * `n_startup_trials` - Number of random trials before TPE sampling.
|
||||
/// * `n_ei_candidates` - Number of candidates to evaluate per sample.
|
||||
/// * `kde_bandwidth` - Optional fixed bandwidth for KDE. If None, uses Scott's rule.
|
||||
/// * `seed` - Optional seed for reproducibility.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::{SqrtGamma, TpeSampler};
|
||||
///
|
||||
/// let sampler = TpeSampler::with_strategy(
|
||||
/// SqrtGamma::default(),
|
||||
/// 10, // n_startup_trials
|
||||
/// 24, // n_ei_candidates
|
||||
/// None, // kde_bandwidth
|
||||
/// Some(42), // seed
|
||||
/// )
|
||||
/// .unwrap();
|
||||
/// ```
|
||||
pub fn with_strategy<G: GammaStrategy + 'static>(
|
||||
gamma_strategy: G,
|
||||
n_startup_trials: usize,
|
||||
n_ei_candidates: usize,
|
||||
kde_bandwidth: Option<f64>,
|
||||
seed: Option<u64>,
|
||||
) -> Result<Self> {
|
||||
if gamma <= 0.0 || gamma >= 1.0 {
|
||||
return Err(Error::InvalidGamma(gamma));
|
||||
}
|
||||
if let Some(bw) = kde_bandwidth
|
||||
&& bw <= 0.0
|
||||
{
|
||||
@@ -251,11 +126,11 @@ impl TpeSampler {
|
||||
|
||||
let rng = match seed {
|
||||
Some(s) => StdRng::seed_from_u64(s),
|
||||
None => rand::make_rng(),
|
||||
None => StdRng::from_os_rng(),
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
gamma_strategy: Arc::new(gamma_strategy),
|
||||
gamma,
|
||||
n_startup_trials,
|
||||
n_ei_candidates,
|
||||
kde_bandwidth,
|
||||
@@ -263,16 +138,8 @@ impl TpeSampler {
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the gamma strategy used by this sampler.
|
||||
#[must_use]
|
||||
pub fn gamma_strategy(&self) -> &dyn GammaStrategy {
|
||||
self.gamma_strategy.as_ref()
|
||||
}
|
||||
|
||||
/// Splits trials into good and bad groups based on the gamma quantile.
|
||||
///
|
||||
/// The gamma value is computed dynamically using the configured [`GammaStrategy`].
|
||||
///
|
||||
/// Returns (`good_trials`, `bad_trials`) where `good_trials` contains trials
|
||||
/// with values below the gamma quantile (for minimization).
|
||||
#[allow(
|
||||
@@ -280,7 +147,6 @@ impl TpeSampler {
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
#[must_use]
|
||||
fn split_trials<'a>(
|
||||
&self,
|
||||
history: &'a [CompletedTrial],
|
||||
@@ -298,15 +164,9 @@ impl TpeSampler {
|
||||
.unwrap_or(core::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
// Compute gamma using the strategy and clamp to valid range
|
||||
let gamma = self
|
||||
.gamma_strategy
|
||||
.gamma(history.len())
|
||||
.clamp(f64::EPSILON, 1.0 - f64::EPSILON);
|
||||
|
||||
// Calculate the split point (gamma quantile)
|
||||
// Ensure at least 1 trial in each group if possible
|
||||
let n_good = ((history.len() as f64 * gamma).ceil() as usize)
|
||||
let n_good = ((history.len() as f64 * self.gamma).ceil() as usize)
|
||||
.max(1)
|
||||
.min(history.len() - 1);
|
||||
|
||||
@@ -561,8 +421,6 @@ impl Default for TpeSampler {
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// Using a fixed gamma value:
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
||||
///
|
||||
@@ -574,23 +432,9 @@ impl Default for TpeSampler {
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// ```
|
||||
///
|
||||
/// Using a custom gamma strategy:
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::{SqrtGamma, TpeSamplerBuilder};
|
||||
///
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .gamma_strategy(SqrtGamma::default())
|
||||
/// .n_startup_trials(20)
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// ```
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TpeSamplerBuilder {
|
||||
gamma_strategy: Box<dyn GammaStrategy>,
|
||||
/// Raw gamma value for deferred validation (Some if `gamma()` was called)
|
||||
raw_gamma: Option<f64>,
|
||||
gamma: f64,
|
||||
n_startup_trials: usize,
|
||||
n_ei_candidates: usize,
|
||||
kde_bandwidth: Option<f64>,
|
||||
@@ -601,7 +445,7 @@ impl TpeSamplerBuilder {
|
||||
/// Creates a new builder with default settings.
|
||||
///
|
||||
/// Default settings:
|
||||
/// - gamma strategy: [`FixedGamma`] with gamma = 0.25
|
||||
/// - gamma: 0.25 (top 25% of trials are considered "good")
|
||||
/// - `n_startup_trials`: 10 (random sampling for first 10 trials)
|
||||
/// - `n_ei_candidates`: 24 (evaluate 24 candidates per sample)
|
||||
/// - `kde_bandwidth`: None (uses Scott's rule for automatic bandwidth)
|
||||
@@ -609,8 +453,7 @@ impl TpeSamplerBuilder {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
gamma_strategy: Box::new(FixedGamma::default()),
|
||||
raw_gamma: None,
|
||||
gamma: 0.25,
|
||||
n_startup_trials: 10,
|
||||
n_ei_candidates: 24,
|
||||
kde_bandwidth: None,
|
||||
@@ -618,10 +461,7 @@ impl TpeSamplerBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets a fixed gamma value for splitting trials into good/bad groups.
|
||||
///
|
||||
/// This is a convenience method that creates a [`FixedGamma`] strategy.
|
||||
/// For more advanced gamma strategies, use [`gamma_strategy`](Self::gamma_strategy).
|
||||
/// Sets the gamma quantile for splitting trials into good/bad groups.
|
||||
///
|
||||
/// A gamma of 0.25 means the top 25% of trials (by objective value) are
|
||||
/// considered "good" and used to build the l(x) distribution.
|
||||
@@ -647,68 +487,7 @@ impl TpeSamplerBuilder {
|
||||
/// `build()` will return `Err(Error::InvalidGamma)`.
|
||||
#[must_use]
|
||||
pub fn gamma(mut self, gamma: f64) -> Self {
|
||||
// We defer validation to build() time for consistency with the existing API
|
||||
// Store the raw value for validation later
|
||||
self.raw_gamma = Some(gamma);
|
||||
self
|
||||
}
|
||||
|
||||
/// Sets a custom gamma strategy for splitting trials into good/bad groups.
|
||||
///
|
||||
/// The gamma strategy determines what fraction of trials are considered
|
||||
/// "good" based on the number of completed trials. This allows the gamma
|
||||
/// value to adapt dynamically during optimization.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `strategy` - A type implementing [`GammaStrategy`].
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// Using built-in strategies:
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::{LinearGamma, SqrtGamma, TpeSamplerBuilder};
|
||||
///
|
||||
/// // Square root strategy (Optuna-style)
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .gamma_strategy(SqrtGamma::default())
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
///
|
||||
/// // Linear interpolation strategy
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .gamma_strategy(LinearGamma::new(0.1, 0.3, 50).unwrap())
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// ```
|
||||
///
|
||||
/// Using a custom strategy:
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::{GammaStrategy, TpeSamplerBuilder};
|
||||
///
|
||||
/// #[derive(Debug, Clone)]
|
||||
/// struct MyGamma;
|
||||
///
|
||||
/// impl GammaStrategy for MyGamma {
|
||||
/// fn gamma(&self, n_trials: usize) -> f64 {
|
||||
/// 0.25 // Always return 0.25
|
||||
/// }
|
||||
/// fn clone_box(&self) -> Box<dyn GammaStrategy> {
|
||||
/// Box::new(self.clone())
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .gamma_strategy(MyGamma)
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn gamma_strategy<G: GammaStrategy + 'static>(mut self, strategy: G) -> Self {
|
||||
self.gamma_strategy = Box::new(strategy);
|
||||
self.raw_gamma = None; // Clear any raw gamma set by gamma()
|
||||
self.gamma = gamma;
|
||||
self
|
||||
}
|
||||
|
||||
@@ -823,7 +602,7 @@ impl TpeSamplerBuilder {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::InvalidGamma` if a fixed gamma value was set and is not in (0.0, 1.0).
|
||||
/// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0).
|
||||
/// Returns `Error::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
|
||||
///
|
||||
/// # Examples
|
||||
@@ -840,33 +619,13 @@ impl TpeSamplerBuilder {
|
||||
/// .unwrap();
|
||||
/// ```
|
||||
pub fn build(self) -> Result<TpeSampler> {
|
||||
// Determine the gamma strategy to use
|
||||
let gamma_strategy: Arc<dyn GammaStrategy> = if let Some(raw) = self.raw_gamma {
|
||||
// Validate and create FixedGamma from raw value
|
||||
Arc::new(FixedGamma::new(raw)?)
|
||||
} else {
|
||||
Arc::from(self.gamma_strategy)
|
||||
};
|
||||
|
||||
// Validate bandwidth
|
||||
if let Some(bw) = self.kde_bandwidth
|
||||
&& bw <= 0.0
|
||||
{
|
||||
return Err(Error::InvalidBandwidth(bw));
|
||||
}
|
||||
|
||||
let rng = match self.seed {
|
||||
Some(s) => StdRng::seed_from_u64(s),
|
||||
None => rand::make_rng(),
|
||||
};
|
||||
|
||||
Ok(TpeSampler {
|
||||
gamma_strategy,
|
||||
n_startup_trials: self.n_startup_trials,
|
||||
n_ei_candidates: self.n_ei_candidates,
|
||||
kde_bandwidth: self.kde_bandwidth,
|
||||
rng: Mutex::new(rng),
|
||||
})
|
||||
TpeSampler::with_config(
|
||||
self.gamma,
|
||||
self.n_startup_trials,
|
||||
self.n_ei_candidates,
|
||||
self.kde_bandwidth,
|
||||
self.seed,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1024,27 +783,25 @@ mod tests {
|
||||
|
||||
use super::*;
|
||||
use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};
|
||||
use crate::parameter::ParamId;
|
||||
|
||||
fn create_trial(
|
||||
id: u64,
|
||||
value: f64,
|
||||
params: Vec<(ParamId, ParamValue, Distribution)>,
|
||||
params: Vec<(&str, ParamValue, Distribution)>,
|
||||
) -> 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);
|
||||
for (name, pv, dist) in params {
|
||||
param_map.insert(name.to_string(), pv);
|
||||
dist_map.insert(name.to_string(), dist);
|
||||
}
|
||||
CompletedTrial::new(id, param_map, dist_map, HashMap::new(), value)
|
||||
CompletedTrial::new(id, param_map, dist_map, value)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_sampler_new() {
|
||||
let sampler = TpeSampler::new();
|
||||
// Default uses FixedGamma with 0.25
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma - 0.25).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 10);
|
||||
assert_eq!(sampler.n_ei_candidates, 24);
|
||||
}
|
||||
@@ -1052,8 +809,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_tpe_sampler_with_config() {
|
||||
let sampler = TpeSampler::with_config(0.15, 20, 32, None, Some(42)).unwrap();
|
||||
// with_config uses FixedGamma
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma - 0.15).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 20);
|
||||
assert_eq!(sampler.n_ei_candidates, 32);
|
||||
}
|
||||
@@ -1105,13 +861,12 @@ mod tests {
|
||||
});
|
||||
|
||||
// Create 20 trials with values 0..20
|
||||
let x_id = ParamId::new();
|
||||
let history: Vec<CompletedTrial> = (0..20)
|
||||
.map(|i| {
|
||||
create_trial(
|
||||
i as u64,
|
||||
f64::from(i),
|
||||
vec![(x_id, ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
||||
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
@@ -1140,7 +895,6 @@ mod tests {
|
||||
});
|
||||
|
||||
// Create history where low values (near 0.2) are "good"
|
||||
let x_id = ParamId::new();
|
||||
let history: Vec<CompletedTrial> = (0..20)
|
||||
.map(|i| {
|
||||
let x = f64::from(i) / 20.0;
|
||||
@@ -1149,7 +903,7 @@ mod tests {
|
||||
create_trial(
|
||||
i as u64,
|
||||
value,
|
||||
vec![(x_id, ParamValue::Float(x), dist.clone())],
|
||||
vec![("x", ParamValue::Float(x), dist.clone())],
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
@@ -1178,7 +932,6 @@ mod tests {
|
||||
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 4 });
|
||||
|
||||
// Create history where category 1 is consistently good
|
||||
let cat_id = ParamId::new();
|
||||
let history: Vec<CompletedTrial> = (0..20)
|
||||
.map(|i| {
|
||||
let category = i % 4;
|
||||
@@ -1188,7 +941,7 @@ mod tests {
|
||||
i as u64,
|
||||
value,
|
||||
vec![(
|
||||
cat_id,
|
||||
"cat",
|
||||
ParamValue::Categorical(category as usize),
|
||||
dist.clone(),
|
||||
)],
|
||||
@@ -1224,7 +977,6 @@ mod tests {
|
||||
});
|
||||
|
||||
// Create history where values near 30 are good
|
||||
let x_id = ParamId::new();
|
||||
let history: Vec<CompletedTrial> = (0..20)
|
||||
.map(|i| {
|
||||
let x = i * 5; // 0, 5, 10, ..., 95
|
||||
@@ -1232,7 +984,7 @@ mod tests {
|
||||
create_trial(
|
||||
i as u64,
|
||||
value,
|
||||
vec![(x_id, ParamValue::Int(x), dist.clone())],
|
||||
vec![("x", ParamValue::Int(x), dist.clone())],
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
@@ -1257,13 +1009,12 @@ mod tests {
|
||||
step: None,
|
||||
});
|
||||
|
||||
let x_id = ParamId::new();
|
||||
let history: Vec<CompletedTrial> = (0..20)
|
||||
.map(|i| {
|
||||
create_trial(
|
||||
i as u64,
|
||||
f64::from(i),
|
||||
vec![(x_id, ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
||||
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
@@ -1282,7 +1033,7 @@ mod tests {
|
||||
fn test_tpe_sampler_builder_default() {
|
||||
let builder = TpeSamplerBuilder::new();
|
||||
let sampler = builder.build().unwrap();
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma - 0.25).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 10);
|
||||
assert_eq!(sampler.n_ei_candidates, 24);
|
||||
}
|
||||
@@ -1296,7 +1047,7 @@ mod tests {
|
||||
.seed(42)
|
||||
.build()
|
||||
.unwrap();
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma - 0.15).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 20);
|
||||
assert_eq!(sampler.n_ei_candidates, 32);
|
||||
}
|
||||
@@ -1309,7 +1060,7 @@ mod tests {
|
||||
.n_ei_candidates(48)
|
||||
.build()
|
||||
.unwrap();
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.10).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma - 0.10).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 15);
|
||||
assert_eq!(sampler.n_ei_candidates, 48);
|
||||
}
|
||||
@@ -1318,7 +1069,7 @@ mod tests {
|
||||
fn test_tpe_sampler_builder_partial() {
|
||||
// Test setting only some options
|
||||
let sampler = TpeSamplerBuilder::new().gamma(0.20).build().unwrap();
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.20).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma - 0.20).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 10); // default
|
||||
assert_eq!(sampler.n_ei_candidates, 24); // default
|
||||
}
|
||||
@@ -1338,13 +1089,12 @@ mod tests {
|
||||
step: None,
|
||||
});
|
||||
|
||||
let x_id = ParamId::new();
|
||||
let history: Vec<CompletedTrial> = (0..20u32)
|
||||
.map(|i| {
|
||||
create_trial(
|
||||
u64::from(i),
|
||||
f64::from(i),
|
||||
vec![(x_id, ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
||||
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
@@ -1,670 +0,0 @@
|
||||
use core::fmt::Debug;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
/// A strategy for computing the gamma quantile in TPE.
|
||||
///
|
||||
/// The gamma value determines what fraction of trials are considered "good"
|
||||
/// when splitting the trial history. Different strategies can adapt this
|
||||
/// fraction based on the number of completed trials.
|
||||
///
|
||||
/// # Implementation Notes
|
||||
///
|
||||
/// - The returned gamma must be in the range (0.0, 1.0)
|
||||
/// - Implementations should be deterministic for reproducibility
|
||||
/// - The `clone_box` method enables trait object cloning
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::GammaStrategy;
|
||||
///
|
||||
/// #[derive(Debug, Clone)]
|
||||
/// struct ConstantGamma(f64);
|
||||
///
|
||||
/// impl GammaStrategy for ConstantGamma {
|
||||
/// fn gamma(&self, _n_trials: usize) -> f64 {
|
||||
/// self.0
|
||||
/// }
|
||||
///
|
||||
/// fn clone_box(&self) -> Box<dyn GammaStrategy> {
|
||||
/// Box::new(self.clone())
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub trait GammaStrategy: Send + Sync + Debug {
|
||||
/// Computes the gamma quantile based on the number of completed trials.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `n_trials` - The number of completed trials in the history.
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// A gamma value in the range (0.0, 1.0). Values outside this range
|
||||
/// will be clamped by the sampler.
|
||||
fn gamma(&self, n_trials: usize) -> f64;
|
||||
|
||||
/// Creates a boxed clone of this strategy.
|
||||
///
|
||||
/// This method enables cloning of trait objects, which is necessary
|
||||
/// for the builder pattern and sampler configuration.
|
||||
fn clone_box(&self) -> Box<dyn GammaStrategy>;
|
||||
}
|
||||
|
||||
impl Clone for Box<dyn GammaStrategy> {
|
||||
fn clone(&self) -> Self {
|
||||
self.clone_box()
|
||||
}
|
||||
}
|
||||
|
||||
/// A fixed gamma strategy that returns a constant value.
|
||||
///
|
||||
/// This is the simplest strategy and the default behavior of TPE.
|
||||
/// The gamma value remains constant regardless of the number of trials.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::{FixedGamma, TpeSampler};
|
||||
///
|
||||
/// // Use 15% of trials as "good"
|
||||
/// let sampler = TpeSampler::builder()
|
||||
/// .gamma_strategy(FixedGamma::new(0.15).unwrap())
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct FixedGamma {
|
||||
gamma: f64,
|
||||
}
|
||||
|
||||
impl FixedGamma {
|
||||
/// Creates a new fixed gamma strategy.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `gamma` - The constant gamma value to use.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0).
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::FixedGamma;
|
||||
///
|
||||
/// let strategy = FixedGamma::new(0.25).unwrap();
|
||||
/// assert!((strategy.value() - 0.25).abs() < f64::EPSILON);
|
||||
/// ```
|
||||
pub fn new(gamma: f64) -> crate::Result<Self> {
|
||||
if gamma <= 0.0 || gamma >= 1.0 {
|
||||
return Err(Error::InvalidGamma(gamma));
|
||||
}
|
||||
Ok(Self { gamma })
|
||||
}
|
||||
|
||||
/// Returns the fixed gamma value.
|
||||
#[must_use]
|
||||
pub fn value(&self) -> f64 {
|
||||
self.gamma
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FixedGamma {
|
||||
/// Creates a fixed gamma strategy with the default value of 0.25.
|
||||
fn default() -> Self {
|
||||
Self { gamma: 0.25 }
|
||||
}
|
||||
}
|
||||
|
||||
impl GammaStrategy for FixedGamma {
|
||||
fn gamma(&self, _n_trials: usize) -> f64 {
|
||||
self.gamma
|
||||
}
|
||||
|
||||
fn clone_box(&self) -> Box<dyn GammaStrategy> {
|
||||
Box::new(*self)
|
||||
}
|
||||
}
|
||||
|
||||
/// A linear gamma strategy that interpolates between min and max values.
|
||||
///
|
||||
/// The gamma value increases linearly from `gamma_min` to `gamma_max` as the
|
||||
/// number of trials grows from 0 to `n_trials_max`. Beyond `n_trials_max`,
|
||||
/// gamma remains at `gamma_max`.
|
||||
///
|
||||
/// This strategy is useful when you want to be more explorative early on
|
||||
/// (smaller gamma = fewer "good" trials) and more exploitative later
|
||||
/// (larger gamma = more "good" trials).
|
||||
///
|
||||
/// # Formula
|
||||
///
|
||||
/// ```text
|
||||
/// gamma = gamma_min + (gamma_max - gamma_min) * min(n_trials / n_trials_max, 1.0)
|
||||
/// ```
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::{GammaStrategy, LinearGamma, TpeSampler};
|
||||
///
|
||||
/// let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap();
|
||||
///
|
||||
/// // At 0 trials: gamma = 0.1
|
||||
/// assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON);
|
||||
///
|
||||
/// // At 50 trials: gamma = 0.25 (midpoint)
|
||||
/// assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON);
|
||||
///
|
||||
/// // At 100+ trials: gamma = 0.4
|
||||
/// assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON);
|
||||
/// assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON);
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct LinearGamma {
|
||||
gamma_min: f64,
|
||||
gamma_max: f64,
|
||||
n_trials_max: usize,
|
||||
}
|
||||
|
||||
impl LinearGamma {
|
||||
/// Creates a new linear gamma strategy.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `gamma_min` - The minimum gamma value (at 0 trials).
|
||||
/// * `gamma_max` - The maximum gamma value (at `n_trials_max` trials).
|
||||
/// * `n_trials_max` - The number of trials at which gamma reaches its maximum.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::InvalidGamma` if:
|
||||
/// - `gamma_min` is not in (0.0, 1.0)
|
||||
/// - `gamma_max` is not in (0.0, 1.0)
|
||||
/// - `gamma_min > gamma_max`
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::LinearGamma;
|
||||
///
|
||||
/// // Gamma goes from 0.1 to 0.3 over 50 trials
|
||||
/// let strategy = LinearGamma::new(0.1, 0.3, 50).unwrap();
|
||||
/// ```
|
||||
pub fn new(gamma_min: f64, gamma_max: f64, n_trials_max: usize) -> crate::Result<Self> {
|
||||
if gamma_min <= 0.0 || gamma_min >= 1.0 {
|
||||
return Err(Error::InvalidGamma(gamma_min));
|
||||
}
|
||||
if gamma_max <= 0.0 || gamma_max >= 1.0 {
|
||||
return Err(Error::InvalidGamma(gamma_max));
|
||||
}
|
||||
if gamma_min > gamma_max {
|
||||
return Err(Error::InvalidGamma(gamma_min));
|
||||
}
|
||||
Ok(Self {
|
||||
gamma_min,
|
||||
gamma_max,
|
||||
n_trials_max,
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the minimum gamma value.
|
||||
#[must_use]
|
||||
pub fn gamma_min(&self) -> f64 {
|
||||
self.gamma_min
|
||||
}
|
||||
|
||||
/// Returns the maximum gamma value.
|
||||
#[must_use]
|
||||
pub fn gamma_max(&self) -> f64 {
|
||||
self.gamma_max
|
||||
}
|
||||
|
||||
/// Returns the number of trials at which gamma reaches its maximum.
|
||||
#[must_use]
|
||||
pub fn n_trials_max(&self) -> usize {
|
||||
self.n_trials_max
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for LinearGamma {
|
||||
/// Creates a linear gamma strategy with default values:
|
||||
/// - `gamma_min`: 0.10
|
||||
/// - `gamma_max`: 0.25
|
||||
/// - `n_trials_max`: 100
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
gamma_min: 0.10,
|
||||
gamma_max: 0.25,
|
||||
n_trials_max: 100,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl GammaStrategy for LinearGamma {
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
fn gamma(&self, n_trials: usize) -> f64 {
|
||||
if self.n_trials_max == 0 {
|
||||
return self.gamma_max;
|
||||
}
|
||||
let t = (n_trials as f64 / self.n_trials_max as f64).min(1.0);
|
||||
self.gamma_min + (self.gamma_max - self.gamma_min) * t
|
||||
}
|
||||
|
||||
fn clone_box(&self) -> Box<dyn GammaStrategy> {
|
||||
Box::new(*self)
|
||||
}
|
||||
}
|
||||
|
||||
/// A square root gamma strategy inspired by Optuna's default behavior.
|
||||
///
|
||||
/// The gamma value is computed based on the inverse square root of the number
|
||||
/// of trials, providing a balance between exploration and exploitation that
|
||||
/// naturally adapts as more data becomes available.
|
||||
///
|
||||
/// # Formula
|
||||
///
|
||||
/// ```text
|
||||
/// n_good = max(1, floor(gamma_factor / sqrt(n_trials)))
|
||||
/// gamma = min(gamma_max, n_good / n_trials)
|
||||
/// ```
|
||||
///
|
||||
/// When `n_trials` is 0, returns `gamma_max`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::{GammaStrategy, SqrtGamma, TpeSampler};
|
||||
///
|
||||
/// let strategy = SqrtGamma::default();
|
||||
///
|
||||
/// // Gamma decreases as trials increase
|
||||
/// let g10 = strategy.gamma(10);
|
||||
/// let g100 = strategy.gamma(100);
|
||||
/// assert!(g10 > g100, "Gamma should decrease with more trials");
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct SqrtGamma {
|
||||
gamma_factor: f64,
|
||||
gamma_max: f64,
|
||||
}
|
||||
|
||||
impl SqrtGamma {
|
||||
/// Creates a new square root gamma strategy.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `gamma_factor` - The factor controlling how quickly gamma decreases.
|
||||
/// Higher values mean more "good" trials at any given point.
|
||||
/// * `gamma_max` - The maximum gamma value (used when `n_trials` is small).
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::InvalidGamma` if:
|
||||
/// - `gamma_factor` is not positive
|
||||
/// - `gamma_max` is not in (0.0, 1.0)
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::SqrtGamma;
|
||||
///
|
||||
/// let strategy = SqrtGamma::new(1.0, 0.25).unwrap();
|
||||
/// ```
|
||||
pub fn new(gamma_factor: f64, gamma_max: f64) -> crate::Result<Self> {
|
||||
if gamma_factor <= 0.0 {
|
||||
return Err(Error::InvalidGamma(gamma_factor));
|
||||
}
|
||||
if gamma_max <= 0.0 || gamma_max >= 1.0 {
|
||||
return Err(Error::InvalidGamma(gamma_max));
|
||||
}
|
||||
Ok(Self {
|
||||
gamma_factor,
|
||||
gamma_max,
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the gamma factor.
|
||||
#[must_use]
|
||||
pub fn gamma_factor(&self) -> f64 {
|
||||
self.gamma_factor
|
||||
}
|
||||
|
||||
/// Returns the maximum gamma value.
|
||||
#[must_use]
|
||||
pub fn gamma_max(&self) -> f64 {
|
||||
self.gamma_max
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SqrtGamma {
|
||||
/// Creates a square root gamma strategy with default values:
|
||||
/// - `gamma_factor`: 1.0
|
||||
/// - `gamma_max`: 0.25
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
gamma_factor: 1.0,
|
||||
gamma_max: 0.25,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl GammaStrategy for SqrtGamma {
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
fn gamma(&self, n_trials: usize) -> f64 {
|
||||
if n_trials == 0 {
|
||||
return self.gamma_max;
|
||||
}
|
||||
let n_good = (self.gamma_factor / (n_trials as f64).sqrt()).max(1.0);
|
||||
(n_good / n_trials as f64).min(self.gamma_max)
|
||||
}
|
||||
|
||||
fn clone_box(&self) -> Box<dyn GammaStrategy> {
|
||||
Box::new(*self)
|
||||
}
|
||||
}
|
||||
|
||||
/// A Hyperopt-style gamma strategy.
|
||||
///
|
||||
/// This strategy computes gamma as `min(gamma_max, (gamma_base + 1) / n_trials)`,
|
||||
/// which is inspired by the original Hyperopt TPE implementation.
|
||||
///
|
||||
/// # Formula
|
||||
///
|
||||
/// ```text
|
||||
/// gamma = min(gamma_max, (gamma_base + 1) / n_trials)
|
||||
/// ```
|
||||
///
|
||||
/// When `n_trials` is 0, returns `gamma_max`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::{GammaStrategy, HyperoptGamma};
|
||||
///
|
||||
/// // With gamma_base=24 and gamma_max=0.5:
|
||||
/// // - At n=25: gamma = min(0.5, 25/25) = 0.5 (capped)
|
||||
/// // - At n=100: gamma = min(0.5, 25/100) = 0.25
|
||||
/// let strategy = HyperoptGamma::new(24.0, 0.5).unwrap();
|
||||
///
|
||||
/// // Early trials have higher gamma
|
||||
/// let g50 = strategy.gamma(50);
|
||||
/// let g200 = strategy.gamma(200);
|
||||
/// assert!(g50 > g200, "Gamma should decrease with more trials");
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct HyperoptGamma {
|
||||
gamma_base: f64,
|
||||
gamma_max: f64,
|
||||
}
|
||||
|
||||
impl HyperoptGamma {
|
||||
/// Creates a new Hyperopt-style gamma strategy.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `gamma_base` - The base value added to 1 in the numerator.
|
||||
/// * `gamma_max` - The maximum gamma value.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::InvalidGamma` if:
|
||||
/// - `gamma_base` is negative
|
||||
/// - `gamma_max` is not in (0.0, 1.0)
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::HyperoptGamma;
|
||||
///
|
||||
/// let strategy = HyperoptGamma::new(24.0, 0.25).unwrap();
|
||||
/// ```
|
||||
pub fn new(gamma_base: f64, gamma_max: f64) -> crate::Result<Self> {
|
||||
if gamma_base < 0.0 {
|
||||
return Err(Error::InvalidGamma(gamma_base));
|
||||
}
|
||||
if gamma_max <= 0.0 || gamma_max >= 1.0 {
|
||||
return Err(Error::InvalidGamma(gamma_max));
|
||||
}
|
||||
Ok(Self {
|
||||
gamma_base,
|
||||
gamma_max,
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the gamma base value.
|
||||
#[must_use]
|
||||
pub fn gamma_base(&self) -> f64 {
|
||||
self.gamma_base
|
||||
}
|
||||
|
||||
/// Returns the maximum gamma value.
|
||||
#[must_use]
|
||||
pub fn gamma_max(&self) -> f64 {
|
||||
self.gamma_max
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for HyperoptGamma {
|
||||
/// Creates a Hyperopt-style gamma strategy with default values:
|
||||
/// - `gamma_base`: 24.0
|
||||
/// - `gamma_max`: 0.25
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
gamma_base: 24.0,
|
||||
gamma_max: 0.25,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl GammaStrategy for HyperoptGamma {
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
fn gamma(&self, n_trials: usize) -> f64 {
|
||||
if n_trials == 0 {
|
||||
return self.gamma_max;
|
||||
}
|
||||
((self.gamma_base + 1.0) / n_trials as f64).min(self.gamma_max)
|
||||
}
|
||||
|
||||
fn clone_box(&self) -> Box<dyn GammaStrategy> {
|
||||
Box::new(*self)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::sampler::tpe::TpeSampler;
|
||||
|
||||
#[test]
|
||||
fn test_fixed_gamma_default() {
|
||||
let strategy = FixedGamma::default();
|
||||
assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON);
|
||||
assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON);
|
||||
assert!((strategy.value() - 0.25).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fixed_gamma_custom() {
|
||||
let strategy = FixedGamma::new(0.15).unwrap();
|
||||
assert!((strategy.gamma(0) - 0.15).abs() < f64::EPSILON);
|
||||
assert!((strategy.gamma(50) - 0.15).abs() < f64::EPSILON);
|
||||
assert!((strategy.gamma(1000) - 0.15).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fixed_gamma_invalid() {
|
||||
assert!(FixedGamma::new(0.0).is_err());
|
||||
assert!(FixedGamma::new(1.0).is_err());
|
||||
assert!(FixedGamma::new(-0.1).is_err());
|
||||
assert!(FixedGamma::new(1.5).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_linear_gamma_default() {
|
||||
let strategy = LinearGamma::default();
|
||||
assert!((strategy.gamma(0) - 0.10).abs() < f64::EPSILON);
|
||||
assert!((strategy.gamma(50) - 0.175).abs() < f64::EPSILON); // midpoint
|
||||
assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON);
|
||||
assert!((strategy.gamma(200) - 0.25).abs() < f64::EPSILON); // capped
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_linear_gamma_custom() {
|
||||
let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap();
|
||||
assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON);
|
||||
assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON);
|
||||
assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON);
|
||||
assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_linear_gamma_invalid() {
|
||||
assert!(LinearGamma::new(0.0, 0.5, 100).is_err());
|
||||
assert!(LinearGamma::new(0.1, 1.0, 100).is_err());
|
||||
assert!(LinearGamma::new(0.5, 0.2, 100).is_err()); // min > max
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sqrt_gamma_default() {
|
||||
let strategy = SqrtGamma::default();
|
||||
// At n=0, returns gamma_max
|
||||
assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON);
|
||||
|
||||
// gamma decreases with more trials
|
||||
let g10 = strategy.gamma(10);
|
||||
let g100 = strategy.gamma(100);
|
||||
assert!(g10 > g100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sqrt_gamma_custom() {
|
||||
let strategy = SqrtGamma::new(2.0, 0.5).unwrap();
|
||||
assert!((strategy.gamma(0) - 0.5).abs() < f64::EPSILON);
|
||||
|
||||
// At n=4: n_good = max(1, 2/2) = 1, gamma = 1/4 = 0.25
|
||||
let g4 = strategy.gamma(4);
|
||||
assert!((g4 - 0.25).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sqrt_gamma_invalid() {
|
||||
assert!(SqrtGamma::new(0.0, 0.25).is_err()); // factor must be positive
|
||||
assert!(SqrtGamma::new(-1.0, 0.25).is_err());
|
||||
assert!(SqrtGamma::new(1.0, 0.0).is_err());
|
||||
assert!(SqrtGamma::new(1.0, 1.0).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hyperopt_gamma_default() {
|
||||
let strategy = HyperoptGamma::default();
|
||||
// At n=0, returns gamma_max
|
||||
assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON);
|
||||
|
||||
// At n=100: (24+1)/100 = 0.25, so capped to 0.25
|
||||
assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON);
|
||||
|
||||
// At n=200: (24+1)/200 = 0.125
|
||||
assert!((strategy.gamma(200) - 0.125).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hyperopt_gamma_custom() {
|
||||
let strategy = HyperoptGamma::new(9.0, 0.5).unwrap();
|
||||
// At n=20: (9+1)/20 = 0.5, capped to 0.5
|
||||
assert!((strategy.gamma(20) - 0.5).abs() < f64::EPSILON);
|
||||
|
||||
// At n=100: (9+1)/100 = 0.1
|
||||
assert!((strategy.gamma(100) - 0.1).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hyperopt_gamma_invalid() {
|
||||
assert!(HyperoptGamma::new(-1.0, 0.25).is_err());
|
||||
assert!(HyperoptGamma::new(24.0, 0.0).is_err());
|
||||
assert!(HyperoptGamma::new(24.0, 1.0).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gamma_strategy_clone_box() {
|
||||
let fixed: Box<dyn GammaStrategy> = Box::new(FixedGamma::new(0.3).unwrap());
|
||||
let cloned = fixed.clone();
|
||||
assert!((cloned.gamma(0) - 0.3).abs() < f64::EPSILON);
|
||||
|
||||
let linear: Box<dyn GammaStrategy> = Box::new(LinearGamma::default());
|
||||
let cloned = linear.clone();
|
||||
assert!((cloned.gamma(0) - 0.10).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_with_linear_gamma_strategy() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.gamma_strategy(LinearGamma::new(0.1, 0.3, 50).unwrap())
|
||||
.n_startup_trials(5)
|
||||
.seed(42)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
// Verify the strategy is applied
|
||||
let g = sampler.gamma_strategy().gamma(25);
|
||||
assert!((g - 0.2).abs() < f64::EPSILON); // midpoint of 0.1 to 0.3
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gamma_overrides_gamma_strategy() {
|
||||
// When gamma() is called after gamma_strategy(), it should take precedence
|
||||
let sampler = TpeSampler::builder()
|
||||
.gamma_strategy(SqrtGamma::default())
|
||||
.gamma(0.15) // This should override
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
// Should use fixed gamma of 0.15
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma_strategy().gamma(100) - 0.15).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gamma_strategy_overrides_gamma() {
|
||||
// When gamma_strategy() is called after gamma(), it should take precedence
|
||||
let sampler = TpeSampler::builder()
|
||||
.gamma(0.15)
|
||||
.gamma_strategy(SqrtGamma::default()) // This should override
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
// Should use SqrtGamma - gamma decreases with trials
|
||||
let g10 = sampler.gamma_strategy().gamma(10);
|
||||
let g100 = sampler.gamma_strategy().gamma(100);
|
||||
assert!(g10 > g100, "SqrtGamma should decrease with more trials");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_custom_gamma_strategy() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct DoubleGamma;
|
||||
|
||||
impl GammaStrategy for DoubleGamma {
|
||||
fn gamma(&self, n_trials: usize) -> f64 {
|
||||
// Double the trial count-based calculation, capped at 0.5
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
(0.01 * n_trials as f64).min(0.5)
|
||||
}
|
||||
|
||||
fn clone_box(&self) -> Box<dyn GammaStrategy> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
}
|
||||
|
||||
let sampler = TpeSampler::builder()
|
||||
.gamma_strategy(DoubleGamma)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
assert!((sampler.gamma_strategy().gamma(10) - 0.1).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma_strategy().gamma(50) - 0.5).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma_strategy().gamma(100) - 0.5).abs() < f64::EPSILON);
|
||||
}
|
||||
}
|
||||
@@ -1,16 +0,0 @@
|
||||
//! Tree-Parzen Estimator (TPE) sampler implementation and utilities.
|
||||
//!
|
||||
//! This module provides TPE-based sampling for Bayesian optimization,
|
||||
//! including support for intersection search space calculation.
|
||||
|
||||
mod gamma;
|
||||
mod multivariate;
|
||||
mod sampler;
|
||||
pub mod search_space;
|
||||
|
||||
pub use gamma::{FixedGamma, GammaStrategy, HyperoptGamma, LinearGamma, SqrtGamma};
|
||||
pub use multivariate::{
|
||||
ConstantLiarStrategy, MultivariateTpeSampler, MultivariateTpeSamplerBuilder,
|
||||
};
|
||||
pub use sampler::{TpeSampler, TpeSamplerBuilder};
|
||||
pub use search_space::{GroupDecomposedSearchSpace, IntersectionSearchSpace};
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+282
-1330
File diff suppressed because it is too large
Load Diff
+746
-206
File diff suppressed because it is too large
Load Diff
@@ -2,7 +2,6 @@
|
||||
|
||||
/// 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,
|
||||
@@ -12,7 +11,6 @@ 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,
|
||||
@@ -20,6 +18,4 @@ pub enum TrialState {
|
||||
Complete,
|
||||
/// The trial failed with an error.
|
||||
Failed,
|
||||
/// The trial was pruned (stopped early).
|
||||
Pruned,
|
||||
}
|
||||
|
||||
+23
-61
@@ -4,7 +4,6 @@
|
||||
|
||||
#![cfg(feature = "async")]
|
||||
|
||||
use optimizer::parameter::{FloatParam, Parameter};
|
||||
use optimizer::sampler::random::RandomSampler;
|
||||
use optimizer::sampler::tpe::TpeSampler;
|
||||
use optimizer::{Direction, Error, Study};
|
||||
@@ -14,15 +13,10 @@ async fn test_optimize_async_basic() {
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_async(10, move |mut trial| {
|
||||
let x_param = x_param.clone();
|
||||
async move {
|
||||
let x = x_param.suggest(&mut trial)?;
|
||||
Ok::<_, Error>((trial, x * x))
|
||||
}
|
||||
.optimize_async(10, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
Ok::<_, Error>((trial, x * x))
|
||||
})
|
||||
.await
|
||||
.expect("async optimization should succeed");
|
||||
@@ -33,7 +27,7 @@ async fn test_optimize_async_basic() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_optimize_async_with_tpe() {
|
||||
async fn test_optimize_async_with_sampler() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
@@ -42,15 +36,10 @@ async fn test_optimize_async_with_tpe() {
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(-5.0, 5.0);
|
||||
|
||||
study
|
||||
.optimize_async(15, move |mut trial| {
|
||||
let x_param = x_param.clone();
|
||||
async move {
|
||||
let x = x_param.suggest(&mut trial)?;
|
||||
Ok::<_, Error>((trial, x * x))
|
||||
}
|
||||
.optimize_async_with_sampler(15, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
Ok::<_, Error>((trial, x * x))
|
||||
})
|
||||
.await
|
||||
.expect("async optimization with sampler should succeed");
|
||||
@@ -65,15 +54,10 @@ async fn test_optimize_parallel() {
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_parallel(20, 4, move |mut trial| {
|
||||
let x_param = x_param.clone();
|
||||
async move {
|
||||
let x = x_param.suggest(&mut trial)?;
|
||||
Ok::<_, Error>((trial, x * x))
|
||||
}
|
||||
.optimize_parallel(20, 4, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
Ok::<_, Error>((trial, x * x))
|
||||
})
|
||||
.await
|
||||
.expect("parallel optimization should succeed");
|
||||
@@ -82,7 +66,7 @@ async fn test_optimize_parallel() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_optimize_parallel_with_tpe() {
|
||||
async fn test_optimize_parallel_with_sampler() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
@@ -91,18 +75,11 @@ async fn test_optimize_parallel_with_tpe() {
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(-5.0, 5.0);
|
||||
let y_param = FloatParam::new(-5.0, 5.0);
|
||||
|
||||
study
|
||||
.optimize_parallel(15, 3, move |mut trial| {
|
||||
let x_param = x_param.clone();
|
||||
let y_param = y_param.clone();
|
||||
async move {
|
||||
let x = x_param.suggest(&mut trial)?;
|
||||
let y = y_param.suggest(&mut trial)?;
|
||||
Ok::<_, Error>((trial, x * x + y * y))
|
||||
}
|
||||
.optimize_parallel_with_sampler(15, 3, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
let y = trial.suggest_float("y", -5.0, 5.0)?;
|
||||
Ok::<_, Error>((trial, x * x + y * y))
|
||||
})
|
||||
.await
|
||||
.expect("parallel optimization with sampler should succeed");
|
||||
@@ -128,7 +105,6 @@ async fn test_optimize_async_all_failures() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(deprecated)]
|
||||
async fn test_optimize_async_with_sampler_all_failures() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
|
||||
@@ -163,7 +139,6 @@ async fn test_optimize_parallel_all_failures() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(deprecated)]
|
||||
async fn test_optimize_parallel_with_sampler_all_failures() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
|
||||
@@ -187,15 +162,12 @@ async fn test_optimize_async_partial_failures() {
|
||||
|
||||
let counter = std::sync::atomic::AtomicUsize::new(0);
|
||||
|
||||
let x_param = FloatParam::new(0.0, 10.0);
|
||||
|
||||
study
|
||||
.optimize_async(10, move |mut trial| {
|
||||
.optimize_async(10, |mut trial| {
|
||||
let count = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
let x_param = x_param.clone();
|
||||
async move {
|
||||
if count.is_multiple_of(2) {
|
||||
let x = x_param.suggest(&mut trial)?;
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, Error>((trial, x))
|
||||
} else {
|
||||
Err(Error::NoCompletedTrials) // Use as error type
|
||||
@@ -214,16 +186,11 @@ async fn test_optimize_parallel_high_concurrency() {
|
||||
let sampler = RandomSampler::with_seed(42);
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(0.0, 10.0);
|
||||
|
||||
// Run with concurrency higher than n_trials
|
||||
study
|
||||
.optimize_parallel(5, 10, move |mut trial| {
|
||||
let x_param = x_param.clone();
|
||||
async move {
|
||||
let x = x_param.suggest(&mut trial)?;
|
||||
Ok::<_, Error>((trial, x))
|
||||
}
|
||||
.optimize_parallel(5, 10, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, Error>((trial, x))
|
||||
})
|
||||
.await
|
||||
.expect("should handle high concurrency");
|
||||
@@ -236,16 +203,11 @@ async fn test_optimize_parallel_single_concurrency() {
|
||||
let sampler = RandomSampler::with_seed(42);
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(0.0, 10.0);
|
||||
|
||||
// Run with concurrency of 1 (sequential)
|
||||
study
|
||||
.optimize_parallel(10, 1, move |mut trial| {
|
||||
let x_param = x_param.clone();
|
||||
async move {
|
||||
let x = x_param.suggest(&mut trial)?;
|
||||
Ok::<_, Error>((trial, x))
|
||||
}
|
||||
.optimize_parallel(10, 1, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, Error>((trial, x))
|
||||
})
|
||||
.await
|
||||
.expect("should work with single concurrency");
|
||||
|
||||
@@ -1,133 +0,0 @@
|
||||
//! 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}"
|
||||
);
|
||||
}
|
||||
@@ -1,259 +0,0 @@
|
||||
#![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
|
||||
);
|
||||
}
|
||||
@@ -1,61 +0,0 @@
|
||||
use optimizer::Trial;
|
||||
use optimizer::parameter::{EnumParam, Parameter};
|
||||
use optimizer_derive::Categorical;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Categorical)]
|
||||
enum Color {
|
||||
Red,
|
||||
Green,
|
||||
Blue,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Categorical)]
|
||||
enum SingleVariant {
|
||||
Only,
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn derive_categorical_n_choices() {
|
||||
use optimizer::parameter::Categorical;
|
||||
assert_eq!(Color::N_CHOICES, 3);
|
||||
assert_eq!(SingleVariant::N_CHOICES, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn derive_categorical_roundtrip() {
|
||||
use optimizer::parameter::Categorical;
|
||||
for i in 0..Color::N_CHOICES {
|
||||
let val = Color::from_index(i);
|
||||
assert_eq!(val.to_index(), i);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn derive_categorical_values() {
|
||||
use optimizer::parameter::Categorical;
|
||||
assert_eq!(Color::from_index(0), Color::Red);
|
||||
assert_eq!(Color::from_index(1), Color::Green);
|
||||
assert_eq!(Color::from_index(2), Color::Blue);
|
||||
assert_eq!(Color::Red.to_index(), 0);
|
||||
assert_eq!(Color::Green.to_index(), 1);
|
||||
assert_eq!(Color::Blue.to_index(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn derive_categorical_with_enum_param() {
|
||||
let mut trial = Trial::new(0);
|
||||
let param = EnumParam::<Color>::new();
|
||||
let color = param.suggest(&mut trial).unwrap();
|
||||
assert!([Color::Red, Color::Green, Color::Blue].contains(&color));
|
||||
|
||||
// Cached (same param id)
|
||||
let color2 = param.suggest(&mut trial).unwrap();
|
||||
assert_eq!(color, color2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn derive_categorical_suggest_via_trial() {
|
||||
let mut trial = Trial::new(0);
|
||||
let color = trial.suggest_param(&EnumParam::<Color>::new()).unwrap();
|
||||
assert!([Color::Red, Color::Green, Color::Blue].contains(&color));
|
||||
}
|
||||
@@ -1,166 +0,0 @@
|
||||
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
|
||||
);
|
||||
}
|
||||
+354
-1188
File diff suppressed because it is too large
Load Diff
@@ -1,228 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use optimizer::Direction;
|
||||
use optimizer::pruner::{MedianPruner, Pruner};
|
||||
use optimizer::sampler::CompletedTrial;
|
||||
|
||||
/// Helper to build a completed trial with given intermediate values.
|
||||
fn trial_with_values(id: u64, intermediate_values: Vec<(u64, f64)>) -> CompletedTrial {
|
||||
CompletedTrial::with_intermediate_values(
|
||||
id,
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
0.0,
|
||||
intermediate_values,
|
||||
HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
// --- Minimize direction ---
|
||||
|
||||
#[test]
|
||||
fn prune_when_worse_than_median_minimize() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize);
|
||||
// 3 completed trials with values at step 2: [1.0, 2.0, 3.0] => median = 2.0
|
||||
let completed = vec![
|
||||
trial_with_values(0, vec![(0, 0.5), (1, 0.8), (2, 1.0)]),
|
||||
trial_with_values(1, vec![(0, 0.6), (1, 1.5), (2, 2.0)]),
|
||||
trial_with_values(2, vec![(0, 0.7), (1, 2.0), (2, 3.0)]),
|
||||
];
|
||||
// Current trial value at step 2 is 2.5 > median 2.0 => prune
|
||||
let current = vec![(0, 0.5), (1, 1.0), (2, 2.5)];
|
||||
assert!(pruner.should_prune(3, 2, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_better_than_median_minimize() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize);
|
||||
let completed = vec![
|
||||
trial_with_values(0, vec![(0, 0.5), (1, 0.8), (2, 1.0)]),
|
||||
trial_with_values(1, vec![(0, 0.6), (1, 1.5), (2, 2.0)]),
|
||||
trial_with_values(2, vec![(0, 0.7), (1, 2.0), (2, 3.0)]),
|
||||
];
|
||||
// Current trial value at step 2 is 1.5 < median 2.0 => don't prune
|
||||
let current = vec![(0, 0.5), (1, 1.0), (2, 1.5)];
|
||||
assert!(!pruner.should_prune(3, 2, ¤t, &completed));
|
||||
}
|
||||
|
||||
// --- Maximize direction ---
|
||||
|
||||
#[test]
|
||||
fn prune_when_worse_than_median_maximize() {
|
||||
let pruner = MedianPruner::new(Direction::Maximize);
|
||||
// Values at step 1: [5.0, 7.0, 9.0] => median = 7.0
|
||||
let completed = vec![
|
||||
trial_with_values(0, vec![(0, 3.0), (1, 5.0)]),
|
||||
trial_with_values(1, vec![(0, 4.0), (1, 7.0)]),
|
||||
trial_with_values(2, vec![(0, 5.0), (1, 9.0)]),
|
||||
];
|
||||
// Current value 6.0 < median 7.0 => prune (worse for maximize)
|
||||
let current = vec![(0, 4.0), (1, 6.0)];
|
||||
assert!(pruner.should_prune(3, 1, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_better_than_median_maximize() {
|
||||
let pruner = MedianPruner::new(Direction::Maximize);
|
||||
let completed = vec![
|
||||
trial_with_values(0, vec![(0, 3.0), (1, 5.0)]),
|
||||
trial_with_values(1, vec![(0, 4.0), (1, 7.0)]),
|
||||
trial_with_values(2, vec![(0, 5.0), (1, 9.0)]),
|
||||
];
|
||||
// Current value 8.0 > median 7.0 => don't prune
|
||||
let current = vec![(0, 4.0), (1, 8.0)];
|
||||
assert!(!pruner.should_prune(3, 1, ¤t, &completed));
|
||||
}
|
||||
|
||||
// --- Warmup steps ---
|
||||
|
||||
#[test]
|
||||
fn no_prune_during_warmup() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize).n_warmup_steps(5);
|
||||
let completed = vec![trial_with_values(0, vec![(0, 1.0), (1, 1.0), (2, 1.0)])];
|
||||
// Step 2 < warmup 5 => never prune, even if value is terrible
|
||||
let current = vec![(0, 100.0), (1, 100.0), (2, 100.0)];
|
||||
assert!(!pruner.should_prune(1, 2, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_after_warmup() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize).n_warmup_steps(2);
|
||||
let completed = vec![trial_with_values(0, vec![(0, 1.0), (1, 1.0), (2, 1.0)])];
|
||||
// Step 2 >= warmup 2 => pruning allowed; current 100.0 > median 1.0
|
||||
let current = vec![(0, 100.0), (1, 100.0), (2, 100.0)];
|
||||
assert!(pruner.should_prune(1, 2, ¤t, &completed));
|
||||
}
|
||||
|
||||
// --- n_min_trials ---
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_fewer_than_n_min_trials() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize).n_min_trials(3);
|
||||
// Only 2 completed trials — below threshold of 3
|
||||
let completed = vec![
|
||||
trial_with_values(0, vec![(0, 1.0)]),
|
||||
trial_with_values(1, vec![(0, 2.0)]),
|
||||
];
|
||||
let current = vec![(0, 100.0)];
|
||||
assert!(!pruner.should_prune(2, 0, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_when_at_least_n_min_trials() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize).n_min_trials(3);
|
||||
// 3 completed trials with step 0: [1.0, 2.0, 3.0] => median 2.0
|
||||
let completed = vec![
|
||||
trial_with_values(0, vec![(0, 1.0)]),
|
||||
trial_with_values(1, vec![(0, 2.0)]),
|
||||
trial_with_values(2, vec![(0, 3.0)]),
|
||||
];
|
||||
// 5.0 > median 2.0 => prune
|
||||
let current = vec![(0, 5.0)];
|
||||
assert!(pruner.should_prune(3, 0, ¤t, &completed));
|
||||
}
|
||||
|
||||
// --- No completed trials with values at step ---
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_no_completed_trials_at_step() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize);
|
||||
// Completed trials only have values at step 0, not step 5
|
||||
let completed = vec![
|
||||
trial_with_values(0, vec![(0, 1.0)]),
|
||||
trial_with_values(1, vec![(0, 2.0)]),
|
||||
];
|
||||
let current = vec![(0, 0.5), (5, 100.0)];
|
||||
assert!(!pruner.should_prune(2, 5, ¤t, &completed));
|
||||
}
|
||||
|
||||
// --- Median calculation edge cases ---
|
||||
|
||||
#[test]
|
||||
fn correct_median_with_even_number_of_trials() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize);
|
||||
// 4 trials at step 0: [1.0, 2.0, 3.0, 4.0] => median = 2.5
|
||||
let completed = vec![
|
||||
trial_with_values(0, vec![(0, 1.0)]),
|
||||
trial_with_values(1, vec![(0, 2.0)]),
|
||||
trial_with_values(2, vec![(0, 3.0)]),
|
||||
trial_with_values(3, vec![(0, 4.0)]),
|
||||
];
|
||||
// 2.6 > 2.5 => prune
|
||||
let current = vec![(0, 2.6)];
|
||||
assert!(pruner.should_prune(4, 0, ¤t, &completed));
|
||||
// 2.4 < 2.5 => don't prune
|
||||
let current = vec![(0, 2.4)];
|
||||
assert!(!pruner.should_prune(4, 0, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn correct_median_with_odd_number_of_trials() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize);
|
||||
// 5 trials at step 0: [1.0, 2.0, 3.0, 4.0, 5.0] => median = 3.0
|
||||
let completed = vec![
|
||||
trial_with_values(0, vec![(0, 1.0)]),
|
||||
trial_with_values(1, vec![(0, 2.0)]),
|
||||
trial_with_values(2, vec![(0, 3.0)]),
|
||||
trial_with_values(3, vec![(0, 4.0)]),
|
||||
trial_with_values(4, vec![(0, 5.0)]),
|
||||
];
|
||||
// 3.5 > 3.0 => prune
|
||||
let current = vec![(0, 3.5)];
|
||||
assert!(pruner.should_prune(5, 0, ¤t, &completed));
|
||||
// 2.5 < 3.0 => don't prune
|
||||
let current = vec![(0, 2.5)];
|
||||
assert!(!pruner.should_prune(5, 0, ¤t, &completed));
|
||||
}
|
||||
|
||||
// --- Non-contiguous step numbers ---
|
||||
|
||||
#[test]
|
||||
fn works_with_non_contiguous_steps() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize);
|
||||
// Steps are 0, 10, 100 — non-contiguous
|
||||
let completed = vec![
|
||||
trial_with_values(0, vec![(0, 1.0), (10, 2.0), (100, 3.0)]),
|
||||
trial_with_values(1, vec![(0, 1.5), (10, 2.5), (100, 4.0)]),
|
||||
trial_with_values(2, vec![(0, 2.0), (10, 3.0), (100, 5.0)]),
|
||||
];
|
||||
// At step 100: [3.0, 4.0, 5.0] => median = 4.0
|
||||
let current = vec![(0, 1.0), (10, 2.0), (100, 4.5)];
|
||||
assert!(pruner.should_prune(3, 100, ¤t, &completed));
|
||||
|
||||
let current = vec![(0, 1.0), (10, 2.0), (100, 3.5)];
|
||||
assert!(!pruner.should_prune(3, 100, ¤t, &completed));
|
||||
}
|
||||
|
||||
// --- No intermediate values for current trial ---
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_no_intermediate_values() {
|
||||
let pruner = MedianPruner::new(Direction::Minimize);
|
||||
let completed = vec![trial_with_values(0, vec![(0, 1.0)])];
|
||||
assert!(!pruner.should_prune(1, 0, &[], &completed));
|
||||
}
|
||||
|
||||
// --- Pruned trials are excluded from median calculation ---
|
||||
|
||||
#[test]
|
||||
fn pruned_trials_excluded_from_median() {
|
||||
use optimizer::TrialState;
|
||||
|
||||
let pruner = MedianPruner::new(Direction::Minimize);
|
||||
|
||||
let mut pruned = trial_with_values(0, vec![(0, 0.1)]);
|
||||
pruned.state = TrialState::Pruned;
|
||||
|
||||
// Only the completed trial (value 5.0) counts. Pruned trial (0.1) is excluded.
|
||||
let completed = vec![pruned, trial_with_values(1, vec![(0, 5.0)])];
|
||||
|
||||
// 3.0 < 5.0 => don't prune (only 1 completed trial with median 5.0)
|
||||
let current = vec![(0, 3.0)];
|
||||
assert!(!pruner.should_prune(2, 0, ¤t, &completed));
|
||||
|
||||
// 6.0 > 5.0 => prune
|
||||
let current = vec![(0, 6.0)];
|
||||
assert!(pruner.should_prune(2, 0, ¤t, &completed));
|
||||
}
|
||||
@@ -1,424 +0,0 @@
|
||||
//! Integration tests for the Multivariate TPE sampler.
|
||||
//!
|
||||
//! These tests compare the performance of `MultivariateTpeSampler` against
|
||||
//! the standard `TpeSampler` on problems with and without parameter correlations.
|
||||
|
||||
#![allow(
|
||||
clippy::cast_sign_loss,
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation
|
||||
)]
|
||||
|
||||
use optimizer::parameter::{CategoricalParam, FloatParam, IntParam, Parameter};
|
||||
use optimizer::sampler::tpe::{MultivariateTpeSampler, TpeSampler};
|
||||
use optimizer::{Direction, Error, Study};
|
||||
|
||||
// =============================================================================
|
||||
// Rosenbrock function: f(x,y) = (a-x)^2 + b*(y-x^2)^2
|
||||
// This is a classic benchmark with strong parameter correlation.
|
||||
// Optimal: (x, y) = (a, a^2) with f(x, y) = 0
|
||||
// Standard parameters: a = 1, b = 100
|
||||
// The "banana-shaped" valley makes this hard for independent samplers.
|
||||
// =============================================================================
|
||||
|
||||
/// Computes the Rosenbrock function value.
|
||||
///
|
||||
/// f(x, y) = (a - x)^2 + b * (y - x^2)^2
|
||||
///
|
||||
/// With standard parameters a = 1, b = 100:
|
||||
/// - Optimal point: (1, 1)
|
||||
/// - Optimal value: 0
|
||||
fn rosenbrock(x: f64, y: f64) -> f64 {
|
||||
let a = 1.0;
|
||||
let b = 100.0;
|
||||
(a - x).powi(2) + b * (y - x * x).powi(2)
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Test: Multivariate TPE on Rosenbrock function (correlated parameters)
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_tpe_rosenbrock_finds_good_solution() {
|
||||
// Multivariate TPE should find a good solution on Rosenbrock
|
||||
// because it can model the correlation between x and y
|
||||
let sampler = MultivariateTpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(10)
|
||||
.n_ei_candidates(24)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(-2.0, 2.0);
|
||||
let y_param = FloatParam::new(-2.0, 4.0);
|
||||
|
||||
study
|
||||
.optimize(100, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let y = y_param.suggest(trial)?;
|
||||
Ok::<_, Error>(rosenbrock(x, y))
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
let best = study.best_trial().expect("should have at least one trial");
|
||||
|
||||
// Multivariate TPE should find a reasonably good solution
|
||||
// The global minimum is 0, but getting close is challenging
|
||||
assert!(
|
||||
best.value < 10.0,
|
||||
"Multivariate TPE should find good Rosenbrock solution: best value {} should be < 10.0",
|
||||
best.value
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_independent_tpe_rosenbrock() {
|
||||
// Independent TPE (standard TpeSampler) on Rosenbrock
|
||||
// This establishes a baseline for comparison
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(10)
|
||||
.n_ei_candidates(24)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(-2.0, 2.0);
|
||||
let y_param = FloatParam::new(-2.0, 4.0);
|
||||
|
||||
study
|
||||
.optimize(100, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let y = y_param.suggest(trial)?;
|
||||
Ok::<_, Error>(rosenbrock(x, y))
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
let best = study.best_trial().expect("should have at least one trial");
|
||||
|
||||
// Independent TPE should still find a decent solution,
|
||||
// but may not be as good as multivariate TPE on this correlated problem
|
||||
assert!(
|
||||
best.value < 50.0,
|
||||
"Independent TPE should find reasonable Rosenbrock solution: best value {} should be < 50.0",
|
||||
best.value
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_tpe_outperforms_on_correlated_problem() {
|
||||
// Run multiple seeds and compare average performance
|
||||
// Multivariate TPE should generally find better solutions on Rosenbrock
|
||||
let n_runs = 5;
|
||||
let n_trials = 80;
|
||||
|
||||
let mut multivariate_best_values = Vec::new();
|
||||
let mut independent_best_values = Vec::new();
|
||||
|
||||
for seed in 0..n_runs {
|
||||
// Multivariate TPE
|
||||
let multivariate_sampler = MultivariateTpeSampler::builder()
|
||||
.seed(seed as u64)
|
||||
.n_startup_trials(10)
|
||||
.n_ei_candidates(24)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, multivariate_sampler);
|
||||
|
||||
let x_param = FloatParam::new(-2.0, 2.0);
|
||||
let y_param = FloatParam::new(-2.0, 4.0);
|
||||
|
||||
study
|
||||
.optimize(n_trials, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let y = y_param.suggest(trial)?;
|
||||
Ok::<_, Error>(rosenbrock(x, y))
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
multivariate_best_values.push(study.best_trial().unwrap().value);
|
||||
|
||||
// Independent TPE
|
||||
let independent_sampler = TpeSampler::builder()
|
||||
.seed(seed as u64)
|
||||
.n_startup_trials(10)
|
||||
.n_ei_candidates(24)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, independent_sampler);
|
||||
|
||||
let x_param = FloatParam::new(-2.0, 2.0);
|
||||
let y_param = FloatParam::new(-2.0, 4.0);
|
||||
|
||||
study
|
||||
.optimize(n_trials, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let y = y_param.suggest(trial)?;
|
||||
Ok::<_, Error>(rosenbrock(x, y))
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
independent_best_values.push(study.best_trial().unwrap().value);
|
||||
}
|
||||
|
||||
let multivariate_mean: f64 = multivariate_best_values.iter().sum::<f64>() / n_runs as f64;
|
||||
let independent_mean: f64 = independent_best_values.iter().sum::<f64>() / n_runs as f64;
|
||||
|
||||
// Log results for debugging (these won't show in normal test runs,
|
||||
// but are useful when running with --nocapture)
|
||||
eprintln!("Multivariate TPE mean best: {multivariate_mean:.4}");
|
||||
eprintln!("Independent TPE mean best: {independent_mean:.4}");
|
||||
eprintln!("Multivariate best values: {multivariate_best_values:?}");
|
||||
eprintln!("Independent best values: {independent_best_values:?}");
|
||||
|
||||
// Both methods should find reasonable solutions
|
||||
assert!(
|
||||
multivariate_mean < 20.0,
|
||||
"Multivariate TPE mean {multivariate_mean:.4} should be < 20.0"
|
||||
);
|
||||
assert!(
|
||||
independent_mean < 100.0,
|
||||
"Independent TPE mean {independent_mean:.4} should be < 100.0"
|
||||
);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Independent parameter problem: f(x,y) = x^2 + y^2
|
||||
// No correlation between parameters - both methods should work equally well.
|
||||
// =============================================================================
|
||||
|
||||
/// Simple sphere function with independent parameters.
|
||||
///
|
||||
/// f(x, y) = x^2 + y^2
|
||||
///
|
||||
/// Optimal point: (0, 0)
|
||||
/// Optimal value: 0
|
||||
fn sphere(x: f64, y: f64) -> f64 {
|
||||
x * x + y * y
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_tpe_independent_problem() {
|
||||
// On an independent problem, multivariate TPE should still work well
|
||||
let sampler = MultivariateTpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(10)
|
||||
.n_ei_candidates(24)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(-5.0, 5.0);
|
||||
let y_param = FloatParam::new(-5.0, 5.0);
|
||||
|
||||
study
|
||||
.optimize(50, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let y = y_param.suggest(trial)?;
|
||||
Ok::<_, Error>(sphere(x, y))
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
let best = study.best_trial().expect("should have at least one trial");
|
||||
|
||||
assert!(
|
||||
best.value < 5.0,
|
||||
"Multivariate TPE should find good solution on sphere: best value {} should be < 5.0",
|
||||
best.value
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_independent_tpe_independent_problem() {
|
||||
// Baseline: independent TPE on sphere function
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(10)
|
||||
.n_ei_candidates(24)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(-5.0, 5.0);
|
||||
let y_param = FloatParam::new(-5.0, 5.0);
|
||||
|
||||
study
|
||||
.optimize(50, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let y = y_param.suggest(trial)?;
|
||||
Ok::<_, Error>(sphere(x, y))
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
let best = study.best_trial().expect("should have at least one trial");
|
||||
|
||||
assert!(
|
||||
best.value < 5.0,
|
||||
"Independent TPE should find good solution on sphere: best value {} should be < 5.0",
|
||||
best.value
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_both_samplers_work_on_independent_problem() {
|
||||
// Run both samplers on the independent sphere function
|
||||
// and verify they both achieve similar performance
|
||||
let n_runs = 5;
|
||||
let n_trials = 50;
|
||||
|
||||
let mut multivariate_results = Vec::new();
|
||||
let mut independent_results = Vec::new();
|
||||
|
||||
for seed in 0..n_runs {
|
||||
// Multivariate TPE
|
||||
let sampler = MultivariateTpeSampler::builder()
|
||||
.seed(seed as u64)
|
||||
.n_startup_trials(10)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(-5.0, 5.0);
|
||||
let y_param = FloatParam::new(-5.0, 5.0);
|
||||
|
||||
study
|
||||
.optimize(n_trials, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let y = y_param.suggest(trial)?;
|
||||
Ok::<_, Error>(sphere(x, y))
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
multivariate_results.push(study.best_trial().unwrap().value);
|
||||
|
||||
// Independent TPE
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(seed as u64)
|
||||
.n_startup_trials(10)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(-5.0, 5.0);
|
||||
let y_param = FloatParam::new(-5.0, 5.0);
|
||||
|
||||
study
|
||||
.optimize(n_trials, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let y = y_param.suggest(trial)?;
|
||||
Ok::<_, Error>(sphere(x, y))
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
independent_results.push(study.best_trial().unwrap().value);
|
||||
}
|
||||
|
||||
let multivariate_mean: f64 = multivariate_results.iter().sum::<f64>() / n_runs as f64;
|
||||
let independent_mean: f64 = independent_results.iter().sum::<f64>() / n_runs as f64;
|
||||
|
||||
eprintln!("Sphere function results:");
|
||||
eprintln!(" Multivariate TPE mean: {multivariate_mean:.4}");
|
||||
eprintln!(" Independent TPE mean: {independent_mean:.4}");
|
||||
|
||||
// Both should find good solutions on this simple problem
|
||||
assert!(
|
||||
multivariate_mean < 5.0,
|
||||
"Multivariate TPE mean {multivariate_mean:.4} should be < 5.0 on sphere"
|
||||
);
|
||||
assert!(
|
||||
independent_mean < 5.0,
|
||||
"Independent TPE mean {independent_mean:.4} should be < 5.0 on sphere"
|
||||
);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Test: Multivariate TPE with group decomposition
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_tpe_with_group_decomposition() {
|
||||
// Test that group decomposition works correctly
|
||||
let sampler = MultivariateTpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(10)
|
||||
.group(true) // Enable group decomposition
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(-5.0, 5.0);
|
||||
let y_param = FloatParam::new(-5.0, 5.0);
|
||||
|
||||
study
|
||||
.optimize(50, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let y = y_param.suggest(trial)?;
|
||||
Ok::<_, Error>(sphere(x, y))
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
let best = study.best_trial().expect("should have at least one trial");
|
||||
|
||||
assert!(
|
||||
best.value < 10.0,
|
||||
"Multivariate TPE with groups should find good solution: best value {} should be < 10.0",
|
||||
best.value
|
||||
);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Test: Multivariate TPE with mixed parameter types
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_multivariate_tpe_mixed_parameter_types() {
|
||||
// Test with float, int, and categorical parameters
|
||||
let sampler = MultivariateTpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(10)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
let x_param = FloatParam::new(-5.0, 5.0);
|
||||
let n_param = IntParam::new(1, 10);
|
||||
let mode_param = CategoricalParam::new(vec!["a", "b", "c"]);
|
||||
|
||||
study
|
||||
.optimize(50, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let n = n_param.suggest(trial)?;
|
||||
let mode = mode_param.suggest(trial)?;
|
||||
|
||||
// Objective depends on all parameters
|
||||
let mode_factor = match mode {
|
||||
"a" => 1.0,
|
||||
"b" => 0.5,
|
||||
"c" => 2.0,
|
||||
_ => unreachable!(),
|
||||
};
|
||||
|
||||
Ok::<_, Error>(x * x + (n as f64 - 5.0).powi(2) * mode_factor)
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
let best = study.best_trial().expect("should have at least one trial");
|
||||
|
||||
// Should find a reasonable solution
|
||||
assert!(
|
||||
best.value < 25.0,
|
||||
"Multivariate TPE should handle mixed types: best value {} should be < 25.0",
|
||||
best.value
|
||||
);
|
||||
}
|
||||
@@ -1,240 +0,0 @@
|
||||
use optimizer::parameter::{
|
||||
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter,
|
||||
};
|
||||
use optimizer::{Direction, Study, Trial};
|
||||
|
||||
#[test]
|
||||
fn suggest_float_param_via_trial() {
|
||||
let param = FloatParam::new(0.0, 1.0);
|
||||
let mut trial = Trial::new(0);
|
||||
let x = trial.suggest_param(¶m).unwrap();
|
||||
assert!((0.0..=1.0).contains(&x));
|
||||
|
||||
// Cached
|
||||
let x2 = trial.suggest_param(¶m).unwrap();
|
||||
assert_eq!(x, x2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggest_float_log_param_via_trial() {
|
||||
let param = FloatParam::new(1e-5, 1e-1).log_scale();
|
||||
let mut trial = Trial::new(0);
|
||||
let lr = trial.suggest_param(¶m).unwrap();
|
||||
assert!((1e-5..=1e-1).contains(&lr));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggest_float_step_param_via_trial() {
|
||||
let param = FloatParam::new(0.0, 1.0).step(0.25);
|
||||
let mut trial = Trial::new(0);
|
||||
let x = trial.suggest_param(¶m).unwrap();
|
||||
assert!((0.0..=1.0).contains(&x));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggest_int_param_via_trial() {
|
||||
let param = IntParam::new(1, 10);
|
||||
let mut trial = Trial::new(0);
|
||||
let n = trial.suggest_param(¶m).unwrap();
|
||||
assert!((1..=10).contains(&n));
|
||||
|
||||
// Cached
|
||||
let n2 = trial.suggest_param(¶m).unwrap();
|
||||
assert_eq!(n, n2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggest_int_log_param_via_trial() {
|
||||
let param = IntParam::new(1, 1024).log_scale();
|
||||
let mut trial = Trial::new(0);
|
||||
let batch = trial.suggest_param(¶m).unwrap();
|
||||
assert!((1..=1024).contains(&batch));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggest_int_step_param_via_trial() {
|
||||
let param = IntParam::new(32, 512).step(32);
|
||||
let mut trial = Trial::new(0);
|
||||
let units = trial.suggest_param(¶m).unwrap();
|
||||
assert!((32..=512).contains(&units));
|
||||
assert_eq!((units - 32) % 32, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggest_categorical_param_via_trial() {
|
||||
let choices = vec!["sgd", "adam", "rmsprop"];
|
||||
let param = CategoricalParam::new(choices.clone());
|
||||
let mut trial = Trial::new(0);
|
||||
let opt = trial.suggest_param(¶m).unwrap();
|
||||
assert!(choices.contains(&opt));
|
||||
|
||||
// Cached
|
||||
let opt2 = trial.suggest_param(¶m).unwrap();
|
||||
assert_eq!(opt, opt2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggest_bool_param_via_trial() {
|
||||
let param = BoolParam::new();
|
||||
let mut trial = Trial::new(0);
|
||||
let val = trial.suggest_param(¶m).unwrap();
|
||||
let _ = val;
|
||||
|
||||
// Cached
|
||||
let val2 = trial.suggest_param(¶m).unwrap();
|
||||
assert_eq!(val, val2);
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
enum Activation {
|
||||
Relu,
|
||||
Sigmoid,
|
||||
Tanh,
|
||||
}
|
||||
|
||||
impl Categorical for Activation {
|
||||
const N_CHOICES: usize = 3;
|
||||
|
||||
fn from_index(index: usize) -> Self {
|
||||
match index {
|
||||
0 => Activation::Relu,
|
||||
1 => Activation::Sigmoid,
|
||||
2 => Activation::Tanh,
|
||||
_ => panic!("invalid index"),
|
||||
}
|
||||
}
|
||||
|
||||
fn to_index(&self) -> usize {
|
||||
match self {
|
||||
Activation::Relu => 0,
|
||||
Activation::Sigmoid => 1,
|
||||
Activation::Tanh => 2,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggest_enum_param_via_trial() {
|
||||
let param = EnumParam::<Activation>::new();
|
||||
let mut trial = Trial::new(0);
|
||||
let act = trial.suggest_param(¶m).unwrap();
|
||||
assert!([Activation::Relu, Activation::Sigmoid, Activation::Tanh].contains(&act));
|
||||
|
||||
// Cached
|
||||
let act2 = trial.suggest_param(¶m).unwrap();
|
||||
assert_eq!(act, act2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parameter_conflict_detection() {
|
||||
let float_param = FloatParam::new(0.0, 1.0);
|
||||
let int_param = IntParam::new(0, 10);
|
||||
let mut trial = Trial::new(0);
|
||||
let _ = trial.suggest_param(&float_param).unwrap();
|
||||
|
||||
// Different param type with different id - no conflict
|
||||
let result = trial.suggest_param(&int_param);
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Different bounds for same param type but different id - no conflict
|
||||
let float_param2 = FloatParam::new(0.0, 2.0);
|
||||
let result = trial.suggest_param(&float_param2);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validation_prevents_suggest() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
assert!(trial.suggest_param(&FloatParam::new(1.0, 0.0)).is_err());
|
||||
assert!(
|
||||
trial
|
||||
.suggest_param(&FloatParam::new(-1.0, 1.0).log_scale())
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
trial
|
||||
.suggest_param(&FloatParam::new(0.0, 1.0).step(-0.1))
|
||||
.is_err()
|
||||
);
|
||||
assert!(trial.suggest_param(&IntParam::new(10, 1)).is_err());
|
||||
assert!(
|
||||
trial
|
||||
.suggest_param(&IntParam::new(0, 10).log_scale())
|
||||
.is_err()
|
||||
);
|
||||
assert!(trial.suggest_param(&IntParam::new(0, 10).step(-1)).is_err());
|
||||
assert!(
|
||||
trial
|
||||
.suggest_param(&CategoricalParam::<&str>::new(vec![]))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parameter_api_with_study() {
|
||||
let x_param = FloatParam::new(-5.0, 5.0);
|
||||
let n_param = IntParam::new(1, 10);
|
||||
let dropout_param = BoolParam::new();
|
||||
let opt_param = CategoricalParam::new(vec!["sgd", "adam"]);
|
||||
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
study
|
||||
.optimize(5, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
let n = n_param.suggest(trial)?;
|
||||
let dropout = dropout_param.suggest(trial)?;
|
||||
let opt = opt_param.suggest(trial)?;
|
||||
let _ = (n, dropout, opt);
|
||||
Ok::<_, optimizer::Error>(x * x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
assert!(best.value >= 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parameter_suggest_method() {
|
||||
let param = FloatParam::new(0.0, 1.0);
|
||||
let mut trial = Trial::new(0);
|
||||
let x = param.suggest(&mut trial).unwrap();
|
||||
assert!((0.0..=1.0).contains(&x));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn existing_suggest_methods_still_work() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
let x_param = FloatParam::new(0.0, 1.0);
|
||||
let x = x_param.suggest(&mut trial).unwrap();
|
||||
assert!((0.0..=1.0).contains(&x));
|
||||
|
||||
let lr_param = FloatParam::new(1e-5, 1e-1).log_scale();
|
||||
let lr = lr_param.suggest(&mut trial).unwrap();
|
||||
assert!((1e-5..=1e-1).contains(&lr));
|
||||
|
||||
let step_param = FloatParam::new(0.0, 1.0).step(0.25);
|
||||
let step = step_param.suggest(&mut trial).unwrap();
|
||||
assert!((0.0..=1.0).contains(&step));
|
||||
|
||||
let n_param = IntParam::new(1, 10);
|
||||
let n = n_param.suggest(&mut trial).unwrap();
|
||||
assert!((1..=10).contains(&n));
|
||||
|
||||
let batch_param = IntParam::new(1, 1024).log_scale();
|
||||
let batch = batch_param.suggest(&mut trial).unwrap();
|
||||
assert!((1..=1024).contains(&batch));
|
||||
|
||||
let units_param = IntParam::new(32, 512).step(32);
|
||||
let units = units_param.suggest(&mut trial).unwrap();
|
||||
assert!((32..=512).contains(&units));
|
||||
|
||||
let opt_param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]);
|
||||
let opt = opt_param.suggest(&mut trial).unwrap();
|
||||
assert!(["sgd", "adam", "rmsprop"].contains(&opt));
|
||||
|
||||
let flag_param = BoolParam::new();
|
||||
let flag = flag_param.suggest(&mut trial).unwrap();
|
||||
let _ = flag;
|
||||
}
|
||||
@@ -1,295 +0,0 @@
|
||||
#![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
|
||||
}
|
||||
@@ -1,45 +0,0 @@
|
||||
#[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);
|
||||
}
|
||||
@@ -1,64 +0,0 @@
|
||||
use optimizer::pruner::{Pruner, ThresholdPruner};
|
||||
|
||||
#[test]
|
||||
fn prune_when_value_exceeds_upper_threshold() {
|
||||
let pruner = ThresholdPruner::new().upper(10.0);
|
||||
let values = vec![(0, 5.0), (1, 8.0), (2, 11.0)];
|
||||
assert!(pruner.should_prune(0, 2, &values, &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_when_value_falls_below_lower_threshold() {
|
||||
let pruner = ThresholdPruner::new().lower(0.0);
|
||||
let values = vec![(0, 5.0), (1, 2.0), (2, -1.0)];
|
||||
assert!(pruner.should_prune(0, 2, &values, &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_value_within_bounds() {
|
||||
let pruner = ThresholdPruner::new().upper(10.0).lower(0.0);
|
||||
let values = vec![(0, 3.0), (1, 5.0), (2, 7.0)];
|
||||
assert!(!pruner.should_prune(0, 2, &values, &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_no_intermediate_values() {
|
||||
let pruner = ThresholdPruner::new().upper(10.0).lower(0.0);
|
||||
assert!(!pruner.should_prune(0, 0, &[], &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn works_with_only_upper_set() {
|
||||
let pruner = ThresholdPruner::new().upper(5.0);
|
||||
let below = vec![(0, 3.0)];
|
||||
let above = vec![(0, 6.0)];
|
||||
assert!(!pruner.should_prune(0, 0, &below, &[]));
|
||||
assert!(pruner.should_prune(0, 0, &above, &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn works_with_only_lower_set() {
|
||||
let pruner = ThresholdPruner::new().lower(2.0);
|
||||
let above = vec![(0, 5.0)];
|
||||
let below = vec![(0, 1.0)];
|
||||
assert!(!pruner.should_prune(0, 0, &above, &[]));
|
||||
assert!(pruner.should_prune(0, 0, &below, &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn works_with_both_thresholds_set() {
|
||||
let pruner = ThresholdPruner::new().upper(10.0).lower(0.0);
|
||||
|
||||
// Within bounds
|
||||
assert!(!pruner.should_prune(0, 0, &[(0, 5.0)], &[]));
|
||||
|
||||
// Exceeds upper
|
||||
assert!(pruner.should_prune(0, 0, &[(0, 15.0)], &[]));
|
||||
|
||||
// Below lower
|
||||
assert!(pruner.should_prune(0, 0, &[(0, -3.0)], &[]));
|
||||
|
||||
// At exact boundary (not pruned — strictly greater/less)
|
||||
assert!(!pruner.should_prune(0, 0, &[(0, 10.0)], &[]));
|
||||
assert!(!pruner.should_prune(0, 0, &[(0, 0.0)], &[]));
|
||||
}
|
||||
@@ -1,148 +0,0 @@
|
||||
use optimizer::parameter::{FloatParam, Parameter};
|
||||
use optimizer::{AttrValue, Direction, Study};
|
||||
|
||||
#[test]
|
||||
fn set_and_get_float_attr() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("score", 42.5);
|
||||
assert_eq!(trial.user_attr("score"), Some(&AttrValue::Float(42.5)));
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_and_get_int_attr() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("epoch", 42_i64);
|
||||
assert_eq!(trial.user_attr("epoch"), Some(&AttrValue::Int(42)));
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_and_get_string_attr() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("model", "resnet50");
|
||||
assert_eq!(
|
||||
trial.user_attr("model"),
|
||||
Some(&AttrValue::String("resnet50".to_owned()))
|
||||
);
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_and_get_bool_attr() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("converged", true);
|
||||
assert_eq!(trial.user_attr("converged"), Some(&AttrValue::Bool(true)));
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attrs_propagate_to_completed_trial() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("time_secs", 1.5);
|
||||
trial.set_user_attr("tag", "baseline");
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
assert_eq!(best.user_attr("time_secs"), Some(&AttrValue::Float(1.5)));
|
||||
assert_eq!(
|
||||
best.user_attr("tag"),
|
||||
Some(&AttrValue::String("baseline".to_owned()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn overwrite_attr_replaces_value() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("key", "old");
|
||||
trial.set_user_attr("key", "new");
|
||||
assert_eq!(
|
||||
trial.user_attr("key"),
|
||||
Some(&AttrValue::String("new".to_owned()))
|
||||
);
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
assert_eq!(
|
||||
best.user_attr("key"),
|
||||
Some(&AttrValue::String("new".to_owned()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_attr_returns_none() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
assert_eq!(trial.user_attr("nonexistent"), None);
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
assert_eq!(best.user_attr("nonexistent"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_attrs_map_returns_all() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("a", 1.0);
|
||||
trial.set_user_attr("b", true);
|
||||
assert_eq!(trial.user_attrs().len(), 2);
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
assert_eq!(best.user_attrs().len(), 2);
|
||||
}
|
||||
@@ -1,3 +1,2 @@
|
||||
[default.extend-words]
|
||||
Tpe = "Tpe"
|
||||
consts = "consts"
|
||||
|
||||
Reference in New Issue
Block a user