gavindiaz/rust-optimizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement Parameters API

Browse Source
This commit is contained in:
Manuel Raimann
2026-02-06 17:15:30 +01:00
parent 2fef1ab3c3
commit 55e50e6afb
19 changed files with 3023 additions and 2089 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.toml
+11
View File
@@ -1,3 +1,6 @@
[workspace]
members = ["optimizer-derive"]
[package]
name = "optimizer"
version = "0.4.0"
@@ -17,13 +20,16 @@ rand = "0.9"
thiserror = "2"
parking_lot = "0.12"
tokio = { version = "1", features = ["sync", "rt-multi-thread"], optional = true }
optimizer-derive = { version = "0.1.0", path = "optimizer-derive", optional = true }
[features]
default = []
async = ["dep:tokio"]
derive = ["dep:optimizer-derive"]
[dev-dependencies]
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
optimizer-derive = { version = "0.1.0", path = "optimizer-derive" }
[[example]]
name = "async_api_optimization"
@@ -33,3 +39,8 @@ required-features = ["async"]
[[example]]
name = "ml_hyperparameter_tuning"
path = "examples/ml_hyperparameter_tuning.rs"
[[example]]
name = "parameter_api"
path = "examples/parameter_api.rs"
required-features = ["derive"]
examples/async_api_optimization.rs
+75 -60
View File
@@ -26,6 +26,7 @@
use std::time::{Duration, Instant};
use optimizer::parameter::{BoolParam, CategoricalParam, IntParam, Parameter};
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{Direction, ParamValue, Study, Trial};
@@ -133,32 +134,27 @@ async fn evaluate_service(config: &ServiceConfig) -> f64 {
/// 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"],
)?;
#[allow(clippy::too_many_arguments)]
async fn objective(
mut trial: Trial,
cache_size_mb_param: &IntParam,
connection_pool_size_param: &IntParam,
request_timeout_ms_param: &IntParam,
retry_count_param: &IntParam,
batch_size_param: &IntParam,
compression_level_param: &IntParam,
use_http2_param: &BoolParam,
load_balancing_param: &CategoricalParam<&str>,
) -> optimizer::Result<(Trial, f64)> {
// Sample configuration parameters using parameter definitions
let cache_size_mb = cache_size_mb_param.suggest(&mut trial)?;
let connection_pool_size = connection_pool_size_param.suggest(&mut trial)?;
let request_timeout_ms = request_timeout_ms_param.suggest(&mut trial)?;
let retry_count = retry_count_param.suggest(&mut trial)?;
let batch_size = batch_size_param.suggest(&mut trial)?;
let compression_level = compression_level_param.suggest(&mut trial)?;
let use_http2 = use_http2_param.suggest(&mut trial)?;
let load_balancing = load_balancing_param.suggest(&mut trial)?;
// Build configuration
let config = ServiceConfig {
@@ -184,20 +180,11 @@ async fn objective(mut trial: Trial) -> optimizer::Result<(Trial, f64)> {
// ============================================================================
/// 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}"),
fn format_param(value: &ParamValue) -> String {
match value {
ParamValue::Float(v) => format!("{v:.4}"),
ParamValue::Int(v) => format!("{v}"),
ParamValue::Categorical(idx) => format!("category_{idx}"),
}
}
@@ -226,23 +213,13 @@ fn print_best_config(study: &Study<f64>) -> optimizer::Result<()> {
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}");
}
for (id, value) in &best.params {
let label = best
.param_labels
.get(id)
.map_or_else(|| format!("{id}"), |l| l.clone());
let display = format_param(value);
println!(" {label}: {display}");
}
Ok(())
@@ -284,7 +261,22 @@ async fn main() -> optimizer::Result<()> {
// Step 2: Create a study to minimize the score
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
// Step 3: Configure optimization
// Step 3: Define parameter search spaces
let cache_size_mb_param = IntParam::new(64, 1024).step(64);
let connection_pool_size_param = IntParam::new(10, 200).step(10);
let request_timeout_ms_param = IntParam::new(1000, 10000).step(500);
let retry_count_param = IntParam::new(0, 5);
let batch_size_param = IntParam::new(1, 256).log_scale();
let compression_level_param = IntParam::new(0, 9);
let use_http2_param = BoolParam::new();
let load_balancing_param = CategoricalParam::new(vec![
"round_robin",
"least_connections",
"random",
"ip_hash",
]);
// Step 4: Configure optimization
let n_trials = 40;
let concurrency = 4; // Run 4 trials in parallel
@@ -292,7 +284,7 @@ async fn main() -> optimizer::Result<()> {
let start = Instant::now();
// Step 4: Run parallel async optimization
// Step 5: Run parallel async optimization
//
// optimize_parallel_with_sampler:
// - Runs up to `concurrency` trials simultaneously
@@ -303,7 +295,30 @@ async fn main() -> optimizer::Result<()> {
// 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)
.optimize_parallel_with_sampler(n_trials, concurrency, move |trial| {
let cache_size_mb_param = cache_size_mb_param.clone();
let connection_pool_size_param = connection_pool_size_param.clone();
let request_timeout_ms_param = request_timeout_ms_param.clone();
let retry_count_param = retry_count_param.clone();
let batch_size_param = batch_size_param.clone();
let compression_level_param = compression_level_param.clone();
let use_http2_param = use_http2_param.clone();
let load_balancing_param = load_balancing_param.clone();
async move {
objective(
trial,
&cache_size_mb_param,
&connection_pool_size_param,
&request_timeout_ms_param,
&retry_count_param,
&batch_size_param,
&compression_level_param,
&use_http2_param,
&load_balancing_param,
)
.await
}
})
.await?;
let elapsed = start.elapsed();
examples/ml_hyperparameter_tuning.rs
+71 -70
View File
@@ -23,6 +23,7 @@
use std::ops::ControlFlow;
use optimizer::parameter::{FloatParam, IntParam, Parameter};
use optimizer::sampler::CompletedTrial;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{Direction, ParamValue, Study, Trial};
@@ -87,34 +88,32 @@ fn evaluate_model(config: &ModelConfig) -> f64 {
/// 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
/// 1. Uses parameter definitions passed as arguments
/// 2. Builds a model configuration from the suggested values
/// 3. Evaluates the model and returns the loss
///
/// The optimizer learns from the results to suggest better parameters
/// in future trials.
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)?;
#[allow(clippy::too_many_arguments)]
fn objective(
trial: &mut Trial,
learning_rate_param: &FloatParam,
max_depth_param: &IntParam,
n_estimators_param: &IntParam,
subsample_param: &FloatParam,
colsample_bytree_param: &FloatParam,
min_child_weight_param: &IntParam,
reg_alpha_param: &FloatParam,
reg_lambda_param: &FloatParam,
) -> optimizer::Result<f64> {
let learning_rate = learning_rate_param.suggest(trial)?;
let max_depth = max_depth_param.suggest(trial)?;
let n_estimators = n_estimators_param.suggest(trial)?;
let subsample = subsample_param.suggest(trial)?;
let colsample_bytree = colsample_bytree_param.suggest(trial)?;
let min_child_weight = min_child_weight_param.suggest(trial)?;
let reg_alpha = reg_alpha_param.suggest(trial)?;
let reg_lambda = reg_lambda_param.suggest(trial)?;
// Build configuration and evaluate
let config = ModelConfig {
@@ -148,35 +147,16 @@ fn objective(trial: &mut Trial) -> optimizer::Result<f64> {
/// 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,
// Print trial number and objective value
print!("{:>5} ", study.n_trials());
for value in trial.params.values() {
match value {
ParamValue::Float(v) => print!("{v:>12.5} "),
ParamValue::Int(v) => print!("{v:>12} "),
ParamValue::Categorical(v) => print!("{v:>12} "),
}
};
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,
);
}
println!("{:>12.6}", trial.value);
// Early stopping: if we find an excellent solution, stop early
if trial.value < 0.16 {
@@ -218,21 +198,22 @@ fn main() -> optimizer::Result<()> {
// 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"
"{:>5} {:>12} (parameters...) {:>12}",
"Trial", "Params", "Loss"
);
println!("{}", "-".repeat(110));
println!("{}", "-".repeat(60));
// Step 3: Run optimization
// Step 3: Define parameter search spaces
let learning_rate_param = FloatParam::new(0.001, 0.3).log_scale();
let max_depth_param = IntParam::new(3, 12);
let n_estimators_param = IntParam::new(50, 500).step(50);
let subsample_param = FloatParam::new(0.5, 1.0);
let colsample_bytree_param = FloatParam::new(0.5, 1.0);
let min_child_weight_param = IntParam::new(1, 10);
let reg_alpha_param = FloatParam::new(1e-3, 10.0).log_scale();
let reg_lambda_param = FloatParam::new(1e-3, 10.0).log_scale();
// Step 4: Run optimization
//
// optimize_with_callback_sampler runs the objective function for up to
// n_trials iterations. After each trial, it calls the callback.
@@ -240,7 +221,23 @@ fn main() -> optimizer::Result<()> {
// history for informed sampling.
let n_trials = 50;
study.optimize_with_callback_sampler(n_trials, objective, on_trial_complete)?;
study.optimize_with_callback_sampler(
n_trials,
|trial| {
objective(
trial,
&learning_rate_param,
&max_depth_param,
&n_estimators_param,
&subsample_param,
&colsample_bytree_param,
&min_child_weight_param,
&reg_alpha_param,
&reg_lambda_param,
)
},
on_trial_complete,
)?;
// Step 4: Get the best result
println!("\n{}", "=".repeat(110));
@@ -252,11 +249,15 @@ fn main() -> optimizer::Result<()> {
println!(" Loss: {:.6}", best.value);
println!(" Parameters:");
for (name, value) in &best.params {
for (id, value) in &best.params {
let label = best
.param_labels
.get(id)
.map_or_else(|| format!("{id}"), |l| l.clone());
match value {
ParamValue::Float(v) => println!(" {name}: {v:.6}"),
ParamValue::Int(v) => println!(" {name}: {v}"),
ParamValue::Categorical(v) => println!(" {name}: category {v}"),
ParamValue::Float(v) => println!(" {label}: {v:.6}"),
ParamValue::Int(v) => println!(" {label}: {v}"),
ParamValue::Categorical(v) => println!(" {label}: category {v}"),
}
}
examples/parameter_api.rs
+55
View File
@@ -0,0 +1,55 @@
use optimizer::parameter::{
BoolParam, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter,
};
use optimizer::{Direction, Study};
use optimizer_derive::Categorical;
#[derive(Clone, Debug, Categorical)]
enum Activation {
Relu,
Sigmoid,
Tanh,
}
fn main() {
let study: Study<f64> = Study::new(Direction::Minimize);
// Define parameters outside the objective function
let lr_param = FloatParam::new(1e-5, 1e-1).log_scale();
let n_layers_param = IntParam::new(1, 5);
let units_param = IntParam::new(32, 512).step(32);
let optimizer_param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]);
let activation_param = EnumParam::<Activation>::new();
let batch_size_param = IntParam::new(16, 256).log_scale();
let use_dropout_param = BoolParam::new();
study
.optimize(20, |trial| {
let lr = lr_param.suggest(trial)?;
let n_layers = n_layers_param.suggest(trial)?;
let units = units_param.suggest(trial)?;
let optimizer = optimizer_param.suggest(trial)?;
let use_dropout = use_dropout_param.suggest(trial)?;
let activation = activation_param.suggest(trial)?;
let batch_size = batch_size_param.suggest(trial)?;
// Simulate a loss function
let loss = lr * (n_layers as f64) + (units as f64) * 0.001
- if use_dropout { 0.1 } else { 0.0 };
println!(
"Trial {}: lr={lr:.6}, layers={n_layers}, units={units}, opt={optimizer}, \
dropout={use_dropout}, activation={activation:?}, batch={batch_size} -> loss={loss:.4}",
trial.id()
);
Ok::<_, optimizer::Error>(loss)
})
.unwrap();
let best = study.best_trial().unwrap();
println!("\nBest trial: value={:.4}", best.value);
for (id, label) in &best.param_labels {
println!(" {}: {:?}", label, best.params[id]);
}
}
optimizer-derive/Cargo.toml
+16
View File
@@ -0,0 +1,16 @@
[package]
name = "optimizer-derive"
version = "0.1.0"
edition = "2024"
rust-version = "1.88"
license = "MIT"
description = "Derive macros for the optimizer crate"
repository = "https://github.com/raimannma/rust-optimizer"
[lib]
proc-macro = true
[dependencies]
syn = { version = "2", features = ["full"] }
quote = "1"
proc-macro2 = "1"
optimizer-derive/src/lib.rs
+71
View File
@@ -0,0 +1,71 @@
use proc_macro::TokenStream;
use quote::quote;
use syn::{Data, DeriveInput, Fields, parse_macro_input};
/// Derive macro for the `Categorical` trait on fieldless enums.
///
/// Generates an implementation of `optimizer::Categorical` that maps
/// enum variants to/from sequential indices.
///
/// # Example
///
/// ```ignore
/// use optimizer::Categorical;
///
/// #[derive(Clone, Categorical)]
/// enum Color {
/// Red,
/// Green,
/// Blue,
/// }
/// ```
#[proc_macro_derive(Categorical)]
pub fn derive_categorical(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let name = &input.ident;
let Data::Enum(data_enum) = &input.data else {
return syn::Error::new_spanned(&input, "Categorical can only be derived for enums")
.to_compile_error()
.into();
};
// Validate all variants are fieldless
for variant in &data_enum.variants {
if !matches!(variant.fields, Fields::Unit) {
return syn::Error::new_spanned(
variant,
"Categorical can only be derived for enums with unit variants (no fields)",
)
.to_compile_error()
.into();
}
}
let n_choices = data_enum.variants.len();
let variant_names: Vec<_> = data_enum.variants.iter().map(|v| &v.ident).collect();
let indices: Vec<usize> = (0..n_choices).collect();
let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
let expanded = quote! {
impl #impl_generics optimizer::Categorical for #name #ty_generics #where_clause {
const N_CHOICES: usize = #n_choices;
fn from_index(index: usize) -> Self {
match index {
#(#indices => #name::#variant_names,)*
_ => panic!("invalid index {} for {} with {} variants", index, stringify!(#name), #n_choices),
}
}
fn to_index(&self) -> usize {
match self {
#(#name::#variant_names => #indices,)*
}
}
}
};
expanded.into()
}
src/lib.rs
+52 -23
View File
@@ -28,6 +28,7 @@
//! # Quick Start
//!
//! ```
//! use optimizer::parameter::{FloatParam, Parameter};
//! use optimizer::sampler::tpe::TpeSampler;
//! use optimizer::{Direction, Study};
//!
@@ -35,17 +36,23 @@
//! let sampler = TpeSampler::builder().seed(42).build().unwrap();
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//!
//! // Define parameter search space
//! let x_param = FloatParam::new(-10.0, 10.0);
//!
//! // Optimize x^2 for 20 trials
//! study
//! .optimize_with_sampler(20, |trial| {
//! let x = trial.suggest_float("x", -10.0, 10.0)?;
//! let x = x_param.suggest(trial)?;
//! Ok::<_, optimizer::Error>(x * x)
//! })
//! .unwrap();
//!
//! // Get the best result
//! let best = study.best_trial().unwrap();
//! println!("Best value: {} at x={:?}", best.value, best.params);
//! println!("Best value: {}", best.value);
//! for (id, label) in &best.param_labels {
//! println!(" {}: {:?}", label, best.params[id]);
//! }
//! ```
//!
//! # Creating a Study
@@ -69,29 +76,35 @@
//!
//! # Suggesting Parameters
//!
//! Within the objective function, use [`Trial`] to suggest parameter values:
//! Within the objective function, use parameter types to suggest values:
//!
//! ```
//! use optimizer::parameter::{BoolParam, CategoricalParam, FloatParam, IntParam, Parameter};
//! use optimizer::{Direction, Study};
//!
//! let study: Study<f64> = Study::new(Direction::Minimize);
//!
//! // Define parameter search spaces
//! let x_param = FloatParam::new(0.0, 1.0);
//! let lr_param = FloatParam::new(1e-5, 1e-1).log_scale();
//! let step_param = FloatParam::new(0.0, 1.0).step(0.1);
//! let n_param = IntParam::new(1, 10);
//! let batch_param = IntParam::new(16, 256).log_scale();
//! let units_param = IntParam::new(32, 512).step(32);
//! let flag_param = BoolParam::new();
//! let optimizer_param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]);
//!
//! study
//! .optimize(10, |trial| {
//! // Float parameters
//! let x = trial.suggest_float("x", 0.0, 1.0)?;
//! let lr = trial.suggest_float_log("learning_rate", 1e-5, 1e-1)?;
//! let step = trial.suggest_float_step("step", 0.0, 1.0, 0.1)?;
//! let x = x_param.suggest(trial)?;
//! let lr = lr_param.suggest(trial)?;
//! let step = step_param.suggest(trial)?;
//! let n = n_param.suggest(trial)?;
//! let batch = batch_param.suggest(trial)?;
//! let units = units_param.suggest(trial)?;
//! let flag = flag_param.suggest(trial)?;
//! let optimizer = optimizer_param.suggest(trial)?;
//!
//! // Integer parameters
//! let n = trial.suggest_int("n_layers", 1, 10)?;
//! let batch = trial.suggest_int_log("batch_size", 16, 256)?;
//! let units = trial.suggest_int_step("units", 32, 512, 32)?;
//!
//! // Categorical parameters
//! let optimizer = trial.suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])?;
//!
//! // Return objective value
//! Ok::<_, optimizer::Error>(x * n as f64)
//! })
//! .unwrap();
@@ -147,35 +160,51 @@
//!
//! ```ignore
//! use optimizer::{Study, Direction};
//! use optimizer::parameter::{FloatParam, Parameter};
//!
//! let x_param = FloatParam::new(0.0, 1.0);
//!
//! // Sequential async
//! study.optimize_async(10, |mut trial| async move {
//! let x = trial.suggest_float("x", 0.0, 1.0)?;
//! Ok((trial, x * x))
//! study.optimize_async(10, |mut trial| {
//! let x_param = x_param.clone();
//! async move {
//! let x = x_param.suggest(&mut trial)?;
//! Ok((trial, x * x))
//! }
//! }).await?;
//!
//! // Parallel with bounded concurrency
//! study.optimize_parallel(10, 4, |mut trial| async move {
//! let x = trial.suggest_float("x", 0.0, 1.0)?;
//! Ok((trial, x * x))
//! study.optimize_parallel(10, 4, |mut trial| {
//! let x_param = x_param.clone();
//! async move {
//! let x = x_param.suggest(&mut trial)?;
//! Ok((trial, x * x))
//! }
//! }).await?;
//! ```
//!
//! # Feature Flags
//!
//! - `async`: Enable async optimization methods (requires tokio)
//! - `derive`: Enable `#[derive(Categorical)]` for enum parameters
mod distribution;
mod error;
mod kde;
mod param;
pub mod parameter;
pub mod sampler;
mod study;
mod trial;
mod types;
pub use error::{Error, Result};
#[cfg(feature = "derive")]
pub use optimizer_derive::Categorical;
pub use param::ParamValue;
pub use parameter::{
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, ParamId, Parameter,
};
pub use study::Study;
pub use trial::{SuggestableRange, Trial};
pub use trial::Trial;
pub use types::{Direction, TrialState};
src/parameter.rs
+913
View File
@@ -0,0 +1,913 @@
//! Central parameter trait and built-in parameter types.
//!
//! The [`Parameter`] trait provides a unified way to define parameter types
//! and suggest values from a [`Trial`]. Built-in implementations
//! cover floats, integers, categoricals, booleans, and enum types.
//!
//! # Example
//!
//! ```
//! use optimizer::Trial;
//! use optimizer::parameter::{BoolParam, FloatParam, IntParam, Parameter};
//!
//! let mut trial = Trial::new(0);
//!
//! let lr = FloatParam::new(1e-5, 1e-1)
//! .log_scale()
//! .suggest(&mut trial)
//! .unwrap();
//! let layers = IntParam::new(1, 10).suggest(&mut trial).unwrap();
//! let dropout = BoolParam::new().suggest(&mut trial).unwrap();
//! ```
use core::fmt::Debug;
use core::sync::atomic::{AtomicU64, Ordering};
use crate::distribution::{
CategoricalDistribution, Distribution, FloatDistribution, IntDistribution,
};
use crate::error::{Error, Result};
use crate::param::ParamValue;
use crate::trial::Trial;
static NEXT_PARAM_ID: AtomicU64 = AtomicU64::new(0);
/// A unique identifier for a parameter instance.
///
/// Each parameter is assigned a unique `ParamId` at creation time. Cloning a parameter
/// copies its `ParamId`, so clones refer to the same logical parameter.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ParamId(u64);
impl ParamId {
/// Creates a new unique `ParamId`.
pub fn new() -> Self {
Self(NEXT_PARAM_ID.fetch_add(1, Ordering::Relaxed))
}
}
impl Default for ParamId {
fn default() -> Self {
Self::new()
}
}
impl core::fmt::Display for ParamId {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "param_{}", self.0)
}
}
/// A trait for defining parameter types that can be suggested by a [`Trial`].
///
/// Implementors specify the distribution to sample from and how to convert
/// the raw [`ParamValue`] back into a typed value.
pub trait Parameter: Debug {
/// The typed value returned after sampling.
type Value;
/// Returns the unique identifier for this parameter.
fn id(&self) -> ParamId;
/// Returns the distribution that this parameter samples from.
fn distribution(&self) -> Distribution;
/// Converts a raw [`ParamValue`] into the typed value.
///
/// # Errors
///
/// Returns an error if the `ParamValue` variant doesn't match what this parameter expects.
fn cast_param_value(&self, param_value: &ParamValue) -> Result<Self::Value>;
/// Validates the parameter configuration.
///
/// Called before sampling. The default implementation accepts all configurations.
///
/// # Errors
///
/// Returns an error if the parameter configuration is invalid.
fn validate(&self) -> Result<()> {
Ok(())
}
/// Returns a human-readable label for this parameter.
///
/// Defaults to the `Debug` output of the parameter.
fn label(&self) -> String {
format!("{self:?}")
}
/// Suggests a value for this parameter from the given trial.
///
/// This is a convenience method that delegates to [`Trial::suggest_param`].
///
/// # Errors
///
/// Returns an error if validation fails, the parameter conflicts with
/// a previously suggested parameter of the same id, or sampling fails.
fn suggest(&self, trial: &mut Trial) -> Result<Self::Value>
where
Self: Sized,
{
trial.suggest_param(self)
}
}
/// A floating-point parameter with optional log-scale and step size.
///
/// # Example
///
/// ```
/// use optimizer::Trial;
/// use optimizer::parameter::{FloatParam, Parameter};
///
/// let mut trial = Trial::new(0);
///
/// // Simple range
/// let x = FloatParam::new(0.0, 1.0).suggest(&mut trial).unwrap();
///
/// // Log-scale
/// let lr = FloatParam::new(1e-5, 1e-1)
/// .log_scale()
/// .suggest(&mut trial)
/// .unwrap();
///
/// // Stepped
/// let step = FloatParam::new(0.0, 1.0)
/// .step(0.25)
/// .suggest(&mut trial)
/// .unwrap();
/// ```
#[derive(Clone, Debug)]
pub struct FloatParam {
id: ParamId,
low: f64,
high: f64,
log_scale: bool,
step: Option<f64>,
}
impl FloatParam {
/// Creates a new float parameter with the given bounds.
#[must_use]
pub fn new(low: f64, high: f64) -> Self {
Self {
id: ParamId::new(),
low,
high,
log_scale: false,
step: None,
}
}
/// Enables log-scale sampling.
#[must_use]
pub fn log_scale(mut self) -> Self {
self.log_scale = true;
self
}
/// Sets a step size for discretized sampling.
#[must_use]
pub fn step(mut self, step: f64) -> Self {
self.step = Some(step);
self
}
}
impl Parameter for FloatParam {
type Value = f64;
fn id(&self) -> ParamId {
self.id
}
fn distribution(&self) -> Distribution {
Distribution::Float(FloatDistribution {
low: self.low,
high: self.high,
log_scale: self.log_scale,
step: self.step,
})
}
fn cast_param_value(&self, param_value: &ParamValue) -> Result<f64> {
match param_value {
ParamValue::Float(v) => Ok(*v),
_ => Err(Error::Internal(
"Float distribution should return Float value",
)),
}
}
fn validate(&self) -> Result<()> {
if self.low > self.high {
return Err(Error::InvalidBounds {
low: self.low,
high: self.high,
});
}
if self.log_scale && self.low <= 0.0 {
return Err(Error::InvalidLogBounds);
}
if let Some(step) = self.step
&& step <= 0.0
{
return Err(Error::InvalidStep);
}
Ok(())
}
}
/// An integer parameter with optional log-scale and step size.
///
/// # Example
///
/// ```
/// use optimizer::Trial;
/// use optimizer::parameter::{IntParam, Parameter};
///
/// let mut trial = Trial::new(0);
///
/// // Simple range
/// let n = IntParam::new(1, 10).suggest(&mut trial).unwrap();
///
/// // Log-scale
/// let batch = IntParam::new(1, 1024)
/// .log_scale()
/// .suggest(&mut trial)
/// .unwrap();
///
/// // Stepped
/// let units = IntParam::new(32, 512).step(32).suggest(&mut trial).unwrap();
/// ```
#[derive(Clone, Debug)]
pub struct IntParam {
id: ParamId,
low: i64,
high: i64,
log_scale: bool,
step: Option<i64>,
}
impl IntParam {
/// Creates a new integer parameter with the given bounds.
#[must_use]
pub fn new(low: i64, high: i64) -> Self {
Self {
id: ParamId::new(),
low,
high,
log_scale: false,
step: None,
}
}
/// Enables log-scale sampling.
#[must_use]
pub fn log_scale(mut self) -> Self {
self.log_scale = true;
self
}
/// Sets a step size for discretized sampling.
#[must_use]
pub fn step(mut self, step: i64) -> Self {
self.step = Some(step);
self
}
}
impl Parameter for IntParam {
type Value = i64;
fn id(&self) -> ParamId {
self.id
}
fn distribution(&self) -> Distribution {
Distribution::Int(IntDistribution {
low: self.low,
high: self.high,
log_scale: self.log_scale,
step: self.step,
})
}
#[allow(clippy::cast_precision_loss)]
fn cast_param_value(&self, param_value: &ParamValue) -> Result<i64> {
match param_value {
ParamValue::Int(v) => Ok(*v),
_ => Err(Error::Internal("Int distribution should return Int value")),
}
}
#[allow(clippy::cast_precision_loss)]
fn validate(&self) -> Result<()> {
if self.low > self.high {
return Err(Error::InvalidBounds {
low: self.low as f64,
high: self.high as f64,
});
}
if self.log_scale && self.low < 1 {
return Err(Error::InvalidLogBounds);
}
if let Some(step) = self.step
&& step <= 0
{
return Err(Error::InvalidStep);
}
Ok(())
}
}
/// A categorical parameter that selects from a list of choices.
///
/// # Example
///
/// ```
/// use optimizer::Trial;
/// use optimizer::parameter::{CategoricalParam, Parameter};
///
/// let mut trial = Trial::new(0);
/// let opt = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"])
/// .suggest(&mut trial)
/// .unwrap();
/// ```
#[derive(Clone, Debug)]
pub struct CategoricalParam<T: Clone> {
id: ParamId,
choices: Vec<T>,
}
impl<T: Clone> CategoricalParam<T> {
/// Creates a new categorical parameter with the given choices.
#[must_use]
pub fn new(choices: Vec<T>) -> Self {
Self {
id: ParamId::new(),
choices,
}
}
}
impl<T: Clone + Debug> Parameter for CategoricalParam<T> {
type Value = T;
fn id(&self) -> ParamId {
self.id
}
fn distribution(&self) -> Distribution {
Distribution::Categorical(CategoricalDistribution {
n_choices: self.choices.len(),
})
}
fn cast_param_value(&self, param_value: &ParamValue) -> Result<T> {
match param_value {
ParamValue::Categorical(index) => Ok(self.choices[*index].clone()),
_ => Err(Error::Internal(
"Categorical distribution should return Categorical value",
)),
}
}
fn validate(&self) -> Result<()> {
if self.choices.is_empty() {
return Err(Error::EmptyChoices);
}
Ok(())
}
}
/// A boolean parameter (equivalent to a categorical with `[false, true]`).
///
/// # Example
///
/// ```
/// use optimizer::Trial;
/// use optimizer::parameter::{BoolParam, Parameter};
///
/// let mut trial = Trial::new(0);
/// let dropout = BoolParam::new().suggest(&mut trial).unwrap();
/// ```
#[derive(Clone, Debug)]
pub struct BoolParam {
id: ParamId,
}
impl BoolParam {
/// Creates a new boolean parameter.
#[must_use]
pub fn new() -> Self {
Self { id: ParamId::new() }
}
}
impl Default for BoolParam {
fn default() -> Self {
Self::new()
}
}
impl Parameter for BoolParam {
type Value = bool;
fn id(&self) -> ParamId {
self.id
}
fn distribution(&self) -> Distribution {
Distribution::Categorical(CategoricalDistribution { n_choices: 2 })
}
fn cast_param_value(&self, param_value: &ParamValue) -> Result<bool> {
match param_value {
ParamValue::Categorical(index) => Ok(*index != 0),
_ => Err(Error::Internal(
"Categorical distribution should return Categorical value",
)),
}
}
}
/// A trait for enum types that can be used as categorical parameters.
///
/// This trait maps enum variants to sequential indices and back. It can be
/// derived automatically for fieldless enums using `#[derive(Categorical)]`
/// when the `derive` feature is enabled.
///
/// # Example
///
/// Manual implementation:
///
/// ```
/// use optimizer::parameter::Categorical;
///
/// #[derive(Clone)]
/// enum Activation {
/// Relu,
/// Sigmoid,
/// Tanh,
/// }
///
/// impl Categorical for Activation {
/// const N_CHOICES: usize = 3;
///
/// fn from_index(index: usize) -> Self {
/// match index {
/// 0 => Activation::Relu,
/// 1 => Activation::Sigmoid,
/// 2 => Activation::Tanh,
/// _ => panic!("invalid index"),
/// }
/// }
///
/// fn to_index(&self) -> usize {
/// match self {
/// Activation::Relu => 0,
/// Activation::Sigmoid => 1,
/// Activation::Tanh => 2,
/// }
/// }
/// }
/// ```
pub trait Categorical: Sized + Clone {
/// The number of variants in the enum.
const N_CHOICES: usize;
/// Creates an instance from a variant index.
///
/// # Panics
///
/// Panics if `index >= N_CHOICES`.
fn from_index(index: usize) -> Self;
/// Returns the index of this variant.
fn to_index(&self) -> usize;
}
/// A parameter that selects from the variants of an enum implementing [`Categorical`].
///
/// # Example
///
/// ```
/// use optimizer::Trial;
/// use optimizer::parameter::{Categorical, EnumParam, Parameter};
///
/// #[derive(Clone, Debug)]
/// enum Optimizer {
/// Sgd,
/// Adam,
/// Rmsprop,
/// }
///
/// impl Categorical for Optimizer {
/// const N_CHOICES: usize = 3;
/// fn from_index(index: usize) -> Self {
/// match index {
/// 0 => Optimizer::Sgd,
/// 1 => Optimizer::Adam,
/// 2 => Optimizer::Rmsprop,
/// _ => panic!("invalid index"),
/// }
/// }
/// fn to_index(&self) -> usize {
/// match self {
/// Optimizer::Sgd => 0,
/// Optimizer::Adam => 1,
/// Optimizer::Rmsprop => 2,
/// }
/// }
/// }
///
/// let mut trial = Trial::new(0);
/// let opt = EnumParam::<Optimizer>::new().suggest(&mut trial).unwrap();
/// ```
#[derive(Clone, Debug)]
pub struct EnumParam<T: Categorical> {
id: ParamId,
_marker: core::marker::PhantomData<T>,
}
impl<T: Categorical> EnumParam<T> {
/// Creates a new enum parameter.
#[must_use]
pub fn new() -> Self {
Self {
id: ParamId::new(),
_marker: core::marker::PhantomData,
}
}
}
impl<T: Categorical> Default for EnumParam<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: Categorical + Debug> Parameter for EnumParam<T> {
type Value = T;
fn id(&self) -> ParamId {
self.id
}
fn distribution(&self) -> Distribution {
Distribution::Categorical(CategoricalDistribution {
n_choices: T::N_CHOICES,
})
}
fn cast_param_value(&self, param_value: &ParamValue) -> Result<T> {
match param_value {
ParamValue::Categorical(index) => Ok(T::from_index(*index)),
_ => Err(Error::Internal(
"Categorical distribution should return Categorical value",
)),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn float_param_distribution() {
let param = FloatParam::new(0.0, 1.0);
assert_eq!(
param.distribution(),
Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
})
);
}
#[test]
fn float_param_log_scale() {
let param = FloatParam::new(1e-5, 1e-1).log_scale();
assert_eq!(
param.distribution(),
Distribution::Float(FloatDistribution {
low: 1e-5,
high: 1e-1,
log_scale: true,
step: None,
})
);
}
#[test]
fn float_param_step() {
let param = FloatParam::new(0.0, 1.0).step(0.25);
assert_eq!(
param.distribution(),
Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: Some(0.25),
})
);
}
#[test]
fn float_param_validate_invalid_bounds() {
let param = FloatParam::new(1.0, 0.0);
assert!(param.validate().is_err());
}
#[test]
fn float_param_validate_invalid_log() {
let param = FloatParam::new(-1.0, 1.0).log_scale();
assert!(param.validate().is_err());
}
#[test]
fn float_param_validate_invalid_step() {
let param = FloatParam::new(0.0, 1.0).step(-0.1);
assert!(param.validate().is_err());
}
#[test]
#[allow(clippy::float_cmp)]
fn float_param_cast_param_value() {
let param = FloatParam::new(0.0, 1.0);
assert_eq!(
param.cast_param_value(&ParamValue::Float(0.5)).unwrap(),
0.5
);
assert!(param.cast_param_value(&ParamValue::Int(1)).is_err());
}
#[test]
fn int_param_distribution() {
let param = IntParam::new(1, 100);
assert_eq!(
param.distribution(),
Distribution::Int(IntDistribution {
low: 1,
high: 100,
log_scale: false,
step: None,
})
);
}
#[test]
fn int_param_log_scale() {
let param = IntParam::new(1, 1024).log_scale();
assert_eq!(
param.distribution(),
Distribution::Int(IntDistribution {
low: 1,
high: 1024,
log_scale: true,
step: None,
})
);
}
#[test]
fn int_param_step() {
let param = IntParam::new(100, 500).step(50);
assert_eq!(
param.distribution(),
Distribution::Int(IntDistribution {
low: 100,
high: 500,
log_scale: false,
step: Some(50),
})
);
}
#[test]
fn int_param_validate_invalid_bounds() {
let param = IntParam::new(10, 1);
assert!(param.validate().is_err());
}
#[test]
fn int_param_validate_invalid_log() {
let param = IntParam::new(0, 10).log_scale();
assert!(param.validate().is_err());
}
#[test]
fn int_param_validate_invalid_step() {
let param = IntParam::new(0, 10).step(-1);
assert!(param.validate().is_err());
}
#[test]
fn int_param_cast_param_value() {
let param = IntParam::new(1, 10);
assert_eq!(param.cast_param_value(&ParamValue::Int(5)).unwrap(), 5);
assert!(param.cast_param_value(&ParamValue::Float(1.0)).is_err());
}
#[test]
fn categorical_param_distribution() {
let param = CategoricalParam::new(vec!["a", "b", "c"]);
assert_eq!(
param.distribution(),
Distribution::Categorical(CategoricalDistribution { n_choices: 3 })
);
}
#[test]
fn categorical_param_validate_empty() {
let param = CategoricalParam::<&str>::new(vec![]);
assert!(param.validate().is_err());
}
#[test]
fn categorical_param_cast_param_value() {
let param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]);
assert_eq!(
param.cast_param_value(&ParamValue::Categorical(1)).unwrap(),
"adam"
);
assert!(param.cast_param_value(&ParamValue::Float(1.0)).is_err());
}
#[test]
fn bool_param_distribution() {
let param = BoolParam::new();
assert_eq!(
param.distribution(),
Distribution::Categorical(CategoricalDistribution { n_choices: 2 })
);
}
#[test]
fn bool_param_cast_param_value() {
let param = BoolParam::new();
assert!(!param.cast_param_value(&ParamValue::Categorical(0)).unwrap());
assert!(param.cast_param_value(&ParamValue::Categorical(1)).unwrap());
assert!(param.cast_param_value(&ParamValue::Float(1.0)).is_err());
}
#[derive(Clone, Debug, PartialEq)]
enum TestEnum {
A,
B,
C,
}
impl Categorical for TestEnum {
const N_CHOICES: usize = 3;
fn from_index(index: usize) -> Self {
match index {
0 => TestEnum::A,
1 => TestEnum::B,
2 => TestEnum::C,
_ => panic!("invalid index"),
}
}
fn to_index(&self) -> usize {
match self {
TestEnum::A => 0,
TestEnum::B => 1,
TestEnum::C => 2,
}
}
}
#[test]
fn enum_param_distribution() {
let param = EnumParam::<TestEnum>::new();
assert_eq!(
param.distribution(),
Distribution::Categorical(CategoricalDistribution { n_choices: 3 })
);
}
#[test]
fn enum_param_cast_param_value() {
let param = EnumParam::<TestEnum>::new();
assert_eq!(
param.cast_param_value(&ParamValue::Categorical(0)).unwrap(),
TestEnum::A
);
assert_eq!(
param.cast_param_value(&ParamValue::Categorical(2)).unwrap(),
TestEnum::C
);
assert!(param.cast_param_value(&ParamValue::Float(1.0)).is_err());
}
#[test]
fn float_param_suggest_via_trial() {
let param = FloatParam::new(0.0, 1.0);
let mut trial = Trial::new(0);
let x = param.suggest(&mut trial).unwrap();
assert!((0.0..=1.0).contains(&x));
// Cached value (same param id) - exact equality expected for cached values
let x2 = param.suggest(&mut trial).unwrap();
assert!((x - x2).abs() < f64::EPSILON);
}
#[test]
fn int_param_suggest_via_trial() {
let param = IntParam::new(1, 10);
let mut trial = Trial::new(0);
let n = param.suggest(&mut trial).unwrap();
assert!((1..=10).contains(&n));
// Cached value
let n2 = param.suggest(&mut trial).unwrap();
assert_eq!(n, n2);
}
#[test]
fn categorical_param_suggest_via_trial() {
let choices = vec!["sgd", "adam", "rmsprop"];
let param = CategoricalParam::new(choices.clone());
let mut trial = Trial::new(0);
let opt = param.suggest(&mut trial).unwrap();
assert!(choices.contains(&opt));
// Cached value
let opt2 = param.suggest(&mut trial).unwrap();
assert_eq!(opt, opt2);
}
#[test]
fn bool_param_suggest_via_trial() {
let param = BoolParam::new();
let mut trial = Trial::new(0);
let val = param.suggest(&mut trial).unwrap();
let _ = val; // just check it doesn't error
// Cached value
let val2 = param.suggest(&mut trial).unwrap();
assert_eq!(val, val2);
}
#[test]
fn enum_param_suggest_via_trial() {
let param = EnumParam::<TestEnum>::new();
let mut trial = Trial::new(0);
let val = param.suggest(&mut trial).unwrap();
assert!([TestEnum::A, TestEnum::B, TestEnum::C].contains(&val));
// Cached value
let val2 = param.suggest(&mut trial).unwrap();
assert_eq!(val, val2);
}
#[test]
fn parameter_conflict_detection() {
let param_float = FloatParam::new(0.0, 1.0);
let param_int = IntParam::new(0, 10);
let mut trial = Trial::new(0);
let _ = param_float.suggest(&mut trial).unwrap();
// Different param with different distribution but same id won't happen
// since each param gets a unique id. But different param object = different id = no conflict.
let result = param_int.suggest(&mut trial);
assert!(result.is_ok()); // Different id, no conflict
}
#[test]
fn float_param_validation_prevents_suggest() {
let param = FloatParam::new(1.0, 0.0);
let mut trial = Trial::new(0);
let result = param.suggest(&mut trial);
assert!(result.is_err());
}
#[test]
fn categorical_trait_roundtrip() {
for i in 0..TestEnum::N_CHOICES {
let val = TestEnum::from_index(i);
assert_eq!(val.to_index(), i);
}
}
#[test]
fn param_id_uniqueness() {
let id1 = ParamId::new();
let id2 = ParamId::new();
assert_ne!(id1, id2);
}
#[test]
fn param_clone_preserves_id() {
let param = FloatParam::new(0.0, 1.0);
let cloned = param.clone();
assert_eq!(param.id(), cloned.id());
}
}
src/sampler/mod.rs
+21 -54
View File
@@ -8,6 +8,7 @@ use std::collections::HashMap;
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::parameter::ParamId;
/// A completed trial with its parameters, distributions, and objective value.
///
@@ -18,10 +19,12 @@ use crate::param::ParamValue;
pub struct CompletedTrial<V = f64> {
/// 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>,
/// The sampled parameter values, keyed by parameter id.
pub params: HashMap<ParamId, ParamValue>,
/// The parameter distributions used, keyed by parameter id.
pub distributions: HashMap<ParamId, Distribution>,
/// Human-readable labels for parameters, keyed by parameter id.
pub param_labels: HashMap<ParamId, String>,
/// The objective value returned by the objective function.
pub value: V,
}
@@ -30,14 +33,16 @@ impl<V> CompletedTrial<V> {
/// Creates a new completed trial.
pub fn new(
id: u64,
params: HashMap<String, ParamValue>,
distributions: HashMap<String, Distribution>,
params: HashMap<ParamId, ParamValue>,
distributions: HashMap<ParamId, Distribution>,
param_labels: HashMap<ParamId, String>,
value: V,
) -> Self {
Self {
id,
params,
distributions,
param_labels,
value,
}
}
@@ -48,34 +53,16 @@ impl<V> CompletedTrial<V> {
/// 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>,
/// The sampled parameter values, keyed by parameter id.
pub params: HashMap<ParamId, ParamValue>,
/// The parameter distributions used, keyed by parameter id.
pub distributions: HashMap<ParamId, Distribution>,
/// Human-readable labels for parameters, keyed by parameter id.
pub param_labels: HashMap<ParamId, String>,
}
impl PendingTrial {
@@ -83,13 +70,15 @@ impl PendingTrial {
#[must_use]
pub fn new(
id: u64,
params: HashMap<String, ParamValue>,
distributions: HashMap<String, Distribution>,
params: HashMap<ParamId, ParamValue>,
distributions: HashMap<ParamId, Distribution>,
param_labels: HashMap<ParamId, String>,
) -> Self {
Self {
id,
params,
distributions,
param_labels,
}
}
}
@@ -99,28 +88,6 @@ impl PendingTrial {
/// Samplers are responsible for generating parameter values based on
/// the distribution and historical trial data. The trait requires
/// `Send + Sync` to support concurrent and async optimization.
///
/// # Examples
///
/// Implementing a custom sampler:
///
/// ```ignore
/// use optimizer::{Sampler, ParamValue, Distribution, CompletedTrial};
///
/// struct MySampler;
///
/// impl Sampler for MySampler {
/// fn sample(
/// &self,
/// distribution: &Distribution,
/// trial_id: u64,
/// history: &[CompletedTrial],
/// ) -> ParamValue {
/// // Custom sampling logic here
/// todo!()
/// }
/// }
/// ```
pub trait Sampler: Send + Sync {
/// Samples a parameter value from the given distribution.
///
src/sampler/tpe/multivariate.rs
+701 -517
View File
File diff suppressed because it is too large Load Diff
src/sampler/tpe/sampler.rs
+18 -11
View File
@@ -1024,19 +1024,20 @@ mod tests {
use super::*;
use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};
use crate::parameter::ParamId;
fn create_trial(
id: u64,
value: f64,
params: Vec<(&str, ParamValue, Distribution)>,
params: Vec<(ParamId, ParamValue, Distribution)>,
) -> CompletedTrial {
let mut param_map = HashMap::new();
let mut dist_map = HashMap::new();
for (name, pv, dist) in params {
param_map.insert(name.to_string(), pv);
dist_map.insert(name.to_string(), dist);
for (param_id, pv, dist) in params {
param_map.insert(param_id, pv);
dist_map.insert(param_id, dist);
}
CompletedTrial::new(id, param_map, dist_map, value)
CompletedTrial::new(id, param_map, dist_map, HashMap::new(), value)
}
#[test]
@@ -1104,12 +1105,13 @@ mod tests {
});
// Create 20 trials with values 0..20
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..20)
.map(|i| {
create_trial(
i as u64,
f64::from(i),
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
vec![(x_id, ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
)
})
.collect();
@@ -1138,6 +1140,7 @@ mod tests {
});
// Create history where low values (near 0.2) are "good"
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..20)
.map(|i| {
let x = f64::from(i) / 20.0;
@@ -1146,7 +1149,7 @@ mod tests {
create_trial(
i as u64,
value,
vec![("x", ParamValue::Float(x), dist.clone())],
vec![(x_id, ParamValue::Float(x), dist.clone())],
)
})
.collect();
@@ -1175,6 +1178,7 @@ mod tests {
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 4 });
// Create history where category 1 is consistently good
let cat_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..20)
.map(|i| {
let category = i % 4;
@@ -1184,7 +1188,7 @@ mod tests {
i as u64,
value,
vec![(
"cat",
cat_id,
ParamValue::Categorical(category as usize),
dist.clone(),
)],
@@ -1220,6 +1224,7 @@ mod tests {
});
// Create history where values near 30 are good
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..20)
.map(|i| {
let x = i * 5; // 0, 5, 10, ..., 95
@@ -1227,7 +1232,7 @@ mod tests {
create_trial(
i as u64,
value,
vec![("x", ParamValue::Int(x), dist.clone())],
vec![(x_id, ParamValue::Int(x), dist.clone())],
)
})
.collect();
@@ -1252,12 +1257,13 @@ mod tests {
step: None,
});
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..20)
.map(|i| {
create_trial(
i as u64,
f64::from(i),
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
vec![(x_id, ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
)
})
.collect();
@@ -1332,12 +1338,13 @@ mod tests {
step: None,
});
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..20u32)
.map(|i| {
create_trial(
u64::from(i),
f64::from(i),
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
vec![(x_id, ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
)
})
.collect();
src/sampler/tpe/search_space.rs
+232 -191
View File
@@ -11,6 +11,7 @@
use std::collections::{HashMap, HashSet, VecDeque};
use crate::distribution::Distribution;
use crate::parameter::ParamId;
use crate::sampler::CompletedTrial;
/// Computes the intersection of search spaces across completed trials.
@@ -87,34 +88,30 @@ impl IntersectionSearchSpace {
/// - For dynamic search spaces where not all parameters are sampled in every trial,
/// this helps identify the "stable" set of parameters that can be modeled jointly.
#[must_use]
pub fn calculate(trials: &[CompletedTrial]) -> HashMap<String, Distribution> {
pub fn calculate(trials: &[CompletedTrial]) -> HashMap<ParamId, Distribution> {
if trials.is_empty() {
return HashMap::new();
}
// Get parameter names from the first trial as the initial candidate set
// Get parameter ids from the first trial as the initial candidate set
let first_trial = &trials[0];
let mut candidate_params: HashSet<&str> = first_trial
.distributions
.keys()
.map(String::as_str)
.collect();
let mut candidate_params: HashSet<ParamId> =
first_trial.distributions.keys().copied().collect();
// Intersect with parameter sets from all other trials
for trial in trials.iter().skip(1) {
let trial_params: HashSet<&str> =
trial.distributions.keys().map(String::as_str).collect();
let trial_params: HashSet<ParamId> = trial.distributions.keys().copied().collect();
candidate_params.retain(|param| trial_params.contains(param));
}
// Build the result map using distributions from the first trial
// that contains each parameter
let mut result = HashMap::new();
for param_name in candidate_params {
for param_id in candidate_params {
// Find the first trial that has this parameter and use its distribution
for trial in trials {
if let Some(dist) = trial.distributions.get(param_name) {
result.insert(param_name.to_string(), dist.clone());
if let Some(dist) = trial.distributions.get(&param_id) {
result.insert(param_id, dist.clone());
break;
}
}
@@ -193,7 +190,7 @@ impl GroupDecomposedSearchSpace {
///
/// # Returns
///
/// A `Vec<HashSet<String>>` where each `HashSet` represents an independent group
/// A `Vec<HashSet<ParamId>>` where each `HashSet` represents an independent group
/// of parameters that co-occur. Parameters within the same group have appeared
/// together in at least one trial (directly or transitively). Parameters in
/// different groups have never appeared in the same trial.
@@ -212,16 +209,16 @@ impl GroupDecomposedSearchSpace {
/// - This is useful for `MultivariateTpeSampler` when `group=true` to sample
/// independent parameter groups separately.
#[must_use]
pub fn calculate(trials: &[CompletedTrial]) -> Vec<HashSet<String>> {
pub fn calculate(trials: &[CompletedTrial]) -> Vec<HashSet<ParamId>> {
if trials.is_empty() {
return Vec::new();
}
// Collect all unique parameter names
let mut all_params: HashSet<String> = HashSet::new();
// Collect all unique parameter ids
let mut all_params: HashSet<ParamId> = HashSet::new();
for trial in trials {
for param_name in trial.distributions.keys() {
all_params.insert(param_name.clone());
for &param_id in trial.distributions.keys() {
all_params.insert(param_id);
}
}
@@ -231,51 +228,51 @@ impl GroupDecomposedSearchSpace {
// Build adjacency list for co-occurrence graph
// Two parameters are adjacent if they appear in the same trial
let mut adjacency: HashMap<String, HashSet<String>> = HashMap::new();
for param in &all_params {
adjacency.insert(param.clone(), HashSet::new());
let mut adjacency: HashMap<ParamId, HashSet<ParamId>> = HashMap::new();
for &param in &all_params {
adjacency.insert(param, HashSet::new());
}
for trial in trials {
let trial_params: Vec<&String> = trial.distributions.keys().collect();
let trial_params: Vec<ParamId> = trial.distributions.keys().copied().collect();
// Connect all pairs of parameters in this trial
for (i, param1) in trial_params.iter().enumerate() {
for param2 in trial_params.iter().skip(i + 1) {
for (i, &param1) in trial_params.iter().enumerate() {
for &param2 in trial_params.iter().skip(i + 1) {
adjacency
.get_mut(*param1)
.get_mut(&param1)
.expect("param should exist in adjacency map")
.insert((*param2).clone());
.insert(param2);
adjacency
.get_mut(*param2)
.get_mut(&param2)
.expect("param should exist in adjacency map")
.insert((*param1).clone());
.insert(param1);
}
}
}
// Find connected components using BFS
let mut visited: HashSet<String> = HashSet::new();
let mut groups: Vec<HashSet<String>> = Vec::new();
let mut visited: HashSet<ParamId> = HashSet::new();
let mut groups: Vec<HashSet<ParamId>> = Vec::new();
for param in &all_params {
if visited.contains(param) {
for &param in &all_params {
if visited.contains(&param) {
continue;
}
// BFS to find all parameters in this component
let mut component: HashSet<String> = HashSet::new();
let mut queue: VecDeque<String> = VecDeque::new();
queue.push_back(param.clone());
visited.insert(param.clone());
let mut component: HashSet<ParamId> = HashSet::new();
let mut queue: VecDeque<ParamId> = VecDeque::new();
queue.push_back(param);
visited.insert(param);
while let Some(current) = queue.pop_front() {
component.insert(current.clone());
component.insert(current);
if let Some(neighbors) = adjacency.get(&current) {
for neighbor in neighbors {
if !visited.contains(neighbor) {
visited.insert(neighbor.clone());
queue.push_back(neighbor.clone());
for &neighbor in neighbors {
if !visited.contains(&neighbor) {
visited.insert(neighbor);
queue.push_back(neighbor);
}
}
}
@@ -293,19 +290,20 @@ mod tests {
use super::*;
use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};
use crate::param::ParamValue;
use crate::parameter::ParamId;
fn create_trial(
id: u64,
params: Vec<(&str, ParamValue, Distribution)>,
params: Vec<(ParamId, ParamValue, Distribution)>,
value: f64,
) -> CompletedTrial {
let mut param_map = HashMap::new();
let mut dist_map = HashMap::new();
for (name, pv, dist) in params {
param_map.insert(name.to_string(), pv);
dist_map.insert(name.to_string(), dist);
for (param_id, pv, dist) in params {
param_map.insert(param_id, pv);
dist_map.insert(param_id, dist);
}
CompletedTrial::new(id, param_map, dist_map, value)
CompletedTrial::new(id, param_map, dist_map, HashMap::new(), value)
}
#[test]
@@ -317,6 +315,8 @@ mod tests {
#[test]
fn test_single_trial() {
let x_id = ParamId::new();
let y_id = ParamId::new();
let dist_x = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
@@ -333,22 +333,24 @@ mod tests {
let trial = create_trial(
0,
vec![
("x", ParamValue::Float(0.5), dist_x.clone()),
("y", ParamValue::Int(5), dist_y.clone()),
(x_id, ParamValue::Float(0.5), dist_x.clone()),
(y_id, ParamValue::Int(5), dist_y.clone()),
],
1.0,
);
let result = IntersectionSearchSpace::calculate(&[trial]);
assert_eq!(result.len(), 2);
assert!(result.contains_key("x"));
assert!(result.contains_key("y"));
assert_eq!(result.get("x"), Some(&dist_x));
assert_eq!(result.get("y"), Some(&dist_y));
assert!(result.contains_key(&x_id));
assert!(result.contains_key(&y_id));
assert_eq!(result.get(&x_id), Some(&dist_x));
assert_eq!(result.get(&y_id), Some(&dist_y));
}
#[test]
fn test_all_trials_same_params() {
let x_id = ParamId::new();
let y_id = ParamId::new();
let dist_x = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
@@ -369,8 +371,8 @@ mod tests {
create_trial(
i,
vec![
("x", ParamValue::Float(val), dist_x.clone()),
("y", ParamValue::Float(val - 0.5), dist_y.clone()),
(x_id, ParamValue::Float(val), dist_x.clone()),
(y_id, ParamValue::Float(val - 0.5), dist_y.clone()),
],
val * val,
)
@@ -379,12 +381,15 @@ mod tests {
let result = IntersectionSearchSpace::calculate(&trials);
assert_eq!(result.len(), 2);
assert!(result.contains_key("x"));
assert!(result.contains_key("y"));
assert!(result.contains_key(&x_id));
assert!(result.contains_key(&y_id));
}
#[test]
fn test_partial_overlap() {
let x_id = ParamId::new();
let y_id = ParamId::new();
let z_id = ParamId::new();
let dist_x = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
@@ -408,8 +413,8 @@ mod tests {
let trial1 = create_trial(
0,
vec![
("x", ParamValue::Float(0.5), dist_x.clone()),
("y", ParamValue::Float(0.3), dist_y.clone()),
(x_id, ParamValue::Float(0.5), dist_x.clone()),
(y_id, ParamValue::Float(0.3), dist_y.clone()),
],
1.0,
);
@@ -418,8 +423,8 @@ mod tests {
let trial2 = create_trial(
1,
vec![
("x", ParamValue::Float(0.7), dist_x.clone()),
("z", ParamValue::Float(0.2), dist_z.clone()),
(x_id, ParamValue::Float(0.7), dist_x.clone()),
(z_id, ParamValue::Float(0.2), dist_z.clone()),
],
0.5,
);
@@ -428,24 +433,26 @@ mod tests {
let trial3 = create_trial(
2,
vec![
("x", ParamValue::Float(0.6), dist_x.clone()),
("y", ParamValue::Float(0.4), dist_y.clone()),
("z", ParamValue::Float(0.1), dist_z.clone()),
(x_id, ParamValue::Float(0.6), dist_x.clone()),
(y_id, ParamValue::Float(0.4), dist_y.clone()),
(z_id, ParamValue::Float(0.1), dist_z.clone()),
],
0.8,
);
let result = IntersectionSearchSpace::calculate(&[trial1, trial2, trial3]);
// Only "x" appears in all three trials
// Only x appears in all three trials
assert_eq!(result.len(), 1);
assert!(result.contains_key("x"));
assert!(!result.contains_key("y"));
assert!(!result.contains_key("z"));
assert!(result.contains_key(&x_id));
assert!(!result.contains_key(&y_id));
assert!(!result.contains_key(&z_id));
}
#[test]
fn test_no_common_params() {
let x_id = ParamId::new();
let y_id = ParamId::new();
let dist_x = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
@@ -460,10 +467,10 @@ mod tests {
});
// Trial 1: only x
let trial1 = create_trial(0, vec![("x", ParamValue::Float(0.5), dist_x.clone())], 1.0);
let trial1 = create_trial(0, vec![(x_id, ParamValue::Float(0.5), dist_x.clone())], 1.0);
// Trial 2: only y
let trial2 = create_trial(1, vec![("y", ParamValue::Float(0.3), dist_y.clone())], 0.5);
let trial2 = create_trial(1, vec![(y_id, ParamValue::Float(0.3), dist_y.clone())], 0.5);
let result = IntersectionSearchSpace::calculate(&[trial1, trial2]);
@@ -473,6 +480,9 @@ mod tests {
#[test]
fn test_mixed_distribution_types() {
let lr_id = ParamId::new();
let n_layers_id = ParamId::new();
let optimizer_id = ParamId::new();
let dist_float = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
@@ -490,9 +500,9 @@ mod tests {
let trial1 = create_trial(
0,
vec![
("learning_rate", ParamValue::Float(0.01), dist_float.clone()),
("n_layers", ParamValue::Int(3), dist_int.clone()),
("optimizer", ParamValue::Categorical(0), dist_cat.clone()),
(lr_id, ParamValue::Float(0.01), dist_float.clone()),
(n_layers_id, ParamValue::Int(3), dist_int.clone()),
(optimizer_id, ParamValue::Categorical(0), dist_cat.clone()),
],
1.0,
);
@@ -500,13 +510,9 @@ mod tests {
let trial2 = create_trial(
1,
vec![
(
"learning_rate",
ParamValue::Float(0.001),
dist_float.clone(),
),
("n_layers", ParamValue::Int(5), dist_int.clone()),
("optimizer", ParamValue::Categorical(1), dist_cat.clone()),
(lr_id, ParamValue::Float(0.001), dist_float.clone()),
(n_layers_id, ParamValue::Int(5), dist_int.clone()),
(optimizer_id, ParamValue::Categorical(1), dist_cat.clone()),
],
0.8,
);
@@ -514,19 +520,20 @@ mod tests {
let result = IntersectionSearchSpace::calculate(&[trial1, trial2]);
assert_eq!(result.len(), 3);
assert!(matches!(result.get(&lr_id), Some(Distribution::Float(_))));
assert!(matches!(
result.get("learning_rate"),
Some(Distribution::Float(_))
result.get(&n_layers_id),
Some(Distribution::Int(_))
));
assert!(matches!(result.get("n_layers"), Some(Distribution::Int(_))));
assert!(matches!(
result.get("optimizer"),
result.get(&optimizer_id),
Some(Distribution::Categorical(_))
));
}
#[test]
fn test_distribution_from_first_trial() {
let x_id = ParamId::new();
// Test that when distributions differ, the first trial's distribution is used
let dist_x_v1 = Distribution::Float(FloatDistribution {
low: 0.0,
@@ -543,13 +550,13 @@ mod tests {
let trial1 = create_trial(
0,
vec![("x", ParamValue::Float(0.5), dist_x_v1.clone())],
vec![(x_id, ParamValue::Float(0.5), dist_x_v1.clone())],
1.0,
);
let trial2 = create_trial(
1,
vec![("x", ParamValue::Float(5.0), dist_x_v2.clone())],
vec![(x_id, ParamValue::Float(5.0), dist_x_v2.clone())],
0.5,
);
@@ -557,13 +564,15 @@ mod tests {
assert_eq!(result.len(), 1);
// Should use the distribution from the first trial
assert_eq!(result.get("x"), Some(&dist_x_v1));
assert_eq!(result.get(&x_id), Some(&dist_x_v1));
}
#[test]
fn test_many_trials_with_conditional_params() {
let lr_id = ParamId::new();
let use_dropout_id = ParamId::new();
let dropout_rate_id = ParamId::new();
// Simulate a scenario with conditional parameters
// e.g., "use_dropout" is a bool, and "dropout_rate" only exists when use_dropout=true
let dist_lr = Distribution::Float(FloatDistribution {
low: 1e-5,
high: 1e-1,
@@ -582,14 +591,14 @@ mod tests {
let trial1 = create_trial(
0,
vec![
("lr", ParamValue::Float(0.01), dist_lr.clone()),
(lr_id, ParamValue::Float(0.01), dist_lr.clone()),
(
"use_dropout",
use_dropout_id,
ParamValue::Categorical(1),
dist_dropout.clone(),
),
(
"dropout_rate",
dropout_rate_id,
ParamValue::Float(0.2),
dist_dropout_rate.clone(),
),
@@ -601,9 +610,9 @@ mod tests {
let trial2 = create_trial(
1,
vec![
("lr", ParamValue::Float(0.001), dist_lr.clone()),
(lr_id, ParamValue::Float(0.001), dist_lr.clone()),
(
"use_dropout",
use_dropout_id,
ParamValue::Categorical(0),
dist_dropout.clone(),
),
@@ -615,14 +624,14 @@ mod tests {
let trial3 = create_trial(
2,
vec![
("lr", ParamValue::Float(0.005), dist_lr.clone()),
(lr_id, ParamValue::Float(0.005), dist_lr.clone()),
(
"use_dropout",
use_dropout_id,
ParamValue::Categorical(1),
dist_dropout.clone(),
),
(
"dropout_rate",
dropout_rate_id,
ParamValue::Float(0.3),
dist_dropout_rate.clone(),
),
@@ -632,11 +641,11 @@ mod tests {
let result = IntersectionSearchSpace::calculate(&[trial1, trial2, trial3]);
// Only "lr" and "use_dropout" appear in all trials
// Only lr and use_dropout appear in all trials
assert_eq!(result.len(), 2);
assert!(result.contains_key("lr"));
assert!(result.contains_key("use_dropout"));
assert!(!result.contains_key("dropout_rate")); // Not in trial2
assert!(result.contains_key(&lr_id));
assert!(result.contains_key(&use_dropout_id));
assert!(!result.contains_key(&dropout_rate_id)); // Not in trial2
}
// ==================== GroupDecomposedSearchSpace Tests ====================
@@ -657,12 +666,13 @@ mod tests {
step: None,
});
let trial = create_trial(0, vec![("x", ParamValue::Float(0.5), dist)], 1.0);
let x_id = ParamId::new();
let trial = create_trial(0, vec![(x_id, ParamValue::Float(0.5), dist)], 1.0);
let groups = GroupDecomposedSearchSpace::calculate(&[trial]);
assert_eq!(groups.len(), 1);
assert!(groups[0].contains("x"));
assert!(groups[0].contains(&x_id));
assert_eq!(groups[0].len(), 1);
}
@@ -675,12 +685,15 @@ mod tests {
step: None,
});
let x_id = ParamId::new();
let y_id = ParamId::new();
let z_id = ParamId::new();
let trial = create_trial(
0,
vec![
("x", ParamValue::Float(0.5), dist.clone()),
("y", ParamValue::Float(0.3), dist.clone()),
("z", ParamValue::Float(0.7), dist),
(x_id, ParamValue::Float(0.5), dist.clone()),
(y_id, ParamValue::Float(0.3), dist.clone()),
(z_id, ParamValue::Float(0.7), dist),
],
1.0,
);
@@ -690,9 +703,9 @@ mod tests {
// All params appear together, so one group
assert_eq!(groups.len(), 1);
assert_eq!(groups[0].len(), 3);
assert!(groups[0].contains("x"));
assert!(groups[0].contains("y"));
assert!(groups[0].contains("z"));
assert!(groups[0].contains(&x_id));
assert!(groups[0].contains(&y_id));
assert!(groups[0].contains(&z_id));
}
#[test]
@@ -704,12 +717,17 @@ mod tests {
step: None,
});
let x_id = ParamId::new();
let y_id = ParamId::new();
let a_id = ParamId::new();
let b_id = ParamId::new();
// Trial 1: x, y (group 1)
let trial1 = create_trial(
0,
vec![
("x", ParamValue::Float(0.1), dist.clone()),
("y", ParamValue::Float(0.2), dist.clone()),
(x_id, ParamValue::Float(0.1), dist.clone()),
(y_id, ParamValue::Float(0.2), dist.clone()),
],
1.0,
);
@@ -718,8 +736,8 @@ mod tests {
let trial2 = create_trial(
1,
vec![
("a", ParamValue::Float(0.3), dist.clone()),
("b", ParamValue::Float(0.4), dist),
(a_id, ParamValue::Float(0.3), dist.clone()),
(b_id, ParamValue::Float(0.4), dist),
],
0.5,
);
@@ -730,8 +748,8 @@ mod tests {
assert_eq!(groups.len(), 2);
// Find which group has x/y and which has a/b
let group_xy = groups.iter().find(|g| g.contains("x"));
let group_ab = groups.iter().find(|g| g.contains("a"));
let group_xy = groups.iter().find(|g| g.contains(&x_id));
let group_ab = groups.iter().find(|g| g.contains(&a_id));
assert!(group_xy.is_some());
assert!(group_ab.is_some());
@@ -740,12 +758,12 @@ mod tests {
let group_ab = group_ab.expect("group with a should exist");
assert_eq!(group_xy.len(), 2);
assert!(group_xy.contains("x"));
assert!(group_xy.contains("y"));
assert!(group_xy.contains(&x_id));
assert!(group_xy.contains(&y_id));
assert_eq!(group_ab.len(), 2);
assert!(group_ab.contains("a"));
assert!(group_ab.contains("b"));
assert!(group_ab.contains(&a_id));
assert!(group_ab.contains(&b_id));
}
#[test]
@@ -757,12 +775,16 @@ mod tests {
step: None,
});
let x_id = ParamId::new();
let y_id = ParamId::new();
let z_id = ParamId::new();
// Trial 1: x, y
let trial1 = create_trial(
0,
vec![
("x", ParamValue::Float(0.1), dist.clone()),
("y", ParamValue::Float(0.2), dist.clone()),
(x_id, ParamValue::Float(0.1), dist.clone()),
(y_id, ParamValue::Float(0.2), dist.clone()),
],
1.0,
);
@@ -771,8 +793,8 @@ mod tests {
let trial2 = create_trial(
1,
vec![
("y", ParamValue::Float(0.3), dist.clone()),
("z", ParamValue::Float(0.4), dist),
(y_id, ParamValue::Float(0.3), dist.clone()),
(z_id, ParamValue::Float(0.4), dist),
],
0.5,
);
@@ -782,9 +804,9 @@ mod tests {
// All should be in one group due to transitive connection via y
assert_eq!(groups.len(), 1);
assert_eq!(groups[0].len(), 3);
assert!(groups[0].contains("x"));
assert!(groups[0].contains("y"));
assert!(groups[0].contains("z"));
assert!(groups[0].contains(&x_id));
assert!(groups[0].contains(&y_id));
assert!(groups[0].contains(&z_id));
}
#[test]
@@ -796,33 +818,35 @@ mod tests {
step: None,
});
let a_id = ParamId::new();
let b_id = ParamId::new();
let c_id = ParamId::new();
let d_id = ParamId::new();
// Create a chain: a-b, b-c, c-d
// Trial 1: a, b
let trial1 = create_trial(
0,
vec![
("a", ParamValue::Float(0.1), dist.clone()),
("b", ParamValue::Float(0.2), dist.clone()),
(a_id, ParamValue::Float(0.1), dist.clone()),
(b_id, ParamValue::Float(0.2), dist.clone()),
],
1.0,
);
// Trial 2: b, c
let trial2 = create_trial(
1,
vec![
("b", ParamValue::Float(0.3), dist.clone()),
("c", ParamValue::Float(0.4), dist.clone()),
(b_id, ParamValue::Float(0.3), dist.clone()),
(c_id, ParamValue::Float(0.4), dist.clone()),
],
0.5,
);
// Trial 3: c, d
let trial3 = create_trial(
2,
vec![
("c", ParamValue::Float(0.5), dist.clone()),
("d", ParamValue::Float(0.6), dist),
(c_id, ParamValue::Float(0.5), dist.clone()),
(d_id, ParamValue::Float(0.6), dist),
],
0.3,
);
@@ -832,10 +856,10 @@ mod tests {
// All should be connected in one group
assert_eq!(groups.len(), 1);
assert_eq!(groups[0].len(), 4);
assert!(groups[0].contains("a"));
assert!(groups[0].contains("b"));
assert!(groups[0].contains("c"));
assert!(groups[0].contains("d"));
assert!(groups[0].contains(&a_id));
assert!(groups[0].contains(&b_id));
assert!(groups[0].contains(&c_id));
assert!(groups[0].contains(&d_id));
}
#[test]
@@ -847,10 +871,14 @@ mod tests {
step: None,
});
let x_id = ParamId::new();
let y_id = ParamId::new();
let z_id = ParamId::new();
// Each trial has exactly one parameter (all isolated)
let trial1 = create_trial(0, vec![("x", ParamValue::Float(0.1), dist.clone())], 1.0);
let trial2 = create_trial(1, vec![("y", ParamValue::Float(0.2), dist.clone())], 0.5);
let trial3 = create_trial(2, vec![("z", ParamValue::Float(0.3), dist)], 0.3);
let trial1 = create_trial(0, vec![(x_id, ParamValue::Float(0.1), dist.clone())], 1.0);
let trial2 = create_trial(1, vec![(y_id, ParamValue::Float(0.2), dist.clone())], 0.5);
let trial3 = create_trial(2, vec![(z_id, ParamValue::Float(0.3), dist)], 0.3);
let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3]);
@@ -861,10 +889,10 @@ mod tests {
}
// Verify all params are covered
let all_params: HashSet<String> = groups.iter().flatten().cloned().collect();
assert!(all_params.contains("x"));
assert!(all_params.contains("y"));
assert!(all_params.contains("z"));
let all_params: HashSet<ParamId> = groups.iter().flatten().copied().collect();
assert!(all_params.contains(&x_id));
assert!(all_params.contains(&y_id));
assert!(all_params.contains(&z_id));
}
#[test]
@@ -876,51 +904,54 @@ mod tests {
step: None,
});
let x_id = ParamId::new();
let y_id = ParamId::new();
let z_id = ParamId::new();
let a_id = ParamId::new();
let b_id = ParamId::new();
let w_id = ParamId::new();
// Group 1: x, y, z connected
// Group 2: a, b connected
// w is isolated
// Trial 1: x, y
let trial1 = create_trial(
0,
vec![
("x", ParamValue::Float(0.1), dist.clone()),
("y", ParamValue::Float(0.2), dist.clone()),
(x_id, ParamValue::Float(0.1), dist.clone()),
(y_id, ParamValue::Float(0.2), dist.clone()),
],
1.0,
);
// Trial 2: y, z
let trial2 = create_trial(
1,
vec![
("y", ParamValue::Float(0.3), dist.clone()),
("z", ParamValue::Float(0.4), dist.clone()),
(y_id, ParamValue::Float(0.3), dist.clone()),
(z_id, ParamValue::Float(0.4), dist.clone()),
],
0.5,
);
// Trial 3: a, b
let trial3 = create_trial(
2,
vec![
("a", ParamValue::Float(0.5), dist.clone()),
("b", ParamValue::Float(0.6), dist.clone()),
(a_id, ParamValue::Float(0.5), dist.clone()),
(b_id, ParamValue::Float(0.6), dist.clone()),
],
0.3,
);
// Trial 4: w (isolated)
let trial4 = create_trial(3, vec![("w", ParamValue::Float(0.7), dist)], 0.2);
let trial4 = create_trial(3, vec![(w_id, ParamValue::Float(0.7), dist)], 0.2);
let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3, trial4]);
// Should have 3 groups: {x,y,z}, {a,b}, {w}
assert_eq!(groups.len(), 3);
let group_xyz = groups.iter().find(|g| g.contains("x"));
let group_ab = groups.iter().find(|g| g.contains("a"));
let group_w = groups.iter().find(|g| g.contains("w"));
let group_xyz = groups.iter().find(|g| g.contains(&x_id));
let group_ab = groups.iter().find(|g| g.contains(&a_id));
let group_w = groups.iter().find(|g| g.contains(&w_id));
assert!(group_xyz.is_some());
assert!(group_ab.is_some());
@@ -928,18 +959,18 @@ mod tests {
let group_xyz = group_xyz.expect("group with x should exist");
assert_eq!(group_xyz.len(), 3);
assert!(group_xyz.contains("x"));
assert!(group_xyz.contains("y"));
assert!(group_xyz.contains("z"));
assert!(group_xyz.contains(&x_id));
assert!(group_xyz.contains(&y_id));
assert!(group_xyz.contains(&z_id));
let group_ab = group_ab.expect("group with a should exist");
assert_eq!(group_ab.len(), 2);
assert!(group_ab.contains("a"));
assert!(group_ab.contains("b"));
assert!(group_ab.contains(&a_id));
assert!(group_ab.contains(&b_id));
let group_w = group_w.expect("group with w should exist");
assert_eq!(group_w.len(), 1);
assert!(group_w.contains("w"));
assert!(group_w.contains(&w_id));
}
#[test]
@@ -951,14 +982,19 @@ mod tests {
step: None,
});
let a_id = ParamId::new();
let b_id = ParamId::new();
let c_id = ParamId::new();
let d_id = ParamId::new();
// All params in single trial
let trial = create_trial(
0,
vec![
("a", ParamValue::Float(0.1), dist.clone()),
("b", ParamValue::Float(0.2), dist.clone()),
("c", ParamValue::Float(0.3), dist.clone()),
("d", ParamValue::Float(0.4), dist),
(a_id, ParamValue::Float(0.1), dist.clone()),
(b_id, ParamValue::Float(0.2), dist.clone()),
(c_id, ParamValue::Float(0.3), dist.clone()),
(d_id, ParamValue::Float(0.4), dist),
],
1.0,
);
@@ -972,6 +1008,9 @@ mod tests {
#[test]
fn test_group_with_mixed_distribution_types() {
let lr_id = ParamId::new();
let n_layers_id = ParamId::new();
let optimizer_id = ParamId::new();
let dist_float = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
@@ -991,23 +1030,23 @@ mod tests {
let trial1 = create_trial(
0,
vec![
("learning_rate", ParamValue::Float(0.01), dist_float.clone()),
("n_layers", ParamValue::Int(3), dist_int.clone()),
(lr_id, ParamValue::Float(0.01), dist_float.clone()),
(n_layers_id, ParamValue::Int(3), dist_int.clone()),
],
1.0,
);
let trial2 = create_trial(
1,
vec![("optimizer", ParamValue::Categorical(1), dist_cat)],
vec![(optimizer_id, ParamValue::Categorical(1), dist_cat)],
0.5,
);
let trial3 = create_trial(
2,
vec![
("learning_rate", ParamValue::Float(0.001), dist_float),
("n_layers", ParamValue::Int(5), dist_int),
(lr_id, ParamValue::Float(0.001), dist_float),
(n_layers_id, ParamValue::Int(5), dist_int),
],
0.8,
);
@@ -1017,20 +1056,20 @@ mod tests {
// Should have 2 groups: {learning_rate, n_layers} and {optimizer}
assert_eq!(groups.len(), 2);
let group_lr = groups.iter().find(|g| g.contains("learning_rate"));
let group_opt = groups.iter().find(|g| g.contains("optimizer"));
let group_lr = groups.iter().find(|g| g.contains(&lr_id));
let group_opt = groups.iter().find(|g| g.contains(&optimizer_id));
assert!(group_lr.is_some());
assert!(group_opt.is_some());
let group_lr = group_lr.expect("group with learning_rate should exist");
assert_eq!(group_lr.len(), 2);
assert!(group_lr.contains("learning_rate"));
assert!(group_lr.contains("n_layers"));
assert!(group_lr.contains(&lr_id));
assert!(group_lr.contains(&n_layers_id));
let group_opt = group_opt.expect("group with optimizer should exist");
assert_eq!(group_opt.len(), 1);
assert!(group_opt.contains("optimizer"));
assert!(group_opt.contains(&optimizer_id));
}
#[test]
@@ -1042,15 +1081,17 @@ mod tests {
step: None,
});
let center_id = ParamId::new();
let a_id = ParamId::new();
let b_id = ParamId::new();
let c_id = ParamId::new();
// Star topology: center connects to all others
// Trial 1: center, a
// Trial 2: center, b
// Trial 3: center, c
let trial1 = create_trial(
0,
vec![
("center", ParamValue::Float(0.1), dist.clone()),
("a", ParamValue::Float(0.2), dist.clone()),
(center_id, ParamValue::Float(0.1), dist.clone()),
(a_id, ParamValue::Float(0.2), dist.clone()),
],
1.0,
);
@@ -1058,8 +1099,8 @@ mod tests {
let trial2 = create_trial(
1,
vec![
("center", ParamValue::Float(0.3), dist.clone()),
("b", ParamValue::Float(0.4), dist.clone()),
(center_id, ParamValue::Float(0.3), dist.clone()),
(b_id, ParamValue::Float(0.4), dist.clone()),
],
0.5,
);
@@ -1067,8 +1108,8 @@ mod tests {
let trial3 = create_trial(
2,
vec![
("center", ParamValue::Float(0.5), dist.clone()),
("c", ParamValue::Float(0.6), dist),
(center_id, ParamValue::Float(0.5), dist.clone()),
(c_id, ParamValue::Float(0.6), dist),
],
0.3,
);
@@ -1078,9 +1119,9 @@ mod tests {
// All connected via center
assert_eq!(groups.len(), 1);
assert_eq!(groups[0].len(), 4);
assert!(groups[0].contains("center"));
assert!(groups[0].contains("a"));
assert!(groups[0].contains("b"));
assert!(groups[0].contains("c"));
assert!(groups[0].contains(&center_id));
assert!(groups[0].contains(&a_id));
assert!(groups[0].contains(&b_id));
assert!(groups[0].contains(&c_id));
}
}
src/study.rs
+83 -32
View File
@@ -166,11 +166,13 @@ where
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x_param = FloatParam::new(0.0, 1.0);
/// let mut trial = study.create_trial();
/// let x = trial.suggest_float("x", 0.0, 1.0).unwrap();
/// let x = x_param.suggest(&mut trial).unwrap();
/// let objective_value = x * x;
/// study.complete_trial(trial, objective_value);
///
@@ -182,6 +184,7 @@ where
trial.id(),
trial.params().clone(),
trial.distributions().clone(),
trial.param_labels().clone(),
value,
);
self.completed_trials.write().push(completed);
@@ -228,11 +231,13 @@ where
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x_param = FloatParam::new(0.0, 1.0);
/// let mut trial = study.create_trial();
/// let _ = trial.suggest_float("x", 0.0, 1.0);
/// let _ = x_param.suggest(&mut trial);
/// study.complete_trial(trial, 0.5);
///
/// for completed in study.trials() {
@@ -253,13 +258,15 @@ where
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// assert_eq!(study.n_trials(), 0);
///
/// let x_param = FloatParam::new(0.0, 1.0);
/// let mut trial = study.create_trial();
/// let _ = trial.suggest_float("x", 0.0, 1.0);
/// let _ = x_param.suggest(&mut trial);
/// study.complete_trial(trial, 0.5);
/// assert_eq!(study.n_trials(), 1);
/// ```
@@ -280,6 +287,7 @@ where
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
@@ -287,12 +295,14 @@ where
/// // Error when no trials completed
/// assert!(study.best_trial().is_err());
///
/// let x_param = FloatParam::new(0.0, 1.0);
///
/// let mut trial1 = study.create_trial();
/// let _ = trial1.suggest_float("x", 0.0, 1.0);
/// let _ = x_param.suggest(&mut trial1);
/// study.complete_trial(trial1, 0.8);
///
/// let mut trial2 = study.create_trial();
/// let _ = trial2.suggest_float("x", 0.0, 1.0);
/// let _ = x_param.suggest(&mut trial2);
/// study.complete_trial(trial2, 0.3);
///
/// let best = study.best_trial().unwrap();
@@ -343,6 +353,7 @@ where
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Maximize);
@@ -350,12 +361,14 @@ where
/// // Error when no trials completed
/// assert!(study.best_value().is_err());
///
/// let x_param = FloatParam::new(0.0, 1.0);
///
/// let mut trial1 = study.create_trial();
/// let _ = trial1.suggest_float("x", 0.0, 1.0);
/// let _ = x_param.suggest(&mut trial1);
/// study.complete_trial(trial1, 0.3);
///
/// let mut trial2 = study.create_trial();
/// let _ = trial2.suggest_float("x", 0.0, 1.0);
/// let _ = x_param.suggest(&mut trial2);
/// study.complete_trial(trial2, 0.8);
///
/// let best = study.best_value().unwrap();
@@ -392,6 +405,7 @@ where
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
@@ -399,9 +413,11 @@ where
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
///
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize(10, |trial| {
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
/// let x = x_param.suggest(trial)?;
/// Ok::<_, optimizer::Error>(x * x)
/// })
/// .unwrap();
@@ -460,6 +476,7 @@ where
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
@@ -469,12 +486,17 @@ where
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
///
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize_async(10, |mut trial| async move {
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
/// // Simulate async work (e.g., network request)
/// let value = x * x;
/// Ok::<_, optimizer::Error>((trial, value))
/// .optimize_async(10, |mut trial| {
/// let x_param = x_param.clone();
/// async move {
/// let x = x_param.suggest(&mut trial)?;
/// // Simulate async work (e.g., network request)
/// let value = x * x;
/// Ok::<_, optimizer::Error>((trial, value))
/// }
/// })
/// .await?;
///
@@ -542,6 +564,7 @@ where
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
@@ -551,12 +574,17 @@ where
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
///
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize_parallel(10, 4, |mut trial| async move {
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
/// // Async objective function (e.g., network request)
/// let value = x * x;
/// Ok::<_, optimizer::Error>((trial, value))
/// .optimize_parallel(10, 4, move |mut trial| {
/// let x_param = x_param.clone();
/// async move {
/// let x = x_param.suggest(&mut trial)?;
/// // Async objective function (e.g., network request)
/// let value = x * x;
/// Ok::<_, optimizer::Error>((trial, value))
/// }
/// })
/// .await?;
///
@@ -652,6 +680,7 @@ where
/// ```
/// use std::ops::ControlFlow;
///
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
@@ -659,11 +688,13 @@ where
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
///
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize_with_callback(
/// 100,
/// |trial| {
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
/// let x = x_param.suggest(trial)?;
/// Ok::<_, optimizer::Error>(x * x)
/// },
/// |_study, completed_trial| {
@@ -746,6 +777,7 @@ impl Study<f64> {
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
@@ -755,7 +787,8 @@ impl Study<f64> {
/// let mut trial = study.create_trial_with_sampler();
///
/// // Parameter suggestions now use the study's sampler and history
/// let x = trial.suggest_float("x", 0.0, 1.0).unwrap();
/// let x_param = FloatParam::new(0.0, 1.0);
/// let x = x_param.suggest(&mut trial).unwrap();
/// ```
pub fn create_trial_with_sampler(&self) -> Trial {
let id = self.next_trial_id();
@@ -788,6 +821,7 @@ impl Study<f64> {
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
@@ -795,9 +829,11 @@ impl Study<f64> {
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
///
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize_with_sampler(10, |trial| {
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
/// let x = x_param.suggest(trial)?;
/// Ok::<_, optimizer::Error>(x * x)
/// })
/// .unwrap();
@@ -859,6 +895,7 @@ impl Study<f64> {
/// ```
/// use std::ops::ControlFlow;
///
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
@@ -866,11 +903,13 @@ impl Study<f64> {
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
///
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize_with_callback_sampler(
/// 100,
/// |trial| {
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
/// let x = x_param.suggest(trial)?;
/// Ok::<_, optimizer::Error>(x * x)
/// },
/// |study, _completed_trial| {
@@ -958,6 +997,7 @@ impl Study<f64> {
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
@@ -967,12 +1007,17 @@ impl Study<f64> {
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
///
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize_async_with_sampler(10, |mut trial| async move {
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
/// // Simulate async work (e.g., network request)
/// let value = x * x;
/// Ok::<_, optimizer::Error>((trial, value))
/// .optimize_async_with_sampler(10, |mut trial| {
/// let x_param = x_param.clone();
/// async move {
/// let x = x_param.suggest(&mut trial)?;
/// // Simulate async work (e.g., network request)
/// let value = x * x;
/// Ok::<_, optimizer::Error>((trial, value))
/// }
/// })
/// .await?;
///
@@ -1040,6 +1085,7 @@ impl Study<f64> {
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
@@ -1049,12 +1095,17 @@ impl Study<f64> {
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
///
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize_parallel_with_sampler(10, 4, |mut trial| async move {
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
/// // Async objective function (e.g., network request)
/// let value = x * x;
/// Ok::<_, optimizer::Error>((trial, value))
/// .optimize_parallel_with_sampler(10, 4, move |mut trial| {
/// let x_param = x_param.clone();
/// async move {
/// let x = x_param.suggest(&mut trial)?;
/// // Async objective function (e.g., network request)
/// let value = x * x;
/// Ok::<_, optimizer::Error>((trial, value))
/// }
/// })
/// .await?;
///
src/trial.rs
+53 -767
View File
@@ -1,81 +1,17 @@
//! Trial implementation for tracking sampled parameters and trial state.
use core::ops::{Range, RangeInclusive};
use std::collections::HashMap;
use std::sync::Arc;
use parking_lot::RwLock;
use crate::distribution::{
CategoricalDistribution, Distribution, FloatDistribution, IntDistribution,
};
use crate::distribution::Distribution;
use crate::error::{Error, Result};
use crate::param::ParamValue;
use crate::parameter::{ParamId, Parameter};
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
@@ -90,10 +26,12 @@ pub struct Trial {
id: u64,
/// Current state of the trial.
state: TrialState,
/// Sampled parameter values, keyed by parameter name.
params: HashMap<String, ParamValue>,
/// Parameter distributions, keyed by parameter name.
distributions: HashMap<String, Distribution>,
/// Sampled parameter values, keyed by parameter id.
params: HashMap<ParamId, ParamValue>,
/// Parameter distributions, keyed by parameter id.
distributions: HashMap<ParamId, Distribution>,
/// Human-readable labels for parameters, keyed by parameter id.
param_labels: HashMap<ParamId, String>,
/// The sampler to use for generating parameter values.
sampler: Option<Arc<dyn Sampler>>,
/// Access to the history of completed trials (shared with Study).
@@ -107,6 +45,7 @@ impl core::fmt::Debug for Trial {
.field("state", &self.state)
.field("params", &self.params)
.field("distributions", &self.distributions)
.field("param_labels", &self.param_labels)
.field("has_sampler", &self.sampler.is_some())
.field("has_history", &self.history.is_some())
.finish()
@@ -118,7 +57,7 @@ impl Trial {
///
/// The trial starts in the `Running` state with no parameters sampled.
/// This constructor creates a trial without a sampler, which will use
/// local random sampling for suggest_* methods.
/// local random sampling for suggest methods.
///
/// For trials that use the study's sampler, use `Trial::with_sampler` instead.
///
@@ -141,6 +80,7 @@ impl Trial {
state: TrialState::Running,
params: HashMap::new(),
distributions: HashMap::new(),
param_labels: HashMap::new(),
sampler: None,
history: None,
}
@@ -166,6 +106,7 @@ impl Trial {
state: TrialState::Running,
params: HashMap::new(),
distributions: HashMap::new(),
param_labels: HashMap::new(),
sampler: Some(sampler),
history: Some(history),
}
@@ -202,16 +143,22 @@ impl Trial {
/// Returns a reference to the sampled parameters.
#[must_use]
pub fn params(&self) -> &HashMap<String, ParamValue> {
pub fn params(&self) -> &HashMap<ParamId, ParamValue> {
&self.params
}
/// Returns a reference to the parameter distributions.
#[must_use]
pub fn distributions(&self) -> &HashMap<String, Distribution> {
pub fn distributions(&self) -> &HashMap<ParamId, Distribution> {
&self.distributions
}
/// Returns a reference to the parameter labels.
#[must_use]
pub fn param_labels(&self) -> &HashMap<ParamId, String> {
&self.param_labels
}
/// Sets the trial state to Complete.
pub(crate) fn set_complete(&mut self) {
self.state = TrialState::Complete;
@@ -222,730 +169,69 @@ impl Trial {
self.state = TrialState::Failed;
}
/// Suggests a float parameter with the given bounds.
/// Suggests a parameter value using a [`Parameter`] definition.
///
/// 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.
/// This is the primary entry point for sampling parameters. It handles
/// validation, caching, conflict detection, sampling, and conversion.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive).
/// * `high` - The upper bound (inclusive).
/// * `param` - The parameter definition.
///
/// # Errors
///
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different bounds.
/// Returns an error if:
/// - The parameter fails validation
/// - The parameter conflicts with a previously suggested parameter of the same id
/// - Sampling or conversion fails
///
/// # Examples
///
/// ```
/// use optimizer::Trial;
/// use optimizer::parameter::{BoolParam, FloatParam, IntParam, Parameter};
///
/// let x_param = FloatParam::new(0.0, 1.0);
/// let n_param = IntParam::new(1, 10);
/// let flag_param = BoolParam::new();
///
/// let mut trial = Trial::new(0);
/// let x = trial.suggest_float("x", 0.0, 1.0).unwrap();
/// assert!(x >= 0.0 && x <= 1.0);
///
/// // Calling again with same bounds returns cached value
/// let x2 = trial.suggest_float("x", 0.0, 1.0).unwrap();
/// assert_eq!(x, x2);
/// let x = trial.suggest_param(&x_param).unwrap();
/// let n = trial.suggest_param(&n_param).unwrap();
/// let flag = trial.suggest_param(&flag_param).unwrap();
/// ```
pub fn suggest_float(&mut self, name: impl Into<String>, low: f64, high: f64) -> Result<f64> {
if low > high {
return Err(Error::InvalidBounds { low, high });
}
pub fn suggest_param<P: Parameter>(&mut self, param: &P) -> Result<P::Value> {
param.validate()?;
let name = name.into();
let distribution = FloatDistribution {
low,
high,
log_scale: false,
step: None,
};
let param_id = param.id();
let distribution = param.distribution();
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Float(existing) = existing_dist
&& (existing.low - low).abs() < f64::EPSILON
&& (existing.high - high).abs() < f64::EPSILON
&& !existing.log_scale
&& existing.step.is_none()
{
if let Some(existing_dist) = self.distributions.get(&param_id) {
if *existing_dist == distribution {
// Same distribution, return cached value
if let Some(ParamValue::Float(value)) = self.params.get(&name) {
return Ok(*value);
if let Some(value) = self.params.get(&param_id) {
return param.cast_param_value(value);
}
}
// Distribution exists but doesn't match
return Err(Error::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
name: param.label(),
reason: "parameter was previously sampled with different configuration or type"
.to_string(),
});
}
// Sample using the sampler
let dist = Distribution::Float(distribution);
let ParamValue::Float(value) = self.sample_value(&dist) else {
return Err(Error::Internal(
"Float distribution should return Float value",
));
};
let value = self.sample_value(&distribution);
let result = param.cast_param_value(&value)?;
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Float(value));
// Store distribution, value, and label
self.distributions.insert(param_id, distribution);
self.params.insert(param_id, value);
self.param_labels.insert(param_id, param.label());
Ok(value)
}
/// Suggests a float parameter sampled on a logarithmic scale.
///
/// The value is sampled uniformly in log space, which is useful for parameters
/// that span multiple orders of magnitude (e.g., learning rates).
///
/// If the parameter has already been sampled with the same bounds and `log_scale=true`,
/// the cached value is returned. If the parameter was sampled with different configuration,
/// a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive, must be positive).
/// * `high` - The upper bound (inclusive).
///
/// # Errors
///
/// Returns `InvalidLogBounds` if `low <= 0`.
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
///
/// # Examples
///
/// ```
/// use optimizer::Trial;
///
/// let mut trial = Trial::new(0);
/// let lr = trial
/// .suggest_float_log("learning_rate", 1e-5, 1e-1)
/// .unwrap();
/// assert!(lr >= 1e-5 && lr <= 1e-1);
///
/// // Calling again with same bounds returns cached value
/// let lr2 = trial
/// .suggest_float_log("learning_rate", 1e-5, 1e-1)
/// .unwrap();
/// assert_eq!(lr, lr2);
/// ```
pub fn suggest_float_log(
&mut self,
name: impl Into<String>,
low: f64,
high: f64,
) -> Result<f64> {
if low <= 0.0 {
return Err(Error::InvalidLogBounds);
}
if low > high {
return Err(Error::InvalidBounds { low, high });
}
let name = name.into();
let distribution = FloatDistribution {
low,
high,
log_scale: true,
step: None,
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Float(existing) = existing_dist
&& (existing.low - low).abs() < f64::EPSILON
&& (existing.high - high).abs() < f64::EPSILON
&& existing.log_scale
&& existing.step.is_none()
{
// Same distribution, return cached value
if let Some(ParamValue::Float(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(Error::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// 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(Error::Internal(
"Float distribution should return Float value",
));
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Float(value));
Ok(value)
}
/// Suggests a float parameter that snaps to a step grid.
///
/// The value is sampled from the discrete set {low, low + step, low + 2*step, ...}
/// where each value is <= high.
///
/// If the parameter has already been sampled with the same configuration,
/// the cached value is returned. If the parameter was sampled with different configuration,
/// a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive).
/// * `high` - The upper bound (inclusive).
/// * `step` - The step size (must be positive).
///
/// # Errors
///
/// Returns `InvalidStep` if `step <= 0`.
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
///
/// # Examples
///
/// ```
/// use optimizer::Trial;
///
/// let mut trial = Trial::new(0);
/// let x = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap();
/// // x will be one of: 0.0, 0.25, 0.5, 0.75, 1.0
/// assert!(x >= 0.0 && x <= 1.0);
/// assert!((x / 0.25).fract().abs() < 1e-10 || (x / 0.25).fract().abs() > 1.0 - 1e-10);
///
/// // Calling again with same bounds returns cached value
/// let x2 = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap();
/// assert_eq!(x, x2);
/// ```
pub fn suggest_float_step(
&mut self,
name: impl Into<String>,
low: f64,
high: f64,
step: f64,
) -> Result<f64> {
if step <= 0.0 {
return Err(Error::InvalidStep);
}
if low > high {
return Err(Error::InvalidBounds { low, high });
}
let name = name.into();
let distribution = FloatDistribution {
low,
high,
log_scale: false,
step: Some(step),
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Float(existing) = existing_dist
&& (existing.low - low).abs() < f64::EPSILON
&& (existing.high - high).abs() < f64::EPSILON
&& !existing.log_scale
&& existing.step == Some(step)
{
// Same distribution, return cached value
if let Some(ParamValue::Float(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(Error::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// Sample using the sampler (sampler handles step-grid)
let dist = Distribution::Float(distribution);
let ParamValue::Float(value) = self.sample_value(&dist) else {
return Err(Error::Internal(
"Float distribution should return Float value",
));
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Float(value));
Ok(value)
}
/// Suggests an integer parameter with the given bounds.
///
/// The value is sampled uniformly from the range [low, high] inclusive.
///
/// 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.
/// * `low` - The lower bound (inclusive).
/// * `high` - The upper bound (inclusive).
///
/// # Errors
///
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different bounds.
///
/// # Examples
///
/// ```
/// use optimizer::Trial;
///
/// let mut trial = Trial::new(0);
/// let n = trial.suggest_int("n_layers", 1, 10).unwrap();
/// assert!(n >= 1 && n <= 10);
///
/// // Calling again with same bounds returns cached value
/// let n2 = trial.suggest_int("n_layers", 1, 10).unwrap();
/// assert_eq!(n, n2);
/// ```
#[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(Error::InvalidBounds {
low: low as f64,
high: high as f64,
});
}
let name = name.into();
let distribution = IntDistribution {
low,
high,
log_scale: false,
step: None,
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Int(existing) = existing_dist
&& existing.low == low
&& existing.high == high
&& !existing.log_scale
&& existing.step.is_none()
{
// Same distribution, return cached value
if let Some(ParamValue::Int(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(Error::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// Sample using the sampler
let dist = Distribution::Int(distribution);
let ParamValue::Int(value) = self.sample_value(&dist) else {
return Err(Error::Internal("Int distribution should return Int value"));
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Int(value));
Ok(value)
}
/// Suggests an integer parameter sampled on a logarithmic scale.
///
/// The value is sampled uniformly in log space, which is useful for parameters
/// that span multiple orders of magnitude (e.g., batch sizes).
///
/// If the parameter has already been sampled with the same bounds and `log_scale=true`,
/// the cached value is returned. If the parameter was sampled with different configuration,
/// a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive, must be >= 1).
/// * `high` - The upper bound (inclusive).
///
/// # Errors
///
/// Returns `InvalidLogBounds` if `low < 1`.
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
///
/// # Examples
///
/// ```
/// use optimizer::Trial;
///
/// let mut trial = Trial::new(0);
/// let batch_size = trial.suggest_int_log("batch_size", 1, 1024).unwrap();
/// assert!(batch_size >= 1 && batch_size <= 1024);
///
/// // Calling again with same bounds returns cached value
/// let batch_size2 = trial.suggest_int_log("batch_size", 1, 1024).unwrap();
/// assert_eq!(batch_size, batch_size2);
/// ```
#[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(Error::InvalidLogBounds);
}
if low > high {
return Err(Error::InvalidBounds {
low: low as f64,
high: high as f64,
});
}
let name = name.into();
let distribution = IntDistribution {
low,
high,
log_scale: true,
step: None,
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Int(existing) = existing_dist
&& existing.low == low
&& existing.high == high
&& existing.log_scale
&& existing.step.is_none()
{
// Same distribution, return cached value
if let Some(ParamValue::Int(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(Error::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// 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(Error::Internal("Int distribution should return Int value"));
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Int(value));
Ok(value)
}
/// Suggests an integer parameter that snaps to a step grid.
///
/// The value is sampled from the discrete set {low, low + step, low + 2*step, ...}
/// where each value is <= high.
///
/// If the parameter has already been sampled with the same configuration,
/// the cached value is returned. If the parameter was sampled with different configuration,
/// a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive).
/// * `high` - The upper bound (inclusive).
/// * `step` - The step size (must be positive).
///
/// # Errors
///
/// Returns `InvalidStep` if `step <= 0`.
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
///
/// # Examples
///
/// ```
/// use optimizer::Trial;
///
/// let mut trial = Trial::new(0);
/// let n = trial
/// .suggest_int_step("n_estimators", 100, 500, 50)
/// .unwrap();
/// // n will be one of: 100, 150, 200, 250, 300, 350, 400, 450, 500
/// assert!(n >= 100 && n <= 500);
/// assert!((n - 100) % 50 == 0);
///
/// // Calling again with same bounds returns cached value
/// let n2 = trial
/// .suggest_int_step("n_estimators", 100, 500, 50)
/// .unwrap();
/// assert_eq!(n, n2);
/// ```
#[allow(clippy::cast_precision_loss)]
pub fn suggest_int_step(
&mut self,
name: impl Into<String>,
low: i64,
high: i64,
step: i64,
) -> Result<i64> {
if step <= 0 {
return Err(Error::InvalidStep);
}
if low > high {
return Err(Error::InvalidBounds {
low: low as f64,
high: high as f64,
});
}
let name = name.into();
let distribution = IntDistribution {
low,
high,
log_scale: false,
step: Some(step),
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Int(existing) = existing_dist
&& existing.low == low
&& existing.high == high
&& !existing.log_scale
&& existing.step == Some(step)
{
// Same distribution, return cached value
if let Some(ParamValue::Int(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(Error::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// Sample using the sampler (sampler handles step-grid)
let dist = Distribution::Int(distribution);
let ParamValue::Int(value) = self.sample_value(&dist) else {
return Err(Error::Internal("Int distribution should return Int value"));
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Int(value));
Ok(value)
}
/// Suggests a categorical parameter from the given choices.
///
/// The value is selected uniformly at random from the provided choices.
/// Internally, the index of the selected choice is stored.
///
/// If the parameter has already been sampled with the same number of choices,
/// the cached value (same index) is returned. If the parameter was sampled with
/// a different number of choices, a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `choices` - A slice of choices to select from.
///
/// # Type Parameters
///
/// * `T` - The type of the choices. Only requires `Clone`.
///
/// # Errors
///
/// Returns `EmptyChoices` if `choices` is empty.
/// Returns `ParameterConflict` if the parameter was previously sampled with a different number of choices.
///
/// # Examples
///
/// ```
/// use optimizer::Trial;
///
/// let mut trial = Trial::new(0);
/// let optimizer = trial
/// .suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])
/// .unwrap();
/// assert!(["sgd", "adam", "rmsprop"].contains(&optimizer));
///
/// // Calling again with same choices returns cached value
/// let optimizer2 = trial
/// .suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])
/// .unwrap();
/// assert_eq!(optimizer, optimizer2);
/// ```
pub fn suggest_categorical<T: Clone>(
&mut self,
name: impl Into<String>,
choices: &[T],
) -> Result<T> {
if choices.is_empty() {
return Err(Error::EmptyChoices);
}
let name = name.into();
let n_choices = choices.len();
let distribution = CategoricalDistribution { n_choices };
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Categorical(existing) = existing_dist
&& existing.n_choices == n_choices
{
// Same distribution, return cached value
if let Some(ParamValue::Categorical(index)) = self.params.get(&name) {
return Ok(choices[*index].clone());
}
}
// Distribution exists but doesn't match
return Err(Error::ParameterConflict {
name,
reason: "parameter was previously sampled with different number of choices or type"
.to_string(),
});
}
// Sample using the sampler
let dist = Distribution::Categorical(distribution);
let ParamValue::Categorical(index) = self.sample_value(&dist) else {
return Err(Error::Internal(
"Categorical distribution should return Categorical value",
));
};
// Store distribution and value (store the index)
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Categorical(index));
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())
Ok(result)
}
}
tests/async_tests.rs
+57 -21
View File
@@ -4,6 +4,7 @@
#![cfg(feature = "async")]
use optimizer::parameter::{FloatParam, Parameter};
use optimizer::sampler::random::RandomSampler;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{Direction, Error, Study};
@@ -13,10 +14,15 @@ async fn test_optimize_async_basic() {
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-10.0, 10.0);
study
.optimize_async(10, |mut trial| async move {
let x = trial.suggest_float("x", -10.0, 10.0)?;
Ok::<_, Error>((trial, x * x))
.optimize_async(10, move |mut trial| {
let x_param = x_param.clone();
async move {
let x = x_param.suggest(&mut trial)?;
Ok::<_, Error>((trial, x * x))
}
})
.await
.expect("async optimization should succeed");
@@ -36,10 +42,15 @@ async fn test_optimize_async_with_sampler() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-5.0, 5.0);
study
.optimize_async_with_sampler(15, |mut trial| async move {
let x = trial.suggest_float("x", -5.0, 5.0)?;
Ok::<_, Error>((trial, x * x))
.optimize_async_with_sampler(15, move |mut trial| {
let x_param = x_param.clone();
async move {
let x = x_param.suggest(&mut trial)?;
Ok::<_, Error>((trial, x * x))
}
})
.await
.expect("async optimization with sampler should succeed");
@@ -54,10 +65,15 @@ async fn test_optimize_parallel() {
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-10.0, 10.0);
study
.optimize_parallel(20, 4, |mut trial| async move {
let x = trial.suggest_float("x", -10.0, 10.0)?;
Ok::<_, Error>((trial, x * x))
.optimize_parallel(20, 4, move |mut trial| {
let x_param = x_param.clone();
async move {
let x = x_param.suggest(&mut trial)?;
Ok::<_, Error>((trial, x * x))
}
})
.await
.expect("parallel optimization should succeed");
@@ -75,11 +91,18 @@ async fn test_optimize_parallel_with_sampler() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-5.0, 5.0);
let y_param = FloatParam::new(-5.0, 5.0);
study
.optimize_parallel_with_sampler(15, 3, |mut trial| async move {
let x = trial.suggest_float("x", -5.0, 5.0)?;
let y = trial.suggest_float("y", -5.0, 5.0)?;
Ok::<_, Error>((trial, x * x + y * y))
.optimize_parallel_with_sampler(15, 3, move |mut trial| {
let x_param = x_param.clone();
let y_param = y_param.clone();
async move {
let x = x_param.suggest(&mut trial)?;
let y = y_param.suggest(&mut trial)?;
Ok::<_, Error>((trial, x * x + y * y))
}
})
.await
.expect("parallel optimization with sampler should succeed");
@@ -162,12 +185,15 @@ async fn test_optimize_async_partial_failures() {
let counter = std::sync::atomic::AtomicUsize::new(0);
let x_param = FloatParam::new(0.0, 10.0);
study
.optimize_async(10, |mut trial| {
.optimize_async(10, move |mut trial| {
let count = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let x_param = x_param.clone();
async move {
if count.is_multiple_of(2) {
let x = trial.suggest_float("x", 0.0, 10.0)?;
let x = x_param.suggest(&mut trial)?;
Ok::<_, Error>((trial, x))
} else {
Err(Error::NoCompletedTrials) // Use as error type
@@ -186,11 +212,16 @@ async fn test_optimize_parallel_high_concurrency() {
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(0.0, 10.0);
// Run with concurrency higher than n_trials
study
.optimize_parallel(5, 10, |mut trial| async move {
let x = trial.suggest_float("x", 0.0, 10.0)?;
Ok::<_, Error>((trial, x))
.optimize_parallel(5, 10, move |mut trial| {
let x_param = x_param.clone();
async move {
let x = x_param.suggest(&mut trial)?;
Ok::<_, Error>((trial, x))
}
})
.await
.expect("should handle high concurrency");
@@ -203,11 +234,16 @@ async fn test_optimize_parallel_single_concurrency() {
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(0.0, 10.0);
// Run with concurrency of 1 (sequential)
study
.optimize_parallel(10, 1, |mut trial| async move {
let x = trial.suggest_float("x", 0.0, 10.0)?;
Ok::<_, Error>((trial, x))
.optimize_parallel(10, 1, move |mut trial| {
let x_param = x_param.clone();
async move {
let x = x_param.suggest(&mut trial)?;
Ok::<_, Error>((trial, x))
}
})
.await
.expect("should work with single concurrency");
tests/derive_tests.rs
+61
View File
@@ -0,0 +1,61 @@
use optimizer::Trial;
use optimizer::parameter::{EnumParam, Parameter};
use optimizer_derive::Categorical;
#[derive(Clone, Debug, PartialEq, Categorical)]
enum Color {
Red,
Green,
Blue,
}
#[derive(Clone, Debug, PartialEq, Categorical)]
enum SingleVariant {
Only,
}
#[test]
fn derive_categorical_n_choices() {
use optimizer::parameter::Categorical;
assert_eq!(Color::N_CHOICES, 3);
assert_eq!(SingleVariant::N_CHOICES, 1);
}
#[test]
fn derive_categorical_roundtrip() {
use optimizer::parameter::Categorical;
for i in 0..Color::N_CHOICES {
let val = Color::from_index(i);
assert_eq!(val.to_index(), i);
}
}
#[test]
fn derive_categorical_values() {
use optimizer::parameter::Categorical;
assert_eq!(Color::from_index(0), Color::Red);
assert_eq!(Color::from_index(1), Color::Green);
assert_eq!(Color::from_index(2), Color::Blue);
assert_eq!(Color::Red.to_index(), 0);
assert_eq!(Color::Green.to_index(), 1);
assert_eq!(Color::Blue.to_index(), 2);
}
#[test]
fn derive_categorical_with_enum_param() {
let mut trial = Trial::new(0);
let param = EnumParam::<Color>::new();
let color = param.suggest(&mut trial).unwrap();
assert!([Color::Red, Color::Green, Color::Blue].contains(&color));
// Cached (same param id)
let color2 = param.suggest(&mut trial).unwrap();
assert_eq!(color, color2);
}
#[test]
fn derive_categorical_suggest_via_trial() {
let mut trial = Trial::new(0);
let color = trial.suggest_param(&EnumParam::<Color>::new()).unwrap();
assert!([Color::Red, Color::Green, Color::Blue].contains(&color));
}
tests/integration.rs
+240 -322
View File
File diff suppressed because it is too large Load Diff
tests/multivariate_tpe_integration.rs
+53 -21
View File
@@ -9,6 +9,7 @@
clippy::cast_possible_truncation
)]
use optimizer::parameter::{CategoricalParam, FloatParam, IntParam, Parameter};
use optimizer::sampler::tpe::{MultivariateTpeSampler, TpeSampler};
use optimizer::{Direction, Error, Study};
@@ -50,10 +51,13 @@ fn test_multivariate_tpe_rosenbrock_finds_good_solution() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-2.0, 2.0);
let y_param = FloatParam::new(-2.0, 4.0);
study
.optimize_with_sampler(100, |trial| {
let x = trial.suggest_float("x", -2.0, 2.0)?;
let y = trial.suggest_float("y", -2.0, 4.0)?;
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(rosenbrock(x, y))
})
.expect("optimization should succeed");
@@ -82,10 +86,13 @@ fn test_independent_tpe_rosenbrock() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-2.0, 2.0);
let y_param = FloatParam::new(-2.0, 4.0);
study
.optimize_with_sampler(100, |trial| {
let x = trial.suggest_float("x", -2.0, 2.0)?;
let y = trial.suggest_float("y", -2.0, 4.0)?;
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(rosenbrock(x, y))
})
.expect("optimization should succeed");
@@ -122,10 +129,13 @@ fn test_multivariate_tpe_outperforms_on_correlated_problem() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, multivariate_sampler);
let x_param = FloatParam::new(-2.0, 2.0);
let y_param = FloatParam::new(-2.0, 4.0);
study
.optimize_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)?;
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(rosenbrock(x, y))
})
.unwrap();
@@ -142,10 +152,13 @@ fn test_multivariate_tpe_outperforms_on_correlated_problem() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, independent_sampler);
let x_param = FloatParam::new(-2.0, 2.0);
let y_param = FloatParam::new(-2.0, 4.0);
study
.optimize_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)?;
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(rosenbrock(x, y))
})
.unwrap();
@@ -201,10 +214,13 @@ fn test_multivariate_tpe_independent_problem() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-5.0, 5.0);
let y_param = FloatParam::new(-5.0, 5.0);
study
.optimize_with_sampler(50, |trial| {
let x = trial.suggest_float("x", -5.0, 5.0)?;
let y = trial.suggest_float("y", -5.0, 5.0)?;
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(sphere(x, y))
})
.expect("optimization should succeed");
@@ -230,10 +246,13 @@ fn test_independent_tpe_independent_problem() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-5.0, 5.0);
let y_param = FloatParam::new(-5.0, 5.0);
study
.optimize_with_sampler(50, |trial| {
let x = trial.suggest_float("x", -5.0, 5.0)?;
let y = trial.suggest_float("y", -5.0, 5.0)?;
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(sphere(x, y))
})
.expect("optimization should succeed");
@@ -267,10 +286,13 @@ fn test_both_samplers_work_on_independent_problem() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-5.0, 5.0);
let y_param = FloatParam::new(-5.0, 5.0);
study
.optimize_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)?;
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(sphere(x, y))
})
.unwrap();
@@ -286,10 +308,13 @@ fn test_both_samplers_work_on_independent_problem() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-5.0, 5.0);
let y_param = FloatParam::new(-5.0, 5.0);
study
.optimize_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)?;
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(sphere(x, y))
})
.unwrap();
@@ -331,10 +356,13 @@ fn test_multivariate_tpe_with_group_decomposition() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-5.0, 5.0);
let y_param = FloatParam::new(-5.0, 5.0);
study
.optimize_with_sampler(50, |trial| {
let x = trial.suggest_float("x", -5.0, 5.0)?;
let y = trial.suggest_float("y", -5.0, 5.0)?;
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(sphere(x, y))
})
.expect("optimization should succeed");
@@ -363,11 +391,15 @@ fn test_multivariate_tpe_mixed_parameter_types() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-5.0, 5.0);
let n_param = IntParam::new(1, 10);
let mode_param = CategoricalParam::new(vec!["a", "b", "c"]);
study
.optimize_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"])?;
let x = x_param.suggest(trial)?;
let n = n_param.suggest(trial)?;
let mode = mode_param.suggest(trial)?;
// Objective depends on all parameters
let mode_factor = match mode {
tests/parameter_tests.rs
+240
View File
@@ -0,0 +1,240 @@
use optimizer::parameter::{
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter,
};
use optimizer::{Direction, Study, Trial};
#[test]
fn suggest_float_param_via_trial() {
let param = FloatParam::new(0.0, 1.0);
let mut trial = Trial::new(0);
let x = trial.suggest_param(&param).unwrap();
assert!((0.0..=1.0).contains(&x));
// Cached
let x2 = trial.suggest_param(&param).unwrap();
assert_eq!(x, x2);
}
#[test]
fn suggest_float_log_param_via_trial() {
let param = FloatParam::new(1e-5, 1e-1).log_scale();
let mut trial = Trial::new(0);
let lr = trial.suggest_param(&param).unwrap();
assert!((1e-5..=1e-1).contains(&lr));
}
#[test]
fn suggest_float_step_param_via_trial() {
let param = FloatParam::new(0.0, 1.0).step(0.25);
let mut trial = Trial::new(0);
let x = trial.suggest_param(&param).unwrap();
assert!((0.0..=1.0).contains(&x));
}
#[test]
fn suggest_int_param_via_trial() {
let param = IntParam::new(1, 10);
let mut trial = Trial::new(0);
let n = trial.suggest_param(&param).unwrap();
assert!((1..=10).contains(&n));
// Cached
let n2 = trial.suggest_param(&param).unwrap();
assert_eq!(n, n2);
}
#[test]
fn suggest_int_log_param_via_trial() {
let param = IntParam::new(1, 1024).log_scale();
let mut trial = Trial::new(0);
let batch = trial.suggest_param(&param).unwrap();
assert!((1..=1024).contains(&batch));
}
#[test]
fn suggest_int_step_param_via_trial() {
let param = IntParam::new(32, 512).step(32);
let mut trial = Trial::new(0);
let units = trial.suggest_param(&param).unwrap();
assert!((32..=512).contains(&units));
assert_eq!((units - 32) % 32, 0);
}
#[test]
fn suggest_categorical_param_via_trial() {
let choices = vec!["sgd", "adam", "rmsprop"];
let param = CategoricalParam::new(choices.clone());
let mut trial = Trial::new(0);
let opt = trial.suggest_param(&param).unwrap();
assert!(choices.contains(&opt));
// Cached
let opt2 = trial.suggest_param(&param).unwrap();
assert_eq!(opt, opt2);
}
#[test]
fn suggest_bool_param_via_trial() {
let param = BoolParam::new();
let mut trial = Trial::new(0);
let val = trial.suggest_param(&param).unwrap();
let _ = val;
// Cached
let val2 = trial.suggest_param(&param).unwrap();
assert_eq!(val, val2);
}
#[derive(Clone, Debug, PartialEq)]
enum Activation {
Relu,
Sigmoid,
Tanh,
}
impl Categorical for Activation {
const N_CHOICES: usize = 3;
fn from_index(index: usize) -> Self {
match index {
0 => Activation::Relu,
1 => Activation::Sigmoid,
2 => Activation::Tanh,
_ => panic!("invalid index"),
}
}
fn to_index(&self) -> usize {
match self {
Activation::Relu => 0,
Activation::Sigmoid => 1,
Activation::Tanh => 2,
}
}
}
#[test]
fn suggest_enum_param_via_trial() {
let param = EnumParam::<Activation>::new();
let mut trial = Trial::new(0);
let act = trial.suggest_param(&param).unwrap();
assert!([Activation::Relu, Activation::Sigmoid, Activation::Tanh].contains(&act));
// Cached
let act2 = trial.suggest_param(&param).unwrap();
assert_eq!(act, act2);
}
#[test]
fn parameter_conflict_detection() {
let float_param = FloatParam::new(0.0, 1.0);
let int_param = IntParam::new(0, 10);
let mut trial = Trial::new(0);
let _ = trial.suggest_param(&float_param).unwrap();
// Different param type with different id - no conflict
let result = trial.suggest_param(&int_param);
assert!(result.is_ok());
// Different bounds for same param type but different id - no conflict
let float_param2 = FloatParam::new(0.0, 2.0);
let result = trial.suggest_param(&float_param2);
assert!(result.is_ok());
}
#[test]
fn validation_prevents_suggest() {
let mut trial = Trial::new(0);
assert!(trial.suggest_param(&FloatParam::new(1.0, 0.0)).is_err());
assert!(
trial
.suggest_param(&FloatParam::new(-1.0, 1.0).log_scale())
.is_err()
);
assert!(
trial
.suggest_param(&FloatParam::new(0.0, 1.0).step(-0.1))
.is_err()
);
assert!(trial.suggest_param(&IntParam::new(10, 1)).is_err());
assert!(
trial
.suggest_param(&IntParam::new(0, 10).log_scale())
.is_err()
);
assert!(trial.suggest_param(&IntParam::new(0, 10).step(-1)).is_err());
assert!(
trial
.suggest_param(&CategoricalParam::<&str>::new(vec![]))
.is_err()
);
}
#[test]
fn parameter_api_with_study() {
let x_param = FloatParam::new(-5.0, 5.0);
let n_param = IntParam::new(1, 10);
let dropout_param = BoolParam::new();
let opt_param = CategoricalParam::new(vec!["sgd", "adam"]);
let study: Study<f64> = Study::new(Direction::Minimize);
study
.optimize(5, |trial| {
let x = x_param.suggest(trial)?;
let n = n_param.suggest(trial)?;
let dropout = dropout_param.suggest(trial)?;
let opt = opt_param.suggest(trial)?;
let _ = (n, dropout, opt);
Ok::<_, optimizer::Error>(x * x)
})
.unwrap();
let best = study.best_trial().unwrap();
assert!(best.value >= 0.0);
}
#[test]
fn parameter_suggest_method() {
let param = FloatParam::new(0.0, 1.0);
let mut trial = Trial::new(0);
let x = param.suggest(&mut trial).unwrap();
assert!((0.0..=1.0).contains(&x));
}
#[test]
fn existing_suggest_methods_still_work() {
let mut trial = Trial::new(0);
let x_param = FloatParam::new(0.0, 1.0);
let x = x_param.suggest(&mut trial).unwrap();
assert!((0.0..=1.0).contains(&x));
let lr_param = FloatParam::new(1e-5, 1e-1).log_scale();
let lr = lr_param.suggest(&mut trial).unwrap();
assert!((1e-5..=1e-1).contains(&lr));
let step_param = FloatParam::new(0.0, 1.0).step(0.25);
let step = step_param.suggest(&mut trial).unwrap();
assert!((0.0..=1.0).contains(&step));
let n_param = IntParam::new(1, 10);
let n = n_param.suggest(&mut trial).unwrap();
assert!((1..=10).contains(&n));
let batch_param = IntParam::new(1, 1024).log_scale();
let batch = batch_param.suggest(&mut trial).unwrap();
assert!((1..=1024).contains(&batch));
let units_param = IntParam::new(32, 512).step(32);
let units = units_param.suggest(&mut trial).unwrap();
assert!((32..=512).contains(&units));
let opt_param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]);
let opt = opt_param.suggest(&mut trial).unwrap();
assert!(["sgd", "adam", "rmsprop"].contains(&opt));
let flag_param = BoolParam::new();
let flag = flag_param.suggest(&mut trial).unwrap();
let _ = flag;
}