From 55e50e6afb2ffcc81d15bd3dd82f17993a45c774 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Fri, 6 Feb 2026 17:15:30 +0100 Subject: [PATCH] Implement Parameters API --- Cargo.toml | 11 + examples/async_api_optimization.rs | 135 +-- examples/ml_hyperparameter_tuning.rs | 141 +-- examples/parameter_api.rs | 55 ++ optimizer-derive/Cargo.toml | 16 + optimizer-derive/src/lib.rs | 71 ++ src/lib.rs | 75 +- src/parameter.rs | 913 ++++++++++++++++++ src/sampler/mod.rs | 75 +- src/sampler/tpe/multivariate.rs | 1218 ++++++++++++++----------- src/sampler/tpe/sampler.rs | 29 +- src/sampler/tpe/search_space.rs | 423 +++++---- src/study.rs | 115 ++- src/trial.rs | 820 ++--------------- tests/async_tests.rs | 78 +- tests/derive_tests.rs | 61 ++ tests/integration.rs | 562 +++++------- tests/multivariate_tpe_integration.rs | 74 +- tests/parameter_tests.rs | 240 +++++ 19 files changed, 3023 insertions(+), 2089 deletions(-) create mode 100644 examples/parameter_api.rs create mode 100644 optimizer-derive/Cargo.toml create mode 100644 optimizer-derive/src/lib.rs create mode 100644 src/parameter.rs create mode 100644 tests/derive_tests.rs create mode 100644 tests/parameter_tests.rs diff --git a/Cargo.toml b/Cargo.toml index 1947bcf..660f64d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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"] diff --git a/examples/async_api_optimization.rs b/examples/async_api_optimization.rs index 7d3fc53..a9017ba 100644 --- a/examples/async_api_optimization.rs +++ b/examples/async_api_optimization.rs @@ -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) -> 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 = 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(); diff --git a/examples/ml_hyperparameter_tuning.rs b/examples/ml_hyperparameter_tuning.rs index 8b8e499..ee1d49f 100644 --- a/examples/ml_hyperparameter_tuning.rs +++ b/examples/ml_hyperparameter_tuning.rs @@ -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 { - // 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 { + 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 { /// Return `ControlFlow::Continue(())` to keep optimizing. /// Return `ControlFlow::Break(())` to stop early. fn on_trial_complete(study: &Study, trial: &CompletedTrial) -> 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, + ®_alpha_param, + ®_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}"), } } diff --git a/examples/parameter_api.rs b/examples/parameter_api.rs new file mode 100644 index 0000000..b76baf8 --- /dev/null +++ b/examples/parameter_api.rs @@ -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 = 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::::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]); + } +} diff --git a/optimizer-derive/Cargo.toml b/optimizer-derive/Cargo.toml new file mode 100644 index 0000000..eb3ed8d --- /dev/null +++ b/optimizer-derive/Cargo.toml @@ -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" diff --git a/optimizer-derive/src/lib.rs b/optimizer-derive/src/lib.rs new file mode 100644 index 0000000..4d019f9 --- /dev/null +++ b/optimizer-derive/src/lib.rs @@ -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 = (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() +} diff --git a/src/lib.rs b/src/lib.rs index a51b2b8..38d470c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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 = 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 = 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}; diff --git a/src/parameter.rs b/src/parameter.rs new file mode 100644 index 0000000..bb979ca --- /dev/null +++ b/src/parameter.rs @@ -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; + + /// 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 + 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, +} + +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 { + 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, +} + +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 { + 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 { + id: ParamId, + choices: Vec, +} + +impl CategoricalParam { + /// Creates a new categorical parameter with the given choices. + #[must_use] + pub fn new(choices: Vec) -> Self { + Self { + id: ParamId::new(), + choices, + } + } +} + +impl Parameter for CategoricalParam { + 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 { + 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 { + 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::::new().suggest(&mut trial).unwrap(); +/// ``` +#[derive(Clone, Debug)] +pub struct EnumParam { + id: ParamId, + _marker: core::marker::PhantomData, +} + +impl EnumParam { + /// Creates a new enum parameter. + #[must_use] + pub fn new() -> Self { + Self { + id: ParamId::new(), + _marker: core::marker::PhantomData, + } + } +} + +impl Default for EnumParam { + fn default() -> Self { + Self::new() + } +} + +impl Parameter for EnumParam { + 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 { + 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::::new(); + assert_eq!( + param.distribution(), + Distribution::Categorical(CategoricalDistribution { n_choices: 3 }) + ); + } + + #[test] + fn enum_param_cast_param_value() { + let param = EnumParam::::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::::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()); + } +} diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index 487b7f4..91d2207 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -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 { /// The unique identifier for this trial. pub id: u64, - /// The sampled parameter values, keyed by parameter name. - pub params: HashMap, - /// The parameter distributions used, keyed by parameter name. - pub distributions: HashMap, + /// The sampled parameter values, keyed by parameter id. + pub params: HashMap, + /// The parameter distributions used, keyed by parameter id. + pub distributions: HashMap, + /// Human-readable labels for parameters, keyed by parameter id. + pub param_labels: HashMap, /// The objective value returned by the objective function. pub value: V, } @@ -30,14 +33,16 @@ impl CompletedTrial { /// Creates a new completed trial. pub fn new( id: u64, - params: HashMap, - distributions: HashMap, + params: HashMap, + distributions: HashMap, + param_labels: HashMap, value: V, ) -> Self { Self { id, params, distributions, + param_labels, value, } } @@ -48,34 +53,16 @@ impl CompletedTrial { /// 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, - /// The parameter distributions used, keyed by parameter name. - pub distributions: HashMap, + /// The sampled parameter values, keyed by parameter id. + pub params: HashMap, + /// The parameter distributions used, keyed by parameter id. + pub distributions: HashMap, + /// Human-readable labels for parameters, keyed by parameter id. + pub param_labels: HashMap, } impl PendingTrial { @@ -83,13 +70,15 @@ impl PendingTrial { #[must_use] pub fn new( id: u64, - params: HashMap, - distributions: HashMap, + params: HashMap, + distributions: HashMap, + param_labels: HashMap, ) -> 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. /// diff --git a/src/sampler/tpe/multivariate.rs b/src/sampler/tpe/multivariate.rs index 7967222..93dcdf9 100644 --- a/src/sampler/tpe/multivariate.rs +++ b/src/sampler/tpe/multivariate.rs @@ -123,6 +123,7 @@ use super::{FixedGamma, GammaStrategy}; use crate::distribution::Distribution; use crate::error::Result; use crate::param::ParamValue; +use crate::parameter::ParamId; use crate::sampler::{CompletedTrial, PendingTrial, Sampler}; /// Strategy for imputing objective values for pending/running trials during parallel optimization. @@ -191,7 +192,7 @@ pub struct MultivariateTpeSampler { rng: Mutex, /// Cache for joint samples to maintain consistency across parameters within the same trial. /// The tuple contains (`trial_id`, cached joint sample). - joint_sample_cache: Mutex)>>, + joint_sample_cache: Mutex)>>, } impl MultivariateTpeSampler { @@ -306,6 +307,7 @@ impl MultivariateTpeSampler { /// ConstantLiarStrategy, MultivariateTpeSampler, CompletedTrial, PendingTrial, /// }; /// use optimizer::param::ParamValue; + /// use optimizer::parameter::ParamId; /// use optimizer::distribution::{Distribution, FloatDistribution}; /// /// // Create a sampler with mean imputation @@ -318,17 +320,20 @@ impl MultivariateTpeSampler { /// let dist = Distribution::Float(FloatDistribution { /// low: 0.0, high: 1.0, log_scale: false, step: None, /// }); + /// let x_id = ParamId::new(); /// let completed = vec![ /// CompletedTrial::new( /// 0, - /// [("x".to_string(), ParamValue::Float(0.2))].into_iter().collect(), - /// [("x".to_string(), dist.clone())].into_iter().collect(), + /// [(x_id, ParamValue::Float(0.2))].into_iter().collect(), + /// [(x_id, dist.clone())].into_iter().collect(), + /// HashMap::new(), /// 1.0, /// ), /// CompletedTrial::new( /// 1, - /// [("x".to_string(), ParamValue::Float(0.8))].into_iter().collect(), - /// [("x".to_string(), dist.clone())].into_iter().collect(), + /// [(x_id, ParamValue::Float(0.8))].into_iter().collect(), + /// [(x_id, dist.clone())].into_iter().collect(), + /// HashMap::new(), /// 3.0, /// ), /// ]; @@ -337,8 +342,9 @@ impl MultivariateTpeSampler { /// let pending = vec![ /// PendingTrial::new( /// 2, - /// [("x".to_string(), ParamValue::Float(0.5))].into_iter().collect(), - /// [("x".to_string(), dist.clone())].into_iter().collect(), + /// [(x_id, ParamValue::Float(0.5))].into_iter().collect(), + /// [(x_id, dist.clone())].into_iter().collect(), + /// HashMap::new(), /// ), /// ]; /// @@ -375,6 +381,7 @@ impl MultivariateTpeSampler { pending.id, pending.params.clone(), pending.distributions.clone(), + HashMap::new(), imputed_value, )); } @@ -457,7 +464,7 @@ impl MultivariateTpeSampler { pub fn filter_trials<'a>( &self, history: &'a [CompletedTrial], - search_space: &HashMap, + search_space: &HashMap, ) -> Vec<&'a CompletedTrial> { history .iter() @@ -465,7 +472,7 @@ impl MultivariateTpeSampler { // Include trial only if it has ALL parameters in the search space search_space .keys() - .all(|param| trial.params.contains_key(param)) + .all(|param_id| trial.params.contains_key(param_id)) }) .collect() } @@ -589,13 +596,16 @@ impl MultivariateTpeSampler { /// ```ignore /// use std::collections::HashMap; /// use optimizer::sampler::tpe::MultivariateTpeSampler; + /// use optimizer::parameter::ParamId; /// /// let sampler = MultivariateTpeSampler::new(); /// let trials = vec![/* ... completed trials ... */]; /// let filtered = sampler.filter_trials(&trials, &search_space); /// /// // Extract observations for x and y in that order - /// let param_order = vec!["x".to_string(), "y".to_string()]; + /// let x_id = ParamId::new(); + /// let y_id = ParamId::new(); + /// let param_order = vec![x_id, y_id]; /// let observations = sampler.extract_observations(&filtered, ¶m_order); /// /// // observations[i][0] is the x value for trial i @@ -606,15 +616,15 @@ impl MultivariateTpeSampler { pub fn extract_observations( &self, trials: &[&CompletedTrial], - param_order: &[String], + param_order: &[ParamId], ) -> Vec> { trials .iter() .map(|trial| { param_order .iter() - .filter_map(|param_name| { - trial.params.get(param_name).and_then(|value| match value { + .filter_map(|param_id| { + trial.params.get(param_id).and_then(|value| match value { crate::param::ParamValue::Float(f) => Some(*f), crate::param::ParamValue::Int(i) => Some(*i as f64), crate::param::ParamValue::Categorical(_) => None, // Skip categorical @@ -758,12 +768,12 @@ impl MultivariateTpeSampler { #[allow(clippy::unused_self)] fn sample_all_uniform( &self, - search_space: &HashMap, + search_space: &HashMap, rng: &mut rand::rngs::StdRng, - ) -> HashMap { + ) -> HashMap { search_space .iter() - .map(|(name, dist)| (name.clone(), Self::sample_uniform_single(dist, rng))) + .map(|(id, dist)| (*id, Self::sample_uniform_single(dist, rng))) .collect() } @@ -830,7 +840,7 @@ impl MultivariateTpeSampler { /// /// # Returns /// - /// A `HashMap` containing sampled values for all parameters + /// A `HashMap` containing sampled values for all parameters /// in the search space. /// /// # Algorithm @@ -855,6 +865,7 @@ impl MultivariateTpeSampler { /// ```ignore /// use std::collections::HashMap; /// use optimizer::sampler::tpe::MultivariateTpeSampler; + /// use optimizer::parameter::ParamId; /// use optimizer::distribution::{Distribution, FloatDistribution}; /// /// let sampler = MultivariateTpeSampler::builder() @@ -863,11 +874,13 @@ impl MultivariateTpeSampler { /// .build() /// .unwrap(); /// + /// let x_id = ParamId::new(); + /// let y_id = ParamId::new(); /// let mut search_space = HashMap::new(); - /// search_space.insert("x".to_string(), Distribution::Float(FloatDistribution { + /// search_space.insert(x_id, Distribution::Float(FloatDistribution { /// low: 0.0, high: 1.0, log_scale: false, step: None, /// })); - /// search_space.insert("y".to_string(), Distribution::Float(FloatDistribution { + /// search_space.insert(y_id, Distribution::Float(FloatDistribution { /// low: 0.0, high: 1.0, log_scale: false, step: None, /// })); /// @@ -878,9 +891,9 @@ impl MultivariateTpeSampler { #[allow(clippy::too_many_lines)] pub fn sample_joint( &self, - search_space: &HashMap, + search_space: &HashMap, history: &[CompletedTrial], - ) -> HashMap { + ) -> HashMap { let mut rng = self.rng.lock(); // Early returns for cases requiring random sampling @@ -915,9 +928,9 @@ impl MultivariateTpeSampler { /// A `HashMap` mapping parameter names to their sampled values. fn sample_with_groups( &self, - search_space: &HashMap, + search_space: &HashMap, history: &[CompletedTrial], - ) -> HashMap { + ) -> HashMap { use std::collections::HashSet; use super::GroupDecomposedSearchSpace; @@ -925,15 +938,15 @@ impl MultivariateTpeSampler { // Decompose the search space into independent parameter groups let groups = GroupDecomposedSearchSpace::calculate(history); - let mut result: HashMap = HashMap::new(); + let mut result: HashMap = HashMap::new(); // Sample each group independently for group in &groups { // Build a sub-search space for this group - let group_search_space: HashMap = search_space + let group_search_space: HashMap = search_space .iter() - .filter(|(name, _)| group.contains(*name)) - .map(|(name, dist)| (name.clone(), dist.clone())) + .filter(|(id, _)| group.contains(id)) + .map(|(id, dist)| (*id, dist.clone())) .collect(); if group_search_space.is_empty() { @@ -947,7 +960,7 @@ impl MultivariateTpeSampler { trial .distributions .keys() - .any(|param| group.contains(param)) + .any(|param_id| group.contains(param_id)) }) .collect(); @@ -963,25 +976,25 @@ impl MultivariateTpeSampler { drop(rng); // Merge group results into the main result - for (name, value) in group_result { - result.insert(name, value); + for (id, value) in group_result { + result.insert(id, value); } } // Handle parameters not in any group (sample independently) - let grouped_params: HashSet = groups.iter().flatten().cloned().collect(); - let ungrouped_params: HashMap = search_space + let grouped_params: HashSet = groups.iter().flatten().copied().collect(); + let ungrouped_params: HashMap = search_space .iter() - .filter(|(name, _)| !grouped_params.contains(*name) && !result.contains_key(*name)) - .map(|(name, dist)| (name.clone(), dist.clone())) + .filter(|(id, _)| !grouped_params.contains(id) && !result.contains_key(id)) + .map(|(id, dist)| (*id, dist.clone())) .collect(); if !ungrouped_params.is_empty() { // Sample ungrouped parameters uniformly (no history for them) let mut rng = self.rng.lock(); - for (name, dist) in &ungrouped_params { + for (id, dist) in &ungrouped_params { let value = Self::sample_uniform_single(dist, &mut rng); - result.insert(name.clone(), value); + result.insert(*id, value); } } @@ -1005,10 +1018,10 @@ impl MultivariateTpeSampler { #[allow(clippy::too_many_lines)] fn sample_single_group( &self, - search_space: &HashMap, + search_space: &HashMap, history: &[CompletedTrial], rng: &mut StdRng, - ) -> HashMap { + ) -> HashMap { use super::IntersectionSearchSpace; use crate::kde::MultivariateKDE; @@ -1030,22 +1043,22 @@ impl MultivariateTpeSampler { let (good, bad) = self.split_trials(&filtered); // Sample categorical parameters using TPE with l(x)/g(x) ratio - let mut result: HashMap = HashMap::new(); - for (name, dist) in &intersection { + let mut result: HashMap = HashMap::new(); + for (param_id, dist) in &intersection { if let Distribution::Categorical(d) = dist { - let good_indices = Self::extract_categorical_indices(&good, name); - let bad_indices = Self::extract_categorical_indices(&bad, name); + let good_indices = Self::extract_categorical_indices(&good, *param_id); + let bad_indices = Self::extract_categorical_indices(&bad, *param_id); let idx = Self::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng); - result.insert(name.clone(), ParamValue::Categorical(idx)); + result.insert(*param_id, ParamValue::Categorical(idx)); } } // Collect continuous parameters - let mut param_order: Vec = intersection + let mut param_order: Vec = intersection .iter() .filter(|(_, dist)| !matches!(dist, Distribution::Categorical(_))) - .map(|(name, _)| name.clone()) + .map(|(id, _)| *id) .collect(); if param_order.is_empty() { @@ -1060,7 +1073,7 @@ impl MultivariateTpeSampler { return result; } - param_order.sort(); + param_order.sort_by_key(|id| format!("{id}")); // Extract observations and validate let good_obs = self.extract_observations(&good, ¶m_order); @@ -1111,13 +1124,13 @@ impl MultivariateTpeSampler { let selected = self.select_candidate_with_rng(&good_kde, &bad_kde, rng); - // Map selected values to parameter names - for (idx, param_name) in param_order.iter().enumerate() { - if let Some(dist) = intersection.get(param_name) { + // Map selected values to parameter ids + for (idx, param_id) in param_order.iter().enumerate() { + if let Some(dist) = intersection.get(param_id) { let value = selected[idx]; let param_value = self.convert_to_param_value(value, dist); if let Some(pv) = param_value { - result.insert(param_name.clone(), pv); + result.insert(*param_id, pv); } } } @@ -1173,15 +1186,15 @@ impl MultivariateTpeSampler { #[allow(dead_code)] fn fill_remaining_independent( &self, - search_space: &HashMap, - _intersection: &HashMap, + search_space: &HashMap, + _intersection: &HashMap, history: &[CompletedTrial], - result: &mut HashMap, + result: &mut HashMap, ) { // Identify parameters not in result (and not in intersection) - let missing_params: Vec<(&String, &Distribution)> = search_space + let missing_params: Vec<(&ParamId, &Distribution)> = search_space .iter() - .filter(|(name, _)| !result.contains_key(*name)) + .filter(|(id, _)| !result.contains_key(id)) .collect(); if missing_params.is_empty() { @@ -1193,10 +1206,10 @@ impl MultivariateTpeSampler { let mut rng = self.rng.lock(); - for (name, dist) in missing_params { + for (param_id, dist) in missing_params { let value = - self.sample_independent_tpe(name, dist, &good_trials, &bad_trials, &mut rng); - result.insert(name.clone(), value); + self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, &mut rng); + result.insert(*param_id, value); } } @@ -1205,16 +1218,16 @@ impl MultivariateTpeSampler { /// This variant accepts an external RNG, used when the caller already holds the lock. fn fill_remaining_independent_with_rng( &self, - search_space: &HashMap, - _intersection: &HashMap, + search_space: &HashMap, + _intersection: &HashMap, history: &[CompletedTrial], - result: &mut HashMap, + result: &mut HashMap, rng: &mut StdRng, ) { // Identify parameters not in result (and not in intersection) - let missing_params: Vec<(&String, &Distribution)> = search_space + let missing_params: Vec<(&ParamId, &Distribution)> = search_space .iter() - .filter(|(name, _)| !result.contains_key(*name)) + .filter(|(id, _)| !result.contains_key(id)) .collect(); if missing_params.is_empty() { @@ -1224,9 +1237,10 @@ impl MultivariateTpeSampler { // Split trials for independent sampling let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); - for (name, dist) in missing_params { - let value = self.sample_independent_tpe(name, dist, &good_trials, &bad_trials, rng); - result.insert(name.clone(), value); + for (param_id, dist) in missing_params { + let value = + self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, rng); + result.insert(*param_id, value); } } @@ -1237,7 +1251,7 @@ impl MultivariateTpeSampler { #[allow(clippy::too_many_lines)] fn sample_independent_tpe( &self, - param_name: &str, + param_id: ParamId, distribution: &Distribution, good_trials: &[&CompletedTrial], bad_trials: &[&CompletedTrial], @@ -1247,7 +1261,7 @@ impl MultivariateTpeSampler { Distribution::Float(d) => { let good_values: Vec = good_trials .iter() - .filter_map(|t| t.params.get(param_name)) + .filter_map(|t| t.params.get(¶m_id)) .filter_map(|v| match v { ParamValue::Float(f) => Some(*f), _ => None, @@ -1257,7 +1271,7 @@ impl MultivariateTpeSampler { let bad_values: Vec = bad_trials .iter() - .filter_map(|t| t.params.get(param_name)) + .filter_map(|t| t.params.get(¶m_id)) .filter_map(|v| match v { ParamValue::Float(f) => Some(*f), _ => None, @@ -1283,7 +1297,7 @@ impl MultivariateTpeSampler { Distribution::Int(d) => { let good_values: Vec = good_trials .iter() - .filter_map(|t| t.params.get(param_name)) + .filter_map(|t| t.params.get(¶m_id)) .filter_map(|v| match v { ParamValue::Int(i) => Some(*i), _ => None, @@ -1293,7 +1307,7 @@ impl MultivariateTpeSampler { let bad_values: Vec = bad_trials .iter() - .filter_map(|t| t.params.get(param_name)) + .filter_map(|t| t.params.get(¶m_id)) .filter_map(|v| match v { ParamValue::Int(i) => Some(*i), _ => None, @@ -1319,7 +1333,7 @@ impl MultivariateTpeSampler { Distribution::Categorical(d) => { let good_indices: Vec = good_trials .iter() - .filter_map(|t| t.params.get(param_name)) + .filter_map(|t| t.params.get(¶m_id)) .filter_map(|v| match v { ParamValue::Categorical(i) => Some(*i), _ => None, @@ -1329,7 +1343,7 @@ impl MultivariateTpeSampler { let bad_indices: Vec = bad_trials .iter() - .filter_map(|t| t.params.get(param_name)) + .filter_map(|t| t.params.get(¶m_id)) .filter_map(|v| match v { ParamValue::Categorical(i) => Some(*i), _ => None, @@ -1483,19 +1497,19 @@ impl MultivariateTpeSampler { #[allow(dead_code)] // Used by tests fn sample_all_independent( &self, - search_space: &HashMap, + search_space: &HashMap, history: &[CompletedTrial], - ) -> HashMap { + ) -> HashMap { // Split trials for independent sampling let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); let mut rng = self.rng.lock(); let mut result = HashMap::new(); - for (name, dist) in search_space { + for (param_id, dist) in search_space { let value = - self.sample_independent_tpe(name, dist, &good_trials, &bad_trials, &mut rng); - result.insert(name.clone(), value); + self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, &mut rng); + result.insert(*param_id, value); } result @@ -1506,18 +1520,19 @@ impl MultivariateTpeSampler { /// This variant accepts an external RNG, used when the caller already holds the lock. fn sample_all_independent_with_rng( &self, - search_space: &HashMap, + search_space: &HashMap, history: &[CompletedTrial], rng: &mut StdRng, - ) -> HashMap { + ) -> HashMap { // Split trials for independent sampling let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); let mut result = HashMap::new(); - for (name, dist) in search_space { - let value = self.sample_independent_tpe(name, dist, &good_trials, &bad_trials, rng); - result.insert(name.clone(), value); + for (param_id, dist) in search_space { + let value = + self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, rng); + result.insert(*param_id, value); } result @@ -1608,11 +1623,11 @@ impl MultivariateTpeSampler { /// # Returns /// /// A vector of category indices from the trials. - fn extract_categorical_indices(trials: &[&CompletedTrial], param_name: &str) -> Vec { + fn extract_categorical_indices(trials: &[&CompletedTrial], param_id: ParamId) -> Vec { trials .iter() .filter_map(|trial| { - trial.params.get(param_name).and_then(|value| { + trial.params.get(¶m_id).and_then(|value| { if let ParamValue::Categorical(idx) = value { Some(*idx) } else { @@ -1703,7 +1718,7 @@ impl MultivariateTpeSampler { /// distribution bounds and types. fn find_matching_param( distribution: &Distribution, - cached_sample: &HashMap, + cached_sample: &HashMap, ) -> Option { // Match by distribution type and value compatibility for value in cached_sample.values() { @@ -1737,21 +1752,21 @@ impl MultivariateTpeSampler { fn build_search_space_from_history( current_distribution: &Distribution, history: &[CompletedTrial], - ) -> HashMap { + ) -> HashMap { let mut search_space = HashMap::new(); // Collect distributions from history for trial in history { - for (name, dist) in &trial.distributions { + for (param_id, dist) in &trial.distributions { search_space - .entry(name.clone()) + .entry(*param_id) .or_insert_with(|| dist.clone()); } } // If the search space is empty, create a placeholder for the current distribution if search_space.is_empty() { - search_space.insert("_current".to_string(), current_distribution.clone()); + search_space.insert(ParamId::new(), current_distribution.clone()); } search_space @@ -2456,6 +2471,7 @@ mod tests { use super::*; use crate::distribution::FloatDistribution; use crate::param::ParamValue; + use crate::parameter::ParamId; use crate::sampler::{CompletedTrial, PendingTrial}; fn float_dist() -> Distribution { @@ -2468,19 +2484,21 @@ mod tests { } fn create_completed_trial(id: u64, x_value: f64, objective: f64) -> CompletedTrial { + let x_id = ParamId::new(); let mut params = HashMap::new(); - params.insert("x".to_string(), ParamValue::Float(x_value)); + params.insert(x_id, ParamValue::Float(x_value)); let mut distributions = HashMap::new(); - distributions.insert("x".to_string(), float_dist()); - CompletedTrial::new(id, params, distributions, objective) + distributions.insert(x_id, float_dist()); + CompletedTrial::new(id, params, distributions, HashMap::new(), objective) } fn create_pending_trial(id: u64, x_value: f64) -> PendingTrial { + let x_id = ParamId::new(); let mut params = HashMap::new(); - params.insert("x".to_string(), ParamValue::Float(x_value)); + params.insert(x_id, ParamValue::Float(x_value)); let mut distributions = HashMap::new(); - distributions.insert("x".to_string(), float_dist()); - PendingTrial::new(id, params, distributions) + distributions.insert(x_id, float_dist()); + PendingTrial::new(id, params, distributions, HashMap::new()) } #[test] @@ -2726,11 +2744,12 @@ mod tests { let completed = vec![create_completed_trial(0, 0.2, 1.0)]; // Create a pending trial with specific parameter value + let x_id = ParamId::new(); let mut params = HashMap::new(); - params.insert("x".to_string(), ParamValue::Float(0.777)); + params.insert(x_id, ParamValue::Float(0.777)); let mut distributions = HashMap::new(); - distributions.insert("x".to_string(), float_dist()); - let pending = vec![PendingTrial::new(1, params, distributions)]; + distributions.insert(x_id, float_dist()); + let pending = vec![PendingTrial::new(1, params, distributions, HashMap::new())]; let result = sampler.impute_pending_trials(&pending, &completed); @@ -2739,14 +2758,14 @@ mod tests { let imputed = result.iter().find(|t| t.id == 1).unwrap(); // Parameter value should be preserved - if let Some(ParamValue::Float(v)) = imputed.params.get("x") { + if let Some(ParamValue::Float(v)) = imputed.params.get(&x_id) { assert!((*v - 0.777).abs() < f64::EPSILON); } else { - panic!("Expected Float parameter 'x'"); + panic!("Expected Float parameter"); } // Distribution should be preserved - assert!(imputed.distributions.contains_key("x")); + assert!(imputed.distributions.contains_key(&x_id)); } #[test] @@ -2800,20 +2819,21 @@ mod tests { use super::*; use crate::distribution::{FloatDistribution, IntDistribution}; use crate::param::ParamValue; + use crate::parameter::ParamId; use crate::sampler::CompletedTrial; 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) } fn float_dist() -> Distribution { @@ -2838,7 +2858,7 @@ mod tests { fn test_filter_trials_empty_history() { let sampler = MultivariateTpeSampler::new(); let history: Vec = vec![]; - let search_space: HashMap = HashMap::new(); + let search_space: HashMap = HashMap::new(); let filtered = sampler.filter_trials(&history, &search_space); assert!(filtered.is_empty()); @@ -2847,11 +2867,13 @@ mod tests { #[test] fn test_filter_trials_empty_search_space() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); let history = vec![ - create_trial(0, vec![("x", ParamValue::Float(0.5), float_dist())], 1.0), - create_trial(1, vec![("y", ParamValue::Float(0.3), float_dist())], 0.5), + create_trial(0, vec![(x_id, ParamValue::Float(0.5), float_dist())], 1.0), + create_trial(1, vec![(y_id, ParamValue::Float(0.3), float_dist())], 0.5), ]; - let search_space: HashMap = HashMap::new(); + let search_space: HashMap = HashMap::new(); // With empty search space, all trials should pass (vacuously true) let filtered = sampler.filter_trials(&history, &search_space); @@ -2861,28 +2883,30 @@ mod tests { #[test] fn test_filter_trials_all_match() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); let history = vec![ create_trial( 0, vec![ - ("x", ParamValue::Float(0.5), float_dist()), - ("y", ParamValue::Float(0.3), float_dist()), + (x_id, ParamValue::Float(0.5), float_dist()), + (y_id, ParamValue::Float(0.3), float_dist()), ], 1.0, ), create_trial( 1, vec![ - ("x", ParamValue::Float(0.7), float_dist()), - ("y", ParamValue::Float(0.2), float_dist()), + (x_id, ParamValue::Float(0.7), float_dist()), + (y_id, ParamValue::Float(0.2), float_dist()), ], 0.5, ), ]; let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist()); - search_space.insert("y".to_string(), float_dist()); + search_space.insert(x_id, float_dist()); + search_space.insert(y_id, float_dist()); let filtered = sampler.filter_trials(&history, &search_space); assert_eq!(filtered.len(), 2); @@ -2893,6 +2917,8 @@ mod tests { #[test] fn test_filter_trials_partial_match() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); // Trial 0: has x and y // Trial 1: has only x @@ -2901,17 +2927,17 @@ mod tests { create_trial( 0, vec![ - ("x", ParamValue::Float(0.5), float_dist()), - ("y", ParamValue::Float(0.3), float_dist()), + (x_id, ParamValue::Float(0.5), float_dist()), + (y_id, ParamValue::Float(0.3), float_dist()), ], 1.0, ), - create_trial(1, vec![("x", ParamValue::Float(0.7), float_dist())], 0.5), + create_trial(1, vec![(x_id, ParamValue::Float(0.7), float_dist())], 0.5), create_trial( 2, vec![ - ("x", ParamValue::Float(0.6), float_dist()), - ("y", ParamValue::Float(0.4), float_dist()), + (x_id, ParamValue::Float(0.6), float_dist()), + (y_id, ParamValue::Float(0.4), float_dist()), ], 0.8, ), @@ -2919,8 +2945,8 @@ mod tests { // Search space requires both x and y let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist()); - search_space.insert("y".to_string(), float_dist()); + search_space.insert(x_id, float_dist()); + search_space.insert(y_id, float_dist()); let filtered = sampler.filter_trials(&history, &search_space); @@ -2933,30 +2959,33 @@ mod tests { #[test] fn test_filter_trials_none_match() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + let z_id = ParamId::new(); // All trials have x, but search space requires both x and z let history = vec![ create_trial( 0, vec![ - ("x", ParamValue::Float(0.5), float_dist()), - ("y", ParamValue::Float(0.3), float_dist()), + (x_id, ParamValue::Float(0.5), float_dist()), + (y_id, ParamValue::Float(0.3), float_dist()), ], 1.0, ), create_trial( 1, vec![ - ("x", ParamValue::Float(0.7), float_dist()), - ("y", ParamValue::Float(0.2), float_dist()), + (x_id, ParamValue::Float(0.7), float_dist()), + (y_id, ParamValue::Float(0.2), float_dist()), ], 0.5, ), ]; let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist()); - search_space.insert("z".to_string(), float_dist()); // No trial has z + search_space.insert(x_id, float_dist()); + search_space.insert(z_id, float_dist()); // No trial has z let filtered = sampler.filter_trials(&history, &search_space); @@ -2967,32 +2996,35 @@ mod tests { #[test] fn test_filter_trials_mixed_param_types() { let sampler = MultivariateTpeSampler::new(); + let lr_id = ParamId::new(); + let layers_id = ParamId::new(); + let dropout_id = ParamId::new(); // Trials with mixed parameter types let history = vec![ create_trial( 0, vec![ - ("learning_rate", ParamValue::Float(0.01), float_dist()), - ("n_layers", ParamValue::Int(3), int_dist()), + (lr_id, ParamValue::Float(0.01), float_dist()), + (layers_id, ParamValue::Int(3), int_dist()), ], 1.0, ), create_trial( 1, vec![ - ("learning_rate", ParamValue::Float(0.001), float_dist()), - ("n_layers", ParamValue::Int(5), int_dist()), - ("dropout", ParamValue::Float(0.2), float_dist()), // Extra param + (lr_id, ParamValue::Float(0.001), float_dist()), + (layers_id, ParamValue::Int(5), int_dist()), + (dropout_id, ParamValue::Float(0.2), float_dist()), // Extra param ], 0.8, ), create_trial( 2, vec![ - ("learning_rate", ParamValue::Float(0.005), float_dist()), + (lr_id, ParamValue::Float(0.005), float_dist()), // Missing n_layers - ("dropout", ParamValue::Float(0.1), float_dist()), + (dropout_id, ParamValue::Float(0.1), float_dist()), ], 0.9, ), @@ -3000,8 +3032,8 @@ mod tests { // Search space requires learning_rate and n_layers let mut search_space = HashMap::new(); - search_space.insert("learning_rate".to_string(), float_dist()); - search_space.insert("n_layers".to_string(), int_dist()); + search_space.insert(lr_id, float_dist()); + search_space.insert(layers_id, int_dist()); let filtered = sampler.filter_trials(&history, &search_space); @@ -3014,24 +3046,28 @@ mod tests { #[test] fn test_filter_trials_superset_params_accepted() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + let z_id = ParamId::new(); + let w_id = ParamId::new(); // All trials have more params than the search space requires let history = vec![ create_trial( 0, vec![ - ("x", ParamValue::Float(0.5), float_dist()), - ("y", ParamValue::Float(0.3), float_dist()), - ("z", ParamValue::Float(0.1), float_dist()), + (x_id, ParamValue::Float(0.5), float_dist()), + (y_id, ParamValue::Float(0.3), float_dist()), + (z_id, ParamValue::Float(0.1), float_dist()), ], 1.0, ), create_trial( 1, vec![ - ("x", ParamValue::Float(0.7), float_dist()), - ("y", ParamValue::Float(0.2), float_dist()), - ("w", ParamValue::Float(0.9), float_dist()), // Different extra param + (x_id, ParamValue::Float(0.7), float_dist()), + (y_id, ParamValue::Float(0.2), float_dist()), + (w_id, ParamValue::Float(0.9), float_dist()), // Different extra param ], 0.5, ), @@ -3039,7 +3075,7 @@ mod tests { // Search space only requires x let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist()); + search_space.insert(x_id, float_dist()); let filtered = sampler.filter_trials(&history, &search_space); @@ -3050,13 +3086,14 @@ mod tests { #[test] fn test_filter_trials_preserves_order() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); let history: Vec = (0..10) - .map(|i| create_trial(i, vec![("x", ParamValue::Float(0.5), float_dist())], 1.0)) + .map(|i| create_trial(i, vec![(x_id, ParamValue::Float(0.5), float_dist())], 1.0)) .collect(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist()); + search_space.insert(x_id, float_dist()); let filtered = sampler.filter_trials(&history, &search_space); @@ -3071,24 +3108,26 @@ mod tests { #[test] fn test_filter_trials_single_param_search_space() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); // Some trials have the required param, some don't let history = vec![ create_trial( 0, vec![ - ("x", ParamValue::Float(0.5), float_dist()), - ("y", ParamValue::Float(0.3), float_dist()), + (x_id, ParamValue::Float(0.5), float_dist()), + (y_id, ParamValue::Float(0.3), float_dist()), ], 1.0, ), - create_trial(1, vec![("y", ParamValue::Float(0.7), float_dist())], 0.5), - create_trial(2, vec![("x", ParamValue::Float(0.6), float_dist())], 0.8), + create_trial(1, vec![(y_id, ParamValue::Float(0.7), float_dist())], 0.5), + create_trial(2, vec![(x_id, ParamValue::Float(0.6), float_dist())], 0.8), ]; // Search space only requires x let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist()); + search_space.insert(x_id, float_dist()); let filtered = sampler.filter_trials(&history, &search_space); @@ -3110,6 +3149,7 @@ mod tests { use super::*; use crate::distribution::FloatDistribution; use crate::param::ParamValue; + use crate::parameter::ParamId; use crate::sampler::CompletedTrial; fn create_trial(id: u64, value: f64) -> CompletedTrial { @@ -3119,11 +3159,12 @@ mod tests { log_scale: false, step: None, }); + let x_id = ParamId::new(); let mut params = HashMap::new(); - params.insert("x".to_string(), ParamValue::Float(0.5)); + params.insert(x_id, ParamValue::Float(0.5)); let mut distributions = HashMap::new(); - distributions.insert("x".to_string(), float_dist); - CompletedTrial::new(id, params, distributions, value) + distributions.insert(x_id, float_dist); + CompletedTrial::new(id, params, distributions, HashMap::new(), value) } #[test] @@ -3387,6 +3428,8 @@ mod tests { fn test_split_trials_integration_with_filter() { // Test that split_trials works correctly with filter_trials output let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); let float_dist = Distribution::Float(FloatDistribution { low: 0.0, @@ -3399,17 +3442,23 @@ mod tests { let mut trials_vec = vec![]; for i in 0..10 { let mut params = HashMap::new(); - params.insert("x".to_string(), ParamValue::Float(i as f64 / 10.0)); - params.insert("y".to_string(), ParamValue::Float(i as f64 / 10.0)); + params.insert(x_id, ParamValue::Float(i as f64 / 10.0)); + params.insert(y_id, ParamValue::Float(i as f64 / 10.0)); let mut distributions = HashMap::new(); - distributions.insert("x".to_string(), float_dist.clone()); - distributions.insert("y".to_string(), float_dist.clone()); - trials_vec.push(CompletedTrial::new(i, params, distributions, i as f64)); + distributions.insert(x_id, float_dist.clone()); + distributions.insert(y_id, float_dist.clone()); + trials_vec.push(CompletedTrial::new( + i, + params, + distributions, + HashMap::new(), + i as f64, + )); } let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist.clone()); - search_space.insert("y".to_string(), float_dist); + search_space.insert(x_id, float_dist.clone()); + search_space.insert(y_id, float_dist); // Filter trials let filtered = sampler.filter_trials(&trials_vec, &search_space); @@ -3440,6 +3489,7 @@ mod tests { use super::*; use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; use crate::param::ParamValue; + use crate::parameter::ParamId; use crate::sampler::CompletedTrial; fn float_dist() -> Distribution { @@ -3466,23 +3516,23 @@ mod tests { 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] fn test_extract_observations_empty_trials() { let sampler = MultivariateTpeSampler::new(); let trials: Vec<&CompletedTrial> = vec![]; - let param_order = vec!["x".to_string(), "y".to_string()]; + let param_order = vec![ParamId::new(), ParamId::new()]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3492,9 +3542,10 @@ mod tests { #[test] fn test_extract_observations_empty_param_order() { let sampler = MultivariateTpeSampler::new(); - let trial = create_trial(0, vec![("x", ParamValue::Float(0.5), float_dist())], 1.0); + let x_id = ParamId::new(); + let trial = create_trial(0, vec![(x_id, ParamValue::Float(0.5), float_dist())], 1.0); let trials: Vec<&CompletedTrial> = vec![&trial]; - let param_order: Vec = vec![]; + let param_order: Vec = vec![]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3506,13 +3557,14 @@ mod tests { #[test] fn test_extract_observations_single_float_param() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); let trial_data = [ - create_trial(0, vec![("x", ParamValue::Float(0.1), float_dist())], 1.0), - create_trial(1, vec![("x", ParamValue::Float(0.5), float_dist())], 0.5), - create_trial(2, vec![("x", ParamValue::Float(0.9), float_dist())], 0.8), + create_trial(0, vec![(x_id, ParamValue::Float(0.1), float_dist())], 1.0), + create_trial(1, vec![(x_id, ParamValue::Float(0.5), float_dist())], 0.5), + create_trial(2, vec![(x_id, ParamValue::Float(0.9), float_dist())], 0.8), ]; let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); - let param_order = vec!["x".to_string()]; + let param_order = vec![x_id]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3526,26 +3578,28 @@ mod tests { #[test] fn test_extract_observations_multiple_float_params() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); let trial_data = [ create_trial( 0, vec![ - ("x", ParamValue::Float(0.1), float_dist()), - ("y", ParamValue::Float(0.2), float_dist()), + (x_id, ParamValue::Float(0.1), float_dist()), + (y_id, ParamValue::Float(0.2), float_dist()), ], 1.0, ), create_trial( 1, vec![ - ("x", ParamValue::Float(0.3), float_dist()), - ("y", ParamValue::Float(0.4), float_dist()), + (x_id, ParamValue::Float(0.3), float_dist()), + (y_id, ParamValue::Float(0.4), float_dist()), ], 0.5, ), ]; let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); - let param_order = vec!["x".to_string(), "y".to_string()]; + let param_order = vec![x_id, y_id]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3560,21 +3614,24 @@ mod tests { #[test] fn test_extract_observations_respects_param_order() { let sampler = MultivariateTpeSampler::new(); + let a_id = ParamId::new(); + let b_id = ParamId::new(); + let c_id = ParamId::new(); let trial = create_trial( 0, vec![ - ("a", ParamValue::Float(1.0), float_dist()), - ("b", ParamValue::Float(2.0), float_dist()), - ("c", ParamValue::Float(3.0), float_dist()), + (a_id, ParamValue::Float(1.0), float_dist()), + (b_id, ParamValue::Float(2.0), float_dist()), + (c_id, ParamValue::Float(3.0), float_dist()), ], 1.0, ); let trials: Vec<&CompletedTrial> = vec![&trial]; // Different orderings - let order_abc = vec!["a".to_string(), "b".to_string(), "c".to_string()]; - let order_cba = vec!["c".to_string(), "b".to_string(), "a".to_string()]; - let order_bac = vec!["b".to_string(), "a".to_string(), "c".to_string()]; + let order_abc = vec![a_id, b_id, c_id]; + let order_cba = vec![c_id, b_id, a_id]; + let order_bac = vec![b_id, a_id, c_id]; let obs_abc = sampler.extract_observations(&trials, &order_abc); let obs_cba = sampler.extract_observations(&trials, &order_cba); @@ -3596,13 +3653,14 @@ mod tests { #[test] fn test_extract_observations_int_conversion() { let sampler = MultivariateTpeSampler::new(); + let n_layers_id = ParamId::new(); let trial_data = [ - create_trial(0, vec![("n_layers", ParamValue::Int(3), int_dist())], 1.0), - create_trial(1, vec![("n_layers", ParamValue::Int(5), int_dist())], 0.5), - create_trial(2, vec![("n_layers", ParamValue::Int(10), int_dist())], 0.8), + create_trial(0, vec![(n_layers_id, ParamValue::Int(3), int_dist())], 1.0), + create_trial(1, vec![(n_layers_id, ParamValue::Int(5), int_dist())], 0.5), + create_trial(2, vec![(n_layers_id, ParamValue::Int(10), int_dist())], 0.8), ]; let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); - let param_order = vec!["n_layers".to_string()]; + let param_order = vec![n_layers_id]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3615,32 +3673,31 @@ mod tests { #[test] fn test_extract_observations_mixed_float_and_int() { let sampler = MultivariateTpeSampler::new(); + let lr_id = ParamId::new(); + let n_layers_id = ParamId::new(); + let batch_size_id = ParamId::new(); let trial_data = [ create_trial( 0, vec![ - ("learning_rate", ParamValue::Float(0.01), float_dist()), - ("n_layers", ParamValue::Int(3), int_dist()), - ("batch_size", ParamValue::Int(32), int_dist()), + (lr_id, ParamValue::Float(0.01), float_dist()), + (n_layers_id, ParamValue::Int(3), int_dist()), + (batch_size_id, ParamValue::Int(32), int_dist()), ], 1.0, ), create_trial( 1, vec![ - ("learning_rate", ParamValue::Float(0.001), float_dist()), - ("n_layers", ParamValue::Int(5), int_dist()), - ("batch_size", ParamValue::Int(64), int_dist()), + (lr_id, ParamValue::Float(0.001), float_dist()), + (n_layers_id, ParamValue::Int(5), int_dist()), + (batch_size_id, ParamValue::Int(64), int_dist()), ], 0.5, ), ]; let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); - let param_order = vec![ - "learning_rate".to_string(), - "n_layers".to_string(), - "batch_size".to_string(), - ]; + let param_order = vec![lr_id, n_layers_id, batch_size_id]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3659,22 +3716,33 @@ mod tests { #[test] fn test_extract_observations_skips_categorical() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let optimizer_id = ParamId::new(); + let y_id = ParamId::new(); let trial_data = [ create_trial( 0, vec![ - ("x", ParamValue::Float(0.5), float_dist()), - ("optimizer", ParamValue::Categorical(1), categorical_dist(3)), - ("y", ParamValue::Float(0.3), float_dist()), + (x_id, ParamValue::Float(0.5), float_dist()), + ( + optimizer_id, + ParamValue::Categorical(1), + categorical_dist(3), + ), + (y_id, ParamValue::Float(0.3), float_dist()), ], 1.0, ), create_trial( 1, vec![ - ("x", ParamValue::Float(0.7), float_dist()), - ("optimizer", ParamValue::Categorical(0), categorical_dist(3)), - ("y", ParamValue::Float(0.2), float_dist()), + (x_id, ParamValue::Float(0.7), float_dist()), + ( + optimizer_id, + ParamValue::Categorical(0), + categorical_dist(3), + ), + (y_id, ParamValue::Float(0.2), float_dist()), ], 0.5, ), @@ -3682,7 +3750,7 @@ mod tests { let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); // Include categorical in param order - it should be skipped - let param_order = vec!["x".to_string(), "optimizer".to_string(), "y".to_string()]; + let param_order = vec![x_id, optimizer_id, y_id]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3703,18 +3771,24 @@ mod tests { #[test] fn test_extract_observations_only_categorical_in_order() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let optimizer_id = ParamId::new(); let trial = create_trial( 0, vec![ - ("x", ParamValue::Float(0.5), float_dist()), - ("optimizer", ParamValue::Categorical(1), categorical_dist(3)), + (x_id, ParamValue::Float(0.5), float_dist()), + ( + optimizer_id, + ParamValue::Categorical(1), + categorical_dist(3), + ), ], 1.0, ); let trials: Vec<&CompletedTrial> = vec![&trial]; // Only request categorical param - let param_order = vec!["optimizer".to_string()]; + let param_order = vec![optimizer_id]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3726,20 +3800,22 @@ mod tests { #[test] fn test_extract_observations_missing_param() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); let trial_data = [ create_trial( 0, vec![ - ("x", ParamValue::Float(0.5), float_dist()), - ("y", ParamValue::Float(0.3), float_dist()), + (x_id, ParamValue::Float(0.5), float_dist()), + (y_id, ParamValue::Float(0.3), float_dist()), ], 1.0, ), - // This trial is missing "y" - create_trial(1, vec![("x", ParamValue::Float(0.7), float_dist())], 0.5), + // This trial is missing y + create_trial(1, vec![(x_id, ParamValue::Float(0.7), float_dist())], 0.5), ]; let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); - let param_order = vec!["x".to_string(), "y".to_string()]; + let param_order = vec![x_id, y_id]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3752,22 +3828,21 @@ mod tests { #[test] fn test_extract_observations_large_int_values() { let sampler = MultivariateTpeSampler::new(); + let small_int_id = ParamId::new(); + let medium_int_id = ParamId::new(); + let negative_int_id = ParamId::new(); // Test with large integer values to verify precision let trial = create_trial( 0, vec![ - ("small_int", ParamValue::Int(1), int_dist()), - ("medium_int", ParamValue::Int(1_000_000), int_dist()), - ("negative_int", ParamValue::Int(-42), int_dist()), + (small_int_id, ParamValue::Int(1), int_dist()), + (medium_int_id, ParamValue::Int(1_000_000), int_dist()), + (negative_int_id, ParamValue::Int(-42), int_dist()), ], 1.0, ); let trials: Vec<&CompletedTrial> = vec![&trial]; - let param_order = vec![ - "small_int".to_string(), - "medium_int".to_string(), - "negative_int".to_string(), - ]; + let param_order = vec![small_int_id, medium_int_id, negative_int_id]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3781,6 +3856,8 @@ mod tests { #[test] fn test_extract_observations_many_trials() { let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); // Create 100 trials with predictable values let trial_data: Vec = (0..100) @@ -3788,15 +3865,15 @@ mod tests { create_trial( i, vec![ - ("x", ParamValue::Float(i as f64 / 100.0), float_dist()), - ("y", ParamValue::Int(i as i64), int_dist()), + (x_id, ParamValue::Float(i as f64 / 100.0), float_dist()), + (y_id, ParamValue::Int(i as i64), int_dist()), ], i as f64, ) }) .collect(); let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); - let param_order = vec!["x".to_string(), "y".to_string()]; + let param_order = vec![x_id, y_id]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3811,20 +3888,24 @@ mod tests { #[test] fn test_extract_observations_subset_of_params() { let sampler = MultivariateTpeSampler::new(); + let a_id = ParamId::new(); + let b_id = ParamId::new(); + let c_id = ParamId::new(); + let d_id = ParamId::new(); let trial = create_trial( 0, vec![ - ("a", ParamValue::Float(1.0), float_dist()), - ("b", ParamValue::Float(2.0), float_dist()), - ("c", ParamValue::Float(3.0), float_dist()), - ("d", ParamValue::Float(4.0), float_dist()), + (a_id, ParamValue::Float(1.0), float_dist()), + (b_id, ParamValue::Float(2.0), float_dist()), + (c_id, ParamValue::Float(3.0), float_dist()), + (d_id, ParamValue::Float(4.0), float_dist()), ], 1.0, ); let trials: Vec<&CompletedTrial> = vec![&trial]; // Only extract a subset of params - let param_order = vec!["b".to_string(), "d".to_string()]; + let param_order = vec![b_id, d_id]; let observations = sampler.extract_observations(&trials, ¶m_order); @@ -3839,24 +3920,26 @@ mod tests { // Test full pipeline: filter -> split -> extract let sampler = MultivariateTpeSampler::new(); + let x_id = ParamId::new(); + let n_id = ParamId::new(); let float_dist_val = float_dist(); let int_dist_val = int_dist(); let trial_data: Vec = (0..20) .map(|i| { let mut params = HashMap::new(); - params.insert("x".to_string(), ParamValue::Float(i as f64 / 20.0)); - params.insert("n".to_string(), ParamValue::Int(i as i64)); + params.insert(x_id, ParamValue::Float(i as f64 / 20.0)); + params.insert(n_id, ParamValue::Int(i as i64)); let mut distributions = HashMap::new(); - distributions.insert("x".to_string(), float_dist_val.clone()); - distributions.insert("n".to_string(), int_dist_val.clone()); - CompletedTrial::new(i, params, distributions, i as f64) + distributions.insert(x_id, float_dist_val.clone()); + distributions.insert(n_id, int_dist_val.clone()); + CompletedTrial::new(i, params, distributions, HashMap::new(), i as f64) }) .collect(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist_val); - search_space.insert("n".to_string(), int_dist_val); + search_space.insert(x_id, float_dist_val); + search_space.insert(n_id, int_dist_val); // Filter let filtered = sampler.filter_trials(&trial_data, &search_space); @@ -3868,7 +3951,7 @@ mod tests { assert_eq!(bad.len(), 15); // Extract from good trials - let param_order = vec!["x".to_string(), "n".to_string()]; + let param_order = vec![x_id, n_id]; let good_obs = sampler.extract_observations(&good, ¶m_order); let bad_obs = sampler.extract_observations(&bad, ¶m_order); @@ -4182,6 +4265,7 @@ mod tests { // Full integration test: extract observations -> fit KDEs -> select candidate use crate::distribution::FloatDistribution; use crate::param::ParamValue; + use crate::parameter::ParamId; let sampler = MultivariateTpeSampler::builder() .gamma(0.25) @@ -4197,6 +4281,9 @@ mod tests { step: None, }); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + // Create trials with objective values equal to x + y // So good trials have low x and y values let trial_data: Vec = (0..20) @@ -4206,25 +4293,25 @@ mod tests { #[allow(clippy::cast_precision_loss)] let y = (i as f64) / 2.0; let mut params = HashMap::new(); - params.insert("x".to_string(), ParamValue::Float(x)); - params.insert("y".to_string(), ParamValue::Float(y)); + params.insert(x_id, ParamValue::Float(x)); + params.insert(y_id, ParamValue::Float(y)); let mut distributions = HashMap::new(); - distributions.insert("x".to_string(), float_dist.clone()); - distributions.insert("y".to_string(), float_dist.clone()); - CompletedTrial::new(i, params, distributions, x + y) + distributions.insert(x_id, float_dist.clone()); + distributions.insert(y_id, float_dist.clone()); + CompletedTrial::new(i, params, distributions, HashMap::new(), x + y) }) .collect(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist.clone()); - search_space.insert("y".to_string(), float_dist); + search_space.insert(x_id, float_dist.clone()); + search_space.insert(y_id, float_dist); // Filter and split let filtered = sampler.filter_trials(&trial_data, &search_space); let (good, bad) = sampler.split_trials(&filtered); // Extract observations - let param_order = vec!["x".to_string(), "y".to_string()]; + let param_order = vec![x_id, y_id]; let good_obs = sampler.extract_observations(&good, ¶m_order); let bad_obs = sampler.extract_observations(&bad, ¶m_order); @@ -4265,6 +4352,7 @@ mod tests { use super::*; use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; use crate::param::ParamValue; + use crate::parameter::ParamId; use crate::sampler::CompletedTrial; fn float_dist(low: f64, high: f64) -> Distribution { @@ -4291,16 +4379,16 @@ mod tests { 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] @@ -4311,17 +4399,19 @@ mod tests { .build() .unwrap(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); - search_space.insert("y".to_string(), float_dist(0.0, 1.0)); + search_space.insert(x_id, float_dist(0.0, 1.0)); + search_space.insert(y_id, float_dist(0.0, 1.0)); let history: Vec = vec![]; let result = sampler.sample_joint(&search_space, &history); 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] @@ -4332,8 +4422,9 @@ mod tests { .build() .unwrap(); + let x_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert(x_id, float_dist(0.0, 1.0)); // Create 5 trials (less than n_startup_trials=10) let history: Vec = (0..5) @@ -4341,7 +4432,7 @@ mod tests { create_trial( i, vec![( - "x", + x_id, ParamValue::Float(i as f64 / 10.0), float_dist(0.0, 1.0), )], @@ -4352,8 +4443,8 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); - assert!(result.contains_key("x")); - if let Some(ParamValue::Float(v)) = result.get("x") { + assert!(result.contains_key(&x_id)); + if let Some(ParamValue::Float(v)) = result.get(&x_id) { assert!(*v >= 0.0 && *v <= 1.0); } else { panic!("Expected Float value for x"); @@ -4369,9 +4460,11 @@ mod tests { .build() .unwrap(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 10.0)); - search_space.insert("y".to_string(), float_dist(0.0, 10.0)); + search_space.insert(x_id, float_dist(0.0, 10.0)); + search_space.insert(y_id, float_dist(0.0, 10.0)); // Create 20 trials with good values near origin let history: Vec = (0..20) @@ -4381,8 +4474,8 @@ mod tests { create_trial( i, vec![ - ("x", ParamValue::Float(x), float_dist(0.0, 10.0)), - ("y", ParamValue::Float(y), float_dist(0.0, 10.0)), + (x_id, ParamValue::Float(x), float_dist(0.0, 10.0)), + (y_id, ParamValue::Float(y), float_dist(0.0, 10.0)), ], x + y, // Objective: lower is better ) @@ -4392,14 +4485,14 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); 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)); // Values should be in the distribution bounds - if let Some(ParamValue::Float(x)) = result.get("x") { + if let Some(ParamValue::Float(x)) = result.get(&x_id) { assert!(*x >= 0.0 && *x <= 10.0, "x={x} should be within [0, 10]"); } - if let Some(ParamValue::Float(y)) = result.get("y") { + if let Some(ParamValue::Float(y)) = result.get(&y_id) { assert!(*y >= 0.0 && *y <= 10.0, "y={y} should be within [0, 10]"); } } @@ -4413,8 +4506,9 @@ mod tests { .build() .unwrap(); + let x_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 10.0)); + search_space.insert(x_id, float_dist(0.0, 10.0)); // Create trials: good ones near 0, bad ones near 10 let mut history: Vec = vec![]; @@ -4423,7 +4517,7 @@ mod tests { history.push(create_trial( i, vec![( - "x", + x_id, ParamValue::Float(i as f64 * 0.5), float_dist(0.0, 10.0), )], @@ -4435,7 +4529,7 @@ mod tests { history.push(create_trial( i, vec![( - "x", + x_id, ParamValue::Float(5.0 + (i as f64 - 5.0) * 0.5), float_dist(0.0, 10.0), )], @@ -4447,7 +4541,7 @@ mod tests { let mut low_count = 0; for _ in 0..20 { let result = sampler.sample_joint(&search_space, &history); - if let Some(ParamValue::Float(x)) = result.get("x") { + if let Some(ParamValue::Float(x)) = result.get(&x_id) { if *x < 5.0 { low_count += 1; } @@ -4469,15 +4563,16 @@ mod tests { .build() .unwrap(); + let n_layers_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("n_layers".to_string(), int_dist(1, 10)); + search_space.insert(n_layers_id, int_dist(1, 10)); // Create history let history: Vec = (0..10) .map(|i| { create_trial( i, - vec![("n_layers", ParamValue::Int(i as i64 + 1), int_dist(1, 10))], + vec![(n_layers_id, ParamValue::Int(i as i64 + 1), int_dist(1, 10))], i as f64, ) }) @@ -4485,8 +4580,8 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); - assert!(result.contains_key("n_layers")); - if let Some(ParamValue::Int(v)) = result.get("n_layers") { + assert!(result.contains_key(&n_layers_id)); + if let Some(ParamValue::Int(v)) = result.get(&n_layers_id) { assert!(*v >= 1 && *v <= 10, "n_layers={v} should be within [1, 10]"); } else { panic!("Expected Int value for n_layers"); @@ -4501,9 +4596,11 @@ mod tests { .build() .unwrap(); + let lr_id = ParamId::new(); + let n_layers_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("lr".to_string(), float_dist(0.0, 1.0)); - search_space.insert("n_layers".to_string(), int_dist(1, 5)); + search_space.insert(lr_id, float_dist(0.0, 1.0)); + search_space.insert(n_layers_id, int_dist(1, 5)); // Create history with both params let history: Vec = (0..15) @@ -4512,12 +4609,12 @@ mod tests { i, vec![ ( - "lr", + lr_id, ParamValue::Float(i as f64 / 15.0), float_dist(0.0, 1.0), ), ( - "n_layers", + n_layers_id, ParamValue::Int((i % 5) as i64 + 1), int_dist(1, 5), ), @@ -4530,8 +4627,8 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); assert_eq!(result.len(), 2); - assert!(result.contains_key("lr")); - assert!(result.contains_key("n_layers")); + assert!(result.contains_key(&lr_id)); + assert!(result.contains_key(&n_layers_id)); } #[test] @@ -4543,9 +4640,11 @@ mod tests { .build() .unwrap(); + let x_id = ParamId::new(); + let optimizer_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); - search_space.insert("optimizer".to_string(), categorical_dist(3)); + search_space.insert(x_id, float_dist(0.0, 1.0)); + search_space.insert(optimizer_id, categorical_dist(3)); let history: Vec = (0..15) .map(|i| { @@ -4553,12 +4652,12 @@ mod tests { i, vec![ ( - "x", + x_id, ParamValue::Float(i as f64 / 15.0), float_dist(0.0, 1.0), ), ( - "optimizer", + optimizer_id, ParamValue::Categorical(i as usize % 3), categorical_dist(3), ), @@ -4571,10 +4670,10 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); assert_eq!(result.len(), 2); - assert!(result.contains_key("x")); - assert!(result.contains_key("optimizer")); + assert!(result.contains_key(&x_id)); + assert!(result.contains_key(&optimizer_id)); - if let Some(ParamValue::Categorical(v)) = result.get("optimizer") { + if let Some(ParamValue::Categorical(v)) = result.get(&optimizer_id) { assert!(*v < 3, "optimizer={v} should be in [0, 3)"); } else { panic!("Expected Categorical value for optimizer"); @@ -4583,10 +4682,12 @@ mod tests { #[test] fn test_sample_joint_reproducibility() { + let x_id = ParamId::new(); + let y_id = ParamId::new(); let search_space = { let mut s = HashMap::new(); - s.insert("x".to_string(), float_dist(0.0, 1.0)); - s.insert("y".to_string(), float_dist(0.0, 1.0)); + s.insert(x_id, float_dist(0.0, 1.0)); + s.insert(y_id, float_dist(0.0, 1.0)); s }; @@ -4596,12 +4697,12 @@ mod tests { i, vec![ ( - "x", + x_id, ParamValue::Float(i as f64 / 15.0), float_dist(0.0, 1.0), ), ( - "y", + y_id, ParamValue::Float(i as f64 / 15.0), float_dist(0.0, 1.0), ), @@ -4638,9 +4739,10 @@ mod tests { .build() .unwrap(); + let x_id = ParamId::new(); let mut search_space = HashMap::new(); // Narrow distribution - search_space.insert("x".to_string(), float_dist(0.0, 0.1)); + search_space.insert(x_id, float_dist(0.0, 0.1)); // Create trials at the edge let history: Vec = (0..10) @@ -4648,7 +4750,7 @@ mod tests { create_trial( i, vec![( - "x", + x_id, ParamValue::Float(i as f64 / 100.0), float_dist(0.0, 0.1), )], @@ -4660,7 +4762,7 @@ mod tests { // Sample multiple times for _ in 0..10 { let result = sampler.sample_joint(&search_space, &history); - if let Some(ParamValue::Float(x)) = result.get("x") { + if let Some(ParamValue::Float(x)) = result.get(&x_id) { assert!( *x >= 0.0 && *x <= 0.1, "x={x} should be clamped to [0.0, 0.1]" @@ -4677,20 +4779,22 @@ mod tests { .build() .unwrap(); + let x_id = ParamId::new(); + let y_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); - search_space.insert("y".to_string(), float_dist(0.0, 1.0)); + search_space.insert(x_id, float_dist(0.0, 1.0)); + search_space.insert(y_id, float_dist(0.0, 1.0)); // Create trials with non-overlapping parameters let history = vec![ create_trial( 0, - vec![("x", ParamValue::Float(0.5), float_dist(0.0, 1.0))], + vec![(x_id, ParamValue::Float(0.5), float_dist(0.0, 1.0))], 1.0, ), create_trial( 1, - vec![("y", ParamValue::Float(0.5), float_dist(0.0, 1.0))], + vec![(y_id, ParamValue::Float(0.5), float_dist(0.0, 1.0))], 1.0, ), ]; @@ -4699,8 +4803,8 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); 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] @@ -4713,9 +4817,10 @@ mod tests { .unwrap(); // 5D search space + let dim_ids: Vec = (0..5).map(|_| ParamId::new()).collect(); let mut search_space = HashMap::new(); - for i in 0..5 { - search_space.insert(format!("x{i}"), float_dist(0.0, 10.0)); + for &id in &dim_ids { + search_space.insert(id, float_dist(0.0, 10.0)); } // Create history @@ -4723,27 +4828,28 @@ mod tests { .map(|trial_id| { let mut param_map = HashMap::new(); let mut dist_map = HashMap::new(); - for dim in 0..5 { - let name = format!("x{dim}"); + for (dim, &id) in dim_ids.iter().enumerate() { let value = (trial_id as f64 + dim as f64) / 3.0; - param_map.insert(name.clone(), ParamValue::Float(value)); - dist_map.insert(name, float_dist(0.0, 10.0)); + param_map.insert(id, ParamValue::Float(value)); + dist_map.insert(id, float_dist(0.0, 10.0)); } - CompletedTrial::new(trial_id, param_map, dist_map, trial_id as f64) + CompletedTrial::new( + trial_id, + param_map, + dist_map, + HashMap::new(), + trial_id as f64, + ) }) .collect(); let result = sampler.sample_joint(&search_space, &history); assert_eq!(result.len(), 5); - for i in 0..5 { - let name = format!("x{i}"); - assert!(result.contains_key(&name), "Missing parameter {name}"); - if let Some(ParamValue::Float(v)) = result.get(&name) { - assert!( - *v >= 0.0 && *v <= 10.0, - "{name}={v} should be within [0, 10]" - ); + for (i, &id) in dim_ids.iter().enumerate() { + assert!(result.contains_key(&id), "Missing parameter x{i}"); + if let Some(ParamValue::Float(v)) = result.get(&id) { + assert!(*v >= 0.0 && *v <= 10.0, "x{i}={v} should be within [0, 10]"); } } } @@ -4866,62 +4972,67 @@ mod tests { #[test] fn test_extract_categorical_indices_basic() { + let cat_id = ParamId::new(); let trials = [ create_trial( 0, - vec![("cat", ParamValue::Categorical(1), categorical_dist(3))], + vec![(cat_id, ParamValue::Categorical(1), categorical_dist(3))], 0.5, ), create_trial( 1, - vec![("cat", ParamValue::Categorical(0), categorical_dist(3))], + vec![(cat_id, ParamValue::Categorical(0), categorical_dist(3))], 1.0, ), create_trial( 2, - vec![("cat", ParamValue::Categorical(2), categorical_dist(3))], + vec![(cat_id, ParamValue::Categorical(2), categorical_dist(3))], 1.5, ), ]; let trial_refs: Vec<&CompletedTrial> = trials.iter().collect(); - let indices = MultivariateTpeSampler::extract_categorical_indices(&trial_refs, "cat"); + let indices = MultivariateTpeSampler::extract_categorical_indices(&trial_refs, cat_id); assert_eq!(indices, vec![1, 0, 2]); } #[test] fn test_extract_categorical_indices_missing_param() { + let cat_id = ParamId::new(); + let other_id = ParamId::new(); let trials = [ create_trial( 0, - vec![("cat", ParamValue::Categorical(1), categorical_dist(3))], + vec![(cat_id, ParamValue::Categorical(1), categorical_dist(3))], 0.5, ), create_trial( 1, - vec![("other", ParamValue::Float(1.0), float_dist(0.0, 2.0))], + vec![(other_id, ParamValue::Float(1.0), float_dist(0.0, 2.0))], 1.0, ), ]; let trial_refs: Vec<&CompletedTrial> = trials.iter().collect(); - let indices = MultivariateTpeSampler::extract_categorical_indices(&trial_refs, "cat"); + let indices = MultivariateTpeSampler::extract_categorical_indices(&trial_refs, cat_id); - // Only the trial with "cat" should be included + // Only the trial with cat should be included assert_eq!(indices, vec![1]); } #[test] fn test_extract_categorical_indices_wrong_type() { + let param_id = ParamId::new(); let trials = [create_trial( 0, - vec![("param", ParamValue::Float(1.0), float_dist(0.0, 2.0))], + vec![(param_id, ParamValue::Float(1.0), float_dist(0.0, 2.0))], 0.5, )]; let trial_refs: Vec<&CompletedTrial> = trials.iter().collect(); - let indices = MultivariateTpeSampler::extract_categorical_indices(&trial_refs, "param"); + let indices = + MultivariateTpeSampler::extract_categorical_indices(&trial_refs, param_id); // Should be empty since param is Float, not Categorical assert!(indices.is_empty()); @@ -4929,8 +5040,9 @@ mod tests { #[test] fn test_sample_joint_categorical_tpe_sampling() { + let cat_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("cat".to_string(), categorical_dist(3)); + search_space.insert(cat_id, categorical_dist(3)); // Create history where category 0 is consistently good let mut history = Vec::new(); @@ -4938,7 +5050,7 @@ mod tests { for i in 0..5 { history.push(create_trial( i, - vec![("cat", ParamValue::Categorical(0), categorical_dist(3))], + vec![(cat_id, ParamValue::Categorical(0), categorical_dist(3))], 0.1, )); } @@ -4947,7 +5059,7 @@ mod tests { let cat = if i.is_multiple_of(2) { 1 } else { 2 }; history.push(create_trial( i, - vec![("cat", ParamValue::Categorical(cat), categorical_dist(3))], + vec![(cat_id, ParamValue::Categorical(cat), categorical_dist(3))], 10.0, )); } @@ -4961,7 +5073,7 @@ mod tests { .build() .unwrap(); let result = sampler_test.sample_joint(&search_space, &history); - if let Some(ParamValue::Categorical(idx)) = result.get("cat") { + if let Some(ParamValue::Categorical(idx)) = result.get(&cat_id) { counts[*idx] += 1; } } @@ -4981,9 +5093,11 @@ mod tests { .build() .unwrap(); + let x_id = ParamId::new(); + let cat_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); - search_space.insert("cat".to_string(), categorical_dist(3)); + search_space.insert(x_id, float_dist(0.0, 1.0)); + search_space.insert(cat_id, categorical_dist(3)); // Create history with both types let mut history = Vec::new(); @@ -4993,11 +5107,11 @@ mod tests { i, vec![ ( - "x", + x_id, ParamValue::Float(f64::from(i as u32) * 0.05), float_dist(0.0, 1.0), ), - ("cat", ParamValue::Categorical(0), categorical_dist(3)), + (cat_id, ParamValue::Categorical(0), categorical_dist(3)), ], f64::from(i as u32) * 0.1, )); @@ -5009,11 +5123,11 @@ mod tests { i, vec![ ( - "x", + x_id, ParamValue::Float(0.5 + f64::from(i as u32 - 5) * 0.05), float_dist(0.0, 1.0), ), - ("cat", ParamValue::Categorical(cat), categorical_dist(3)), + (cat_id, ParamValue::Categorical(cat), categorical_dist(3)), ], 5.0 + f64::from(i as u32), )); @@ -5022,18 +5136,18 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); // Both parameters should be present - assert!(result.contains_key("x")); - assert!(result.contains_key("cat")); + assert!(result.contains_key(&x_id)); + assert!(result.contains_key(&cat_id)); // x should be Float assert!( - matches!(result.get("x"), Some(ParamValue::Float(_))), + matches!(result.get(&x_id), Some(ParamValue::Float(_))), "x should be Float" ); // cat should be Categorical assert!( - matches!(result.get("cat"), Some(ParamValue::Categorical(_))), + matches!(result.get(&cat_id), Some(ParamValue::Categorical(_))), "cat should be Categorical" ); } @@ -5046,9 +5160,11 @@ mod tests { .build() .unwrap(); + let cat1_id = ParamId::new(); + let cat2_id = ParamId::new(); let mut search_space = HashMap::new(); - search_space.insert("cat1".to_string(), categorical_dist(2)); - search_space.insert("cat2".to_string(), categorical_dist(3)); + search_space.insert(cat1_id, categorical_dist(2)); + search_space.insert(cat2_id, categorical_dist(3)); // Create history with only categorical parameters let mut history = Vec::new(); @@ -5059,8 +5175,8 @@ mod tests { history.push(create_trial( i, vec![ - ("cat1", ParamValue::Categorical(cat1), categorical_dist(2)), - ("cat2", ParamValue::Categorical(cat2), categorical_dist(3)), + (cat1_id, ParamValue::Categorical(cat1), categorical_dist(2)), + (cat2_id, ParamValue::Categorical(cat2), categorical_dist(3)), ], value, )); @@ -5070,11 +5186,11 @@ mod tests { assert_eq!(result.len(), 2); assert!(matches!( - result.get("cat1"), + result.get(&cat1_id), Some(ParamValue::Categorical(_)) )); assert!(matches!( - result.get("cat2"), + result.get(&cat2_id), Some(ParamValue::Categorical(_)) )); } @@ -5086,6 +5202,7 @@ mod tests { #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)] mod sampler_trait_tests { use super::*; + use crate::parameter::ParamId; fn float_dist(low: f64, high: f64) -> Distribution { Distribution::Float(crate::distribution::FloatDistribution { @@ -5111,16 +5228,16 @@ mod tests { 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, param, dist) in params { - param_map.insert(name.to_string(), param); - dist_map.insert(name.to_string(), dist); + for (param_id, param, dist) in params { + param_map.insert(param_id, param); + 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] @@ -5156,12 +5273,13 @@ mod tests { let dist = float_dist(0.0, 1.0); // Create enough history to trigger TPE sampling + let x_id = ParamId::new(); let history: Vec = (0..10) .map(|i| { create_trial( i, vec![( - "x", + x_id, ParamValue::Float(f64::from(i as u32) / 10.0), float_dist(0.0, 1.0), )], @@ -5285,18 +5403,20 @@ mod tests { let dist_y = float_dist(0.0, 1.0); // Create history with two parameters + let x_id = ParamId::new(); + let y_id = ParamId::new(); let history: Vec = (0..10) .map(|i| { create_trial( i, vec![ ( - "x", + x_id, ParamValue::Float(f64::from(i as u32) / 10.0), float_dist(0.0, 1.0), ), ( - "y", + y_id, ParamValue::Float(f64::from((10 - i) as u32) / 10.0), float_dist(0.0, 1.0), ), @@ -5322,9 +5442,11 @@ mod tests { #[test] fn test_find_matching_param_float() { + let x_id = ParamId::new(); + let y_id = ParamId::new(); let mut cached = HashMap::new(); - cached.insert("x".to_string(), ParamValue::Float(0.5)); - cached.insert("y".to_string(), ParamValue::Float(0.8)); + cached.insert(x_id, ParamValue::Float(0.5)); + cached.insert(y_id, ParamValue::Float(0.8)); let dist = float_dist(0.0, 1.0); let result = MultivariateTpeSampler::find_matching_param(&dist, &cached); @@ -5337,8 +5459,9 @@ mod tests { #[test] fn test_find_matching_param_int() { + let n_id = ParamId::new(); let mut cached = HashMap::new(); - cached.insert("n".to_string(), ParamValue::Int(5)); + cached.insert(n_id, ParamValue::Int(5)); let dist = int_dist(0, 10); let result = MultivariateTpeSampler::find_matching_param(&dist, &cached); @@ -5351,8 +5474,9 @@ mod tests { #[test] fn test_find_matching_param_categorical() { + let choice_id = ParamId::new(); let mut cached = HashMap::new(); - cached.insert("choice".to_string(), ParamValue::Categorical(1)); + cached.insert(choice_id, ParamValue::Categorical(1)); let dist = categorical_dist(3); let result = MultivariateTpeSampler::find_matching_param(&dist, &cached); @@ -5365,8 +5489,9 @@ mod tests { #[test] fn test_find_matching_param_no_match() { + let x_id = ParamId::new(); let mut cached = HashMap::new(); - cached.insert("x".to_string(), ParamValue::Float(0.5)); + cached.insert(x_id, ParamValue::Float(0.5)); // Looking for Int, but only Float in cache let dist = int_dist(0, 10); @@ -5377,8 +5502,9 @@ mod tests { #[test] fn test_find_matching_param_out_of_bounds() { + let x_id = ParamId::new(); let mut cached = HashMap::new(); - cached.insert("x".to_string(), ParamValue::Float(5.0)); // Out of bounds + cached.insert(x_id, ParamValue::Float(5.0)); // Out of bounds let dist = float_dist(0.0, 1.0); let result = MultivariateTpeSampler::find_matching_param(&dist, &cached); @@ -5388,20 +5514,22 @@ mod tests { #[test] fn test_build_search_space_from_history() { + let x_id = ParamId::new(); + let y_id = ParamId::new(); let history = vec![ create_trial( 0, vec![ - ("x", ParamValue::Float(0.5), float_dist(0.0, 1.0)), - ("y", ParamValue::Int(5), int_dist(0, 10)), + (x_id, ParamValue::Float(0.5), float_dist(0.0, 1.0)), + (y_id, ParamValue::Int(5), int_dist(0, 10)), ], 1.0, ), create_trial( 1, vec![ - ("x", ParamValue::Float(0.3), float_dist(0.0, 1.0)), - ("y", ParamValue::Int(3), int_dist(0, 10)), + (x_id, ParamValue::Float(0.3), float_dist(0.0, 1.0)), + (y_id, ParamValue::Int(3), int_dist(0, 10)), ], 2.0, ), @@ -5411,8 +5539,8 @@ mod tests { let search_space = MultivariateTpeSampler::build_search_space_from_history(¤t_dist, &history); - assert!(search_space.contains_key("x")); - assert!(search_space.contains_key("y")); + assert!(search_space.contains_key(&x_id)); + assert!(search_space.contains_key(&y_id)); } #[test] @@ -5425,7 +5553,8 @@ mod tests { // Should have a placeholder for current distribution assert!(!search_space.is_empty()); - assert!(search_space.contains_key("_current")); + // The placeholder uses ParamId::new(), so just check it has exactly 1 entry + assert_eq!(search_space.len(), 1); } } @@ -5442,19 +5571,20 @@ mod tests { mod independent_fallback_tests { use super::*; use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; + 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) } fn float_dist(low: f64, high: f64) -> Distribution { @@ -5481,29 +5611,32 @@ mod tests { #[test] fn test_fallback_on_empty_intersection() { + let x_id = ParamId::new(); + let y_id = ParamId::new(); + let z_id = ParamId::new(); // Create trials with completely different parameters let history = vec![ create_trial( 0, - vec![("x", ParamValue::Float(0.1), float_dist(0.0, 1.0))], + vec![(x_id, ParamValue::Float(0.1), float_dist(0.0, 1.0))], 1.0, ), create_trial( 1, - vec![("y", ParamValue::Float(0.2), float_dist(0.0, 1.0))], + vec![(y_id, ParamValue::Float(0.2), float_dist(0.0, 1.0))], 2.0, ), create_trial( 2, - vec![("z", ParamValue::Float(0.3), float_dist(0.0, 1.0))], + vec![(z_id, ParamValue::Float(0.3), float_dist(0.0, 1.0))], 3.0, ), ]; // Request parameters that none of the trials have in common let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); - search_space.insert("y".to_string(), float_dist(0.0, 1.0)); + search_space.insert(x_id, float_dist(0.0, 1.0)); + search_space.insert(y_id, float_dist(0.0, 1.0)); let sampler = MultivariateTpeSampler::builder() .n_startup_trials(1) @@ -5515,16 +5648,16 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); // Should have all requested parameters - assert!(result.contains_key("x")); - assert!(result.contains_key("y")); + assert!(result.contains_key(&x_id)); + assert!(result.contains_key(&y_id)); // Values should be within bounds - if let ParamValue::Float(x) = result.get("x").unwrap() { + if let ParamValue::Float(x) = result.get(&x_id).unwrap() { assert!((0.0..=1.0).contains(x)); } else { panic!("Expected Float for x"); } - if let ParamValue::Float(y) = result.get("y").unwrap() { + if let ParamValue::Float(y) = result.get(&y_id).unwrap() { assert!((0.0..=1.0).contains(y)); } else { panic!("Expected Float for y"); @@ -5533,29 +5666,32 @@ mod tests { #[test] fn test_fallback_fills_missing_params() { + let x_id = ParamId::new(); + let y_id = ParamId::new(); + let z_id = ParamId::new(); // Trials have x and y, but we request x, y, and z let history = vec![ create_trial( 0, vec![ - ("x", ParamValue::Float(0.2), float_dist(0.0, 1.0)), - ("y", ParamValue::Float(0.3), float_dist(0.0, 1.0)), + (x_id, ParamValue::Float(0.2), float_dist(0.0, 1.0)), + (y_id, ParamValue::Float(0.3), float_dist(0.0, 1.0)), ], 1.0, ), create_trial( 1, vec![ - ("x", ParamValue::Float(0.4), float_dist(0.0, 1.0)), - ("y", ParamValue::Float(0.5), float_dist(0.0, 1.0)), + (x_id, ParamValue::Float(0.4), float_dist(0.0, 1.0)), + (y_id, ParamValue::Float(0.5), float_dist(0.0, 1.0)), ], 2.0, ), create_trial( 2, vec![ - ("x", ParamValue::Float(0.6), float_dist(0.0, 1.0)), - ("y", ParamValue::Float(0.7), float_dist(0.0, 1.0)), + (x_id, ParamValue::Float(0.6), float_dist(0.0, 1.0)), + (y_id, ParamValue::Float(0.7), float_dist(0.0, 1.0)), ], 3.0, ), @@ -5563,9 +5699,9 @@ mod tests { // Request x, y (which are in intersection) and z (which is not) let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); - search_space.insert("y".to_string(), float_dist(0.0, 1.0)); - search_space.insert("z".to_string(), float_dist(0.0, 1.0)); + search_space.insert(x_id, float_dist(0.0, 1.0)); + search_space.insert(y_id, float_dist(0.0, 1.0)); + search_space.insert(z_id, float_dist(0.0, 1.0)); let sampler = MultivariateTpeSampler::builder() .n_startup_trials(1) @@ -5576,16 +5712,16 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); // Should have all three parameters - 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)); // All values should be within bounds - for (name, value) in &result { + for value in result.values() { if let ParamValue::Float(v) = value { assert!( (0.0..=1.0).contains(v), - "Parameter {name} has value {v} out of bounds" + "Parameter has value {v} out of bounds" ); } } @@ -5593,20 +5729,22 @@ mod tests { #[test] fn test_independent_tpe_sampling_with_int() { + let n_id = ParamId::new(); + let m_id = ParamId::new(); // Create trials with int parameters let history: Vec = (0..20) .map(|i| { create_trial( i, - vec![("n", ParamValue::Int((i % 10) as i64), int_dist(0, 10))], + vec![(n_id, ParamValue::Int((i % 10) as i64), int_dist(0, 10))], (i as f64) * 0.1, ) }) .collect(); let mut search_space = HashMap::new(); - search_space.insert("n".to_string(), int_dist(0, 10)); - search_space.insert("m".to_string(), int_dist(0, 5)); // Not in history + search_space.insert(n_id, int_dist(0, 10)); + search_space.insert(m_id, int_dist(0, 5)); // Not in history let sampler = MultivariateTpeSampler::builder() .n_startup_trials(5) @@ -5616,15 +5754,15 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); - assert!(result.contains_key("n")); - assert!(result.contains_key("m")); + assert!(result.contains_key(&n_id)); + assert!(result.contains_key(&m_id)); - if let ParamValue::Int(n) = result.get("n").unwrap() { + if let ParamValue::Int(n) = result.get(&n_id).unwrap() { assert!((0..=10).contains(n)); } else { panic!("Expected Int for n"); } - if let ParamValue::Int(m) = result.get("m").unwrap() { + if let ParamValue::Int(m) = result.get(&m_id).unwrap() { assert!((0..=5).contains(m)); } else { panic!("Expected Int for m"); @@ -5633,13 +5771,15 @@ mod tests { #[test] fn test_independent_tpe_sampling_with_categorical() { + let cat_id = ParamId::new(); + let other_cat_id = ParamId::new(); // Create trials with categorical parameters let history: Vec = (0..20) .map(|i| { create_trial( i, vec![( - "cat", + cat_id, ParamValue::Categorical(i as usize % 3), categorical_dist(3), )], @@ -5649,8 +5789,8 @@ mod tests { .collect(); let mut search_space = HashMap::new(); - search_space.insert("cat".to_string(), categorical_dist(3)); - search_space.insert("other_cat".to_string(), categorical_dist(4)); // Not in history + search_space.insert(cat_id, categorical_dist(3)); + search_space.insert(other_cat_id, categorical_dist(4)); // Not in history let sampler = MultivariateTpeSampler::builder() .n_startup_trials(5) @@ -5660,15 +5800,15 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); - assert!(result.contains_key("cat")); - assert!(result.contains_key("other_cat")); + assert!(result.contains_key(&cat_id)); + assert!(result.contains_key(&other_cat_id)); - if let ParamValue::Categorical(cat) = result.get("cat").unwrap() { + if let ParamValue::Categorical(cat) = result.get(&cat_id).unwrap() { assert!(*cat < 3); } else { panic!("Expected Categorical for cat"); } - if let ParamValue::Categorical(other) = result.get("other_cat").unwrap() { + if let ParamValue::Categorical(other) = result.get(&other_cat_id).unwrap() { assert!(*other < 4); } else { panic!("Expected Categorical for other_cat"); @@ -5677,6 +5817,7 @@ mod tests { #[test] fn test_sample_all_independent_uses_tpe() { + let x_id = ParamId::new(); // Create trials with a clear pattern - good values clustered low let history: Vec = (0..30) .map(|i| { @@ -5688,14 +5829,14 @@ mod tests { let value = if i < 10 { 0.0 } else { 1.0 }; create_trial( 0, - vec![("x", ParamValue::Float(x), float_dist(0.0, 1.0))], + vec![(x_id, ParamValue::Float(x), float_dist(0.0, 1.0))], value, ) }) .collect(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert(x_id, float_dist(0.0, 1.0)); let sampler = MultivariateTpeSampler::builder() .n_startup_trials(5) @@ -5709,7 +5850,7 @@ mod tests { let n_samples = 100; for _ in 0..n_samples { let result = sampler.sample_all_independent(&search_space, &history); - if let Some(ParamValue::Float(x)) = result.get("x") + if let Some(ParamValue::Float(x)) = result.get(&x_id) && *x < 0.5 { low_count += 1; @@ -5726,32 +5867,34 @@ mod tests { #[test] fn test_fallback_with_few_filtered_trials() { + let x_id = ParamId::new(); + let y_id = ParamId::new(); // Create trials where only 1 has all the required parameters let history = vec![ create_trial( 0, vec![ - ("x", ParamValue::Float(0.1), float_dist(0.0, 1.0)), - ("y", ParamValue::Float(0.2), float_dist(0.0, 1.0)), + (x_id, ParamValue::Float(0.1), float_dist(0.0, 1.0)), + (y_id, ParamValue::Float(0.2), float_dist(0.0, 1.0)), ], 1.0, ), create_trial( 1, - vec![("x", ParamValue::Float(0.3), float_dist(0.0, 1.0))], + vec![(x_id, ParamValue::Float(0.3), float_dist(0.0, 1.0))], 2.0, ), create_trial( 2, - vec![("x", ParamValue::Float(0.5), float_dist(0.0, 1.0))], + vec![(x_id, ParamValue::Float(0.5), float_dist(0.0, 1.0))], 3.0, ), ]; // Request both x and y - only trial 0 has both let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); - search_space.insert("y".to_string(), float_dist(0.0, 1.0)); + search_space.insert(x_id, float_dist(0.0, 1.0)); + search_space.insert(y_id, float_dist(0.0, 1.0)); let sampler = MultivariateTpeSampler::builder() .n_startup_trials(1) @@ -5762,17 +5905,18 @@ mod tests { // Should fall back since only 1 trial has both params (need at least 2) let result = sampler.sample_joint(&search_space, &history); - 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_fill_remaining_uniform_fallback() { + let x_id = ParamId::new(); // During startup phase, should use uniform sampling let history: Vec = vec![]; // No history let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert(x_id, float_dist(0.0, 1.0)); let sampler = MultivariateTpeSampler::builder() .n_startup_trials(10) // High startup means we'll use random @@ -5783,14 +5927,18 @@ mod tests { // With no history and high startup threshold, should sample uniformly let result = sampler.sample_joint(&search_space, &history); - assert!(result.contains_key("x")); - if let ParamValue::Float(x) = result.get("x").unwrap() { + assert!(result.contains_key(&x_id)); + if let ParamValue::Float(x) = result.get(&x_id).unwrap() { assert!((0.0..=1.0).contains(x)); } } #[test] fn test_mixed_params_intersection_and_not() { + let x_id = ParamId::new(); + let y_id = ParamId::new(); + let z_id = ParamId::new(); + let cat_id = ParamId::new(); // Create history with x,y - both will be in intersection let history: Vec = (0..15) .map(|i| { @@ -5798,11 +5946,11 @@ mod tests { i, vec![ ( - "x", + x_id, ParamValue::Float((i as f64) * 0.05), float_dist(0.0, 1.0), ), - ("y", ParamValue::Int(i as i64 % 5), int_dist(0, 10)), + (y_id, ParamValue::Int(i as i64 % 5), int_dist(0, 10)), ], (i as f64) * 0.1, ) @@ -5811,10 +5959,10 @@ mod tests { // Request x, y (in intersection) and z, cat (not in intersection) let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), float_dist(0.0, 1.0)); - search_space.insert("y".to_string(), int_dist(0, 10)); - search_space.insert("z".to_string(), float_dist(-1.0, 1.0)); - search_space.insert("cat".to_string(), categorical_dist(5)); + search_space.insert(x_id, float_dist(0.0, 1.0)); + search_space.insert(y_id, int_dist(0, 10)); + search_space.insert(z_id, float_dist(-1.0, 1.0)); + search_space.insert(cat_id, categorical_dist(5)); let sampler = MultivariateTpeSampler::builder() .n_startup_trials(5) @@ -5825,22 +5973,22 @@ mod tests { let result = sampler.sample_joint(&search_space, &history); // All parameters should be present - assert!(result.contains_key("x")); - assert!(result.contains_key("y")); - assert!(result.contains_key("z")); - assert!(result.contains_key("cat")); + assert!(result.contains_key(&x_id)); + assert!(result.contains_key(&y_id)); + assert!(result.contains_key(&z_id)); + assert!(result.contains_key(&cat_id)); // Validate bounds - if let ParamValue::Float(x) = result.get("x").unwrap() { + if let ParamValue::Float(x) = result.get(&x_id).unwrap() { assert!((0.0..=1.0).contains(x)); } - if let ParamValue::Int(y) = result.get("y").unwrap() { + if let ParamValue::Int(y) = result.get(&y_id).unwrap() { assert!((0..=10).contains(y)); } - if let ParamValue::Float(z) = result.get("z").unwrap() { + if let ParamValue::Float(z) = result.get(&z_id).unwrap() { assert!((-1.0..=1.0).contains(z)); } - if let ParamValue::Categorical(cat) = result.get("cat").unwrap() { + if let ParamValue::Categorical(cat) = result.get(&cat_id).unwrap() { assert!(*cat < 5); } } @@ -5850,19 +5998,20 @@ mod tests { mod group_sampling_tests { use super::*; use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; + 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] @@ -5882,6 +6031,9 @@ mod tests { step: None, }); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + // Create history with all params together let history: Vec = (0..10) .map(|i| { @@ -5890,8 +6042,8 @@ mod tests { create_trial( i, vec![ - ("x", ParamValue::Float(v), dist.clone()), - ("y", ParamValue::Float(v + 0.05), dist.clone()), + (x_id, ParamValue::Float(v), dist.clone()), + (y_id, ParamValue::Float(v + 0.05), dist.clone()), ], v * v, ) @@ -5899,13 +6051,13 @@ mod tests { .collect(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), dist.clone()); - search_space.insert("y".to_string(), dist); + search_space.insert(x_id, dist.clone()); + search_space.insert(y_id, dist); let result = sampler.sample_joint(&search_space, &history); - assert!(result.contains_key("x")); - assert!(result.contains_key("y")); + assert!(result.contains_key(&x_id)); + assert!(result.contains_key(&y_id)); assert_eq!(result.len(), 2); } @@ -5926,6 +6078,11 @@ mod tests { step: None, }); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + let a_id = ParamId::new(); + let b_id = ParamId::new(); + // Create history with two independent groups: // Group 1: x, y appear together // Group 2: a, b appear together @@ -5937,8 +6094,8 @@ mod tests { history.push(create_trial( i, vec![ - ("x", ParamValue::Float(v), dist.clone()), - ("y", ParamValue::Float(v + 0.05), dist.clone()), + (x_id, ParamValue::Float(v), dist.clone()), + (y_id, ParamValue::Float(v + 0.05), dist.clone()), ], v * v, )); @@ -5949,8 +6106,8 @@ mod tests { history.push(create_trial( i, vec![ - ("a", ParamValue::Float(v), dist.clone()), - ("b", ParamValue::Float(v + 0.1), dist.clone()), + (a_id, ParamValue::Float(v), dist.clone()), + (b_id, ParamValue::Float(v + 0.1), dist.clone()), ], v + 0.5, )); @@ -5958,18 +6115,18 @@ mod tests { // Search space contains params from both groups let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), dist.clone()); - search_space.insert("y".to_string(), dist.clone()); - search_space.insert("a".to_string(), dist.clone()); - search_space.insert("b".to_string(), dist); + search_space.insert(x_id, dist.clone()); + search_space.insert(y_id, dist.clone()); + search_space.insert(a_id, dist.clone()); + search_space.insert(b_id, dist); let result = sampler.sample_joint(&search_space, &history); // All params should be sampled - assert!(result.contains_key("x")); - assert!(result.contains_key("y")); - assert!(result.contains_key("a")); - assert!(result.contains_key("b")); + assert!(result.contains_key(&x_id)); + assert!(result.contains_key(&y_id)); + assert!(result.contains_key(&a_id)); + assert!(result.contains_key(&b_id)); assert_eq!(result.len(), 4); // Values should be within bounds @@ -6004,6 +6161,10 @@ mod tests { step: None, }); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + let z_id = ParamId::new(); + // All params appear together in all trials let history: Vec = (0..10) .map(|i| { @@ -6012,9 +6173,9 @@ mod tests { create_trial( i, vec![ - ("x", ParamValue::Float(v), dist.clone()), - ("y", ParamValue::Float(v + 0.05), dist.clone()), - ("z", ParamValue::Float(v + 0.1), dist.clone()), + (x_id, ParamValue::Float(v), dist.clone()), + (y_id, ParamValue::Float(v + 0.05), dist.clone()), + (z_id, ParamValue::Float(v + 0.1), dist.clone()), ], v * v, ) @@ -6022,9 +6183,9 @@ mod tests { .collect(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), dist.clone()); - search_space.insert("y".to_string(), dist.clone()); - search_space.insert("z".to_string(), dist); + search_space.insert(x_id, dist.clone()); + search_space.insert(y_id, dist.clone()); + search_space.insert(z_id, dist); let result_grouped = sampler_grouped.sample_joint(&search_space, &history); let result_ungrouped = sampler_ungrouped.sample_joint(&search_space, &history); @@ -6060,27 +6221,31 @@ mod tests { step: None, }); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + let z_id = ParamId::new(); + // Each param appears alone (forms its own isolated group) let history = vec![ - create_trial(0, vec![("x", ParamValue::Float(0.3), dist.clone())], 1.0), - create_trial(1, vec![("y", ParamValue::Float(0.5), dist.clone())], 0.5), - create_trial(2, vec![("z", ParamValue::Float(0.7), dist.clone())], 0.8), - create_trial(3, vec![("x", ParamValue::Float(0.2), dist.clone())], 1.2), - create_trial(4, vec![("y", ParamValue::Float(0.6), dist.clone())], 0.4), - create_trial(5, vec![("z", ParamValue::Float(0.8), dist.clone())], 0.7), + create_trial(0, vec![(x_id, ParamValue::Float(0.3), dist.clone())], 1.0), + create_trial(1, vec![(y_id, ParamValue::Float(0.5), dist.clone())], 0.5), + create_trial(2, vec![(z_id, ParamValue::Float(0.7), dist.clone())], 0.8), + create_trial(3, vec![(x_id, ParamValue::Float(0.2), dist.clone())], 1.2), + create_trial(4, vec![(y_id, ParamValue::Float(0.6), dist.clone())], 0.4), + create_trial(5, vec![(z_id, ParamValue::Float(0.8), dist.clone())], 0.7), ]; let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), dist.clone()); - search_space.insert("y".to_string(), dist.clone()); - search_space.insert("z".to_string(), dist); + search_space.insert(x_id, dist.clone()); + search_space.insert(y_id, dist.clone()); + search_space.insert(z_id, dist); let result = sampler.sample_joint(&search_space, &history); // All params should be sampled - 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)); // Values should be within bounds for value in result.values() { @@ -6107,24 +6272,27 @@ mod tests { step: None, }); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + // Only 5 trials (less than n_startup_trials=10) let history: Vec = (0..5) .map(|i| { #[allow(clippy::cast_precision_loss)] let v = (i as f64) * 0.1; - create_trial(i, vec![("x", ParamValue::Float(v), dist.clone())], v * v) + create_trial(i, vec![(x_id, ParamValue::Float(v), dist.clone())], v * v) }) .collect(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), dist.clone()); - search_space.insert("y".to_string(), dist); + search_space.insert(x_id, dist.clone()); + search_space.insert(y_id, dist); let result = sampler.sample_joint(&search_space, &history); // Should sample uniformly for all params - assert!(result.contains_key("x")); - assert!(result.contains_key("y")); + assert!(result.contains_key(&x_id)); + assert!(result.contains_key(&y_id)); } #[test] @@ -6151,6 +6319,10 @@ mod tests { }); let dist_cat = Distribution::Categorical(CategoricalDistribution { n_choices: 3 }); + let lr_id = ParamId::new(); + let layers_id = ParamId::new(); + let opt_id = ParamId::new(); + // Group 1: float, int co-occur // Group 2: categorical alone let mut history = Vec::new(); @@ -6162,8 +6334,8 @@ mod tests { history.push(create_trial( i, vec![ - ("lr", ParamValue::Float(v), dist_float.clone()), - ("layers", ParamValue::Int(int_v), dist_int.clone()), + (lr_id, ParamValue::Float(v), dist_float.clone()), + (layers_id, ParamValue::Int(int_v), dist_int.clone()), ], v * v, )); @@ -6172,7 +6344,7 @@ mod tests { history.push(create_trial( i, vec![( - "opt", + opt_id, ParamValue::Categorical((i % 3) as usize), dist_cat.clone(), )], @@ -6181,21 +6353,21 @@ mod tests { } let mut search_space = HashMap::new(); - search_space.insert("lr".to_string(), dist_float); - search_space.insert("layers".to_string(), dist_int); - search_space.insert("opt".to_string(), dist_cat); + search_space.insert(lr_id, dist_float); + search_space.insert(layers_id, dist_int); + search_space.insert(opt_id, dist_cat); let result = sampler.sample_joint(&search_space, &history); - assert!(result.contains_key("lr")); - assert!(result.contains_key("layers")); - assert!(result.contains_key("opt")); + assert!(result.contains_key(&lr_id)); + assert!(result.contains_key(&layers_id)); + assert!(result.contains_key(&opt_id)); // Check types - assert!(matches!(result.get("lr"), Some(ParamValue::Float(_)))); - assert!(matches!(result.get("layers"), Some(ParamValue::Int(_)))); + assert!(matches!(result.get(&lr_id), Some(ParamValue::Float(_)))); + assert!(matches!(result.get(&layers_id), Some(ParamValue::Int(_)))); assert!(matches!( - result.get("opt"), + result.get(&opt_id), Some(ParamValue::Categorical(_)) )); } @@ -6217,16 +6389,19 @@ mod tests { step: None, }); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + let history: Vec = vec![]; let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), dist.clone()); - search_space.insert("y".to_string(), dist); + search_space.insert(x_id, dist.clone()); + search_space.insert(y_id, dist); let result = sampler.sample_joint(&search_space, &history); - assert!(result.contains_key("x")); - assert!(result.contains_key("y")); + assert!(result.contains_key(&x_id)); + assert!(result.contains_key(&y_id)); } #[test] @@ -6247,6 +6422,9 @@ mod tests { step: None, }); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + // Two non-overlapping groups let mut history = Vec::new(); for i in 0..5 { @@ -6254,7 +6432,7 @@ mod tests { let v = (i as f64) * 0.1; history.push(create_trial( i, - vec![("x", ParamValue::Float(v), dist.clone())], + vec![(x_id, ParamValue::Float(v), dist.clone())], v, )); } @@ -6263,20 +6441,20 @@ mod tests { let v = (i as f64) * 0.05; history.push(create_trial( i, - vec![("y", ParamValue::Float(v), dist.clone())], + vec![(y_id, ParamValue::Float(v), dist.clone())], v, )); } let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), dist.clone()); - search_space.insert("y".to_string(), dist); + search_space.insert(x_id, dist.clone()); + search_space.insert(y_id, dist); let result = sampler.sample_joint(&search_space, &history); // Should still produce valid results via fallback - assert!(result.contains_key("x")); - assert!(result.contains_key("y")); + assert!(result.contains_key(&x_id)); + assert!(result.contains_key(&y_id)); } #[test] @@ -6289,6 +6467,9 @@ mod tests { step: None, }); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + let history: Vec = (0..10) .map(|i| { #[allow(clippy::cast_precision_loss)] @@ -6296,8 +6477,8 @@ mod tests { create_trial( i, vec![ - ("x", ParamValue::Float(v), dist.clone()), - ("y", ParamValue::Float(v + 0.05), dist.clone()), + (x_id, ParamValue::Float(v), dist.clone()), + (y_id, ParamValue::Float(v + 0.05), dist.clone()), ], v * v, ) @@ -6305,8 +6486,8 @@ mod tests { .collect(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), dist.clone()); - search_space.insert("y".to_string(), dist); + search_space.insert(x_id, dist.clone()); + search_space.insert(y_id, dist); // Same seed should produce same results let sampler1 = MultivariateTpeSamplerBuilder::new() @@ -6345,26 +6526,29 @@ mod tests { step: None, }); + let x_id = ParamId::new(); + let y_id = ParamId::new(); + // History only has x, but search space has x and y let history: Vec = (0..10) .map(|i| { let v = (i as f64) * 0.1; - create_trial(i, vec![("x", ParamValue::Float(v), dist.clone())], v * v) + create_trial(i, vec![(x_id, ParamValue::Float(v), dist.clone())], v * v) }) .collect(); let mut search_space = HashMap::new(); - search_space.insert("x".to_string(), dist.clone()); - search_space.insert("y".to_string(), dist); // Not in history + search_space.insert(x_id, dist.clone()); + search_space.insert(y_id, dist); // Not in history let result = sampler.sample_joint(&search_space, &history); // Both should be sampled - assert!(result.contains_key("x")); - assert!(result.contains_key("y")); + assert!(result.contains_key(&x_id)); + assert!(result.contains_key(&y_id)); // y should be sampled uniformly (not in any group) - if let ParamValue::Float(v) = result.get("y").unwrap() { + if let ParamValue::Float(v) = result.get(&y_id).unwrap() { assert!((0.0..=1.0).contains(v)); } } diff --git a/src/sampler/tpe/sampler.rs b/src/sampler/tpe/sampler.rs index 434967a..7ead399 100644 --- a/src/sampler/tpe/sampler.rs +++ b/src/sampler/tpe/sampler.rs @@ -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 = (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 = (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 = (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 = (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 = (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 = (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(); diff --git a/src/sampler/tpe/search_space.rs b/src/sampler/tpe/search_space.rs index 197eeea..216c143 100644 --- a/src/sampler/tpe/search_space.rs +++ b/src/sampler/tpe/search_space.rs @@ -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 { + pub fn calculate(trials: &[CompletedTrial]) -> HashMap { 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 = + 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 = 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(¶m_id) { + result.insert(param_id, dist.clone()); break; } } @@ -193,7 +190,7 @@ impl GroupDecomposedSearchSpace { /// /// # Returns /// - /// A `Vec>` where each `HashSet` represents an independent group + /// A `Vec>` 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> { + pub fn calculate(trials: &[CompletedTrial]) -> Vec> { if trials.is_empty() { return Vec::new(); } - // Collect all unique parameter names - let mut all_params: HashSet = HashSet::new(); + // Collect all unique parameter ids + let mut all_params: HashSet = HashSet::new(); for trial in trials { - for param_name in trial.distributions.keys() { - all_params.insert(param_name.clone()); + for ¶m_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> = HashMap::new(); - for param in &all_params { - adjacency.insert(param.clone(), HashSet::new()); + let mut adjacency: HashMap> = HashMap::new(); + for ¶m 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 = 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, ¶m1) in trial_params.iter().enumerate() { + for ¶m2 in trial_params.iter().skip(i + 1) { adjacency - .get_mut(*param1) + .get_mut(¶m1) .expect("param should exist in adjacency map") - .insert((*param2).clone()); + .insert(param2); adjacency - .get_mut(*param2) + .get_mut(¶m2) .expect("param should exist in adjacency map") - .insert((*param1).clone()); + .insert(param1); } } } // Find connected components using BFS - let mut visited: HashSet = HashSet::new(); - let mut groups: Vec> = Vec::new(); + let mut visited: HashSet = HashSet::new(); + let mut groups: Vec> = Vec::new(); - for param in &all_params { - if visited.contains(param) { + for ¶m in &all_params { + if visited.contains(¶m) { continue; } // BFS to find all parameters in this component - let mut component: HashSet = HashSet::new(); - let mut queue: VecDeque = VecDeque::new(); - queue.push_back(param.clone()); - visited.insert(param.clone()); + let mut component: HashSet = HashSet::new(); + let mut queue: VecDeque = 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(¤t) { - 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 = groups.iter().flatten().cloned().collect(); - assert!(all_params.contains("x")); - assert!(all_params.contains("y")); - assert!(all_params.contains("z")); + let all_params: HashSet = 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(¢er_id)); + assert!(groups[0].contains(&a_id)); + assert!(groups[0].contains(&b_id)); + assert!(groups[0].contains(&c_id)); } } diff --git a/src/study.rs b/src/study.rs index e9fa0e0..52a6626 100644 --- a/src/study.rs +++ b/src/study.rs @@ -166,11 +166,13 @@ where /// # Examples /// /// ``` + /// use optimizer::parameter::{FloatParam, Parameter}; /// use optimizer::{Direction, Study}; /// /// let study: Study = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 { /// # Examples /// /// ``` + /// use optimizer::parameter::{FloatParam, Parameter}; /// use optimizer::sampler::random::RandomSampler; /// use optimizer::{Direction, Study}; /// @@ -755,7 +787,8 @@ impl Study { /// 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 { /// # Examples /// /// ``` + /// use optimizer::parameter::{FloatParam, Parameter}; /// use optimizer::sampler::random::RandomSampler; /// use optimizer::{Direction, Study}; /// @@ -795,9 +829,11 @@ impl Study { /// let sampler = RandomSampler::with_seed(42); /// let study: Study = 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 { /// ``` /// use std::ops::ControlFlow; /// + /// use optimizer::parameter::{FloatParam, Parameter}; /// use optimizer::sampler::random::RandomSampler; /// use optimizer::{Direction, Study}; /// @@ -866,11 +903,13 @@ impl Study { /// let sampler = RandomSampler::with_seed(42); /// let study: Study = 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 { /// # Examples /// /// ``` + /// use optimizer::parameter::{FloatParam, Parameter}; /// use optimizer::sampler::random::RandomSampler; /// use optimizer::{Direction, Study}; /// @@ -967,12 +1007,17 @@ impl Study { /// let sampler = RandomSampler::with_seed(42); /// let study: Study = 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 { /// # Examples /// /// ``` + /// use optimizer::parameter::{FloatParam, Parameter}; /// use optimizer::sampler::random::RandomSampler; /// use optimizer::{Direction, Study}; /// @@ -1049,12 +1095,17 @@ impl Study { /// let sampler = RandomSampler::with_seed(42); /// let study: Study = 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?; /// diff --git a/src/trial.rs b/src/trial.rs index 824ddaa..29cbfd2 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -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` | `0.0..1.0` | Float range (end-exclusive, treated as inclusive for continuous sampling) | -/// | `RangeInclusive` | `0.0..=1.0` | Float range (end-inclusive) | -/// | `Range` | `1..10` | Integer range from 1 to 9 (end-exclusive) | -/// | `RangeInclusive` | `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; -} - -impl SuggestableRange for Range { - type Output = f64; - - fn suggest(self, trial: &mut Trial, name: String) -> Result { - trial.suggest_float(name, self.start, self.end) - } -} - -impl SuggestableRange for RangeInclusive { - type Output = f64; - - fn suggest(self, trial: &mut Trial, name: String) -> Result { - trial.suggest_float(name, *self.start(), *self.end()) - } -} - -impl SuggestableRange for Range { - type Output = i64; - - fn suggest(self, trial: &mut Trial, name: String) -> Result { - // 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 { - type Output = i64; - - fn suggest(self, trial: &mut Trial, name: String) -> Result { - 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, - /// Parameter distributions, keyed by parameter name. - distributions: HashMap, + /// Sampled parameter values, keyed by parameter id. + params: HashMap, + /// Parameter distributions, keyed by parameter id. + distributions: HashMap, + /// Human-readable labels for parameters, keyed by parameter id. + param_labels: HashMap, /// The sampler to use for generating parameter values. sampler: Option>, /// 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 { + pub fn params(&self) -> &HashMap { &self.params } /// Returns a reference to the parameter distributions. #[must_use] - pub fn distributions(&self) -> &HashMap { + pub fn distributions(&self) -> &HashMap { &self.distributions } + /// Returns a reference to the parameter labels. + #[must_use] + pub fn param_labels(&self) -> &HashMap { + &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, low: f64, high: f64) -> Result { - if low > high { - return Err(Error::InvalidBounds { low, high }); - } + pub fn suggest_param(&mut self, param: &P) -> Result { + 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(¶m_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(¶m_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, - low: f64, - high: f64, - ) -> Result { - 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, - low: f64, - high: f64, - step: f64, - ) -> Result { - 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, low: i64, high: i64) -> Result { - 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, low: i64, high: i64) -> Result { - 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, - low: i64, - high: i64, - step: i64, - ) -> Result { - 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( - &mut self, - name: impl Into, - choices: &[T], - ) -> Result { - 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) -> Result { - 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( - &mut self, - name: impl Into, - range: R, - ) -> Result { - range.suggest(self, name.into()) + Ok(result) } } diff --git a/tests/async_tests.rs b/tests/async_tests.rs index dae3862..a978a6e 100644 --- a/tests/async_tests.rs +++ b/tests/async_tests.rs @@ -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 = 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 = 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 = 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 = 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 = 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 = 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"); diff --git a/tests/derive_tests.rs b/tests/derive_tests.rs new file mode 100644 index 0000000..a4d4058 --- /dev/null +++ b/tests/derive_tests.rs @@ -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::::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::::new()).unwrap(); + assert!([Color::Red, Color::Green, Color::Blue].contains(&color)); +} diff --git a/tests/integration.rs b/tests/integration.rs index 00063db..9d8056e 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -6,6 +6,7 @@ clippy::cast_possible_truncation )] +use optimizer::parameter::{BoolParam, CategoricalParam, FloatParam, IntParam, Parameter}; use optimizer::sampler::random::RandomSampler; use optimizer::sampler::tpe::TpeSampler; use optimizer::{Direction, Error, Study, Trial}; @@ -16,7 +17,7 @@ use optimizer::{Direction, Error, Study, Trial}; #[test] fn test_tpe_optimizes_quadratic_function() { - // Minimize f(x) = (x - 3)^2 where x ∈ [-10, 10] + // Minimize f(x) = (x - 3)^2 where x in [-10, 10] // Optimal: x = 3, f(3) = 0 let sampler = TpeSampler::builder() .seed(42) @@ -27,16 +28,18 @@ fn test_tpe_optimizes_quadratic_function() { let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(-10.0, 10.0); + study .optimize_with_sampler(50, |trial| { - let x = trial.suggest_float("x", -10.0, 10.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>((x - 3.0).powi(2)) }) .expect("optimization should succeed"); let best = study.best_trial().expect("should have at least one trial"); - // TPE should find a value close to optimal (x ≈ 3) + // TPE should find a value close to optimal (x ~ 3) // We expect the best value to be small (close to 0) assert!( best.value < 1.0, @@ -47,7 +50,7 @@ fn test_tpe_optimizes_quadratic_function() { #[test] fn test_tpe_optimizes_multivariate_function() { - // Minimize f(x, y) = x^2 + y^2 where x, y ∈ [-5, 5] + // Minimize f(x, y) = x^2 + y^2 where x, y in [-5, 5] // Optimal: (0, 0), f(0, 0) = 0 let sampler = TpeSampler::builder() .seed(123) @@ -57,10 +60,13 @@ fn test_tpe_optimizes_multivariate_function() { let study: Study = 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(100, |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>(x * x + y * y) }) .expect("optimization should succeed"); @@ -77,7 +83,7 @@ fn test_tpe_optimizes_multivariate_function() { #[test] fn test_tpe_maximization() { - // Maximize f(x) = -(x - 2)^2 + 10 where x ∈ [-10, 10] + // Maximize f(x) = -(x - 2)^2 + 10 where x in [-10, 10] // Optimal: x = 2, f(2) = 10 let sampler = TpeSampler::builder() .seed(456) @@ -87,18 +93,17 @@ fn test_tpe_maximization() { let study: Study = Study::with_sampler(Direction::Maximize, sampler); + let x_param = FloatParam::new(-10.0, 10.0); + study .optimize_with_sampler(50, |trial| { - let x = trial.suggest_float("x", -10.0, 10.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(-(x - 2.0).powi(2) + 10.0) }) .expect("optimization should succeed"); let best = study.best_trial().expect("should have at least one trial"); - // For maximization, best value should be better than a random baseline - // The function ranges from -90 (at x=-10 or x=10, when far from x=2) to 10 (at x=2) - // A random approach would average around 0, so finding >5 is a reasonable check assert!( best.value > 5.0, "TPE should find reasonably good solution: best value {} should be > 5.0", @@ -112,16 +117,16 @@ fn test_tpe_maximization() { #[test] fn test_random_sampler_uniform_float_distribution() { - // Test that RandomSampler samples uniformly by running multiple trials - // and checking the distribution of sampled values let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); let n_samples = 1000; let mut samples = Vec::with_capacity(n_samples); + let x_param = FloatParam::new(0.0, 1.0); + study .optimize(n_samples, |trial| { - let x = trial.suggest_float("x", 0.0, 1.0)?; + let x = x_param.suggest(trial)?; samples.push(x); Ok::<_, Error>(x) }) @@ -139,7 +144,6 @@ fn test_random_sampler_uniform_float_distribution() { let q2 = samples[n_samples / 2]; let q3 = samples[3 * n_samples / 4]; - // For uniform distribution, quartiles should be approximately 0.25, 0.5, 0.75 assert!((q1 - 0.25).abs() < 0.1, "Q1 {q1} should be close to 0.25"); assert!( (q2 - 0.5).abs() < 0.1, @@ -155,18 +159,17 @@ fn test_random_sampler_uniform_int_distribution() { let n_samples = 5000; let mut counts = [0u32; 10]; // counts for values 1-10 + let n_param = IntParam::new(1, 10); + study .optimize(n_samples, |trial| { - let n = trial.suggest_int("n", 1, 10)?; + let n = n_param.suggest(trial)?; assert!((1..=10).contains(&n), "sample {n} out of range [1, 10]"); counts[(n - 1) as usize] += 1; Ok::<_, Error>(n as f64) }) .unwrap(); - // Each value should appear roughly n_samples / 10 times - // With 5000 samples, expected ~500 per bucket, std dev ~21 - // 20% tolerance allows for ~4.5 std devs which is very safe let expected = n_samples as f64 / 10.0; for (i, &count) in counts.iter().enumerate() { let diff = (count as f64 - expected).abs() / expected; @@ -189,18 +192,17 @@ fn test_random_sampler_uniform_categorical_distribution() { let mut counts = [0u32; 4]; let choices = ["a", "b", "c", "d"]; + let cat_param = CategoricalParam::new(choices.to_vec()); + study .optimize(n_samples, |trial| { - let choice = trial.suggest_categorical("cat", &choices)?; + let choice = cat_param.suggest(trial)?; let idx = choices.iter().position(|&c| c == choice).unwrap(); counts[idx] += 1; Ok::<_, Error>(idx as f64) }) .unwrap(); - // Each category should appear roughly n_samples / 4 times - // With 2000 samples, expected ~500 per bucket, std dev ~19 - // 15% tolerance allows for ~4 std devs which is very safe let expected = n_samples as f64 / 4.0; for (i, &count) in counts.iter().enumerate() { let diff = (count as f64 - expected).abs() / expected; @@ -217,8 +219,6 @@ fn test_random_sampler_uniform_categorical_distribution() { #[test] fn test_random_sampler_reproducibility() { - // Two studies with the same seed should produce the same sequence - // NOTE: We must use optimize_with_sampler() for Study to get sampler integration let study1: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(999)); let study2: Study = @@ -227,9 +227,12 @@ fn test_random_sampler_reproducibility() { let mut values1 = Vec::new(); let mut values2 = Vec::new(); + let x_param1 = FloatParam::new(0.0, 100.0); + let x_param2 = FloatParam::new(0.0, 100.0); + study1 .optimize_with_sampler(100, |trial| { - let x = trial.suggest_float("x", 0.0, 100.0)?; + let x = x_param1.suggest(trial)?; values1.push(x); Ok::<_, Error>(x) }) @@ -237,7 +240,7 @@ fn test_random_sampler_reproducibility() { study2 .optimize_with_sampler(100, |trial| { - let x = trial.suggest_float("x", 0.0, 100.0)?; + let x = x_param2.suggest(trial)?; values2.push(x); Ok::<_, Error>(x) }) @@ -252,112 +255,122 @@ fn test_random_sampler_reproducibility() { } // ============================================================================= -// Test: suggest_* methods return cached values on repeated calls +// Test: suggest_param returns cached values on repeated calls // ============================================================================= #[test] fn test_suggest_float_caching() { + let param = FloatParam::new(0.0, 10.0); let mut trial = Trial::new(0); - let x1 = trial.suggest_float("x", 0.0, 10.0).unwrap(); - let x2 = trial.suggest_float("x", 0.0, 10.0).unwrap(); - let x3 = trial.suggest_float("x", 0.0, 10.0).unwrap(); + let x1 = param.suggest(&mut trial).unwrap(); + let x2 = param.suggest(&mut trial).unwrap(); + let x3 = param.suggest(&mut trial).unwrap(); - assert_eq!(x1, x2, "repeated suggest_float should return cached value"); - assert_eq!(x2, x3, "repeated suggest_float should return cached value"); + assert_eq!(x1, x2, "repeated suggest should return cached value"); + assert_eq!(x2, x3, "repeated suggest should return cached value"); } #[test] fn test_suggest_float_log_caching() { + let param = FloatParam::new(1e-5, 1e-1).log_scale(); let mut trial = Trial::new(0); - let x1 = trial.suggest_float_log("lr", 1e-5, 1e-1).unwrap(); - let x2 = trial.suggest_float_log("lr", 1e-5, 1e-1).unwrap(); + let x1 = param.suggest(&mut trial).unwrap(); + let x2 = param.suggest(&mut trial).unwrap(); assert_eq!( x1, x2, - "repeated suggest_float_log should return cached value" + "repeated suggest float log should return cached value" ); } #[test] fn test_suggest_float_step_caching() { + let param = FloatParam::new(0.0, 1.0).step(0.1); let mut trial = Trial::new(0); - let x1 = trial.suggest_float_step("step", 0.0, 1.0, 0.1).unwrap(); - let x2 = trial.suggest_float_step("step", 0.0, 1.0, 0.1).unwrap(); + let x1 = param.suggest(&mut trial).unwrap(); + let x2 = param.suggest(&mut trial).unwrap(); assert_eq!( x1, x2, - "repeated suggest_float_step should return cached value" + "repeated suggest float step should return cached value" ); } #[test] fn test_suggest_int_caching() { + let param = IntParam::new(1, 100); let mut trial = Trial::new(0); - let n1 = trial.suggest_int("n", 1, 100).unwrap(); - let n2 = trial.suggest_int("n", 1, 100).unwrap(); + let n1 = param.suggest(&mut trial).unwrap(); + let n2 = param.suggest(&mut trial).unwrap(); - assert_eq!(n1, n2, "repeated suggest_int should return cached value"); + assert_eq!(n1, n2, "repeated suggest int should return cached value"); } #[test] fn test_suggest_int_log_caching() { + let param = IntParam::new(1, 1024).log_scale(); let mut trial = Trial::new(0); - let n1 = trial.suggest_int_log("batch", 1, 1024).unwrap(); - let n2 = trial.suggest_int_log("batch", 1, 1024).unwrap(); + let n1 = param.suggest(&mut trial).unwrap(); + let n2 = param.suggest(&mut trial).unwrap(); assert_eq!( n1, n2, - "repeated suggest_int_log should return cached value" + "repeated suggest int log should return cached value" ); } #[test] fn test_suggest_int_step_caching() { + let param = IntParam::new(32, 512).step(32); let mut trial = Trial::new(0); - let n1 = trial.suggest_int_step("units", 32, 512, 32).unwrap(); - let n2 = trial.suggest_int_step("units", 32, 512, 32).unwrap(); + let n1 = param.suggest(&mut trial).unwrap(); + let n2 = param.suggest(&mut trial).unwrap(); assert_eq!( n1, n2, - "repeated suggest_int_step should return cached value" + "repeated suggest int step should return cached value" ); } #[test] fn test_suggest_categorical_caching() { + let param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]); let mut trial = Trial::new(0); - let choices = ["sgd", "adam", "rmsprop"]; - let c1 = trial.suggest_categorical("optimizer", &choices).unwrap(); - let c2 = trial.suggest_categorical("optimizer", &choices).unwrap(); + let c1 = param.suggest(&mut trial).unwrap(); + let c2 = param.suggest(&mut trial).unwrap(); assert_eq!( c1, c2, - "repeated suggest_categorical should return cached value" + "repeated suggest categorical should return cached value" ); } #[test] fn test_multiple_parameters_independent_caching() { + let x_param = FloatParam::new(0.0, 1.0); + let y_param = FloatParam::new(0.0, 1.0); + let n_param = IntParam::new(1, 10); + let opt_param = CategoricalParam::new(vec!["a", "b"]); let mut trial = Trial::new(0); // Suggest multiple parameters - let x = trial.suggest_float("x", 0.0, 1.0).unwrap(); - let y = trial.suggest_float("y", 0.0, 1.0).unwrap(); - let n = trial.suggest_int("n", 1, 10).unwrap(); - let opt = trial.suggest_categorical("opt", &["a", "b"]).unwrap(); + let x = x_param.suggest(&mut trial).unwrap(); + let y = y_param.suggest(&mut trial).unwrap(); + let n = n_param.suggest(&mut trial).unwrap(); + let opt = opt_param.suggest(&mut trial).unwrap(); // All should be cached independently - assert_eq!(x, trial.suggest_float("x", 0.0, 1.0).unwrap()); - assert_eq!(y, trial.suggest_float("y", 0.0, 1.0).unwrap()); - assert_eq!(n, trial.suggest_int("n", 1, 10).unwrap()); - assert_eq!(opt, trial.suggest_categorical("opt", &["a", "b"]).unwrap()); + assert_eq!(x, x_param.suggest(&mut trial).unwrap()); + assert_eq!(y, y_param.suggest(&mut trial).unwrap()); + assert_eq!(n, n_param.suggest(&mut trial).unwrap()); + assert_eq!(opt, opt_param.suggest(&mut trial).unwrap()); } // ============================================================================= @@ -365,111 +378,30 @@ fn test_multiple_parameters_independent_caching() { // ============================================================================= #[test] -fn test_parameter_conflict_float_different_bounds() { +fn test_parameter_conflict_same_param_different_distribution() { + // With ParamId-based API, conflict happens when the same ParamId is used + // with a different distribution. This can happen via suggest_param with + // a param that has a mismatched distribution for an already-stored id. + // Since each FloatParam::new() gets a unique id, conflicts only happen + // when the same param object is reused with different internal state, + // which is not possible with the immutable API. + // We test that different param objects don't conflict (they have different ids). + let param1 = FloatParam::new(0.0, 1.0); + let param2 = FloatParam::new(0.0, 2.0); let mut trial = Trial::new(0); - trial.suggest_float("x", 0.0, 1.0).unwrap(); - let result = trial.suggest_float("x", 0.0, 2.0); // Different upper bound - - assert!(matches!(result, Err(Error::ParameterConflict { .. }))); + trial.suggest_param(¶m1).unwrap(); + // Different param object = different id = no conflict + let result = trial.suggest_param(¶m2); + assert!(result.is_ok()); } -#[test] -fn test_parameter_conflict_float_vs_log() { - let mut trial = Trial::new(0); - - trial.suggest_float("x", 0.1, 1.0).unwrap(); - let result = trial.suggest_float_log("x", 0.1, 1.0); // Same bounds but log scale - - assert!(matches!(result, Err(Error::ParameterConflict { .. }))); -} - -#[test] -fn test_parameter_conflict_float_vs_step() { - let mut trial = Trial::new(0); - - trial.suggest_float("x", 0.0, 1.0).unwrap(); - let result = trial.suggest_float_step("x", 0.0, 1.0, 0.1); // Same bounds but with step - - assert!(matches!(result, Err(Error::ParameterConflict { .. }))); -} - -#[test] -fn test_parameter_conflict_int_different_bounds() { - let mut trial = Trial::new(0); - - trial.suggest_int("n", 1, 10).unwrap(); - let result = trial.suggest_int("n", 1, 20); // Different upper bound - - assert!(matches!(result, Err(Error::ParameterConflict { .. }))); -} - -#[test] -fn test_parameter_conflict_int_vs_log() { - let mut trial = Trial::new(0); - - trial.suggest_int("n", 1, 100).unwrap(); - let result = trial.suggest_int_log("n", 1, 100); // Same bounds but log scale - - assert!(matches!(result, Err(Error::ParameterConflict { .. }))); -} - -#[test] -fn test_parameter_conflict_categorical_different_n_choices() { - let mut trial = Trial::new(0); - - trial.suggest_categorical("opt", &["a", "b", "c"]).unwrap(); - let result = trial.suggest_categorical("opt", &["x", "y"]); // Different number of choices - - assert!(matches!(result, Err(Error::ParameterConflict { .. }))); -} - -#[test] -fn test_parameter_conflict_float_vs_int() { - let mut trial = Trial::new(0); - - trial.suggest_float("x", 0.0, 10.0).unwrap(); - let result = trial.suggest_int("x", 0, 10); // Different type - - assert!(matches!(result, Err(Error::ParameterConflict { .. }))); -} - -#[test] -fn test_parameter_conflict_returns_name() { - let mut trial = Trial::new(0); - - trial.suggest_float("my_param", 0.0, 1.0).unwrap(); - let result = trial.suggest_float("my_param", 0.0, 2.0); - - match result { - Err(Error::ParameterConflict { name, .. }) => { - assert_eq!(name, "my_param"); - } - _ => panic!("expected ParameterConflict error"), - } -} - -// ============================================================================= -// Test: empty categorical returns error -// ============================================================================= - #[test] fn test_empty_categorical_returns_error() { + let param = CategoricalParam::<&str>::new(vec![]); let mut trial = Trial::new(0); - let empty: &[&str] = &[]; - - let result = trial.suggest_categorical("opt", empty); - - assert!(matches!(result, Err(Error::EmptyChoices))); -} - -#[test] -fn test_empty_categorical_vec_returns_error() { - let mut trial = Trial::new(0); - let empty: Vec = vec![]; - - let result = trial.suggest_categorical("numbers", &empty); + let result = trial.suggest_param(¶m); assert!(matches!(result, Err(Error::EmptyChoices))); } @@ -480,10 +412,11 @@ fn test_empty_categorical_vec_returns_error() { #[test] fn test_study_basic_workflow() { let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(-5.0, 5.0); study .optimize(10, |trial| { - let x = trial.suggest_float("x", -5.0, 5.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x * x) }) .expect("optimization should succeed"); @@ -496,6 +429,7 @@ fn test_study_basic_workflow() { #[test] fn test_study_with_failures() { let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(-5.0, 5.0); // Every other trial fails let mut counter = 0; @@ -505,9 +439,7 @@ fn test_study_with_failures() { if counter % 2 == 0 { return Err::("intentional failure"); } - let x = trial - .suggest_float("x", -5.0, 5.0) - .map_err(|_| "param error")?; + let x = x_param.suggest(trial).map_err(|_| "param error")?; Ok(x * x) }) .expect("optimization should succeed with some failures"); @@ -529,11 +461,11 @@ fn test_invalid_bounds_errors() { let mut trial = Trial::new(0); // low > high for float - let result = trial.suggest_float("x", 10.0, 5.0); + let result = trial.suggest_param(&FloatParam::new(10.0, 5.0)); assert!(matches!(result, Err(Error::InvalidBounds { .. }))); // low > high for int - let result = trial.suggest_int("n", 100, 50); + let result = trial.suggest_param(&IntParam::new(100, 50)); assert!(matches!(result, Err(Error::InvalidBounds { .. }))); } @@ -542,14 +474,14 @@ fn test_invalid_log_bounds_errors() { let mut trial = Trial::new(0); // low <= 0 for log float - let result = trial.suggest_float_log("x", 0.0, 1.0); + let result = trial.suggest_param(&FloatParam::new(0.0, 1.0).log_scale()); assert!(matches!(result, Err(Error::InvalidLogBounds))); - let result = trial.suggest_float_log("y", -1.0, 1.0); + let result = trial.suggest_param(&FloatParam::new(-1.0, 1.0).log_scale()); assert!(matches!(result, Err(Error::InvalidLogBounds))); // low < 1 for log int - let result = trial.suggest_int_log("n", 0, 100); + let result = trial.suggest_param(&IntParam::new(0, 100).log_scale()); assert!(matches!(result, Err(Error::InvalidLogBounds))); } @@ -558,14 +490,14 @@ fn test_invalid_step_errors() { let mut trial = Trial::new(0); // step <= 0 for float - let result = trial.suggest_float_step("x", 0.0, 1.0, 0.0); + let result = trial.suggest_param(&FloatParam::new(0.0, 1.0).step(0.0)); assert!(matches!(result, Err(Error::InvalidStep))); - let result = trial.suggest_float_step("y", 0.0, 1.0, -0.1); + let result = trial.suggest_param(&FloatParam::new(0.0, 1.0).step(-0.1)); assert!(matches!(result, Err(Error::InvalidStep))); // step <= 0 for int - let result = trial.suggest_int_step("n", 0, 100, 0); + let result = trial.suggest_param(&IntParam::new(0, 100).step(0)); assert!(matches!(result, Err(Error::InvalidStep))); } @@ -579,11 +511,14 @@ fn test_tpe_with_categorical_parameter() { let study: Study = Study::with_sampler(Direction::Maximize, sampler); + let model_param = CategoricalParam::new(vec!["linear", "quadratic", "cubic"]); + let x_param = FloatParam::new(0.0, 2.0); + // Optimization where the best choice depends on the categorical study .optimize_with_sampler(30, |trial| { - let choice = trial.suggest_categorical("model", &["linear", "quadratic", "cubic"])?; - let x = trial.suggest_float("x", 0.0, 2.0)?; + let choice = model_param.suggest(trial)?; + let x = x_param.suggest(trial)?; // cubic model is best at x=1 let value = match choice { @@ -597,7 +532,6 @@ fn test_tpe_with_categorical_parameter() { .expect("optimization should succeed"); let best = study.best_trial().expect("should have best trial"); - // The optimizer should find that "cubic" with x≈1 is best assert!( best.value > 5.0, "should find good solution, got {}", @@ -615,17 +549,18 @@ fn test_tpe_with_integer_parameters() { let study: Study = Study::with_sampler(Direction::Minimize, sampler); - // Minimize (n - 7)^2 where n ∈ [1, 10] + let n_param = IntParam::new(1, 10); + + // Minimize (n - 7)^2 where n in [1, 10] study .optimize_with_sampler(30, |trial| { - let n = trial.suggest_int("n", 1, 10)?; + let n = n_param.suggest(trial)?; Ok::<_, Error>(((n - 7) as f64).powi(2)) }) .expect("optimization should succeed"); let best = study.best_trial().expect("should have best trial"); - // Best value should be small (n close to 7) assert!( best.value < 5.0, "should find n close to 7, best value = {}", @@ -640,13 +575,14 @@ fn test_callback_early_stopping() { let study: Study = Study::new(Direction::Minimize); let trials_run = Cell::new(0); + let x_param = FloatParam::new(0.0, 10.0); study .optimize_with_callback( 100, |trial| { trials_run.set(trials_run.get() + 1); - let x = trial.suggest_float("x", 0.0, 10.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x) }, |_study, _trial| { @@ -666,10 +602,11 @@ fn test_callback_early_stopping() { #[test] fn test_study_trials_iteration() { let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(0.0, 1.0); study .optimize(5, |trial| { - let x = trial.suggest_float("x", 0.0, 1.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x) }) .unwrap(); @@ -682,10 +619,6 @@ fn test_study_trials_iteration() { !trial.params.is_empty(), "each trial should have parameters" ); - assert!( - trial.params.contains_key("x"), - "each trial should have parameter 'x'" - ); } } @@ -708,22 +641,23 @@ fn test_trial_state() { #[test] fn test_trial_params_access() { + let x_param = FloatParam::new(0.0, 1.0); + let n_param = IntParam::new(1, 10); let mut trial = Trial::new(0); - trial.suggest_float("x", 0.0, 1.0).unwrap(); - trial.suggest_int("n", 1, 10).unwrap(); + x_param.suggest(&mut trial).unwrap(); + n_param.suggest(&mut trial).unwrap(); let params = trial.params(); assert_eq!(params.len(), 2); - assert!(params.contains_key("x")); - assert!(params.contains_key("n")); } #[test] fn test_log_scale_float_range() { + let param = FloatParam::new(1e-5, 1e-1).log_scale(); let mut trial = Trial::new(0); - let lr = trial.suggest_float_log("lr", 1e-5, 1e-1).unwrap(); + let lr = param.suggest(&mut trial).unwrap(); assert!( (1e-5..=1e-1).contains(&lr), "log-scale value {lr} out of range" @@ -732,9 +666,10 @@ fn test_log_scale_float_range() { #[test] fn test_step_float_snaps_to_grid() { + let param = FloatParam::new(0.0, 1.0).step(0.25); let mut trial = Trial::new(0); - let x = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap(); + let x = param.suggest(&mut trial).unwrap(); // x should be one of: 0.0, 0.25, 0.5, 0.75, 1.0 let valid_values = [0.0, 0.25, 0.5, 0.75, 1.0]; @@ -744,9 +679,10 @@ fn test_step_float_snaps_to_grid() { #[test] fn test_step_int_snaps_to_grid() { + let param = IntParam::new(0, 100).step(25); let mut trial = Trial::new(0); - let n = trial.suggest_int_step("n", 0, 100, 25).unwrap(); + let n = param.suggest(&mut trial).unwrap(); // n should be one of: 0, 25, 50, 75, 100 assert!( @@ -758,10 +694,11 @@ fn test_step_int_snaps_to_grid() { #[test] fn test_best_value() { let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(0.0, 10.0); study .optimize(10, |trial| { - let x = trial.suggest_float("x", 0.0, 10.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x) }) .unwrap(); @@ -781,10 +718,8 @@ fn test_best_value() { #[test] fn test_study_set_sampler() { - // Test that set_sampler allows changing the sampler after study creation let mut study: Study = Study::new(Direction::Minimize); - // Initially uses RandomSampler, now switch to TPE let tpe = TpeSampler::builder() .seed(42) .n_startup_trials(5) @@ -792,10 +727,11 @@ fn test_study_set_sampler() { .unwrap(); study.set_sampler(tpe); - // Should work with the new sampler + let x_param = FloatParam::new(-5.0, 5.0); + study .optimize_with_sampler(10, |trial| { - let x = trial.suggest_float("x", -5.0, 5.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x * x) }) .expect("optimization should succeed with new sampler"); @@ -805,12 +741,12 @@ fn test_study_set_sampler() { #[test] fn test_study_with_i32_value_type() { - // Test Study with non-f64 value type let study: Study = Study::new(Direction::Minimize); + let x_param = IntParam::new(-10, 10); study .optimize(10, |trial| { - let x = trial.suggest_int("x", -10, 10)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x.abs() as i32) }) .expect("optimization should succeed"); @@ -824,7 +760,6 @@ fn test_study_with_i32_value_type() { fn test_optimize_all_trials_fail() { let study: Study = Study::new(Direction::Minimize); - // All trials fail let result = study.optimize(5, |_trial| Err::("always fails")); assert!( @@ -883,12 +818,12 @@ fn test_optimize_with_callback_sampler_all_trials_fail() { #[test] fn test_trial_debug_format() { + let param = FloatParam::new(0.0, 1.0); let mut trial = Trial::new(42); - trial.suggest_float("x", 0.0, 1.0).unwrap(); + param.suggest(&mut trial).unwrap(); let debug_str = format!("{trial:?}"); - // Should contain trial id and other fields assert!(debug_str.contains("Trial")); assert!(debug_str.contains("42")); assert!(debug_str.contains("has_sampler")); @@ -901,11 +836,12 @@ fn test_tpe_sampler_builder_default_trait() { let builder = TpeSamplerBuilder::default(); let sampler = builder.build().unwrap(); - // Should have default values let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(0.0, 1.0); + study .optimize_with_sampler(5, |trial| { - let x = trial.suggest_float("x", 0.0, 1.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x) }) .unwrap(); @@ -918,9 +854,11 @@ fn test_tpe_sampler_default_trait() { let sampler = TpeSampler::default(); let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(0.0, 1.0); + study .optimize_with_sampler(5, |trial| { - let x = trial.suggest_float("x", 0.0, 1.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x) }) .unwrap(); @@ -938,10 +876,11 @@ fn test_tpe_with_fixed_kde_bandwidth() { .unwrap(); let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(-5.0, 5.0); study .optimize_with_sampler(20, |trial| { - let x = trial.suggest_float("x", -5.0, 5.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x * x) }) .expect("optimization should succeed"); @@ -958,18 +897,18 @@ fn test_tpe_sampler_invalid_kde_bandwidth() { #[test] fn test_tpe_split_trials_with_two_trials() { - // Edge case: exactly 2 trials in history let sampler = TpeSampler::builder() .seed(42) - .n_startup_trials(2) // TPE kicks in after 2 trials + .n_startup_trials(2) .build() .unwrap(); let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(0.0, 10.0); study .optimize_with_sampler(5, |trial| { - let x = trial.suggest_float("x", 0.0, 10.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x) }) .expect("optimization should succeed with small history"); @@ -986,11 +925,11 @@ fn test_tpe_with_log_scale_int() { .unwrap(); let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let batch_param = IntParam::new(1, 1024).log_scale(); study .optimize_with_sampler(20, |trial| { - let batch_size = trial.suggest_int_log("batch_size", 1, 1024)?; - // Optimal around batch_size = 32 + let batch_size = batch_param.suggest(trial)?; Ok::<_, Error>(((batch_size as f64).log2() - 5.0).powi(2)) }) .expect("optimization should succeed"); @@ -1008,11 +947,13 @@ fn test_tpe_with_step_distributions() { .unwrap(); let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(0.0, 10.0).step(0.5); + let n_param = IntParam::new(0, 100).step(10); study .optimize_with_sampler(20, |trial| { - let x = trial.suggest_float_step("x", 0.0, 10.0, 0.5)?; - let n = trial.suggest_int_step("n", 0, 100, 10)?; + let x = x_param.suggest(trial)?; + let n = n_param.suggest(trial)?; Ok::<_, Error>((x - 5.0).powi(2) + ((n - 50) as f64).powi(2)) }) .expect("optimization should succeed"); @@ -1035,22 +976,24 @@ fn test_create_trial_vs_create_trial_with_sampler() { assert_eq!(trial2.id(), 1); // Both should work for suggesting parameters + let x_param = FloatParam::new(0.0, 1.0); let mut trial3 = study.create_trial(); - let x = trial3.suggest_float("x", 0.0, 1.0).unwrap(); + let x = x_param.suggest(&mut trial3).unwrap(); assert!((0.0..=1.0).contains(&x)); } #[test] fn test_manual_trial_completion() { let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(0.0, 10.0); // Manually create and complete trials let mut trial = study.create_trial(); - let x = trial.suggest_float("x", 0.0, 10.0).unwrap(); + let x = x_param.suggest(&mut trial).unwrap(); study.complete_trial(trial, x * x); let mut trial2 = study.create_trial(); - let y = trial2.suggest_float("x", 0.0, 10.0).unwrap(); + let y = x_param.suggest(&mut trial2).unwrap(); study.complete_trial(trial2, y * y); // Manually fail a trial @@ -1063,35 +1006,36 @@ fn test_manual_trial_completion() { #[test] fn test_distributions_access() { + let x_param = FloatParam::new(0.0, 1.0); + let n_param = IntParam::new(1, 10); + let opt_param = CategoricalParam::new(vec!["a", "b", "c"]); let mut trial = Trial::new(0); - trial.suggest_float("x", 0.0, 1.0).unwrap(); - trial.suggest_int("n", 1, 10).unwrap(); - trial.suggest_categorical("opt", &["a", "b", "c"]).unwrap(); + x_param.suggest(&mut trial).unwrap(); + n_param.suggest(&mut trial).unwrap(); + opt_param.suggest(&mut trial).unwrap(); let dists = trial.distributions(); assert_eq!(dists.len(), 3); - assert!(dists.contains_key("x")); - assert!(dists.contains_key("n")); - assert!(dists.contains_key("opt")); } #[test] fn test_tpe_empty_good_or_bad_values_fallback() { - // When TPE can't find values in the good/bad groups, it falls back to random let sampler = TpeSampler::builder() .seed(42) .n_startup_trials(5) - .gamma(0.1) // Very small gamma means few "good" trials + .gamma(0.1) .build() .unwrap(); let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(0.0, 10.0); + let y_param = FloatParam::new(0.0, 10.0); // First optimize with one parameter study .optimize_with_sampler(10, |trial| { - let x = trial.suggest_float("x", 0.0, 10.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x) }) .unwrap(); @@ -1099,7 +1043,7 @@ fn test_tpe_empty_good_or_bad_values_fallback() { // Now try with a different parameter - TPE won't have history for "y" study .optimize_with_sampler(5, |trial| { - let y = trial.suggest_float("y", 0.0, 10.0)?; + let y = y_param.suggest(trial)?; Ok::<_, Error>(y) }) .unwrap(); @@ -1112,12 +1056,13 @@ fn test_callback_early_stopping_on_first_trial() { use std::ops::ControlFlow; let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(0.0, 10.0); study .optimize_with_callback( 100, |trial| { - let x = trial.suggest_float("x", 0.0, 10.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x) }, |_study, _trial| { @@ -1136,12 +1081,13 @@ fn test_callback_sampler_early_stopping() { let sampler = RandomSampler::with_seed(42); let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(0.0, 10.0); study .optimize_with_callback_sampler( 100, |trial| { - let x = trial.suggest_float("x", 0.0, 10.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x) }, |study, _trial| { @@ -1162,69 +1108,74 @@ fn test_int_bounds_with_low_equals_high() { let mut trial = Trial::new(0); // When low == high, should return that exact value - let n = trial.suggest_int("n", 5, 5).unwrap(); + let n_param = IntParam::new(5, 5); + let n = n_param.suggest(&mut trial).unwrap(); assert_eq!(n, 5); - let x = trial.suggest_float("x", 3.0, 3.0).unwrap(); + let x_param = FloatParam::new(3.0, 3.0); + let x = x_param.suggest(&mut trial).unwrap(); assert_eq!(x, 3.0); } #[test] fn test_best_trial_with_nan_values() { - // Test behavior when comparing with NaN values (PartialOrd edge case) let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(0.0, 10.0); - // Complete some normal trials study .optimize(5, |trial| { - let x = trial.suggest_float("x", 0.0, 10.0)?; + let x = x_param.suggest(trial)?; Ok::<_, Error>(x) }) .unwrap(); - // best_trial should still work let best = study.best_trial(); assert!(best.is_ok()); } // ============================================================================= -// Tests for suggest_bool +// Tests for BoolParam // ============================================================================= #[test] fn test_suggest_bool_caching() { + let param = BoolParam::new(); let mut trial = Trial::new(0); - let b1 = trial.suggest_bool("flag").unwrap(); - let b2 = trial.suggest_bool("flag").unwrap(); + let b1 = param.suggest(&mut trial).unwrap(); + let b2 = param.suggest(&mut trial).unwrap(); - assert_eq!(b1, b2, "repeated suggest_bool should return cached value"); + assert_eq!(b1, b2, "repeated suggest bool should return cached value"); } #[test] fn test_suggest_bool_multiple_parameters() { + let dropout_param = BoolParam::new(); + let batchnorm_param = BoolParam::new(); + let skip_param = BoolParam::new(); let mut trial = Trial::new(0); - let a = trial.suggest_bool("use_dropout").unwrap(); - let b = trial.suggest_bool("use_batchnorm").unwrap(); - let c = trial.suggest_bool("use_skip_connections").unwrap(); + let a = dropout_param.suggest(&mut trial).unwrap(); + let b = batchnorm_param.suggest(&mut trial).unwrap(); + let c = skip_param.suggest(&mut trial).unwrap(); // All should be cached independently - assert_eq!(a, trial.suggest_bool("use_dropout").unwrap()); - assert_eq!(b, trial.suggest_bool("use_batchnorm").unwrap()); - assert_eq!(c, trial.suggest_bool("use_skip_connections").unwrap()); + assert_eq!(a, dropout_param.suggest(&mut trial).unwrap()); + assert_eq!(b, batchnorm_param.suggest(&mut trial).unwrap()); + assert_eq!(c, skip_param.suggest(&mut trial).unwrap()); } #[test] fn test_suggest_bool_in_optimization() { let study: Study = Study::new(Direction::Minimize); + let use_feature_param = BoolParam::new(); + let x_param = FloatParam::new(0.0, 10.0); study .optimize(10, |trial| { - let use_feature = trial.suggest_bool("use_feature")?; - let x = trial.suggest_float("x", 0.0, 10.0)?; + let use_feature = use_feature_param.suggest(trial)?; + let x = x_param.suggest(trial)?; - // Objective depends on boolean flag let value = if use_feature { x } else { x * 2.0 }; Ok::<_, Error>(value) }) @@ -1242,109 +1193,86 @@ fn test_suggest_bool_with_tpe() { .unwrap(); let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let use_large_param = BoolParam::new(); + let x_param = FloatParam::new(0.0, 10.0); study .optimize_with_sampler(20, |trial| { - let use_large = trial.suggest_bool("use_large")?; - let base = if use_large { 10.0 } else { 1.0 }; - let x = trial.suggest_float("x", 0.0, base)?; - Ok::<_, Error>(x) + let use_large = use_large_param.suggest(trial)?; + let x = x_param.suggest(trial)?; + // The value depends on use_large flag + let base = if use_large { x * 2.0 } else { x }; + Ok::<_, Error>(base) }) .unwrap(); let best = study.best_trial().unwrap(); - // Best should prefer use_large=false for smaller range - assert!(best.value < 5.0); + assert!(best.value < 10.0); } // ============================================================================= -// Tests for suggest_range +// Tests for FloatParam and IntParam ranges // ============================================================================= #[test] -fn test_suggest_range_float_exclusive() { +fn test_float_param_exclusive_range() { + let param = FloatParam::new(0.0, 1.0); let mut trial = Trial::new(0); - let x = trial.suggest_range("x", 0.0..1.0).unwrap(); + let x = param.suggest(&mut trial).unwrap(); assert!((0.0..=1.0).contains(&x), "value {x} out of range 0.0..1.0"); } #[test] -fn test_suggest_range_float_inclusive() { +fn test_float_param_inclusive_range() { + let param = FloatParam::new(0.0, 1.0); let mut trial = Trial::new(0); - let x = trial.suggest_range("x", 0.0..=1.0).unwrap(); + let x = param.suggest(&mut trial).unwrap(); assert!((0.0..=1.0).contains(&x), "value {x} out of range 0.0..=1.0"); } #[test] -fn test_suggest_range_int_exclusive() { +fn test_int_param_range() { + let param = IntParam::new(1, 10); let mut trial = Trial::new(0); - // 1..10 means 1 to 9 inclusive - let n = trial.suggest_range("n", 1_i64..10).unwrap(); - assert!( - (1..=9).contains(&n), - "value {n} out of range 1..10 (should be 1-9)" - ); + let n = param.suggest(&mut trial).unwrap(); + assert!((1..=10).contains(&n), "value {n} out of range 1..=10"); } #[test] -fn test_suggest_range_int_inclusive() { +fn test_param_caching_float() { + let param = FloatParam::new(0.0, 1.0); let mut trial = Trial::new(0); - // 1..=10 means 1 to 10 inclusive - let n = trial.suggest_range("n", 1_i64..=10).unwrap(); - assert!( - (1..=10).contains(&n), - "value {n} out of range 1..=10 (should be 1-10)" - ); + let x1 = param.suggest(&mut trial).unwrap(); + let x2 = param.suggest(&mut trial).unwrap(); + + assert_eq!(x1, x2, "repeated suggest should return cached value"); } #[test] -fn test_suggest_range_caching_float() { +fn test_param_caching_int() { + let param = IntParam::new(1, 100); let mut trial = Trial::new(0); - let x1 = trial.suggest_range("x", 0.0..1.0).unwrap(); - let x2 = trial.suggest_range("x", 0.0..1.0).unwrap(); + let n1 = param.suggest(&mut trial).unwrap(); + let n2 = param.suggest(&mut trial).unwrap(); - assert_eq!(x1, x2, "repeated suggest_range should return cached value"); + assert_eq!(n1, n2, "repeated suggest should return cached value"); } #[test] -fn test_suggest_range_caching_int() { - let mut trial = Trial::new(0); - - let n1 = trial.suggest_range("n", 1_i64..=100).unwrap(); - let n2 = trial.suggest_range("n", 1_i64..=100).unwrap(); - - assert_eq!(n1, n2, "repeated suggest_range should return cached value"); -} - -#[test] -fn test_suggest_range_multiple_parameters() { - let mut trial = Trial::new(0); - - let x = trial.suggest_range("x", 0.0..1.0).unwrap(); - let y = trial.suggest_range("y", -5.0..=5.0).unwrap(); - let n = trial.suggest_range("n", 1_i64..10).unwrap(); - let m = trial.suggest_range("m", 100_i64..=200).unwrap(); - - // All should be cached independently - assert_eq!(x, trial.suggest_range("x", 0.0..1.0).unwrap()); - assert_eq!(y, trial.suggest_range("y", -5.0..=5.0).unwrap()); - assert_eq!(n, trial.suggest_range("n", 1_i64..10).unwrap()); - assert_eq!(m, trial.suggest_range("m", 100_i64..=200).unwrap()); -} - -#[test] -fn test_suggest_range_in_optimization() { +fn test_multiple_params_in_optimization() { let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(-10.0, 10.0); + let n_param = IntParam::new(1, 5); study .optimize(10, |trial| { - let x = trial.suggest_range("x", -10.0..10.0)?; - let n = trial.suggest_range("n", 1_i64..=5)?; + let x = x_param.suggest(trial)?; + let n = n_param.suggest(trial)?; Ok::<_, Error>(x * x + n as f64) }) .unwrap(); @@ -1353,7 +1281,7 @@ fn test_suggest_range_in_optimization() { } #[test] -fn test_suggest_range_with_tpe() { +fn test_params_with_tpe() { let sampler = TpeSampler::builder() .seed(42) .n_startup_trials(5) @@ -1361,46 +1289,36 @@ fn test_suggest_range_with_tpe() { .unwrap(); let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(-5.0, 5.0); + let n_param = IntParam::new(1, 10); study .optimize_with_sampler(30, |trial| { - let x = trial.suggest_range("x", -5.0..=5.0)?; - let n = trial.suggest_range("n", 1_i64..=10)?; + let x = x_param.suggest(trial)?; + let n = n_param.suggest(trial)?; Ok::<_, Error>(x * x + (n as f64 - 5.0).powi(2)) }) .unwrap(); let best = study.best_trial().unwrap(); - // TPE should find near-optimal solution assert!(best.value < 10.0, "TPE should find good solution"); } #[test] -fn test_suggest_range_empty_int_range_error() { +fn test_single_value_int_range() { + let param = IntParam::new(5, 5); let mut trial = Trial::new(0); - // 5..5 is empty (no valid integers) - let result = trial.suggest_range("n", 5_i64..5); - assert!( - matches!(result, Err(Error::InvalidBounds { .. })), - "empty range should return InvalidBounds error" - ); -} - -#[test] -fn test_suggest_range_single_value_int() { - let mut trial = Trial::new(0); - - // 5..=5 has exactly one value: 5 - let n = trial.suggest_range("n", 5_i64..=5).unwrap(); + let n = param.suggest(&mut trial).unwrap(); assert_eq!(n, 5, "single-value range should return that value"); } #[test] -fn test_suggest_range_single_value_float() { +fn test_single_value_float_range() { + let param = FloatParam::new(4.2, 4.2); let mut trial = Trial::new(0); - let x = trial.suggest_range("x", 4.2..=4.2).unwrap(); + let x = param.suggest(&mut trial).unwrap(); assert!( (x - 4.2).abs() < f64::EPSILON, "single-value range should return that value" diff --git a/tests/multivariate_tpe_integration.rs b/tests/multivariate_tpe_integration.rs index 21b54f9..13e1d08 100644 --- a/tests/multivariate_tpe_integration.rs +++ b/tests/multivariate_tpe_integration.rs @@ -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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 { diff --git a/tests/parameter_tests.rs b/tests/parameter_tests.rs new file mode 100644 index 0000000..37e205f --- /dev/null +++ b/tests/parameter_tests.rs @@ -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(¶m).unwrap(); + assert!((0.0..=1.0).contains(&x)); + + // Cached + let x2 = trial.suggest_param(¶m).unwrap(); + assert_eq!(x, x2); +} + +#[test] +fn suggest_float_log_param_via_trial() { + let param = FloatParam::new(1e-5, 1e-1).log_scale(); + let mut trial = Trial::new(0); + let lr = trial.suggest_param(¶m).unwrap(); + assert!((1e-5..=1e-1).contains(&lr)); +} + +#[test] +fn suggest_float_step_param_via_trial() { + let param = FloatParam::new(0.0, 1.0).step(0.25); + let mut trial = Trial::new(0); + let x = trial.suggest_param(¶m).unwrap(); + assert!((0.0..=1.0).contains(&x)); +} + +#[test] +fn suggest_int_param_via_trial() { + let param = IntParam::new(1, 10); + let mut trial = Trial::new(0); + let n = trial.suggest_param(¶m).unwrap(); + assert!((1..=10).contains(&n)); + + // Cached + let n2 = trial.suggest_param(¶m).unwrap(); + assert_eq!(n, n2); +} + +#[test] +fn suggest_int_log_param_via_trial() { + let param = IntParam::new(1, 1024).log_scale(); + let mut trial = Trial::new(0); + let batch = trial.suggest_param(¶m).unwrap(); + assert!((1..=1024).contains(&batch)); +} + +#[test] +fn suggest_int_step_param_via_trial() { + let param = IntParam::new(32, 512).step(32); + let mut trial = Trial::new(0); + let units = trial.suggest_param(¶m).unwrap(); + assert!((32..=512).contains(&units)); + assert_eq!((units - 32) % 32, 0); +} + +#[test] +fn suggest_categorical_param_via_trial() { + let choices = vec!["sgd", "adam", "rmsprop"]; + let param = CategoricalParam::new(choices.clone()); + let mut trial = Trial::new(0); + let opt = trial.suggest_param(¶m).unwrap(); + assert!(choices.contains(&opt)); + + // Cached + let opt2 = trial.suggest_param(¶m).unwrap(); + assert_eq!(opt, opt2); +} + +#[test] +fn suggest_bool_param_via_trial() { + let param = BoolParam::new(); + let mut trial = Trial::new(0); + let val = trial.suggest_param(¶m).unwrap(); + let _ = val; + + // Cached + let val2 = trial.suggest_param(¶m).unwrap(); + assert_eq!(val, val2); +} + +#[derive(Clone, Debug, PartialEq)] +enum Activation { + Relu, + Sigmoid, + Tanh, +} + +impl Categorical for Activation { + const N_CHOICES: usize = 3; + + fn from_index(index: usize) -> Self { + match index { + 0 => Activation::Relu, + 1 => Activation::Sigmoid, + 2 => Activation::Tanh, + _ => panic!("invalid index"), + } + } + + fn to_index(&self) -> usize { + match self { + Activation::Relu => 0, + Activation::Sigmoid => 1, + Activation::Tanh => 2, + } + } +} + +#[test] +fn suggest_enum_param_via_trial() { + let param = EnumParam::::new(); + let mut trial = Trial::new(0); + let act = trial.suggest_param(¶m).unwrap(); + assert!([Activation::Relu, Activation::Sigmoid, Activation::Tanh].contains(&act)); + + // Cached + let act2 = trial.suggest_param(¶m).unwrap(); + assert_eq!(act, act2); +} + +#[test] +fn parameter_conflict_detection() { + let float_param = FloatParam::new(0.0, 1.0); + let int_param = IntParam::new(0, 10); + let mut trial = Trial::new(0); + let _ = trial.suggest_param(&float_param).unwrap(); + + // Different param type with different id - no conflict + let result = trial.suggest_param(&int_param); + assert!(result.is_ok()); + + // Different bounds for same param type but different id - no conflict + let float_param2 = FloatParam::new(0.0, 2.0); + let result = trial.suggest_param(&float_param2); + assert!(result.is_ok()); +} + +#[test] +fn validation_prevents_suggest() { + let mut trial = Trial::new(0); + + assert!(trial.suggest_param(&FloatParam::new(1.0, 0.0)).is_err()); + assert!( + trial + .suggest_param(&FloatParam::new(-1.0, 1.0).log_scale()) + .is_err() + ); + assert!( + trial + .suggest_param(&FloatParam::new(0.0, 1.0).step(-0.1)) + .is_err() + ); + assert!(trial.suggest_param(&IntParam::new(10, 1)).is_err()); + assert!( + trial + .suggest_param(&IntParam::new(0, 10).log_scale()) + .is_err() + ); + assert!(trial.suggest_param(&IntParam::new(0, 10).step(-1)).is_err()); + assert!( + trial + .suggest_param(&CategoricalParam::<&str>::new(vec![])) + .is_err() + ); +} + +#[test] +fn parameter_api_with_study() { + let x_param = FloatParam::new(-5.0, 5.0); + let n_param = IntParam::new(1, 10); + let dropout_param = BoolParam::new(); + let opt_param = CategoricalParam::new(vec!["sgd", "adam"]); + + let study: Study = 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; +}