Compare commits
20 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| b9f8844e60 | |||
| d1c8925fc3 | |||
| 4ce4311c68 | |||
| 293b992e75 | |||
| af8a9f7638 | |||
| b9a90ebec3 | |||
| 24e921b7b9 | |||
| 6363540269 | |||
| 9f3c0f0d3b | |||
| 747e1928b9 | |||
| 15a43d9485 | |||
| 1091429697 | |||
| b482d56e89 | |||
| 90bf73a39f | |||
| 473b973408 | |||
| d8ef4352c3 | |||
| eb57519506 | |||
| cb5ecf33f0 | |||
| 3898136341 | |||
| 6a8a938b6e |
@@ -6,6 +6,9 @@ on:
|
||||
pull_request:
|
||||
branches: [main, master]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
|
||||
@@ -58,6 +61,21 @@ 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
|
||||
@@ -240,6 +258,7 @@ jobs:
|
||||
- fmt
|
||||
- clippy
|
||||
- test
|
||||
- examples
|
||||
- docs
|
||||
- feature-check
|
||||
- coverage
|
||||
@@ -275,7 +294,14 @@ jobs:
|
||||
fi
|
||||
|
||||
- name: Publish
|
||||
if: steps.check.outputs.skip == 'false' && github.ref == 'refs/heads/master'
|
||||
run: cargo publish
|
||||
if: steps.check.outputs.skip == 'false'
|
||||
run: |
|
||||
cargo publish --allow-dirty 2>&1 | tee publish_output.txt || {
|
||||
if grep -q "already uploaded" publish_output.txt; then
|
||||
echo "Version already published, skipping"
|
||||
exit 0
|
||||
fi
|
||||
exit 1
|
||||
}
|
||||
env:
|
||||
CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
|
||||
|
||||
@@ -5,6 +5,9 @@ on:
|
||||
- cron: "0 6 * * *" # Daily at 6:00 UTC
|
||||
workflow_dispatch: # Allow manual trigger
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
|
||||
|
||||
+16
-3
@@ -1,15 +1,19 @@
|
||||
[package]
|
||||
name = "optimizer"
|
||||
version = "0.3.0"
|
||||
version = "0.4.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.88"
|
||||
license = "MIT"
|
||||
authors = ["Manuel Raimann <raimannma@outlook.de"]
|
||||
description = "A Rust library for optimization algorithms."
|
||||
repository = "https://github.com/raimannma/rust-optimizer"
|
||||
|
||||
documentation = "https://docs.rs/optimizer"
|
||||
keywords = ["optimization", "hyperparameter", "tpe", "grid-search", "bayesian"]
|
||||
categories = ["algorithm", "science", "data-structures"]
|
||||
readme = "README.md"
|
||||
|
||||
[dependencies]
|
||||
log = "0.4"
|
||||
rand = "0.9"
|
||||
thiserror = "2"
|
||||
parking_lot = "0.12"
|
||||
@@ -20,4 +24,13 @@ default = []
|
||||
async = ["dep:tokio"]
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
|
||||
|
||||
[[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"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# optimizer
|
||||
|
||||
A Rust library for black-box optimization using Tree-Parzen Estimator (TPE).
|
||||
A Rust library for black-box optimization with multiple sampling strategies.
|
||||
|
||||
[](https://docs.rs/optimizer)
|
||||
[](https://crates.io/crates/optimizer)
|
||||
@@ -9,6 +9,10 @@ A Rust library for black-box optimization using Tree-Parzen Estimator (TPE).
|
||||
## Features
|
||||
|
||||
- Optuna-like API for hyperparameter optimization
|
||||
- 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, and categorical parameter types
|
||||
- Log-scale and stepped parameter sampling
|
||||
- Sync and async optimization with parallel trial evaluation
|
||||
@@ -16,15 +20,16 @@ A Rust library for black-box optimization using Tree-Parzen Estimator (TPE).
|
||||
## Quick Start
|
||||
|
||||
```rust
|
||||
use optimizer::{Direction, Study, TpeSampler};
|
||||
use optimizer::{Direction, Study};
|
||||
use optimizer::sampler::tpe::TpeSampler;
|
||||
|
||||
let sampler = TpeSampler::builder().seed(42).build();
|
||||
let sampler = TpeSampler::builder().seed(42).build().unwrap();
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(20, |trial| {
|
||||
let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
Ok::<_, optimizer::TpeError>(x * x)
|
||||
Ok::<_, optimizer::Error>(x * x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -32,6 +37,100 @@ let best = study.best_trial().unwrap();
|
||||
println!("Best value: {} at x={:?}", best.value, best.params);
|
||||
```
|
||||
|
||||
## Samplers
|
||||
|
||||
### Random Search
|
||||
|
||||
```rust
|
||||
use optimizer::{Direction, Study};
|
||||
use optimizer::sampler::random::RandomSampler;
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(
|
||||
Direction::Minimize,
|
||||
RandomSampler::with_seed(42),
|
||||
);
|
||||
```
|
||||
|
||||
### TPE (Tree-Parzen Estimator)
|
||||
|
||||
```rust
|
||||
use optimizer::{Direction, Study};
|
||||
use optimizer::sampler::tpe::TpeSampler;
|
||||
|
||||
let sampler = TpeSampler::builder()
|
||||
.gamma(0.15) // Quantile for good/bad split
|
||||
.n_startup_trials(20) // Random trials before TPE kicks in
|
||||
.n_ei_candidates(32) // Candidates to evaluate
|
||||
.seed(42)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
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
|
||||
use optimizer::{Direction, Study};
|
||||
use optimizer::sampler::grid::GridSearchSampler;
|
||||
|
||||
let sampler = GridSearchSampler::builder()
|
||||
.n_points_per_param(10) // Number of points per parameter dimension
|
||||
.build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
```
|
||||
|
||||
## Feature Flags
|
||||
|
||||
- `async` - Enable async optimization methods (requires tokio)
|
||||
|
||||
@@ -0,0 +1,317 @@
|
||||
//! 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_with_sampler`
|
||||
//! - 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::sampler::tpe::TpeSampler;
|
||||
use optimizer::{Direction, ParamValue, Study, Trial};
|
||||
|
||||
// ============================================================================
|
||||
// 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.
|
||||
async fn objective(mut trial: Trial) -> optimizer::Result<(Trial, f64)> {
|
||||
// Sample configuration parameters
|
||||
|
||||
// Stepped integers: only sample multiples of the step value
|
||||
let cache_size_mb = trial.suggest_int_step("cache_size_mb", 64, 1024, 64)?;
|
||||
let connection_pool_size = trial.suggest_int_step("connection_pool_size", 10, 200, 10)?;
|
||||
let request_timeout_ms = trial.suggest_int_step("request_timeout_ms", 1000, 10000, 500)?;
|
||||
|
||||
// Regular integer
|
||||
let retry_count = trial.suggest_int("retry_count", 0, 5)?;
|
||||
|
||||
// Log-scale integer: good for parameters like batch sizes
|
||||
// that might vary from 1 to 256
|
||||
let batch_size = trial.suggest_int_log("batch_size", 1, 256)?;
|
||||
|
||||
// Regular integer for compression level
|
||||
let compression_level = trial.suggest_int("compression_level", 0, 9)?;
|
||||
|
||||
// Boolean: internally uses categorical with [false, true]
|
||||
let use_http2 = trial.suggest_bool("use_http2")?;
|
||||
|
||||
// Categorical: choose from a list of options
|
||||
let load_balancing = trial.suggest_categorical(
|
||||
"load_balancing",
|
||||
&["round_robin", "least_connections", "random", "ip_hash"],
|
||||
)?;
|
||||
|
||||
// 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
|
||||
// ============================================================================
|
||||
|
||||
/// Formats a parameter value for display.
|
||||
fn format_param(name: &str, value: &ParamValue) -> String {
|
||||
match (name, value) {
|
||||
(_, ParamValue::Float(v)) => format!("{v:.4}"),
|
||||
(_, ParamValue::Int(v)) => format!("{v}"),
|
||||
("use_http2", ParamValue::Categorical(idx)) => {
|
||||
if *idx == 1 { "true" } else { "false" }.to_string()
|
||||
}
|
||||
("load_balancing", ParamValue::Categorical(idx)) => {
|
||||
["round_robin", "least_connections", "random", "ip_hash"]
|
||||
.get(*idx)
|
||||
.unwrap_or(&"unknown")
|
||||
.to_string()
|
||||
}
|
||||
(_, ParamValue::Categorical(idx)) => format!("category_{idx}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
fn print_best_config(study: &Study<f64>) -> optimizer::Result<()> {
|
||||
let best = study.best_trial()?;
|
||||
|
||||
println!("\nBest configuration found:");
|
||||
println!(" Score: {:.6}", best.value);
|
||||
println!("\n Parameters:");
|
||||
|
||||
// Print parameters in a logical order
|
||||
let param_order = [
|
||||
"cache_size_mb",
|
||||
"connection_pool_size",
|
||||
"request_timeout_ms",
|
||||
"retry_count",
|
||||
"batch_size",
|
||||
"compression_level",
|
||||
"use_http2",
|
||||
"load_balancing",
|
||||
];
|
||||
|
||||
for name in param_order {
|
||||
if let Some(value) = best.params.get(name) {
|
||||
let display = format_param(name, value);
|
||||
println!(" {name}: {display}");
|
||||
}
|
||||
}
|
||||
|
||||
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: 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 4: Run parallel async optimization
|
||||
//
|
||||
// optimize_parallel_with_sampler:
|
||||
// - 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 "_with_sampler" suffix means the TPE sampler gets access to
|
||||
// trial history for informed sampling.
|
||||
study
|
||||
.optimize_parallel_with_sampler(n_trials, concurrency, objective)
|
||||
.await?;
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
// Step 5: Print results
|
||||
print_results(&study, elapsed, n_trials);
|
||||
print_best_config(&study)?;
|
||||
print_top_trials(&study, 5);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,269 @@
|
||||
//! 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::sampler::CompletedTrial;
|
||||
use optimizer::sampler::tpe::TpeSampler;
|
||||
use optimizer::{Direction, ParamValue, Study, Trial};
|
||||
|
||||
// ============================================================================
|
||||
// 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 `trial.suggest_*()` methods to sample hyperparameter values
|
||||
/// 2. Builds a model configuration from those values
|
||||
/// 3. Evaluates the model and returns the loss
|
||||
///
|
||||
/// The optimizer learns from the results to suggest better parameters
|
||||
/// in future trials.
|
||||
fn objective(trial: &mut Trial) -> optimizer::Result<f64> {
|
||||
// Sample hyperparameters using different strategies:
|
||||
|
||||
// Log-scale: Good for parameters spanning multiple orders of magnitude
|
||||
// The learning rate might be 0.001, 0.01, or 0.1 - log-scale samples evenly across these
|
||||
let learning_rate = trial.suggest_float_log("learning_rate", 0.001, 0.3)?;
|
||||
|
||||
// Regular integer: Uniformly samples from the range [3, 12]
|
||||
let max_depth = trial.suggest_int("max_depth", 3, 12)?;
|
||||
|
||||
// Stepped integer: Only samples multiples of 50 (50, 100, 150, ..., 500)
|
||||
// Useful when you only want to test specific values
|
||||
let n_estimators = trial.suggest_int_step("n_estimators", 50, 500, 50)?;
|
||||
|
||||
// Regular float: Uniformly samples from [0.5, 1.0]
|
||||
let subsample = trial.suggest_float("subsample", 0.5, 1.0)?;
|
||||
let colsample_bytree = trial.suggest_float("colsample_bytree", 0.5, 1.0)?;
|
||||
|
||||
// More parameters
|
||||
let min_child_weight = trial.suggest_int("min_child_weight", 1, 10)?;
|
||||
let reg_alpha = trial.suggest_float_log("reg_alpha", 1e-3, 10.0)?;
|
||||
let reg_lambda = trial.suggest_float_log("reg_lambda", 1e-3, 10.0)?;
|
||||
|
||||
// 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<()> {
|
||||
// Helper to extract parameter values
|
||||
let get_float = |name: &str| -> f64 {
|
||||
match trial.params.get(name) {
|
||||
Some(ParamValue::Float(v)) => *v,
|
||||
_ => 0.0,
|
||||
}
|
||||
};
|
||||
|
||||
let get_int = |name: &str| -> i64 {
|
||||
match trial.params.get(name) {
|
||||
Some(ParamValue::Int(v)) => *v,
|
||||
_ => 0,
|
||||
}
|
||||
};
|
||||
|
||||
// Print progress
|
||||
println!(
|
||||
"{:>5} {:>10.5} {:>10} {:>12} {:>10.3} {:>12.3} {:>8} {:>10.4} {:>10.4} {:>12.6}",
|
||||
study.n_trials(),
|
||||
get_float("learning_rate"),
|
||||
get_int("max_depth"),
|
||||
get_int("n_estimators"),
|
||||
get_float("subsample"),
|
||||
get_float("colsample_bytree"),
|
||||
get_int("min_child_weight"),
|
||||
get_float("reg_alpha"),
|
||||
get_float("reg_lambda"),
|
||||
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} {:>10} {:>10} {:>12} {:>10} {:>12} {:>8} {:>10} {:>10} {:>12}",
|
||||
"Trial",
|
||||
"LR",
|
||||
"MaxDepth",
|
||||
"Estimators",
|
||||
"Subsample",
|
||||
"ColSample",
|
||||
"MinCW",
|
||||
"Alpha",
|
||||
"Lambda",
|
||||
"Loss"
|
||||
);
|
||||
println!("{}", "-".repeat(110));
|
||||
|
||||
// Step 3: Run optimization
|
||||
//
|
||||
// optimize_with_callback_sampler runs the objective function for up to
|
||||
// n_trials iterations. After each trial, it calls the callback.
|
||||
// The "_sampler" suffix means the TPE sampler gets access to trial
|
||||
// history for informed sampling.
|
||||
let n_trials = 50;
|
||||
|
||||
study.optimize_with_callback_sampler(n_trials, objective, 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:");
|
||||
|
||||
for (name, value) in &best.params {
|
||||
match value {
|
||||
ParamValue::Float(v) => println!(" {name}: {v:.6}"),
|
||||
ParamValue::Int(v) => println!(" {name}: {v}"),
|
||||
ParamValue::Categorical(v) => println!(" {name}: category {v}"),
|
||||
}
|
||||
}
|
||||
|
||||
// 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(())
|
||||
}
|
||||
+29
-9
@@ -1,10 +1,5 @@
|
||||
//! Error types for the optimizer library.
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
/// The error type for TPE operations.
|
||||
#[derive(Debug, Error)]
|
||||
pub enum TpeError {
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
/// Returned when the lower bound is greater than the upper bound.
|
||||
#[error("invalid bounds: low ({low}) must be less than or equal to high ({high})")]
|
||||
InvalidBounds {
|
||||
@@ -51,6 +46,32 @@ pub enum TpeError {
|
||||
#[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 an internal invariant is violated.
|
||||
#[error("internal error: {0}")]
|
||||
Internal(&'static str),
|
||||
@@ -61,5 +82,4 @@ pub enum TpeError {
|
||||
TaskError(String),
|
||||
}
|
||||
|
||||
/// A specialized Result type for TPE operations.
|
||||
pub type Result<T> = core::result::Result<T, TpeError>;
|
||||
pub type Result<T> = core::result::Result<T, Error>;
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
//! 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;
|
||||
@@ -0,0 +1,921 @@
|
||||
//! 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;
|
||||
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
use rand::Rng;
|
||||
|
||||
use crate::error::{Result, TpeError};
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// A Gaussian kernel density estimator for continuous distributions.
|
||||
///
|
||||
@@ -45,10 +45,10 @@ impl KernelDensityEstimator {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::EmptySamples` if `samples` is empty.
|
||||
/// Returns `Error::EmptySamples` if `samples` is empty.
|
||||
pub(crate) fn new(samples: Vec<f64>) -> Result<Self> {
|
||||
if samples.is_empty() {
|
||||
return Err(TpeError::EmptySamples);
|
||||
return Err(Error::EmptySamples);
|
||||
}
|
||||
|
||||
let bandwidth = Self::scotts_rule(&samples);
|
||||
@@ -61,14 +61,14 @@ impl KernelDensityEstimator {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::EmptySamples` if `samples` is empty.
|
||||
/// Returns `TpeError::InvalidBandwidth` if `bandwidth` is not positive.
|
||||
/// Returns `Error::EmptySamples` if `samples` is empty.
|
||||
/// Returns `Error::InvalidBandwidth` if `bandwidth` is not positive.
|
||||
pub(crate) fn with_bandwidth(samples: Vec<f64>, bandwidth: f64) -> Result<Self> {
|
||||
if samples.is_empty() {
|
||||
return Err(TpeError::EmptySamples);
|
||||
return Err(Error::EmptySamples);
|
||||
}
|
||||
if bandwidth <= 0.0 {
|
||||
return Err(TpeError::InvalidBandwidth(bandwidth));
|
||||
return Err(Error::InvalidBandwidth(bandwidth));
|
||||
}
|
||||
|
||||
Ok(Self { samples, bandwidth })
|
||||
@@ -261,20 +261,20 @@ mod tests {
|
||||
fn test_kde_empty_samples() {
|
||||
let samples: Vec<f64> = vec![];
|
||||
let result = KernelDensityEstimator::new(samples);
|
||||
assert!(matches!(result, Err(TpeError::EmptySamples)));
|
||||
assert!(matches!(result, Err(Error::EmptySamples)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kde_zero_bandwidth() {
|
||||
let samples = vec![1.0, 2.0, 3.0];
|
||||
let result = KernelDensityEstimator::with_bandwidth(samples, 0.0);
|
||||
assert!(matches!(result, Err(TpeError::InvalidBandwidth(_))));
|
||||
assert!(matches!(result, Err(Error::InvalidBandwidth(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kde_negative_bandwidth() {
|
||||
let samples = vec![1.0, 2.0, 3.0];
|
||||
let result = KernelDensityEstimator::with_bandwidth(samples, -1.0);
|
||||
assert!(matches!(result, Err(TpeError::InvalidBandwidth(_))));
|
||||
assert!(matches!(result, Err(Error::InvalidBandwidth(_))));
|
||||
}
|
||||
}
|
||||
+42
-8
@@ -9,10 +9,16 @@
|
||||
#![deny(clippy::pedantic)]
|
||||
#![deny(clippy::std_instead_of_core)]
|
||||
|
||||
//! A Tree-Parzen Estimator (TPE) library for black-box optimization.
|
||||
//! A black-box optimization library with multiple sampling strategies.
|
||||
//!
|
||||
//! This library provides an Optuna-like API for hyperparameter optimization
|
||||
//! using the Tree-Parzen Estimator algorithm. It supports:
|
||||
//! with support for multiple sampling algorithms:
|
||||
//!
|
||||
//! - **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
|
||||
//!
|
||||
//! Additional features include:
|
||||
//!
|
||||
//! - Float, integer, and categorical parameter types
|
||||
//! - Log-scale and stepped parameter sampling
|
||||
@@ -33,7 +39,7 @@
|
||||
//! study
|
||||
//! .optimize_with_sampler(20, |trial| {
|
||||
//! let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
//! Ok::<_, optimizer::TpeError>(x * x)
|
||||
//! Ok::<_, optimizer::Error>(x * x)
|
||||
//! })
|
||||
//! .unwrap();
|
||||
//!
|
||||
@@ -86,14 +92,27 @@
|
||||
//! let optimizer = trial.suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])?;
|
||||
//!
|
||||
//! // Return objective value
|
||||
//! Ok::<_, optimizer::TpeError>(x * n as f64)
|
||||
//! Ok::<_, optimizer::Error>(x * n as f64)
|
||||
//! })
|
||||
//! .unwrap();
|
||||
//! ```
|
||||
//!
|
||||
//! # Configuring TPE
|
||||
//! # Available Samplers
|
||||
//!
|
||||
//! The [`sampler::tpe::TpeSampler`] can be configured using the builder pattern:
|
||||
//! ## Random Search
|
||||
//!
|
||||
//! The simplest sampling strategy, useful for baselines:
|
||||
//!
|
||||
//! ```
|
||||
//! use optimizer::sampler::random::RandomSampler;
|
||||
//! use optimizer::{Direction, Study};
|
||||
//!
|
||||
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
|
||||
//! ```
|
||||
//!
|
||||
//! ## TPE (Tree-Parzen Estimator)
|
||||
//!
|
||||
//! Bayesian optimization that learns from previous trials:
|
||||
//!
|
||||
//! ```
|
||||
//! use optimizer::sampler::tpe::TpeSampler;
|
||||
@@ -107,6 +126,21 @@
|
||||
//! .unwrap();
|
||||
//! ```
|
||||
//!
|
||||
//! ## Grid Search
|
||||
//!
|
||||
//! Exhaustive search over a discretized parameter space:
|
||||
//!
|
||||
//! ```
|
||||
//! use optimizer::sampler::grid::GridSearchSampler;
|
||||
//! use optimizer::{Direction, Study};
|
||||
//!
|
||||
//! let sampler = GridSearchSampler::builder()
|
||||
//! .n_points_per_param(10) // Points per parameter dimension
|
||||
//! .build();
|
||||
//!
|
||||
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
//! ```
|
||||
//!
|
||||
//! # Async and Parallel Optimization
|
||||
//!
|
||||
//! With the `async` feature enabled, you can run trials asynchronously:
|
||||
@@ -140,8 +174,8 @@ mod study;
|
||||
mod trial;
|
||||
mod types;
|
||||
|
||||
pub use error::{Result, TpeError};
|
||||
pub use error::{Error, Result};
|
||||
pub use param::ParamValue;
|
||||
pub use study::Study;
|
||||
pub use trial::Trial;
|
||||
pub use trial::{SuggestableRange, Trial};
|
||||
pub use types::{Direction, TrialState};
|
||||
|
||||
+1157
File diff suppressed because it is too large
Load Diff
@@ -1,10 +1,16 @@
|
||||
//! Sampler trait and implementations for parameter sampling.
|
||||
|
||||
pub mod grid;
|
||||
pub mod multivariate_tpe;
|
||||
pub mod random;
|
||||
pub mod tpe;
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
pub use multivariate_tpe::{
|
||||
ConstantLiarStrategy, MultivariateTpeSampler, MultivariateTpeSamplerBuilder,
|
||||
};
|
||||
|
||||
use crate::distribution::Distribution;
|
||||
use crate::param::ParamValue;
|
||||
|
||||
@@ -42,6 +48,57 @@ impl<V> CompletedTrial<V> {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```ignore
|
||||
/// use std::collections::HashMap;
|
||||
/// use optimizer::sampler::PendingTrial;
|
||||
/// use optimizer::param::ParamValue;
|
||||
/// use optimizer::distribution::{Distribution, FloatDistribution};
|
||||
///
|
||||
/// let mut params = HashMap::new();
|
||||
/// params.insert("x".to_string(), ParamValue::Float(0.5));
|
||||
///
|
||||
/// let mut distributions = HashMap::new();
|
||||
/// distributions.insert("x".to_string(), Distribution::Float(FloatDistribution {
|
||||
/// low: 0.0, high: 1.0, log_scale: false, step: None,
|
||||
/// }));
|
||||
///
|
||||
/// let pending = PendingTrial::new(1, params, distributions);
|
||||
/// assert_eq!(pending.id, 1);
|
||||
/// ```
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PendingTrial {
|
||||
/// The unique identifier for this trial.
|
||||
pub id: u64,
|
||||
/// 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>,
|
||||
}
|
||||
|
||||
impl PendingTrial {
|
||||
/// Creates a new pending trial.
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
id: u64,
|
||||
params: HashMap<String, ParamValue>,
|
||||
distributions: HashMap<String, Distribution>,
|
||||
) -> Self {
|
||||
Self {
|
||||
id,
|
||||
params,
|
||||
distributions,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for pluggable parameter sampling strategies.
|
||||
///
|
||||
/// Samplers are responsible for generating parameter values based on
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,670 @@
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
//! 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 sampler;
|
||||
pub mod search_space;
|
||||
|
||||
pub use gamma::{FixedGamma, GammaStrategy, HyperoptGamma, LinearGamma, SqrtGamma};
|
||||
pub use sampler::{TpeSampler, TpeSamplerBuilder};
|
||||
pub use search_space::{GroupDecomposedSearchSpace, IntersectionSearchSpace};
|
||||
@@ -3,17 +3,80 @@
|
||||
//! 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::{Rng, SeedableRng};
|
||||
|
||||
use crate::distribution::Distribution;
|
||||
use crate::error::{Result, TpeError};
|
||||
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
|
||||
@@ -25,26 +88,42 @@ 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
|
||||
/// // Create with default settings (FixedGamma at 0.25)
|
||||
/// let sampler = TpeSampler::new();
|
||||
///
|
||||
/// // Create with custom settings using the builder
|
||||
/// let sampler = TpeSampler::builder()
|
||||
/// .gamma(0.15)
|
||||
/// .gamma(0.15) // Shorthand for Fixednew(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 {
|
||||
/// Fraction of trials to consider as "good" (gamma quantile).
|
||||
gamma: f64,
|
||||
/// Strategy for computing the gamma quantile.
|
||||
gamma_strategy: Arc<dyn GammaStrategy>,
|
||||
/// 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.
|
||||
@@ -59,14 +138,14 @@ impl TpeSampler {
|
||||
/// Creates a new TPE sampler with default settings.
|
||||
///
|
||||
/// Default settings:
|
||||
/// - gamma: 0.25 (top 25% of trials are considered "good")
|
||||
/// - gamma strategy: [`FixedGamma`] with gamma = 0.25
|
||||
/// - `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: 0.25,
|
||||
gamma_strategy: Arc::new(FixedGamma::default()),
|
||||
n_startup_trials: 10,
|
||||
n_ei_candidates: 24,
|
||||
kde_bandwidth: None,
|
||||
@@ -96,6 +175,10 @@ 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).
|
||||
@@ -106,8 +189,8 @@ impl TpeSampler {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::InvalidGamma` if gamma is not in (0.0, 1.0).
|
||||
/// Returns `TpeError::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
|
||||
/// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0).
|
||||
/// Returns `Error::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
|
||||
pub fn with_config(
|
||||
gamma: f64,
|
||||
n_startup_trials: usize,
|
||||
@@ -115,13 +198,55 @@ impl TpeSampler {
|
||||
kde_bandwidth: Option<f64>,
|
||||
seed: Option<u64>,
|
||||
) -> Result<Self> {
|
||||
if gamma <= 0.0 || gamma >= 1.0 {
|
||||
return Err(TpeError::InvalidGamma(gamma));
|
||||
}
|
||||
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 let Some(bw) = kde_bandwidth
|
||||
&& bw <= 0.0
|
||||
{
|
||||
return Err(TpeError::InvalidBandwidth(bw));
|
||||
return Err(Error::InvalidBandwidth(bw));
|
||||
}
|
||||
|
||||
let rng = match seed {
|
||||
@@ -130,7 +255,7 @@ impl TpeSampler {
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
gamma,
|
||||
gamma_strategy: Arc::new(gamma_strategy),
|
||||
n_startup_trials,
|
||||
n_ei_candidates,
|
||||
kde_bandwidth,
|
||||
@@ -138,8 +263,16 @@ 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(
|
||||
@@ -164,9 +297,15 @@ 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 * self.gamma).ceil() as usize)
|
||||
let n_good = ((history.len() as f64 * gamma).ceil() as usize)
|
||||
.max(1)
|
||||
.min(history.len() - 1);
|
||||
|
||||
@@ -421,6 +560,8 @@ impl Default for TpeSampler {
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// Using a fixed gamma value:
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
||||
///
|
||||
@@ -432,9 +573,23 @@ 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: f64,
|
||||
gamma_strategy: Box<dyn GammaStrategy>,
|
||||
/// Raw gamma value for deferred validation (Some if `gamma()` was called)
|
||||
raw_gamma: Option<f64>,
|
||||
n_startup_trials: usize,
|
||||
n_ei_candidates: usize,
|
||||
kde_bandwidth: Option<f64>,
|
||||
@@ -445,7 +600,7 @@ impl TpeSamplerBuilder {
|
||||
/// Creates a new builder with default settings.
|
||||
///
|
||||
/// Default settings:
|
||||
/// - gamma: 0.25 (top 25% of trials are considered "good")
|
||||
/// - gamma strategy: [`FixedGamma`] with gamma = 0.25
|
||||
/// - `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)
|
||||
@@ -453,7 +608,8 @@ impl TpeSamplerBuilder {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
gamma: 0.25,
|
||||
gamma_strategy: Box::new(FixedGamma::default()),
|
||||
raw_gamma: None,
|
||||
n_startup_trials: 10,
|
||||
n_ei_candidates: 24,
|
||||
kde_bandwidth: None,
|
||||
@@ -461,7 +617,10 @@ impl TpeSamplerBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets the gamma quantile for splitting trials into good/bad groups.
|
||||
/// 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).
|
||||
///
|
||||
/// 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.
|
||||
@@ -484,10 +643,71 @@ impl TpeSamplerBuilder {
|
||||
/// # Note
|
||||
///
|
||||
/// Validation happens at `build()` time. If gamma is not in (0.0, 1.0),
|
||||
/// `build()` will return `Err(TpeError::InvalidGamma)`.
|
||||
/// `build()` will return `Err(Error::InvalidGamma)`.
|
||||
#[must_use]
|
||||
pub fn gamma(mut self, gamma: f64) -> Self {
|
||||
self.gamma = gamma;
|
||||
// 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
|
||||
}
|
||||
|
||||
@@ -569,7 +789,7 @@ impl TpeSamplerBuilder {
|
||||
/// # Note
|
||||
///
|
||||
/// Validation happens at `build()` time. If bandwidth is not positive,
|
||||
/// `build()` will return `Err(TpeError::InvalidBandwidth)`.
|
||||
/// `build()` will return `Err(Error::InvalidBandwidth)`.
|
||||
#[must_use]
|
||||
pub fn kde_bandwidth(mut self, bandwidth: f64) -> Self {
|
||||
self.kde_bandwidth = Some(bandwidth);
|
||||
@@ -602,8 +822,8 @@ impl TpeSamplerBuilder {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::InvalidGamma` if gamma is not in (0.0, 1.0).
|
||||
/// Returns `TpeError::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
|
||||
/// Returns `Error::InvalidGamma` if a fixed gamma value was set and is not in (0.0, 1.0).
|
||||
/// Returns `Error::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -619,13 +839,33 @@ impl TpeSamplerBuilder {
|
||||
/// .unwrap();
|
||||
/// ```
|
||||
pub fn build(self) -> Result<TpeSampler> {
|
||||
TpeSampler::with_config(
|
||||
self.gamma,
|
||||
self.n_startup_trials,
|
||||
self.n_ei_candidates,
|
||||
self.kde_bandwidth,
|
||||
self.seed,
|
||||
)
|
||||
// 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 => StdRng::from_os_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),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -801,7 +1041,8 @@ mod tests {
|
||||
#[test]
|
||||
fn test_tpe_sampler_new() {
|
||||
let sampler = TpeSampler::new();
|
||||
assert!((sampler.gamma - 0.25).abs() < f64::EPSILON);
|
||||
// Default uses FixedGamma with 0.25
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 10);
|
||||
assert_eq!(sampler.n_ei_candidates, 24);
|
||||
}
|
||||
@@ -809,7 +1050,8 @@ mod tests {
|
||||
#[test]
|
||||
fn test_tpe_sampler_with_config() {
|
||||
let sampler = TpeSampler::with_config(0.15, 20, 32, None, Some(42)).unwrap();
|
||||
assert!((sampler.gamma - 0.15).abs() < f64::EPSILON);
|
||||
// with_config uses FixedGamma
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 20);
|
||||
assert_eq!(sampler.n_ei_candidates, 32);
|
||||
}
|
||||
@@ -817,13 +1059,13 @@ mod tests {
|
||||
#[test]
|
||||
fn test_tpe_sampler_invalid_gamma_zero() {
|
||||
let result = TpeSampler::with_config(0.0, 10, 24, None, None);
|
||||
assert!(matches!(result, Err(TpeError::InvalidGamma(_))));
|
||||
assert!(matches!(result, Err(Error::InvalidGamma(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_sampler_invalid_gamma_one() {
|
||||
let result = TpeSampler::with_config(1.0, 10, 24, None, None);
|
||||
assert!(matches!(result, Err(TpeError::InvalidGamma(_))));
|
||||
assert!(matches!(result, Err(Error::InvalidGamma(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1033,7 +1275,7 @@ mod tests {
|
||||
fn test_tpe_sampler_builder_default() {
|
||||
let builder = TpeSamplerBuilder::new();
|
||||
let sampler = builder.build().unwrap();
|
||||
assert!((sampler.gamma - 0.25).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 10);
|
||||
assert_eq!(sampler.n_ei_candidates, 24);
|
||||
}
|
||||
@@ -1047,7 +1289,7 @@ mod tests {
|
||||
.seed(42)
|
||||
.build()
|
||||
.unwrap();
|
||||
assert!((sampler.gamma - 0.15).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 20);
|
||||
assert_eq!(sampler.n_ei_candidates, 32);
|
||||
}
|
||||
@@ -1060,7 +1302,7 @@ mod tests {
|
||||
.n_ei_candidates(48)
|
||||
.build()
|
||||
.unwrap();
|
||||
assert!((sampler.gamma - 0.10).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.10).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 15);
|
||||
assert_eq!(sampler.n_ei_candidates, 48);
|
||||
}
|
||||
@@ -1069,7 +1311,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 - 0.20).abs() < f64::EPSILON);
|
||||
assert!((sampler.gamma_strategy().gamma(0) - 0.20).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.n_startup_trials, 10); // default
|
||||
assert_eq!(sampler.n_ei_candidates, 24); // default
|
||||
}
|
||||
@@ -1077,7 +1319,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_tpe_sampler_builder_invalid_gamma() {
|
||||
let result = TpeSamplerBuilder::new().gamma(1.5).build();
|
||||
assert!(matches!(result, Err(TpeError::InvalidGamma(_))));
|
||||
assert!(matches!(result, Err(Error::InvalidGamma(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
File diff suppressed because it is too large
Load Diff
+38
-38
@@ -275,7 +275,7 @@ where
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if no trials have been completed.
|
||||
/// Returns `Error::NoCompletedTrials` if no trials have been completed.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -305,7 +305,7 @@ where
|
||||
let trials = self.completed_trials.read();
|
||||
|
||||
if trials.is_empty() {
|
||||
return Err(crate::TpeError::NoCompletedTrials);
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
let best = trials
|
||||
@@ -325,7 +325,7 @@ where
|
||||
}
|
||||
}
|
||||
})
|
||||
.ok_or(crate::TpeError::NoCompletedTrials)?;
|
||||
.ok_or(crate::Error::NoCompletedTrials)?;
|
||||
|
||||
Ok(best.clone())
|
||||
}
|
||||
@@ -338,7 +338,7 @@ where
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if no trials have been completed.
|
||||
/// Returns `Error::NoCompletedTrials` if no trials have been completed.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -387,7 +387,7 @@ where
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
/// Returns `Error::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -402,7 +402,7 @@ where
|
||||
/// study
|
||||
/// .optimize(10, |trial| {
|
||||
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
/// Ok::<_, optimizer::TpeError>(x * x)
|
||||
/// Ok::<_, optimizer::Error>(x * x)
|
||||
/// })
|
||||
/// .unwrap();
|
||||
///
|
||||
@@ -431,7 +431,7 @@ where
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
return Err(crate::TpeError::NoCompletedTrials);
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -455,7 +455,7 @@ where
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
/// Returns `Error::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -474,7 +474,7 @@ where
|
||||
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
/// // Simulate async work (e.g., network request)
|
||||
/// let value = x * x;
|
||||
/// Ok::<_, optimizer::TpeError>((trial, value))
|
||||
/// Ok::<_, optimizer::Error>((trial, value))
|
||||
/// })
|
||||
/// .await?;
|
||||
///
|
||||
@@ -511,7 +511,7 @@ where
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
return Err(crate::TpeError::NoCompletedTrials);
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -536,8 +536,8 @@ where
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
/// Returns `TpeError::TaskError` if the semaphore is closed or a spawned task panics.
|
||||
/// Returns `Error::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
/// Returns `Error::TaskError` if the semaphore is closed or a spawned task panics.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -556,7 +556,7 @@ where
|
||||
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
/// // Async objective function (e.g., network request)
|
||||
/// let value = x * x;
|
||||
/// Ok::<_, optimizer::TpeError>((trial, value))
|
||||
/// Ok::<_, optimizer::Error>((trial, value))
|
||||
/// })
|
||||
/// .await?;
|
||||
///
|
||||
@@ -590,7 +590,7 @@ where
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?;
|
||||
.map_err(|e| crate::Error::TaskError(e.to_string()))?;
|
||||
let trial = self.create_trial();
|
||||
let objective = Arc::clone(&objective);
|
||||
|
||||
@@ -607,7 +607,7 @@ where
|
||||
for handle in handles {
|
||||
match handle
|
||||
.await
|
||||
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?
|
||||
.map_err(|e| crate::Error::TaskError(e.to_string()))?
|
||||
{
|
||||
Ok((trial, value)) => {
|
||||
self.complete_trial(trial, value);
|
||||
@@ -620,7 +620,7 @@ where
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
return Err(crate::TpeError::NoCompletedTrials);
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -643,9 +643,9 @@ where
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if no trials completed successfully
|
||||
/// Returns `Error::NoCompletedTrials` if no trials completed successfully
|
||||
/// before optimization stopped (either by completing all trials or early stopping).
|
||||
/// Returns `TpeError::Internal` if a completed trial is not found after adding (internal invariant violation).
|
||||
/// Returns `Error::Internal` if a completed trial is not found after adding (internal invariant violation).
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -664,7 +664,7 @@ where
|
||||
/// 100,
|
||||
/// |trial| {
|
||||
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
/// Ok::<_, optimizer::TpeError>(x * x)
|
||||
/// Ok::<_, optimizer::Error>(x * x)
|
||||
/// },
|
||||
/// |_study, completed_trial| {
|
||||
/// // Stop early if we find a value less than 1.0
|
||||
@@ -702,7 +702,7 @@ where
|
||||
// Get the just-completed trial for the callback
|
||||
let trials = self.completed_trials.read();
|
||||
let Some(completed) = trials.last() else {
|
||||
return Err(crate::TpeError::Internal(
|
||||
return Err(crate::Error::Internal(
|
||||
"completed trial not found after adding",
|
||||
));
|
||||
};
|
||||
@@ -725,7 +725,7 @@ where
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
return Err(crate::TpeError::NoCompletedTrials);
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -783,7 +783,7 @@ impl Study<f64> {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
/// Returns `Error::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -798,7 +798,7 @@ impl Study<f64> {
|
||||
/// study
|
||||
/// .optimize_with_sampler(10, |trial| {
|
||||
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
/// Ok::<_, optimizer::TpeError>(x * x)
|
||||
/// Ok::<_, optimizer::Error>(x * x)
|
||||
/// })
|
||||
/// .unwrap();
|
||||
///
|
||||
@@ -829,7 +829,7 @@ impl Study<f64> {
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
return Err(crate::TpeError::NoCompletedTrials);
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -851,8 +851,8 @@ impl Study<f64> {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if no trials completed successfully.
|
||||
/// Returns `TpeError::Internal` if a completed trial is not found after adding (internal invariant violation).
|
||||
/// Returns `Error::NoCompletedTrials` if no trials completed successfully.
|
||||
/// Returns `Error::Internal` if a completed trial is not found after adding (internal invariant violation).
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -871,7 +871,7 @@ impl Study<f64> {
|
||||
/// 100,
|
||||
/// |trial| {
|
||||
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
/// Ok::<_, optimizer::TpeError>(x * x)
|
||||
/// Ok::<_, optimizer::Error>(x * x)
|
||||
/// },
|
||||
/// |study, _completed_trial| {
|
||||
/// // Stop after finding 5 good trials
|
||||
@@ -907,7 +907,7 @@ impl Study<f64> {
|
||||
// Get the just-completed trial for the callback
|
||||
let trials = self.completed_trials.read();
|
||||
let Some(completed) = trials.last() else {
|
||||
return Err(crate::TpeError::Internal(
|
||||
return Err(crate::Error::Internal(
|
||||
"completed trial not found after adding",
|
||||
));
|
||||
};
|
||||
@@ -930,7 +930,7 @@ impl Study<f64> {
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
return Err(crate::TpeError::NoCompletedTrials);
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -953,7 +953,7 @@ impl Study<f64> {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
/// Returns `Error::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -972,7 +972,7 @@ impl Study<f64> {
|
||||
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
/// // Simulate async work (e.g., network request)
|
||||
/// let value = x * x;
|
||||
/// Ok::<_, optimizer::TpeError>((trial, value))
|
||||
/// Ok::<_, optimizer::Error>((trial, value))
|
||||
/// })
|
||||
/// .await?;
|
||||
///
|
||||
@@ -1009,7 +1009,7 @@ impl Study<f64> {
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
return Err(crate::TpeError::NoCompletedTrials);
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -1034,8 +1034,8 @@ impl Study<f64> {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
/// Returns `TpeError::TaskError` if the semaphore is closed or a spawned task panics.
|
||||
/// Returns `Error::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
/// Returns `Error::TaskError` if the semaphore is closed or a spawned task panics.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -1054,7 +1054,7 @@ impl Study<f64> {
|
||||
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
/// // Async objective function (e.g., network request)
|
||||
/// let value = x * x;
|
||||
/// Ok::<_, optimizer::TpeError>((trial, value))
|
||||
/// Ok::<_, optimizer::Error>((trial, value))
|
||||
/// })
|
||||
/// .await?;
|
||||
///
|
||||
@@ -1087,7 +1087,7 @@ impl Study<f64> {
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?;
|
||||
.map_err(|e| crate::Error::TaskError(e.to_string()))?;
|
||||
let trial = self.create_trial_with_sampler();
|
||||
let objective = Arc::clone(&objective);
|
||||
|
||||
@@ -1104,7 +1104,7 @@ impl Study<f64> {
|
||||
for handle in handles {
|
||||
match handle
|
||||
.await
|
||||
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?
|
||||
.map_err(|e| crate::Error::TaskError(e.to_string()))?
|
||||
{
|
||||
Ok((trial, value)) => {
|
||||
self.complete_trial(trial, value);
|
||||
@@ -1117,7 +1117,7 @@ impl Study<f64> {
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
return Err(crate::TpeError::NoCompletedTrials);
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
||||
+180
-32
@@ -1,5 +1,6 @@
|
||||
//! Trial implementation for tracking sampled parameters and trial state.
|
||||
|
||||
use core::ops::{Range, RangeInclusive};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -8,11 +9,73 @@ use parking_lot::RwLock;
|
||||
use crate::distribution::{
|
||||
CategoricalDistribution, Distribution, FloatDistribution, IntDistribution,
|
||||
};
|
||||
use crate::error::{Result, TpeError};
|
||||
use crate::error::{Error, Result};
|
||||
use crate::param::ParamValue;
|
||||
use crate::sampler::{CompletedTrial, Sampler};
|
||||
use crate::types::TrialState;
|
||||
|
||||
/// A trait for types that can be used with [`Trial::suggest_range`].
|
||||
///
|
||||
/// This trait is implemented for [`Range`] and [`RangeInclusive`] over `f64` and `i64`.
|
||||
/// It allows using Rust's range syntax directly with the optimizer.
|
||||
///
|
||||
/// # Supported Range Types
|
||||
///
|
||||
/// | Range Type | Example | Description |
|
||||
/// |------------|---------|-------------|
|
||||
/// | `Range<f64>` | `0.0..1.0` | Float range (end-exclusive, treated as inclusive for continuous sampling) |
|
||||
/// | `RangeInclusive<f64>` | `0.0..=1.0` | Float range (end-inclusive) |
|
||||
/// | `Range<i64>` | `1..10` | Integer range from 1 to 9 (end-exclusive) |
|
||||
/// | `RangeInclusive<i64>` | `1..=10` | Integer range from 1 to 10 (end-inclusive) |
|
||||
pub trait SuggestableRange {
|
||||
/// The output type when suggesting from this range.
|
||||
type Output;
|
||||
|
||||
/// Suggests a value from this range using the given trial.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if the range is invalid (e.g., empty or low > high).
|
||||
fn suggest(self, trial: &mut Trial, name: String) -> Result<Self::Output>;
|
||||
}
|
||||
|
||||
impl SuggestableRange for Range<f64> {
|
||||
type Output = f64;
|
||||
|
||||
fn suggest(self, trial: &mut Trial, name: String) -> Result<f64> {
|
||||
trial.suggest_float(name, self.start, self.end)
|
||||
}
|
||||
}
|
||||
|
||||
impl SuggestableRange for RangeInclusive<f64> {
|
||||
type Output = f64;
|
||||
|
||||
fn suggest(self, trial: &mut Trial, name: String) -> Result<f64> {
|
||||
trial.suggest_float(name, *self.start(), *self.end())
|
||||
}
|
||||
}
|
||||
|
||||
impl SuggestableRange for Range<i64> {
|
||||
type Output = i64;
|
||||
|
||||
fn suggest(self, trial: &mut Trial, name: String) -> Result<i64> {
|
||||
// Range is exclusive on the end, so subtract 1
|
||||
trial.suggest_int(
|
||||
name,
|
||||
self.start,
|
||||
self.end.checked_sub(1).unwrap_or(self.end),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl SuggestableRange for RangeInclusive<i64> {
|
||||
type Output = i64;
|
||||
|
||||
fn suggest(self, trial: &mut Trial, name: String) -> Result<i64> {
|
||||
trial.suggest_int(name, *self.start(), *self.end())
|
||||
}
|
||||
}
|
||||
|
||||
/// A trial represents a single evaluation of the objective function.
|
||||
///
|
||||
/// Each trial has a unique ID and stores the sampled parameters along with
|
||||
@@ -190,7 +253,7 @@ impl Trial {
|
||||
/// ```
|
||||
pub fn suggest_float(&mut self, name: impl Into<String>, low: f64, high: f64) -> Result<f64> {
|
||||
if low > high {
|
||||
return Err(TpeError::InvalidBounds { low, high });
|
||||
return Err(Error::InvalidBounds { low, high });
|
||||
}
|
||||
|
||||
let name = name.into();
|
||||
@@ -216,7 +279,7 @@ impl Trial {
|
||||
}
|
||||
}
|
||||
// Distribution exists but doesn't match
|
||||
return Err(TpeError::ParameterConflict {
|
||||
return Err(Error::ParameterConflict {
|
||||
name,
|
||||
reason: "parameter was previously sampled with different bounds or type"
|
||||
.to_string(),
|
||||
@@ -226,7 +289,7 @@ impl Trial {
|
||||
// Sample using the sampler
|
||||
let dist = Distribution::Float(distribution);
|
||||
let ParamValue::Float(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
return Err(Error::Internal(
|
||||
"Float distribution should return Float value",
|
||||
));
|
||||
};
|
||||
@@ -283,11 +346,11 @@ impl Trial {
|
||||
high: f64,
|
||||
) -> Result<f64> {
|
||||
if low <= 0.0 {
|
||||
return Err(TpeError::InvalidLogBounds);
|
||||
return Err(Error::InvalidLogBounds);
|
||||
}
|
||||
|
||||
if low > high {
|
||||
return Err(TpeError::InvalidBounds { low, high });
|
||||
return Err(Error::InvalidBounds { low, high });
|
||||
}
|
||||
|
||||
let name = name.into();
|
||||
@@ -313,7 +376,7 @@ impl Trial {
|
||||
}
|
||||
}
|
||||
// Distribution exists but doesn't match
|
||||
return Err(TpeError::ParameterConflict {
|
||||
return Err(Error::ParameterConflict {
|
||||
name,
|
||||
reason: "parameter was previously sampled with different bounds or type"
|
||||
.to_string(),
|
||||
@@ -323,7 +386,7 @@ impl Trial {
|
||||
// Sample using the sampler (sampler handles log-scale transformation)
|
||||
let dist = Distribution::Float(distribution);
|
||||
let ParamValue::Float(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
return Err(Error::Internal(
|
||||
"Float distribution should return Float value",
|
||||
));
|
||||
};
|
||||
@@ -380,11 +443,11 @@ impl Trial {
|
||||
step: f64,
|
||||
) -> Result<f64> {
|
||||
if step <= 0.0 {
|
||||
return Err(TpeError::InvalidStep);
|
||||
return Err(Error::InvalidStep);
|
||||
}
|
||||
|
||||
if low > high {
|
||||
return Err(TpeError::InvalidBounds { low, high });
|
||||
return Err(Error::InvalidBounds { low, high });
|
||||
}
|
||||
|
||||
let name = name.into();
|
||||
@@ -410,7 +473,7 @@ impl Trial {
|
||||
}
|
||||
}
|
||||
// Distribution exists but doesn't match
|
||||
return Err(TpeError::ParameterConflict {
|
||||
return Err(Error::ParameterConflict {
|
||||
name,
|
||||
reason: "parameter was previously sampled with different bounds or type"
|
||||
.to_string(),
|
||||
@@ -420,7 +483,7 @@ impl Trial {
|
||||
// Sample using the sampler (sampler handles step-grid)
|
||||
let dist = Distribution::Float(distribution);
|
||||
let ParamValue::Float(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
return Err(Error::Internal(
|
||||
"Float distribution should return Float value",
|
||||
));
|
||||
};
|
||||
@@ -466,7 +529,7 @@ impl Trial {
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
pub fn suggest_int(&mut self, name: impl Into<String>, low: i64, high: i64) -> Result<i64> {
|
||||
if low > high {
|
||||
return Err(TpeError::InvalidBounds {
|
||||
return Err(Error::InvalidBounds {
|
||||
low: low as f64,
|
||||
high: high as f64,
|
||||
});
|
||||
@@ -495,7 +558,7 @@ impl Trial {
|
||||
}
|
||||
}
|
||||
// Distribution exists but doesn't match
|
||||
return Err(TpeError::ParameterConflict {
|
||||
return Err(Error::ParameterConflict {
|
||||
name,
|
||||
reason: "parameter was previously sampled with different bounds or type"
|
||||
.to_string(),
|
||||
@@ -505,9 +568,7 @@ impl Trial {
|
||||
// Sample using the sampler
|
||||
let dist = Distribution::Int(distribution);
|
||||
let ParamValue::Int(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
"Int distribution should return Int value",
|
||||
));
|
||||
return Err(Error::Internal("Int distribution should return Int value"));
|
||||
};
|
||||
|
||||
// Store distribution and value
|
||||
@@ -554,11 +615,11 @@ impl Trial {
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
pub fn suggest_int_log(&mut self, name: impl Into<String>, low: i64, high: i64) -> Result<i64> {
|
||||
if low < 1 {
|
||||
return Err(TpeError::InvalidLogBounds);
|
||||
return Err(Error::InvalidLogBounds);
|
||||
}
|
||||
|
||||
if low > high {
|
||||
return Err(TpeError::InvalidBounds {
|
||||
return Err(Error::InvalidBounds {
|
||||
low: low as f64,
|
||||
high: high as f64,
|
||||
});
|
||||
@@ -587,7 +648,7 @@ impl Trial {
|
||||
}
|
||||
}
|
||||
// Distribution exists but doesn't match
|
||||
return Err(TpeError::ParameterConflict {
|
||||
return Err(Error::ParameterConflict {
|
||||
name,
|
||||
reason: "parameter was previously sampled with different bounds or type"
|
||||
.to_string(),
|
||||
@@ -597,9 +658,7 @@ impl Trial {
|
||||
// Sample using the sampler (sampler handles log-scale transformation)
|
||||
let dist = Distribution::Int(distribution);
|
||||
let ParamValue::Int(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
"Int distribution should return Int value",
|
||||
));
|
||||
return Err(Error::Internal("Int distribution should return Int value"));
|
||||
};
|
||||
|
||||
// Store distribution and value
|
||||
@@ -659,11 +718,11 @@ impl Trial {
|
||||
step: i64,
|
||||
) -> Result<i64> {
|
||||
if step <= 0 {
|
||||
return Err(TpeError::InvalidStep);
|
||||
return Err(Error::InvalidStep);
|
||||
}
|
||||
|
||||
if low > high {
|
||||
return Err(TpeError::InvalidBounds {
|
||||
return Err(Error::InvalidBounds {
|
||||
low: low as f64,
|
||||
high: high as f64,
|
||||
});
|
||||
@@ -692,7 +751,7 @@ impl Trial {
|
||||
}
|
||||
}
|
||||
// Distribution exists but doesn't match
|
||||
return Err(TpeError::ParameterConflict {
|
||||
return Err(Error::ParameterConflict {
|
||||
name,
|
||||
reason: "parameter was previously sampled with different bounds or type"
|
||||
.to_string(),
|
||||
@@ -702,9 +761,7 @@ impl Trial {
|
||||
// Sample using the sampler (sampler handles step-grid)
|
||||
let dist = Distribution::Int(distribution);
|
||||
let ParamValue::Int(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
"Int distribution should return Int value",
|
||||
));
|
||||
return Err(Error::Internal("Int distribution should return Int value"));
|
||||
};
|
||||
|
||||
// Store distribution and value
|
||||
@@ -760,7 +817,7 @@ impl Trial {
|
||||
choices: &[T],
|
||||
) -> Result<T> {
|
||||
if choices.is_empty() {
|
||||
return Err(TpeError::EmptyChoices);
|
||||
return Err(Error::EmptyChoices);
|
||||
}
|
||||
|
||||
let name = name.into();
|
||||
@@ -779,7 +836,7 @@ impl Trial {
|
||||
}
|
||||
}
|
||||
// Distribution exists but doesn't match
|
||||
return Err(TpeError::ParameterConflict {
|
||||
return Err(Error::ParameterConflict {
|
||||
name,
|
||||
reason: "parameter was previously sampled with different number of choices or type"
|
||||
.to_string(),
|
||||
@@ -789,7 +846,7 @@ impl Trial {
|
||||
// Sample using the sampler
|
||||
let dist = Distribution::Categorical(distribution);
|
||||
let ParamValue::Categorical(index) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
return Err(Error::Internal(
|
||||
"Categorical distribution should return Categorical value",
|
||||
));
|
||||
};
|
||||
@@ -800,4 +857,95 @@ impl Trial {
|
||||
|
||||
Ok(choices[index].clone())
|
||||
}
|
||||
|
||||
/// Suggests a boolean parameter.
|
||||
///
|
||||
/// The value is selected uniformly at random from `{false, true}`.
|
||||
/// This is equivalent to calling `suggest_categorical(name, &[false, true])`.
|
||||
///
|
||||
/// If the parameter has already been sampled, the cached value is returned.
|
||||
/// If the parameter was sampled with a different type, a `ParameterConflict` error is returned.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `name` - The name of the parameter.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `ParameterConflict` if the parameter was previously sampled with a different type.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Trial;
|
||||
///
|
||||
/// let mut trial = Trial::new(0);
|
||||
/// let use_dropout = trial.suggest_bool("use_dropout").unwrap();
|
||||
/// assert!(use_dropout == true || use_dropout == false);
|
||||
///
|
||||
/// // Calling again returns cached value
|
||||
/// let use_dropout2 = trial.suggest_bool("use_dropout").unwrap();
|
||||
/// assert_eq!(use_dropout, use_dropout2);
|
||||
/// ```
|
||||
pub fn suggest_bool(&mut self, name: impl Into<String>) -> Result<bool> {
|
||||
self.suggest_categorical(name, &[false, true])
|
||||
}
|
||||
|
||||
/// Suggests a parameter value from a range.
|
||||
///
|
||||
/// This method accepts both [`Range`] (`..`) and [`RangeInclusive`] (`..=`)
|
||||
/// for both `f64` and `i64` types, allowing natural Rust range syntax.
|
||||
///
|
||||
/// For integer ranges, note that `Range` (`..`) is end-exclusive while
|
||||
/// `RangeInclusive` (`..=`) is end-inclusive, matching Rust's semantics.
|
||||
///
|
||||
/// If the parameter has already been sampled with the same bounds, the cached value is returned.
|
||||
/// If the parameter was sampled with different bounds, a `ParameterConflict` error is returned.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `name` - The name of the parameter.
|
||||
/// * `range` - The range to sample from.
|
||||
///
|
||||
/// # Type Parameters
|
||||
///
|
||||
/// * `R` - A range type implementing [`SuggestableRange`].
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `InvalidBounds` if the range is invalid (e.g., low > high or empty integer range).
|
||||
/// Returns `ParameterConflict` if the parameter was previously sampled with different bounds.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Trial;
|
||||
///
|
||||
/// let mut trial = Trial::new(0);
|
||||
///
|
||||
/// // Float ranges
|
||||
/// let x = trial.suggest_range("x", 0.0..1.0).unwrap();
|
||||
/// assert!(x >= 0.0 && x <= 1.0);
|
||||
///
|
||||
/// let y = trial.suggest_range("y", 0.0..=1.0).unwrap();
|
||||
/// assert!(y >= 0.0 && y <= 1.0);
|
||||
///
|
||||
/// // Integer ranges
|
||||
/// let n = trial.suggest_range("n", 1_i64..10).unwrap(); // 1 to 9 inclusive
|
||||
/// assert!(n >= 1 && n <= 9);
|
||||
///
|
||||
/// let m = trial.suggest_range("m", 1_i64..=10).unwrap(); // 1 to 10 inclusive
|
||||
/// assert!(m >= 1 && m <= 10);
|
||||
///
|
||||
/// // Calling again with same range returns cached value
|
||||
/// let x2 = trial.suggest_range("x", 0.0..1.0).unwrap();
|
||||
/// assert_eq!(x, x2);
|
||||
/// ```
|
||||
pub fn suggest_range<R: SuggestableRange>(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
range: R,
|
||||
) -> Result<R::Output> {
|
||||
range.suggest(self, name.into())
|
||||
}
|
||||
}
|
||||
|
||||
+13
-13
@@ -6,7 +6,7 @@
|
||||
|
||||
use optimizer::sampler::random::RandomSampler;
|
||||
use optimizer::sampler::tpe::TpeSampler;
|
||||
use optimizer::{Direction, Study, TpeError};
|
||||
use optimizer::{Direction, Error, Study};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_optimize_async_basic() {
|
||||
@@ -16,7 +16,7 @@ async fn test_optimize_async_basic() {
|
||||
study
|
||||
.optimize_async(10, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
Ok::<_, TpeError>((trial, x * x))
|
||||
Ok::<_, Error>((trial, x * x))
|
||||
})
|
||||
.await
|
||||
.expect("async optimization should succeed");
|
||||
@@ -39,7 +39,7 @@ async fn test_optimize_async_with_sampler() {
|
||||
study
|
||||
.optimize_async_with_sampler(15, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
Ok::<_, TpeError>((trial, x * x))
|
||||
Ok::<_, Error>((trial, x * x))
|
||||
})
|
||||
.await
|
||||
.expect("async optimization with sampler should succeed");
|
||||
@@ -57,7 +57,7 @@ async fn test_optimize_parallel() {
|
||||
study
|
||||
.optimize_parallel(20, 4, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
Ok::<_, TpeError>((trial, x * x))
|
||||
Ok::<_, Error>((trial, x * x))
|
||||
})
|
||||
.await
|
||||
.expect("parallel optimization should succeed");
|
||||
@@ -79,7 +79,7 @@ async fn test_optimize_parallel_with_sampler() {
|
||||
.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::<_, TpeError>((trial, x * x + y * y))
|
||||
Ok::<_, Error>((trial, x * x + y * y))
|
||||
})
|
||||
.await
|
||||
.expect("parallel optimization with sampler should succeed");
|
||||
@@ -99,7 +99,7 @@ async fn test_optimize_async_all_failures() {
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(TpeError::NoCompletedTrials)),
|
||||
matches!(result, Err(Error::NoCompletedTrials)),
|
||||
"should return NoCompletedTrials when all trials fail"
|
||||
);
|
||||
}
|
||||
@@ -116,7 +116,7 @@ async fn test_optimize_async_with_sampler_all_failures() {
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(TpeError::NoCompletedTrials)),
|
||||
matches!(result, Err(Error::NoCompletedTrials)),
|
||||
"should return NoCompletedTrials when all trials fail"
|
||||
);
|
||||
}
|
||||
@@ -133,7 +133,7 @@ async fn test_optimize_parallel_all_failures() {
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(TpeError::NoCompletedTrials)),
|
||||
matches!(result, Err(Error::NoCompletedTrials)),
|
||||
"should return NoCompletedTrials when all trials fail"
|
||||
);
|
||||
}
|
||||
@@ -150,7 +150,7 @@ async fn test_optimize_parallel_with_sampler_all_failures() {
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(TpeError::NoCompletedTrials)),
|
||||
matches!(result, Err(Error::NoCompletedTrials)),
|
||||
"should return NoCompletedTrials when all trials fail"
|
||||
);
|
||||
}
|
||||
@@ -168,9 +168,9 @@ async fn test_optimize_async_partial_failures() {
|
||||
async move {
|
||||
if count.is_multiple_of(2) {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>((trial, x))
|
||||
Ok::<_, Error>((trial, x))
|
||||
} else {
|
||||
Err(TpeError::NoCompletedTrials) // Use as error type
|
||||
Err(Error::NoCompletedTrials) // Use as error type
|
||||
}
|
||||
}
|
||||
})
|
||||
@@ -190,7 +190,7 @@ async fn test_optimize_parallel_high_concurrency() {
|
||||
study
|
||||
.optimize_parallel(5, 10, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>((trial, x))
|
||||
Ok::<_, Error>((trial, x))
|
||||
})
|
||||
.await
|
||||
.expect("should handle high concurrency");
|
||||
@@ -207,7 +207,7 @@ async fn test_optimize_parallel_single_concurrency() {
|
||||
study
|
||||
.optimize_parallel(10, 1, |mut trial| async move {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>((trial, x))
|
||||
Ok::<_, Error>((trial, x))
|
||||
})
|
||||
.await
|
||||
.expect("should work with single concurrency");
|
||||
|
||||
+281
-57
@@ -8,7 +8,7 @@
|
||||
|
||||
use optimizer::sampler::random::RandomSampler;
|
||||
use optimizer::sampler::tpe::TpeSampler;
|
||||
use optimizer::{Direction, Study, TpeError, Trial};
|
||||
use optimizer::{Direction, Error, Study, Trial};
|
||||
|
||||
// =============================================================================
|
||||
// Test: optimize simple quadratic function with TPE, finds near-optimal
|
||||
@@ -30,7 +30,7 @@ fn test_tpe_optimizes_quadratic_function() {
|
||||
study
|
||||
.optimize_with_sampler(50, |trial| {
|
||||
let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
Ok::<_, TpeError>((x - 3.0).powi(2))
|
||||
Ok::<_, Error>((x - 3.0).powi(2))
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
@@ -61,7 +61,7 @@ fn test_tpe_optimizes_multivariate_function() {
|
||||
.optimize_with_sampler(100, |trial| {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
let y = trial.suggest_float("y", -5.0, 5.0)?;
|
||||
Ok::<_, TpeError>(x * x + y * y)
|
||||
Ok::<_, Error>(x * x + y * y)
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
@@ -90,7 +90,7 @@ fn test_tpe_maximization() {
|
||||
study
|
||||
.optimize_with_sampler(50, |trial| {
|
||||
let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
Ok::<_, TpeError>(-(x - 2.0).powi(2) + 10.0)
|
||||
Ok::<_, Error>(-(x - 2.0).powi(2) + 10.0)
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
@@ -123,7 +123,7 @@ fn test_random_sampler_uniform_float_distribution() {
|
||||
.optimize(n_samples, |trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 1.0)?;
|
||||
samples.push(x);
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -152,7 +152,7 @@ fn test_random_sampler_uniform_float_distribution() {
|
||||
fn test_random_sampler_uniform_int_distribution() {
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(123));
|
||||
|
||||
let n_samples = 1000;
|
||||
let n_samples = 5000;
|
||||
let mut counts = [0u32; 10]; // counts for values 1-10
|
||||
|
||||
study
|
||||
@@ -160,16 +160,18 @@ fn test_random_sampler_uniform_int_distribution() {
|
||||
let n = trial.suggest_int("n", 1, 10)?;
|
||||
assert!((1..=10).contains(&n), "sample {n} out of range [1, 10]");
|
||||
counts[(n - 1) as usize] += 1;
|
||||
Ok::<_, TpeError>(n as f64)
|
||||
Ok::<_, Error>(n as f64)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// Each value should appear roughly n_samples / 10 times
|
||||
// With 5000 samples, expected ~500 per bucket, std dev ~21
|
||||
// 20% tolerance allows for ~4.5 std devs which is very safe
|
||||
let expected = n_samples as f64 / 10.0;
|
||||
for (i, &count) in counts.iter().enumerate() {
|
||||
let diff = (count as f64 - expected).abs() / expected;
|
||||
assert!(
|
||||
diff < 0.3,
|
||||
diff < 0.2,
|
||||
"value {} appeared {} times, expected ~{}, diff = {:.1}%",
|
||||
i + 1,
|
||||
count,
|
||||
@@ -183,7 +185,7 @@ fn test_random_sampler_uniform_int_distribution() {
|
||||
fn test_random_sampler_uniform_categorical_distribution() {
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(456));
|
||||
|
||||
let n_samples = 1000;
|
||||
let n_samples = 2000;
|
||||
let mut counts = [0u32; 4];
|
||||
let choices = ["a", "b", "c", "d"];
|
||||
|
||||
@@ -192,16 +194,18 @@ fn test_random_sampler_uniform_categorical_distribution() {
|
||||
let choice = trial.suggest_categorical("cat", &choices)?;
|
||||
let idx = choices.iter().position(|&c| c == choice).unwrap();
|
||||
counts[idx] += 1;
|
||||
Ok::<_, TpeError>(idx as f64)
|
||||
Ok::<_, Error>(idx as f64)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// Each category should appear roughly n_samples / 4 times
|
||||
// With 2000 samples, expected ~500 per bucket, std dev ~19
|
||||
// 15% tolerance allows for ~4 std devs which is very safe
|
||||
let expected = n_samples as f64 / 4.0;
|
||||
for (i, &count) in counts.iter().enumerate() {
|
||||
let diff = (count as f64 - expected).abs() / expected;
|
||||
assert!(
|
||||
diff < 0.25,
|
||||
diff < 0.15,
|
||||
"category {} appeared {} times, expected ~{}, diff = {:.1}%",
|
||||
i,
|
||||
count,
|
||||
@@ -227,7 +231,7 @@ fn test_random_sampler_reproducibility() {
|
||||
.optimize_with_sampler(100, |trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 100.0)?;
|
||||
values1.push(x);
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -235,7 +239,7 @@ fn test_random_sampler_reproducibility() {
|
||||
.optimize_with_sampler(100, |trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 100.0)?;
|
||||
values2.push(x);
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -367,7 +371,7 @@ fn test_parameter_conflict_float_different_bounds() {
|
||||
trial.suggest_float("x", 0.0, 1.0).unwrap();
|
||||
let result = trial.suggest_float("x", 0.0, 2.0); // Different upper bound
|
||||
|
||||
assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
|
||||
assert!(matches!(result, Err(Error::ParameterConflict { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -377,7 +381,7 @@ fn test_parameter_conflict_float_vs_log() {
|
||||
trial.suggest_float("x", 0.1, 1.0).unwrap();
|
||||
let result = trial.suggest_float_log("x", 0.1, 1.0); // Same bounds but log scale
|
||||
|
||||
assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
|
||||
assert!(matches!(result, Err(Error::ParameterConflict { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -387,7 +391,7 @@ fn test_parameter_conflict_float_vs_step() {
|
||||
trial.suggest_float("x", 0.0, 1.0).unwrap();
|
||||
let result = trial.suggest_float_step("x", 0.0, 1.0, 0.1); // Same bounds but with step
|
||||
|
||||
assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
|
||||
assert!(matches!(result, Err(Error::ParameterConflict { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -397,7 +401,7 @@ fn test_parameter_conflict_int_different_bounds() {
|
||||
trial.suggest_int("n", 1, 10).unwrap();
|
||||
let result = trial.suggest_int("n", 1, 20); // Different upper bound
|
||||
|
||||
assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
|
||||
assert!(matches!(result, Err(Error::ParameterConflict { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -407,7 +411,7 @@ fn test_parameter_conflict_int_vs_log() {
|
||||
trial.suggest_int("n", 1, 100).unwrap();
|
||||
let result = trial.suggest_int_log("n", 1, 100); // Same bounds but log scale
|
||||
|
||||
assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
|
||||
assert!(matches!(result, Err(Error::ParameterConflict { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -417,7 +421,7 @@ fn test_parameter_conflict_categorical_different_n_choices() {
|
||||
trial.suggest_categorical("opt", &["a", "b", "c"]).unwrap();
|
||||
let result = trial.suggest_categorical("opt", &["x", "y"]); // Different number of choices
|
||||
|
||||
assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
|
||||
assert!(matches!(result, Err(Error::ParameterConflict { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -427,7 +431,7 @@ fn test_parameter_conflict_float_vs_int() {
|
||||
trial.suggest_float("x", 0.0, 10.0).unwrap();
|
||||
let result = trial.suggest_int("x", 0, 10); // Different type
|
||||
|
||||
assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
|
||||
assert!(matches!(result, Err(Error::ParameterConflict { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -438,7 +442,7 @@ fn test_parameter_conflict_returns_name() {
|
||||
let result = trial.suggest_float("my_param", 0.0, 2.0);
|
||||
|
||||
match result {
|
||||
Err(TpeError::ParameterConflict { name, .. }) => {
|
||||
Err(Error::ParameterConflict { name, .. }) => {
|
||||
assert_eq!(name, "my_param");
|
||||
}
|
||||
_ => panic!("expected ParameterConflict error"),
|
||||
@@ -456,7 +460,7 @@ fn test_empty_categorical_returns_error() {
|
||||
|
||||
let result = trial.suggest_categorical("opt", empty);
|
||||
|
||||
assert!(matches!(result, Err(TpeError::EmptyChoices)));
|
||||
assert!(matches!(result, Err(Error::EmptyChoices)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -466,7 +470,7 @@ fn test_empty_categorical_vec_returns_error() {
|
||||
|
||||
let result = trial.suggest_categorical("numbers", &empty);
|
||||
|
||||
assert!(matches!(result, Err(TpeError::EmptyChoices)));
|
||||
assert!(matches!(result, Err(Error::EmptyChoices)));
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
@@ -480,7 +484,7 @@ fn test_study_basic_workflow() {
|
||||
study
|
||||
.optimize(10, |trial| {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
Ok::<_, TpeError>(x * x)
|
||||
Ok::<_, Error>(x * x)
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
@@ -517,7 +521,7 @@ fn test_no_completed_trials_error() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
|
||||
let result = study.best_trial();
|
||||
assert!(matches!(result, Err(TpeError::NoCompletedTrials)));
|
||||
assert!(matches!(result, Err(Error::NoCompletedTrials)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -526,11 +530,11 @@ fn test_invalid_bounds_errors() {
|
||||
|
||||
// low > high for float
|
||||
let result = trial.suggest_float("x", 10.0, 5.0);
|
||||
assert!(matches!(result, Err(TpeError::InvalidBounds { .. })));
|
||||
assert!(matches!(result, Err(Error::InvalidBounds { .. })));
|
||||
|
||||
// low > high for int
|
||||
let result = trial.suggest_int("n", 100, 50);
|
||||
assert!(matches!(result, Err(TpeError::InvalidBounds { .. })));
|
||||
assert!(matches!(result, Err(Error::InvalidBounds { .. })));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -539,14 +543,14 @@ fn test_invalid_log_bounds_errors() {
|
||||
|
||||
// low <= 0 for log float
|
||||
let result = trial.suggest_float_log("x", 0.0, 1.0);
|
||||
assert!(matches!(result, Err(TpeError::InvalidLogBounds)));
|
||||
assert!(matches!(result, Err(Error::InvalidLogBounds)));
|
||||
|
||||
let result = trial.suggest_float_log("y", -1.0, 1.0);
|
||||
assert!(matches!(result, Err(TpeError::InvalidLogBounds)));
|
||||
assert!(matches!(result, Err(Error::InvalidLogBounds)));
|
||||
|
||||
// low < 1 for log int
|
||||
let result = trial.suggest_int_log("n", 0, 100);
|
||||
assert!(matches!(result, Err(TpeError::InvalidLogBounds)));
|
||||
assert!(matches!(result, Err(Error::InvalidLogBounds)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -555,14 +559,14 @@ fn test_invalid_step_errors() {
|
||||
|
||||
// step <= 0 for float
|
||||
let result = trial.suggest_float_step("x", 0.0, 1.0, 0.0);
|
||||
assert!(matches!(result, Err(TpeError::InvalidStep)));
|
||||
assert!(matches!(result, Err(Error::InvalidStep)));
|
||||
|
||||
let result = trial.suggest_float_step("y", 0.0, 1.0, -0.1);
|
||||
assert!(matches!(result, Err(TpeError::InvalidStep)));
|
||||
assert!(matches!(result, Err(Error::InvalidStep)));
|
||||
|
||||
// step <= 0 for int
|
||||
let result = trial.suggest_int_step("n", 0, 100, 0);
|
||||
assert!(matches!(result, Err(TpeError::InvalidStep)));
|
||||
assert!(matches!(result, Err(Error::InvalidStep)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -588,7 +592,7 @@ fn test_tpe_with_categorical_parameter() {
|
||||
"cubic" => -((x - 1.0).powi(2)) + 10.0, // peak at x=1, max value 10
|
||||
_ => unreachable!(),
|
||||
};
|
||||
Ok::<_, TpeError>(value)
|
||||
Ok::<_, Error>(value)
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
@@ -615,7 +619,7 @@ fn test_tpe_with_integer_parameters() {
|
||||
study
|
||||
.optimize_with_sampler(30, |trial| {
|
||||
let n = trial.suggest_int("n", 1, 10)?;
|
||||
Ok::<_, TpeError>(((n - 7) as f64).powi(2))
|
||||
Ok::<_, Error>(((n - 7) as f64).powi(2))
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
@@ -643,7 +647,7 @@ fn test_callback_early_stopping() {
|
||||
|trial| {
|
||||
trials_run.set(trials_run.get() + 1);
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
},
|
||||
|_study, _trial| {
|
||||
// Stop after 5 trials
|
||||
@@ -666,7 +670,7 @@ fn test_study_trials_iteration() {
|
||||
study
|
||||
.optimize(5, |trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 1.0)?;
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -758,7 +762,7 @@ fn test_best_value() {
|
||||
study
|
||||
.optimize(10, |trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -792,7 +796,7 @@ fn test_study_set_sampler() {
|
||||
study
|
||||
.optimize_with_sampler(10, |trial| {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
Ok::<_, TpeError>(x * x)
|
||||
Ok::<_, Error>(x * x)
|
||||
})
|
||||
.expect("optimization should succeed with new sampler");
|
||||
|
||||
@@ -807,7 +811,7 @@ fn test_study_with_i32_value_type() {
|
||||
study
|
||||
.optimize(10, |trial| {
|
||||
let x = trial.suggest_int("x", -10, 10)?;
|
||||
Ok::<_, TpeError>(x.abs() as i32)
|
||||
Ok::<_, Error>(x.abs() as i32)
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
@@ -824,7 +828,7 @@ fn test_optimize_all_trials_fail() {
|
||||
let result = study.optimize(5, |_trial| Err::<f64, &str>("always fails"));
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(TpeError::NoCompletedTrials)),
|
||||
matches!(result, Err(Error::NoCompletedTrials)),
|
||||
"should return NoCompletedTrials when all trials fail"
|
||||
);
|
||||
}
|
||||
@@ -842,7 +846,7 @@ fn test_optimize_with_callback_all_trials_fail() {
|
||||
);
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(TpeError::NoCompletedTrials)),
|
||||
matches!(result, Err(Error::NoCompletedTrials)),
|
||||
"should return NoCompletedTrials when all trials fail"
|
||||
);
|
||||
}
|
||||
@@ -854,7 +858,7 @@ fn test_optimize_with_sampler_all_trials_fail() {
|
||||
let result = study.optimize_with_sampler(5, |_trial| Err::<f64, &str>("always fails"));
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(TpeError::NoCompletedTrials)),
|
||||
matches!(result, Err(Error::NoCompletedTrials)),
|
||||
"should return NoCompletedTrials when all trials fail"
|
||||
);
|
||||
}
|
||||
@@ -872,7 +876,7 @@ fn test_optimize_with_callback_sampler_all_trials_fail() {
|
||||
);
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(TpeError::NoCompletedTrials)),
|
||||
matches!(result, Err(Error::NoCompletedTrials)),
|
||||
"should return NoCompletedTrials when all trials fail"
|
||||
);
|
||||
}
|
||||
@@ -882,7 +886,7 @@ fn test_trial_debug_format() {
|
||||
let mut trial = Trial::new(42);
|
||||
trial.suggest_float("x", 0.0, 1.0).unwrap();
|
||||
|
||||
let debug_str = format!("{:?}", trial);
|
||||
let debug_str = format!("{trial:?}");
|
||||
|
||||
// Should contain trial id and other fields
|
||||
assert!(debug_str.contains("Trial"));
|
||||
@@ -902,7 +906,7 @@ fn test_tpe_sampler_builder_default_trait() {
|
||||
study
|
||||
.optimize_with_sampler(5, |trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 1.0)?;
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -917,7 +921,7 @@ fn test_tpe_sampler_default_trait() {
|
||||
study
|
||||
.optimize_with_sampler(5, |trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 1.0)?;
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -938,7 +942,7 @@ fn test_tpe_with_fixed_kde_bandwidth() {
|
||||
study
|
||||
.optimize_with_sampler(20, |trial| {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
Ok::<_, TpeError>(x * x)
|
||||
Ok::<_, Error>(x * x)
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
@@ -949,7 +953,7 @@ fn test_tpe_with_fixed_kde_bandwidth() {
|
||||
#[test]
|
||||
fn test_tpe_sampler_invalid_kde_bandwidth() {
|
||||
let result = TpeSampler::with_config(0.25, 10, 24, Some(-1.0), None);
|
||||
assert!(matches!(result, Err(TpeError::InvalidBandwidth(_))));
|
||||
assert!(matches!(result, Err(Error::InvalidBandwidth(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -966,7 +970,7 @@ fn test_tpe_split_trials_with_two_trials() {
|
||||
study
|
||||
.optimize_with_sampler(5, |trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.expect("optimization should succeed with small history");
|
||||
|
||||
@@ -987,7 +991,7 @@ fn test_tpe_with_log_scale_int() {
|
||||
.optimize_with_sampler(20, |trial| {
|
||||
let batch_size = trial.suggest_int_log("batch_size", 1, 1024)?;
|
||||
// Optimal around batch_size = 32
|
||||
Ok::<_, TpeError>(((batch_size as f64).log2() - 5.0).powi(2))
|
||||
Ok::<_, Error>(((batch_size as f64).log2() - 5.0).powi(2))
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
@@ -1009,7 +1013,7 @@ fn test_tpe_with_step_distributions() {
|
||||
.optimize_with_sampler(20, |trial| {
|
||||
let x = trial.suggest_float_step("x", 0.0, 10.0, 0.5)?;
|
||||
let n = trial.suggest_int_step("n", 0, 100, 10)?;
|
||||
Ok::<_, TpeError>((x - 5.0).powi(2) + ((n - 50) as f64).powi(2))
|
||||
Ok::<_, Error>((x - 5.0).powi(2) + ((n - 50) as f64).powi(2))
|
||||
})
|
||||
.expect("optimization should succeed");
|
||||
|
||||
@@ -1088,7 +1092,7 @@ fn test_tpe_empty_good_or_bad_values_fallback() {
|
||||
study
|
||||
.optimize_with_sampler(10, |trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -1096,7 +1100,7 @@ fn test_tpe_empty_good_or_bad_values_fallback() {
|
||||
study
|
||||
.optimize_with_sampler(5, |trial| {
|
||||
let y = trial.suggest_float("y", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>(y)
|
||||
Ok::<_, Error>(y)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -1114,7 +1118,7 @@ fn test_callback_early_stopping_on_first_trial() {
|
||||
100,
|
||||
|trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
},
|
||||
|_study, _trial| {
|
||||
// Stop immediately after first trial
|
||||
@@ -1138,7 +1142,7 @@ fn test_callback_sampler_early_stopping() {
|
||||
100,
|
||||
|trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
},
|
||||
|study, _trial| {
|
||||
if study.n_trials() >= 3 {
|
||||
@@ -1174,7 +1178,7 @@ fn test_best_trial_with_nan_values() {
|
||||
study
|
||||
.optimize(5, |trial| {
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
Ok::<_, TpeError>(x)
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@@ -1182,3 +1186,223 @@ fn test_best_trial_with_nan_values() {
|
||||
let best = study.best_trial();
|
||||
assert!(best.is_ok());
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests for suggest_bool
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_suggest_bool_caching() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
let b1 = trial.suggest_bool("flag").unwrap();
|
||||
let b2 = trial.suggest_bool("flag").unwrap();
|
||||
|
||||
assert_eq!(b1, b2, "repeated suggest_bool should return cached value");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_bool_multiple_parameters() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
let a = trial.suggest_bool("use_dropout").unwrap();
|
||||
let b = trial.suggest_bool("use_batchnorm").unwrap();
|
||||
let c = trial.suggest_bool("use_skip_connections").unwrap();
|
||||
|
||||
// All should be cached independently
|
||||
assert_eq!(a, trial.suggest_bool("use_dropout").unwrap());
|
||||
assert_eq!(b, trial.suggest_bool("use_batchnorm").unwrap());
|
||||
assert_eq!(c, trial.suggest_bool("use_skip_connections").unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_bool_in_optimization() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
|
||||
study
|
||||
.optimize(10, |trial| {
|
||||
let use_feature = trial.suggest_bool("use_feature")?;
|
||||
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
||||
|
||||
// Objective depends on boolean flag
|
||||
let value = if use_feature { x } else { x * 2.0 };
|
||||
Ok::<_, Error>(value)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(study.n_trials(), 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_bool_with_tpe() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(20, |trial| {
|
||||
let use_large = trial.suggest_bool("use_large")?;
|
||||
let base = if use_large { 10.0 } else { 1.0 };
|
||||
let x = trial.suggest_float("x", 0.0, base)?;
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
// Best should prefer use_large=false for smaller range
|
||||
assert!(best.value < 5.0);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests for suggest_range
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_float_exclusive() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
let x = trial.suggest_range("x", 0.0..1.0).unwrap();
|
||||
assert!((0.0..=1.0).contains(&x), "value {x} out of range 0.0..1.0");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_float_inclusive() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
let x = trial.suggest_range("x", 0.0..=1.0).unwrap();
|
||||
assert!((0.0..=1.0).contains(&x), "value {x} out of range 0.0..=1.0");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_int_exclusive() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
// 1..10 means 1 to 9 inclusive
|
||||
let n = trial.suggest_range("n", 1_i64..10).unwrap();
|
||||
assert!(
|
||||
(1..=9).contains(&n),
|
||||
"value {n} out of range 1..10 (should be 1-9)"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_int_inclusive() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
// 1..=10 means 1 to 10 inclusive
|
||||
let n = trial.suggest_range("n", 1_i64..=10).unwrap();
|
||||
assert!(
|
||||
(1..=10).contains(&n),
|
||||
"value {n} out of range 1..=10 (should be 1-10)"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_caching_float() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
let x1 = trial.suggest_range("x", 0.0..1.0).unwrap();
|
||||
let x2 = trial.suggest_range("x", 0.0..1.0).unwrap();
|
||||
|
||||
assert_eq!(x1, x2, "repeated suggest_range should return cached value");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_caching_int() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
let n1 = trial.suggest_range("n", 1_i64..=100).unwrap();
|
||||
let n2 = trial.suggest_range("n", 1_i64..=100).unwrap();
|
||||
|
||||
assert_eq!(n1, n2, "repeated suggest_range should return cached value");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_multiple_parameters() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
let x = trial.suggest_range("x", 0.0..1.0).unwrap();
|
||||
let y = trial.suggest_range("y", -5.0..=5.0).unwrap();
|
||||
let n = trial.suggest_range("n", 1_i64..10).unwrap();
|
||||
let m = trial.suggest_range("m", 100_i64..=200).unwrap();
|
||||
|
||||
// All should be cached independently
|
||||
assert_eq!(x, trial.suggest_range("x", 0.0..1.0).unwrap());
|
||||
assert_eq!(y, trial.suggest_range("y", -5.0..=5.0).unwrap());
|
||||
assert_eq!(n, trial.suggest_range("n", 1_i64..10).unwrap());
|
||||
assert_eq!(m, trial.suggest_range("m", 100_i64..=200).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_in_optimization() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
|
||||
study
|
||||
.optimize(10, |trial| {
|
||||
let x = trial.suggest_range("x", -10.0..10.0)?;
|
||||
let n = trial.suggest_range("n", 1_i64..=5)?;
|
||||
Ok::<_, Error>(x * x + n as f64)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(study.n_trials(), 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_with_tpe() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(30, |trial| {
|
||||
let x = trial.suggest_range("x", -5.0..=5.0)?;
|
||||
let n = trial.suggest_range("n", 1_i64..=10)?;
|
||||
Ok::<_, Error>(x * x + (n as f64 - 5.0).powi(2))
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
// TPE should find near-optimal solution
|
||||
assert!(best.value < 10.0, "TPE should find good solution");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_empty_int_range_error() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
// 5..5 is empty (no valid integers)
|
||||
let result = trial.suggest_range("n", 5_i64..5);
|
||||
assert!(
|
||||
matches!(result, Err(Error::InvalidBounds { .. })),
|
||||
"empty range should return InvalidBounds error"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_single_value_int() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
// 5..=5 has exactly one value: 5
|
||||
let n = trial.suggest_range("n", 5_i64..=5).unwrap();
|
||||
assert_eq!(n, 5, "single-value range should return that value");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_suggest_range_single_value_float() {
|
||||
let mut trial = Trial::new(0);
|
||||
|
||||
let x = trial.suggest_range("x", 4.2..=4.2).unwrap();
|
||||
assert!(
|
||||
(x - 4.2).abs() < f64::EPSILON,
|
||||
"single-value range should return that value"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,393 @@
|
||||
//! 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::sampler::MultivariateTpeSampler;
|
||||
use optimizer::sampler::tpe::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);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(100, |trial| {
|
||||
let x = trial.suggest_float("x", -2.0, 2.0)?;
|
||||
let y = trial.suggest_float("y", -2.0, 4.0)?;
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(100, |trial| {
|
||||
let x = trial.suggest_float("x", -2.0, 2.0)?;
|
||||
let y = trial.suggest_float("y", -2.0, 4.0)?;
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(n_trials, |trial| {
|
||||
let x = trial.suggest_float("x", -2.0, 2.0)?;
|
||||
let y = trial.suggest_float("y", -2.0, 4.0)?;
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(n_trials, |trial| {
|
||||
let x = trial.suggest_float("x", -2.0, 2.0)?;
|
||||
let y = trial.suggest_float("y", -2.0, 4.0)?;
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(50, |trial| {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
let y = trial.suggest_float("y", -5.0, 5.0)?;
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(50, |trial| {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
let y = trial.suggest_float("y", -5.0, 5.0)?;
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(n_trials, |trial| {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
let y = trial.suggest_float("y", -5.0, 5.0)?;
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(n_trials, |trial| {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
let y = trial.suggest_float("y", -5.0, 5.0)?;
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(50, |trial| {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
let y = trial.suggest_float("y", -5.0, 5.0)?;
|
||||
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);
|
||||
|
||||
study
|
||||
.optimize_with_sampler(50, |trial| {
|
||||
let x = trial.suggest_float("x", -5.0, 5.0)?;
|
||||
let n = trial.suggest_int("n", 1, 10)?;
|
||||
let mode = trial.suggest_categorical("mode", &["a", "b", "c"])?;
|
||||
|
||||
// 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
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user