From d81d1de4ff543ce8f6c4da9ac4085655d25064e1 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Thu, 12 Feb 2026 12:43:44 +0100 Subject: [PATCH] refactor(tests): split integration.rs into focused subfolders - Delete 20 duplicate tests already covered by parameter_tests.rs - Move 11 pure Trial unit tests into src/trial.rs - Split remaining 84 integration tests into tests/study/ (9 modules) - Group sampler tests into tests/sampler/ (7 modules) - Group pruner tests into tests/pruner/ (2 modules) --- src/trial.rs | 161 ++ tests/integration.rs | 2252 ----------------- tests/pruner/main.rs | 2 + .../median.rs} | 0 .../threshold.rs} | 0 .../{bohb_integration.rs => sampler/bohb.rs} | 0 tests/{cma_es_tests.rs => sampler/cma_es.rs} | 2 - .../differential_evolution.rs} | 0 tests/{gp_tests.rs => sampler/gp.rs} | 2 - tests/sampler/main.rs | 15 + .../multivariate_tpe.rs} | 0 tests/sampler/random.rs | 142 ++ tests/sampler/tpe.rs | 381 +++ tests/study/ask_tell.rs | 98 + tests/study/builder.rs | 92 + tests/study/constraints.rs | 101 + tests/study/enqueue.rs | 189 ++ tests/study/iterator.rs | 43 + tests/study/main.rs | 15 + tests/study/objective.rs | 346 +++ tests/study/summary.rs | 71 + tests/study/top_trials.rs | 76 + tests/study/workflow.rs | 259 ++ 23 files changed, 1991 insertions(+), 2256 deletions(-) delete mode 100644 tests/integration.rs create mode 100644 tests/pruner/main.rs rename tests/{median_pruner_tests.rs => pruner/median.rs} (100%) rename tests/{threshold_pruner_tests.rs => pruner/threshold.rs} (100%) rename tests/{bohb_integration.rs => sampler/bohb.rs} (100%) rename tests/{cma_es_tests.rs => sampler/cma_es.rs} (99%) rename tests/{differential_evolution_tests.rs => sampler/differential_evolution.rs} (100%) rename tests/{gp_tests.rs => sampler/gp.rs} (99%) create mode 100644 tests/sampler/main.rs rename tests/{multivariate_tpe_integration.rs => sampler/multivariate_tpe.rs} (100%) create mode 100644 tests/sampler/random.rs create mode 100644 tests/sampler/tpe.rs create mode 100644 tests/study/ask_tell.rs create mode 100644 tests/study/builder.rs create mode 100644 tests/study/constraints.rs create mode 100644 tests/study/enqueue.rs create mode 100644 tests/study/iterator.rs create mode 100644 tests/study/main.rs create mode 100644 tests/study/objective.rs create mode 100644 tests/study/summary.rs create mode 100644 tests/study/top_trials.rs create mode 100644 tests/study/workflow.rs diff --git a/src/trial.rs b/src/trial.rs index e46077e..1ef939e 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -483,3 +483,164 @@ impl Trial { Ok(result) } } + +#[cfg(test)] +#[allow(clippy::float_cmp)] +mod tests { + use crate::parameter::{BoolParam, CategoricalParam, FloatParam, IntParam, Parameter}; + use crate::types::TrialState; + + #[test] + fn trial_state() { + // from test_trial_state (L643 of integration.rs) + let trial = super::Trial::new(0); + assert_eq!(trial.state(), TrialState::Running); + } + + #[test] + fn trial_params_access() { + // from test_trial_params_access (L651) + let x_param = FloatParam::new(0.0, 1.0); + let n_param = IntParam::new(1, 10); + let mut trial = super::Trial::new(0); + + x_param.suggest(&mut trial).unwrap(); + n_param.suggest(&mut trial).unwrap(); + + let params = trial.params(); + assert_eq!(params.len(), 2); + } + + #[test] + fn trial_debug_format() { + // from test_trial_debug_format (L792) + let param = FloatParam::new(0.0, 1.0); + let mut trial = super::Trial::new(42); + param.suggest(&mut trial).unwrap(); + + let debug_str = format!("{trial:?}"); + + assert!(debug_str.contains("Trial")); + assert!(debug_str.contains("42")); + assert!(debug_str.contains("has_sampler")); + } + + #[test] + fn distributions_access() { + // from test_distributions_access (L960) + 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 = super::Trial::new(0); + + 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); + } + + #[test] + fn multiple_parameters_independent_caching() { + // from test_multiple_parameters_independent_caching (L356) + 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 = super::Trial::new(0); + + 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(); + + 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()); + } + + #[test] + fn suggest_bool_multiple_parameters() { + // from test_suggest_bool_multiple_parameters (L1131) + let dropout_param = BoolParam::new(); + let batchnorm_param = BoolParam::new(); + let skip_param = BoolParam::new(); + let mut trial = super::Trial::new(0); + + let a = dropout_param.suggest(&mut trial).unwrap(); + let b = batchnorm_param.suggest(&mut trial).unwrap(); + let c = skip_param.suggest(&mut trial).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 param_name() { + // from test_param_name (L1312) + let param = FloatParam::new(0.0, 1.0).name("learning_rate"); + let mut trial = super::Trial::new(0); + param.suggest(&mut trial).unwrap(); + + let labels = trial.param_labels(); + let label = labels.values().next().unwrap(); + assert_eq!(label, "learning_rate"); + } + + #[test] + fn step_float_snaps_to_grid() { + // from test_step_float_snaps_to_grid (L676) + let param = FloatParam::new(0.0, 1.0).step(0.25); + let mut trial = super::Trial::new(0); + + let x = param.suggest(&mut trial).unwrap(); + + let valid_values = [0.0, 0.25, 0.5, 0.75, 1.0]; + let is_valid = valid_values.iter().any(|&v| (x - v).abs() < 1e-10); + assert!(is_valid, "stepped float {x} should snap to grid"); + } + + #[test] + fn step_int_snaps_to_grid() { + // from test_step_int_snaps_to_grid (L689) + let param = IntParam::new(0, 100).step(25); + let mut trial = super::Trial::new(0); + + let n = param.suggest(&mut trial).unwrap(); + + assert!( + n % 25 == 0 && (0..=100).contains(&n), + "stepped int {n} should snap to grid" + ); + } + + #[test] + fn int_bounds_with_low_equals_high() { + // from test_int_bounds_with_low_equals_high (L1086) + let mut trial = super::Trial::new(0); + + let n_param = IntParam::new(5, 5); + let n = n_param.suggest(&mut trial).unwrap(); + assert_eq!(n, 5); + + let x_param = FloatParam::new(3.0, 3.0); + let x = x_param.suggest(&mut trial).unwrap(); + assert_eq!(x, 3.0); + } + + #[test] + fn single_value_float_range() { + // from test_single_value_float_range (L1296) + let param = FloatParam::new(4.2, 4.2); + let mut trial = super::Trial::new(0); + + 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/integration.rs b/tests/integration.rs deleted file mode 100644 index a231230..0000000 --- a/tests/integration.rs +++ /dev/null @@ -1,2252 +0,0 @@ -//! Integration tests for the optimizer library. - -#![allow( - clippy::cast_sign_loss, - clippy::cast_precision_loss, - 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}; - -// ============================================================================= -// Test: optimize simple quadratic function with TPE, finds near-optimal -// ============================================================================= - -#[test] -fn test_tpe_optimizes_quadratic_function() { - // Minimize f(x) = (x - 3)^2 where x in [-10, 10] - // Optimal: x = 3, f(3) = 0 - let sampler = TpeSampler::builder() - .seed(42) - .n_startup_trials(10) - .n_ei_candidates(24) - .build() - .unwrap(); - - let study: Study = Study::with_sampler(Direction::Minimize, sampler); - - let x_param = FloatParam::new(-10.0, 10.0); - - study - .optimize(100, |trial| { - 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 reasonable value over 100 trials - // With random startup + TPE, we expect to get within a few units of optimal - assert!( - best.value < 5.0, - "TPE should find near-optimal: best value {} should be < 5.0", - best.value - ); -} - -#[test] -fn test_tpe_optimizes_multivariate_function() { - // 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) - .n_startup_trials(10) - .build() - .unwrap(); - - 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(100, |trial| { - let x = x_param.suggest(trial)?; - let y = y_param.suggest(trial)?; - Ok::<_, Error>(x * x + y * y) - }) - .expect("optimization should succeed"); - - let best = study.best_trial().expect("should have at least one trial"); - - // TPE should find a reasonably good solution - assert!( - best.value < 5.0, - "TPE should find near-optimal: best value {} should be < 5.0", - best.value - ); -} - -#[test] -fn test_tpe_maximization() { - // Maximize f(x) = -(x - 2)^2 + 10 where x in [-10, 10] - // Optimal: x = 2, f(2) = 10 - let sampler = TpeSampler::builder() - .seed(456) - .n_startup_trials(5) - .build() - .unwrap(); - - let study: Study = Study::with_sampler(Direction::Maximize, sampler); - - let x_param = FloatParam::new(-10.0, 10.0); - - study - .optimize(50, |trial| { - 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"); - - assert!( - best.value > 5.0, - "TPE should find reasonably good solution: best value {} should be > 5.0", - best.value - ); -} - -// ============================================================================= -// Test: RandomSampler samples uniformly across range -// ============================================================================= - -#[test] -fn test_random_sampler_uniform_float_distribution() { - 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 = x_param.suggest(trial)?; - samples.push(x); - Ok::<_, Error>(x) - }) - .unwrap(); - - // All samples should be in range - for &s in &samples { - assert!((0.0..=1.0).contains(&s), "sample {s} out of range [0, 1]"); - } - - // Check distribution is roughly uniform by looking at quartiles - samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); - - let q1 = samples[n_samples / 4]; - let q2 = samples[n_samples / 2]; - let q3 = samples[3 * n_samples / 4]; - - assert!((q1 - 0.25).abs() < 0.1, "Q1 {q1} should be close to 0.25"); - assert!( - (q2 - 0.5).abs() < 0.1, - "Q2 (median) {q2} should be close to 0.5" - ); - assert!((q3 - 0.75).abs() < 0.1, "Q3 {q3} should be close to 0.75"); -} - -#[test] -fn test_random_sampler_uniform_int_distribution() { - let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(123)); - - 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 = 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(); - - let expected = n_samples as f64 / 10.0; - for (i, &count) in counts.iter().enumerate() { - let diff = (count as f64 - expected).abs() / expected; - assert!( - diff < 0.2, - "value {} appeared {} times, expected ~{}, diff = {:.1}%", - i + 1, - count, - expected, - diff * 100.0 - ); - } -} - -#[test] -fn test_random_sampler_uniform_categorical_distribution() { - let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(456)); - - let n_samples = 2000; - 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 = cat_param.suggest(trial)?; - let idx = choices.iter().position(|&c| c == choice).unwrap(); - counts[idx] += 1; - Ok::<_, Error>(idx as f64) - }) - .unwrap(); - - let expected = n_samples as f64 / 4.0; - for (i, &count) in counts.iter().enumerate() { - let diff = (count as f64 - expected).abs() / expected; - assert!( - diff < 0.15, - "category {} appeared {} times, expected ~{}, diff = {:.1}%", - i, - count, - expected, - diff * 100.0 - ); - } -} - -#[test] -fn test_random_sampler_reproducibility() { - let study1: Study = - Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(999)); - let study2: Study = - Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(999)); - - 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(100, |trial| { - let x = x_param1.suggest(trial)?; - values1.push(x); - Ok::<_, Error>(x) - }) - .unwrap(); - - study2 - .optimize(100, |trial| { - let x = x_param2.suggest(trial)?; - values2.push(x); - Ok::<_, Error>(x) - }) - .unwrap(); - - for (i, (v1, v2)) in values1.iter().zip(values2.iter()).enumerate() { - assert_eq!( - v1, v2, - "values at trial {i} should be identical with same seed: {v1} vs {v2}" - ); - } -} - -// ============================================================================= -// 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 = 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 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 = param.suggest(&mut trial).unwrap(); - let x2 = param.suggest(&mut trial).unwrap(); - - assert_eq!( - x1, x2, - "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 = param.suggest(&mut trial).unwrap(); - let x2 = param.suggest(&mut trial).unwrap(); - - assert_eq!( - x1, x2, - "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 = param.suggest(&mut trial).unwrap(); - let n2 = param.suggest(&mut trial).unwrap(); - - 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 = param.suggest(&mut trial).unwrap(); - let n2 = param.suggest(&mut trial).unwrap(); - - assert_eq!( - n1, n2, - "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 = param.suggest(&mut trial).unwrap(); - let n2 = param.suggest(&mut trial).unwrap(); - - assert_eq!( - n1, n2, - "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 c1 = param.suggest(&mut trial).unwrap(); - let c2 = param.suggest(&mut trial).unwrap(); - - assert_eq!( - c1, c2, - "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 = 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, 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()); -} - -// ============================================================================= -// Test: parameter conflict returns error -// ============================================================================= - -#[test] -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_param(¶m1).unwrap(); - // Different param object = different id = no conflict - let result = trial.suggest_param(¶m2); - assert!(result.is_ok()); -} - -#[test] -fn test_empty_categorical_returns_error() { - let param = CategoricalParam::<&str>::new(vec![]); - let mut trial = Trial::new(0); - - let result = trial.suggest_param(¶m); - assert!(matches!(result, Err(Error::EmptyChoices))); -} - -// ============================================================================= -// Additional integration tests -// ============================================================================= - -#[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 = x_param.suggest(trial)?; - Ok::<_, Error>(x * x) - }) - .expect("optimization should succeed"); - - assert_eq!(study.n_trials(), 10); - let best = study.best_trial().expect("should have best trial"); - assert!(best.value >= 0.0, "x^2 should be non-negative"); -} - -#[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; - study - .optimize(10, |trial| { - counter += 1; - if counter % 2 == 0 { - return Err::("intentional failure"); - } - let x = x_param.suggest(trial).map_err(|_| "param error")?; - Ok(x * x) - }) - .expect("optimization should succeed with some failures"); - - // Only half the trials should have succeeded - assert_eq!(study.n_trials(), 5, "only 5 trials should have completed"); -} - -#[test] -fn test_no_completed_trials_error() { - let study: Study = Study::new(Direction::Minimize); - - let result = study.best_trial(); - assert!(matches!(result, Err(Error::NoCompletedTrials))); -} - -#[test] -fn test_invalid_bounds_errors() { - let mut trial = Trial::new(0); - - // low > high for float - 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_param(&IntParam::new(100, 50)); - assert!(matches!(result, Err(Error::InvalidBounds { .. }))); -} - -#[test] -fn test_invalid_log_bounds_errors() { - let mut trial = Trial::new(0); - - // low <= 0 for log float - let result = trial.suggest_param(&FloatParam::new(0.0, 1.0).log_scale()); - assert!(matches!(result, Err(Error::InvalidLogBounds))); - - 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_param(&IntParam::new(0, 100).log_scale()); - assert!(matches!(result, Err(Error::InvalidLogBounds))); -} - -#[test] -fn test_invalid_step_errors() { - let mut trial = Trial::new(0); - - // step <= 0 for float - let result = trial.suggest_param(&FloatParam::new(0.0, 1.0).step(0.0)); - assert!(matches!(result, Err(Error::InvalidStep))); - - 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_param(&IntParam::new(0, 100).step(0)); - assert!(matches!(result, Err(Error::InvalidStep))); -} - -#[test] -fn test_tpe_with_categorical_parameter() { - let sampler = TpeSampler::builder() - .seed(42) - .n_startup_trials(5) - .build() - .unwrap(); - - 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(30, |trial| { - let choice = model_param.suggest(trial)?; - let x = x_param.suggest(trial)?; - - // cubic model is best at x=1 - let value = match choice { - "linear" => x, - "quadratic" => x * x, - "cubic" => -((x - 1.0).powi(2)) + 10.0, // peak at x=1, max value 10 - _ => unreachable!(), - }; - Ok::<_, Error>(value) - }) - .expect("optimization should succeed"); - - let best = study.best_trial().expect("should have best trial"); - assert!( - best.value > 5.0, - "should find good solution, got {}", - best.value - ); -} - -#[test] -fn test_tpe_with_integer_parameters() { - let sampler = TpeSampler::builder() - .seed(789) - .n_startup_trials(5) - .build() - .unwrap(); - - let study: Study = Study::with_sampler(Direction::Minimize, sampler); - - let n_param = IntParam::new(1, 10); - - // Minimize (n - 7)^2 where n in [1, 10] - study - .optimize(30, |trial| { - 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"); - - assert!( - best.value < 5.0, - "should find n close to 7, best value = {}", - best.value - ); -} - -#[test] -fn test_callback_early_stopping() { - use std::ops::ControlFlow; - - use optimizer::Objective; - use optimizer::sampler::CompletedTrial; - - struct EarlyStopAfter5 { - x_param: FloatParam, - } - - impl Objective for EarlyStopAfter5 { - type Error = Error; - fn evaluate(&self, trial: &mut Trial) -> Result { - let x = self.x_param.suggest(trial)?; - Ok(x) - } - fn after_trial(&self, study: &Study, _trial: &CompletedTrial) -> ControlFlow<()> { - if study.n_trials() >= 5 { - ControlFlow::Break(()) - } else { - ControlFlow::Continue(()) - } - } - } - - let study: Study = Study::new(Direction::Minimize); - study - .optimize_with( - 100, - EarlyStopAfter5 { - x_param: FloatParam::new(0.0, 10.0), - }, - ) - .expect("optimization should succeed"); - - assert_eq!(study.n_trials(), 5, "should have stopped after 5 trials"); -} - -#[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 = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }) - .unwrap(); - - let trials = study.trials(); - assert_eq!(trials.len(), 5); - - for trial in &trials { - assert!( - !trial.params.is_empty(), - "each trial should have parameters" - ); - } -} - -#[test] -fn test_study_direction() { - let study_min: Study = Study::new(Direction::Minimize); - assert_eq!(study_min.direction(), Direction::Minimize); - - let study_max: Study = Study::new(Direction::Maximize); - assert_eq!(study_max.direction(), Direction::Maximize); -} - -#[test] -fn test_trial_state() { - use optimizer::TrialState; - - let trial = Trial::new(0); - assert_eq!(trial.state(), TrialState::Running); -} - -#[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); - - x_param.suggest(&mut trial).unwrap(); - n_param.suggest(&mut trial).unwrap(); - - let params = trial.params(); - assert_eq!(params.len(), 2); -} - -#[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 = param.suggest(&mut trial).unwrap(); - assert!( - (1e-5..=1e-1).contains(&lr), - "log-scale value {lr} out of 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 = 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]; - let is_valid = valid_values.iter().any(|&v| (x - v).abs() < 1e-10); - assert!(is_valid, "stepped float {x} should snap 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 = param.suggest(&mut trial).unwrap(); - - // n should be one of: 0, 25, 50, 75, 100 - assert!( - n % 25 == 0 && (0..=100).contains(&n), - "stepped int {n} should snap 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 = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }) - .unwrap(); - - let best_value = study.best_value().expect("should have best value"); - let best_trial = study.best_trial().expect("should have best trial"); - - assert_eq!( - best_value, best_trial.value, - "best_value should match best_trial.value" - ); -} - -// ============================================================================= -// Additional coverage tests -// ============================================================================= - -#[test] -fn test_study_set_sampler() { - let mut study: Study = Study::new(Direction::Minimize); - - let tpe = TpeSampler::builder() - .seed(42) - .n_startup_trials(5) - .build() - .unwrap(); - study.set_sampler(tpe); - - let x_param = FloatParam::new(-5.0, 5.0); - - study - .optimize(10, |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x * x) - }) - .expect("optimization should succeed with new sampler"); - - assert_eq!(study.n_trials(), 10); -} - -#[test] -fn test_study_with_i32_value_type() { - let study: Study = Study::new(Direction::Minimize); - let x_param = IntParam::new(-10, 10); - - study - .optimize(10, |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x.abs() as i32) - }) - .expect("optimization should succeed"); - - assert_eq!(study.n_trials(), 10); - let best = study.best_trial().expect("should have best trial"); - assert!(best.value >= 0, "absolute value should be non-negative"); -} - -#[test] -fn test_optimize_all_trials_fail() { - let study: Study = Study::new(Direction::Minimize); - - let result = study.optimize(5, |_trial| Err::("always fails")); - - assert!( - matches!(result, Err(Error::NoCompletedTrials)), - "should return NoCompletedTrials when all trials fail" - ); -} - -#[test] -fn test_optimize_with_all_trials_fail() { - let study: Study = Study::new(Direction::Minimize); - - let result = study.optimize(5, |_trial| Err::("always fails")); - - assert!( - matches!(result, Err(Error::NoCompletedTrials)), - "should return NoCompletedTrials when all trials fail" - ); -} - -#[test] -fn test_trial_debug_format() { - let param = FloatParam::new(0.0, 1.0); - let mut trial = Trial::new(42); - param.suggest(&mut trial).unwrap(); - - let debug_str = format!("{trial:?}"); - - assert!(debug_str.contains("Trial")); - assert!(debug_str.contains("42")); - assert!(debug_str.contains("has_sampler")); -} - -#[test] -fn test_tpe_sampler_builder_default_trait() { - use optimizer::sampler::tpe::TpeSamplerBuilder; - - let builder = TpeSamplerBuilder::default(); - let sampler = builder.build().unwrap(); - - let study: Study = Study::with_sampler(Direction::Minimize, sampler); - let x_param = FloatParam::new(0.0, 1.0); - - study - .optimize(5, |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }) - .unwrap(); - - assert_eq!(study.n_trials(), 5); -} - -#[test] -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(5, |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }) - .unwrap(); - - assert_eq!(study.n_trials(), 5); -} - -#[test] -fn test_tpe_with_fixed_kde_bandwidth() { - let sampler = TpeSampler::builder() - .seed(42) - .n_startup_trials(5) - .kde_bandwidth(0.5) - .build() - .unwrap(); - - let study: Study = Study::with_sampler(Direction::Minimize, sampler); - let x_param = FloatParam::new(-5.0, 5.0); - - study - .optimize(20, |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x * x) - }) - .expect("optimization should succeed"); - - let best = study.best_trial().unwrap(); - assert!(best.value < 10.0, "should find reasonable solution"); -} - -#[test] -fn test_tpe_sampler_invalid_kde_bandwidth() { - let result = TpeSampler::with_config(0.25, 10, 24, Some(-1.0), None); - assert!(matches!(result, Err(Error::InvalidBandwidth(_)))); -} - -#[test] -fn test_tpe_split_trials_with_two_trials() { - let sampler = TpeSampler::builder() - .seed(42) - .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(5, |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }) - .expect("optimization should succeed with small history"); - - assert_eq!(study.n_trials(), 5); -} - -#[test] -fn test_tpe_with_log_scale_int() { - let sampler = TpeSampler::builder() - .seed(42) - .n_startup_trials(5) - .build() - .unwrap(); - - let study: Study = Study::with_sampler(Direction::Minimize, sampler); - let batch_param = IntParam::new(1, 1024).log_scale(); - - study - .optimize(20, |trial| { - let batch_size = batch_param.suggest(trial)?; - Ok::<_, Error>(((batch_size as f64).log2() - 5.0).powi(2)) - }) - .expect("optimization should succeed"); - - let best = study.best_trial().unwrap(); - assert!(best.value < 10.0, "should find reasonable solution"); -} - -#[test] -fn test_tpe_with_step_distributions() { - let sampler = TpeSampler::builder() - .seed(42) - .n_startup_trials(5) - .build() - .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(20, |trial| { - 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"); - - let best = study.best_trial().unwrap(); - assert!(best.value < 100.0, "should find reasonable solution"); -} - -#[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 = x_param.suggest(&mut trial).unwrap(); - study.complete_trial(trial, x * x); - - let mut trial2 = study.create_trial(); - let y = x_param.suggest(&mut trial2).unwrap(); - study.complete_trial(trial2, y * y); - - // Manually fail a trial - let trial3 = study.create_trial(); - study.fail_trial(trial3, "test failure"); - - // Only 2 completed trials - assert_eq!(study.n_trials(), 2); -} - -#[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); - - 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); -} - -#[test] -fn test_tpe_empty_good_or_bad_values_fallback() { - let sampler = TpeSampler::builder() - .seed(42) - .n_startup_trials(5) - .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(10, |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }) - .unwrap(); - - // Now try with a different parameter - TPE won't have history for "y" - study - .optimize(5, |trial| { - let y = y_param.suggest(trial)?; - Ok::<_, Error>(y) - }) - .unwrap(); - - assert_eq!(study.n_trials(), 15); -} - -#[test] -fn test_callback_early_stopping_on_first_trial() { - use std::ops::ControlFlow; - - use optimizer::Objective; - use optimizer::sampler::CompletedTrial; - - struct StopImmediately { - x_param: FloatParam, - } - - impl Objective for StopImmediately { - type Error = Error; - fn evaluate(&self, trial: &mut Trial) -> Result { - let x = self.x_param.suggest(trial)?; - Ok(x) - } - fn after_trial( - &self, - _study: &Study, - _trial: &CompletedTrial, - ) -> ControlFlow<()> { - ControlFlow::Break(()) - } - } - - let study: Study = Study::new(Direction::Minimize); - study - .optimize_with( - 100, - StopImmediately { - x_param: FloatParam::new(0.0, 10.0), - }, - ) - .expect("optimization should succeed"); - - assert_eq!(study.n_trials(), 1, "should have stopped after 1 trial"); -} - -#[test] -fn test_callback_sampler_early_stopping() { - use std::ops::ControlFlow; - - use optimizer::Objective; - use optimizer::sampler::CompletedTrial; - - struct StopAfter3 { - x_param: FloatParam, - } - - impl Objective for StopAfter3 { - type Error = Error; - fn evaluate(&self, trial: &mut Trial) -> Result { - let x = self.x_param.suggest(trial)?; - Ok(x) - } - fn after_trial(&self, study: &Study, _trial: &CompletedTrial) -> ControlFlow<()> { - if study.n_trials() >= 3 { - ControlFlow::Break(()) - } else { - ControlFlow::Continue(()) - } - } - } - - let sampler = RandomSampler::with_seed(42); - let study: Study = Study::with_sampler(Direction::Minimize, sampler); - study - .optimize_with( - 100, - StopAfter3 { - x_param: FloatParam::new(0.0, 10.0), - }, - ) - .expect("optimization should succeed"); - - assert_eq!(study.n_trials(), 3); -} - -#[test] -fn test_int_bounds_with_low_equals_high() { - let mut trial = Trial::new(0); - - // When low == high, should return that exact value - let n_param = IntParam::new(5, 5); - let n = n_param.suggest(&mut trial).unwrap(); - assert_eq!(n, 5); - - 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() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); - - study - .optimize(5, |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }) - .unwrap(); - - let best = study.best_trial(); - assert!(best.is_ok()); -} - -// ============================================================================= -// Tests for BoolParam -// ============================================================================= - -#[test] -fn test_suggest_bool_caching() { - let param = BoolParam::new(); - let mut trial = Trial::new(0); - - 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"); -} - -#[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 = 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, 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 = use_feature_param.suggest(trial)?; - let x = x_param.suggest(trial)?; - - let value = if use_feature { x } else { x * 2.0 }; - Ok::<_, Error>(value) - }) - .unwrap(); - - assert_eq!(study.n_trials(), 10); -} - -#[test] -fn test_suggest_bool_with_tpe() { - let sampler = TpeSampler::builder() - .seed(42) - .n_startup_trials(5) - .build() - .unwrap(); - - let study: Study = Study::with_sampler(Direction::Minimize, sampler); - let use_large_param = BoolParam::new(); - let x_param = FloatParam::new(0.0, 10.0); - - study - .optimize(20, |trial| { - 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(); - assert!(best.value < 10.0); -} - -// ============================================================================= -// Tests for FloatParam and IntParam ranges -// ============================================================================= - -#[test] -fn test_float_param_exclusive_range() { - 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), "value {x} out of range 0.0..1.0"); -} - -#[test] -fn test_float_param_inclusive_range() { - 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), "value {x} out of range 0.0..=1.0"); -} - -#[test] -fn test_int_param_range() { - let param = IntParam::new(1, 10); - let mut trial = Trial::new(0); - - let n = param.suggest(&mut trial).unwrap(); - assert!((1..=10).contains(&n), "value {n} out of range 1..=10"); -} - -#[test] -fn test_param_caching_float() { - let param = FloatParam::new(0.0, 1.0); - let mut trial = Trial::new(0); - - 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_param_caching_int() { - let param = IntParam::new(1, 100); - let mut trial = Trial::new(0); - - let n1 = param.suggest(&mut trial).unwrap(); - let n2 = param.suggest(&mut trial).unwrap(); - - assert_eq!(n1, n2, "repeated suggest should return cached value"); -} - -#[test] -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 = x_param.suggest(trial)?; - let n = n_param.suggest(trial)?; - Ok::<_, Error>(x * x + n as f64) - }) - .unwrap(); - - assert_eq!(study.n_trials(), 10); -} - -#[test] -fn test_params_with_tpe() { - let sampler = TpeSampler::builder() - .seed(42) - .n_startup_trials(5) - .build() - .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(30, |trial| { - 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(); - assert!(best.value < 10.0, "TPE should find good solution"); -} - -#[test] -fn test_single_value_int_range() { - let param = IntParam::new(5, 5); - let mut trial = Trial::new(0); - - let n = param.suggest(&mut trial).unwrap(); - assert_eq!(n, 5, "single-value range should return that value"); -} - -#[test] -fn test_single_value_float_range() { - let param = FloatParam::new(4.2, 4.2); - let mut trial = Trial::new(0); - - let x = param.suggest(&mut trial).unwrap(); - assert!( - (x - 4.2).abs() < f64::EPSILON, - "single-value range should return that value" - ); -} - -// ============================================================================= -// Tests for new API features -// ============================================================================= - -#[test] -fn test_param_name() { - let param = FloatParam::new(0.0, 1.0).name("learning_rate"); - let mut trial = Trial::new(0); - param.suggest(&mut trial).unwrap(); - - let labels = trial.param_labels(); - let label = labels.values().next().unwrap(); - assert_eq!(label, "learning_rate"); -} - -#[test] -fn test_completed_trial_get() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(-10.0, 10.0).name("x"); - let n_param = IntParam::new(1, 10).name("n"); - - study - .optimize(5, |trial| { - let x = x_param.suggest(trial)?; - let n = n_param.suggest(trial)?; - Ok::<_, Error>(x * x + n as f64) - }) - .unwrap(); - - let best = study.best_trial().unwrap(); - let x_val: f64 = best.get(&x_param).unwrap(); - let n_val: i64 = best.get(&n_param).unwrap(); - assert!((-10.0..=10.0).contains(&x_val)); - assert!((1..=10).contains(&n_val)); -} - -// ============================================================================= -// Tests for top_trials -// ============================================================================= - -#[test] -fn test_top_trials_minimize() { - let study: Study = Study::new(Direction::Minimize); - - // Manually complete trials with known values - for &val in &[5.0, 1.0, 3.0, 2.0, 4.0] { - let trial = study.create_trial(); - study.complete_trial(trial, val); - } - - let top3 = study.top_trials(3); - assert_eq!(top3.len(), 3); - assert_eq!(top3[0].value, 1.0); - assert_eq!(top3[1].value, 2.0); - assert_eq!(top3[2].value, 3.0); -} - -#[test] -fn test_top_trials_maximize() { - let study: Study = Study::new(Direction::Maximize); - - for &val in &[5.0, 1.0, 3.0, 2.0, 4.0] { - let trial = study.create_trial(); - study.complete_trial(trial, val); - } - - let top3 = study.top_trials(3); - assert_eq!(top3.len(), 3); - assert_eq!(top3[0].value, 5.0); - assert_eq!(top3[1].value, 4.0); - assert_eq!(top3[2].value, 3.0); -} - -#[test] -fn test_top_trials_n_greater_than_total() { - let study: Study = Study::new(Direction::Minimize); - - for &val in &[3.0, 1.0] { - let trial = study.create_trial(); - study.complete_trial(trial, val); - } - - let top = study.top_trials(10); - assert_eq!(top.len(), 2); - assert_eq!(top[0].value, 1.0); - assert_eq!(top[1].value, 3.0); -} - -#[test] -fn test_top_trials_empty() { - let study: Study = Study::new(Direction::Minimize); - - let top = study.top_trials(5); - assert!(top.is_empty()); -} - -#[test] -fn test_top_trials_excludes_pruned() { - let study: Study = Study::new(Direction::Minimize); - - // Complete some trials - for &val in &[5.0, 1.0, 3.0] { - let trial = study.create_trial(); - study.complete_trial(trial, val); - } - - // Prune a trial (it gets a default value of 0.0 but should be excluded) - let trial = study.create_trial(); - study.prune_trial(trial); - - let top = study.top_trials(5); - assert_eq!(top.len(), 3, "pruned trial should be excluded"); - assert_eq!(top[0].value, 1.0); -} - -// ============================================================================= -// Test: ask-and-tell interface -// ============================================================================= - -#[test] -fn test_ask_and_tell_basic() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); - - for _ in 0..10 { - let mut trial = study.ask(); - let x = x_param.suggest(&mut trial).unwrap(); - let value = x * x; - study.tell(trial, Ok::<_, &str>(value)); - } - - assert_eq!(study.n_trials(), 10); - assert!(study.best_value().unwrap() >= 0.0); -} - -#[test] -fn test_ask_and_tell_with_failures() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(-5.0, 5.0); - - // Alternate success and failure - for i in 0..10 { - let mut trial = study.ask(); - let x = x_param.suggest(&mut trial).unwrap(); - if i % 2 == 0 { - study.tell(trial, Ok::<_, &str>(x * x)); - } else { - study.tell(trial, Err::("simulated failure")); - } - } - - // Only successful trials are counted - assert_eq!(study.n_trials(), 5); -} - -#[test] -fn test_ask_and_tell_with_tpe_sampler() { - let sampler = TpeSampler::builder() - .seed(42) - .n_startup_trials(5) - .build() - .unwrap(); - let study: Study = Study::minimize(sampler); - let x_param = FloatParam::new(-10.0, 10.0); - - for _ in 0..30 { - let mut trial = study.ask(); - let x = x_param.suggest(&mut trial).unwrap(); - study.tell(trial, Ok::<_, &str>((x - 3.0).powi(2))); - } - - assert_eq!(study.n_trials(), 30); - assert!( - study.best_value().unwrap() < 5.0, - "TPE ask-and-tell should find a reasonable value" - ); -} - -#[test] -fn test_ask_and_tell_batch() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); - - // Ask a batch of trials - let batch: Vec<_> = (0..5) - .map(|_| { - let mut t = study.ask(); - let x = x_param.suggest(&mut t).unwrap(); - (t, x) - }) - .collect(); - - // Tell results for the batch - for (trial, x) in batch { - study.tell(trial, Ok::<_, &str>(x * x)); - } - - assert_eq!(study.n_trials(), 5); -} - -#[test] -fn test_ask_and_tell_with_custom_value_type() { - // Ask-and-tell works with non-f64 value types too - let study: Study = Study::new(Direction::Maximize); - - for i in 0..5 { - let trial = study.ask(); - study.tell(trial, Ok::<_, &str>(i * 10)); - } - - assert_eq!(study.n_trials(), 5); - assert_eq!(study.best_value().unwrap(), 40); -} - -// ============================================================================= -// Tests: enqueue trials -// ============================================================================= - -use std::collections::HashMap; - -use optimizer::parameter::ParamValue; - -#[test] -fn test_enqueue_params_evaluated_first() { - let study: Study = Study::new(Direction::Minimize); - let x = FloatParam::new(0.0, 10.0); - let y = IntParam::new(1, 100); - - // Enqueue a specific configuration - study.enqueue(HashMap::from([ - (x.id(), ParamValue::Float(5.0)), - (y.id(), ParamValue::Int(42)), - ])); - - // The first trial should use the enqueued params - let mut trial = study.ask(); - let x_val = x.suggest(&mut trial).unwrap(); - let y_val = y.suggest(&mut trial).unwrap(); - - assert_eq!(x_val, 5.0); - assert_eq!(y_val, 42); -} - -#[test] -fn test_enqueue_fifo_order() { - let study: Study = Study::new(Direction::Minimize); - let x = FloatParam::new(0.0, 10.0); - - // Enqueue two configs - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))])); - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))])); - - // First trial gets first enqueued value - let mut trial1 = study.ask(); - assert_eq!(x.suggest(&mut trial1).unwrap(), 1.0); - - // Second trial gets second enqueued value - let mut trial2 = study.ask(); - assert_eq!(x.suggest(&mut trial2).unwrap(), 2.0); -} - -#[test] -fn test_enqueue_then_normal_sampling_resumes() { - let sampler = RandomSampler::with_seed(42); - let study: Study = Study::with_sampler(Direction::Minimize, sampler); - let x = FloatParam::new(0.0, 10.0); - - // Enqueue one config - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(5.0))])); - - // First trial uses enqueued value - let mut trial1 = study.ask(); - assert_eq!(x.suggest(&mut trial1).unwrap(), 5.0); - study.tell(trial1, Ok::<_, &str>(25.0)); - - // Second trial uses normal sampling (not 5.0) - let mut trial2 = study.ask(); - let x_val = x.suggest(&mut trial2).unwrap(); - // The sampled value should be in [0, 10] but extremely unlikely to be exactly 5.0 - assert!((0.0..=10.0).contains(&x_val)); -} - -#[test] -fn test_enqueue_with_optimize() { - let study: Study = Study::new(Direction::Minimize); - let x = FloatParam::new(0.0, 10.0); - - // Enqueue two specific configs - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))])); - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))])); - - let mut values = Vec::new(); - - study - .optimize(5, |trial| { - let x_val = x.suggest(trial)?; - values.push(x_val); - Ok::<_, Error>(x_val * x_val) - }) - .unwrap(); - - // First two trials should use enqueued values - assert_eq!(values[0], 1.0); - assert_eq!(values[1], 2.0); - // All 5 trials should have completed - assert_eq!(study.n_trials(), 5); -} - -#[test] -fn test_enqueue_partial_params_fall_back_to_sampling() { - let study: Study = Study::new(Direction::Minimize); - let x = FloatParam::new(0.0, 10.0); - let y = IntParam::new(1, 100); - - // Enqueue only x, not y - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(3.0))])); - - let mut trial = study.ask(); - let x_val = x.suggest(&mut trial).unwrap(); - let y_val = y.suggest(&mut trial).unwrap(); - - // x should be the enqueued value - assert_eq!(x_val, 3.0); - // y should be sampled (within range) - assert!((1..=100).contains(&y_val)); -} - -#[test] -fn test_enqueue_trials_appear_in_completed_trials() { - let study: Study = Study::new(Direction::Minimize); - let x = FloatParam::new(0.0, 10.0); - - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(7.0))])); - - study - .optimize(1, |trial| { - let x_val = x.suggest(trial)?; - Ok::<_, Error>(x_val) - }) - .unwrap(); - - let trials = study.trials(); - assert_eq!(trials.len(), 1); - assert_eq!(trials[0].value, 7.0); - assert_eq!( - *trials[0].params.get(&x.id()).unwrap(), - ParamValue::Float(7.0) - ); -} - -#[test] -fn test_enqueue_with_ask_and_tell() { - let study: Study = Study::new(Direction::Minimize); - let x = FloatParam::new(0.0, 10.0); - - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(4.0))])); - - let mut trial = study.ask(); - let x_val = x.suggest(&mut trial).unwrap(); - assert_eq!(x_val, 4.0); - - study.tell(trial, Ok::<_, &str>(x_val * x_val)); - assert_eq!(study.n_trials(), 1); - assert_eq!(study.best_value().unwrap(), 16.0); -} - -#[test] -fn test_n_enqueued() { - let study: Study = Study::new(Direction::Minimize); - let x = FloatParam::new(0.0, 10.0); - - assert_eq!(study.n_enqueued(), 0); - - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))])); - assert_eq!(study.n_enqueued(), 1); - - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))])); - assert_eq!(study.n_enqueued(), 2); - - // Creating a trial dequeues one - let _ = study.ask(); - assert_eq!(study.n_enqueued(), 1); - - let _ = study.ask(); - assert_eq!(study.n_enqueued(), 0); -} - -#[test] -fn test_enqueue_counted_in_n_trials() { - let study: Study = Study::new(Direction::Minimize); - let x = FloatParam::new(0.0, 10.0); - - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))])); - study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))])); - - study - .optimize(5, |trial| { - let x_val = x.suggest(trial)?; - Ok::<_, Error>(x_val) - }) - .unwrap(); - - // All 5 trials count, including the 2 enqueued ones - assert_eq!(study.n_trials(), 5); -} - -// ============================================================================= -// Test: Study summary and Display -// ============================================================================= - -#[test] -fn test_summary_with_completed_trials() { - let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(1)); - let x = FloatParam::new(0.0, 10.0).name("x"); - - study - .optimize(5, |trial| { - let val = x.suggest(trial)?; - Ok::<_, Error>(val * val) - }) - .unwrap(); - - let summary = study.summary(); - assert!(summary.contains("Minimize")); - assert!(summary.contains("5 trials")); - assert!(summary.contains("Best value:")); - assert!(summary.contains("x = ")); -} - -#[test] -fn test_summary_no_completed_trials() { - let study: Study = Study::new(Direction::Maximize); - let summary = study.summary(); - assert!(summary.contains("Maximize")); - assert!(summary.contains("0 trials")); - assert!(!summary.contains("Best value:")); -} - -#[test] -fn test_summary_with_pruned_trials() { - let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(1)); - let x = FloatParam::new(0.0, 10.0).name("x"); - - // Manually create some complete and pruned trials - for _ in 0..3 { - let mut trial = study.create_trial(); - let val = x.suggest(&mut trial).unwrap(); - study.complete_trial(trial, val); - } - for _ in 0..2 { - let mut trial = study.create_trial(); - let _ = x.suggest(&mut trial).unwrap(); - study.prune_trial(trial); - } - - let summary = study.summary(); - // Should show breakdown when there are pruned trials - if study.n_pruned_trials() > 0 { - assert!(summary.contains("complete")); - assert!(summary.contains("pruned")); - } -} - -#[test] -fn test_display_matches_summary() { - let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(1)); - let x = FloatParam::new(0.0, 10.0).name("x"); - - study - .optimize(3, |trial| { - let val = x.suggest(trial)?; - Ok::<_, Error>(val) - }) - .unwrap(); - - assert_eq!(format!("{study}"), study.summary()); -} - -// ============================================================================= -// Tests: optimize_with retries via Objective trait -// ============================================================================= - -#[test] -fn test_retries_successful_trials_not_retried() { - use std::sync::Arc; - use std::sync::atomic::{AtomicU32, Ordering}; - - use optimizer::Objective; - - struct SuccessObj { - x_param: FloatParam, - call_count: Arc, - } - - impl Objective for SuccessObj { - type Error = Error; - fn evaluate(&self, trial: &mut Trial) -> Result { - let x = self.x_param.suggest(trial)?; - self.call_count.fetch_add(1, Ordering::Relaxed); - Ok(x * x) - } - fn max_retries(&self) -> usize { - 3 - } - } - - let study: Study = Study::new(Direction::Minimize); - let call_count = Arc::new(AtomicU32::new(0)); - let obj = SuccessObj { - x_param: FloatParam::new(0.0, 10.0), - call_count: Arc::clone(&call_count), - }; - - study.optimize_with(5, obj).unwrap(); - - // All trials succeed on first try — exactly 5 calls - assert_eq!(call_count.load(Ordering::Relaxed), 5); - assert_eq!(study.n_trials(), 5); -} - -#[test] -fn test_retries_failed_trials_retried_up_to_max() { - use std::sync::Arc; - use std::sync::atomic::{AtomicU32, Ordering}; - - use optimizer::Objective; - - struct AlwaysFailObj { - x_param: FloatParam, - call_count: Arc, - } - - impl Objective for AlwaysFailObj { - type Error = String; - fn evaluate(&self, trial: &mut Trial) -> Result { - let _ = self.x_param.suggest(trial).map_err(|e| e.to_string())?; - self.call_count.fetch_add(1, Ordering::Relaxed); - Err("always fails".to_string()) - } - fn max_retries(&self) -> usize { - 3 - } - } - - let study: Study = Study::new(Direction::Minimize); - let call_count = Arc::new(AtomicU32::new(0)); - let obj = AlwaysFailObj { - x_param: FloatParam::new(0.0, 10.0), - call_count: Arc::clone(&call_count), - }; - - let result = study.optimize_with(1, obj); - - // 1 initial attempt + 3 retries = 4 total calls - assert_eq!(call_count.load(Ordering::Relaxed), 4); - // No trials completed - assert!(matches!(result, Err(Error::NoCompletedTrials))); -} - -#[test] -fn test_retries_permanently_failed_after_exhaustion() { - use optimizer::Objective; - - struct AlwaysFailObj { - x_param: FloatParam, - } - - impl Objective for AlwaysFailObj { - type Error = String; - fn evaluate(&self, trial: &mut Trial) -> Result { - let _ = self.x_param.suggest(trial).map_err(|e| e.to_string())?; - Err("transient error".to_string()) - } - fn max_retries(&self) -> usize { - 2 - } - } - - let study: Study = Study::new(Direction::Minimize); - let obj = AlwaysFailObj { - x_param: FloatParam::new(0.0, 10.0), - }; - - let result = study.optimize_with(3, obj); - - assert!( - matches!(result, Err(Error::NoCompletedTrials)), - "all trials should permanently fail" - ); - assert_eq!( - study.n_trials(), - 0, - "no completed trials should be recorded" - ); -} - -#[test] -fn test_retries_uses_same_parameters() { - use std::sync::atomic::{AtomicU32, Ordering}; - use std::sync::{Arc, Mutex}; - - use optimizer::Objective; - - struct RetryObj { - x_param: FloatParam, - seen_values: Arc>>, - call_count: Arc, - } - - impl Objective for RetryObj { - type Error = String; - fn evaluate(&self, trial: &mut Trial) -> Result { - let x = self.x_param.suggest(trial).map_err(|e| e.to_string())?; - self.seen_values.lock().unwrap().push(x); - let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1; - // Fail first two attempts, succeed on third - if count < 3 { - Err("transient".to_string()) - } else { - Ok(x * x) - } - } - fn max_retries(&self) -> usize { - 2 - } - } - - let study: Study = Study::new(Direction::Minimize); - let seen_values = Arc::new(Mutex::new(Vec::new())); - let call_count = Arc::new(AtomicU32::new(0)); - let obj = RetryObj { - x_param: FloatParam::new(0.0, 10.0), - seen_values: Arc::clone(&seen_values), - call_count: Arc::clone(&call_count), - }; - - study.optimize_with(1, obj).unwrap(); - - let values = seen_values.lock().unwrap(); - assert_eq!(values.len(), 3, "should be called 3 times (1 + 2 retries)"); - // All three calls should have gotten the same parameter value - assert_eq!(values[0], values[1]); - assert_eq!(values[1], values[2]); -} - -#[test] -fn test_retries_n_trials_counts_unique_configs() { - use std::sync::Arc; - use std::sync::atomic::{AtomicU32, Ordering}; - - use optimizer::Objective; - - struct FailFirstObj { - x_param: FloatParam, - call_count: Arc, - } - - impl Objective for FailFirstObj { - type Error = String; - fn evaluate(&self, trial: &mut Trial) -> Result { - let x = self.x_param.suggest(trial).map_err(|e| e.to_string())?; - let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1; - // Fail first attempt of each config, succeed on retry - if count % 2 == 1 { - Err("transient".to_string()) - } else { - Ok(x * x) - } - } - fn max_retries(&self) -> usize { - 2 - } - } - - let study: Study = Study::new(Direction::Minimize); - let call_count = Arc::new(AtomicU32::new(0)); - let obj = FailFirstObj { - x_param: FloatParam::new(0.0, 10.0), - call_count: Arc::clone(&call_count), - }; - - study.optimize_with(3, obj).unwrap(); - - // 3 unique configs, each needing 2 calls = 6 total calls - assert_eq!(call_count.load(Ordering::Relaxed), 6); - // But only 3 completed trials - assert_eq!(study.n_trials(), 3); -} - -#[test] -fn test_retries_with_zero_max_retries_same_as_optimize() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); - let call_count = std::cell::Cell::new(0u32); - - study - .optimize(5, |trial| { - let x = x_param.suggest(trial)?; - call_count.set(call_count.get() + 1); - Ok::<_, Error>(x * x) - }) - .unwrap(); - - assert_eq!(call_count.get(), 5); - assert_eq!(study.n_trials(), 5); -} - -// ============================================================================= -// Tests: IntoIterator for &Study -// ============================================================================= - -#[test] -fn test_into_iterator_iterates_all_trials() { - let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); - let x_param = FloatParam::new(0.0, 10.0); - - for _ in 0..5 { - let mut trial = study.create_trial(); - let x = x_param.suggest(&mut trial).unwrap(); - study.complete_trial(trial, x * x); - } - - let mut count = 0; - for trial in &study { - assert_eq!(trial.state, optimizer::TrialState::Complete); - count += 1; - } - assert_eq!(count, 5); -} - -#[test] -fn test_into_iterator_empty_study() { - let study: Study = Study::new(Direction::Minimize); - - let count = (&study).into_iter().count(); - assert_eq!(count, 0); -} - -#[test] -fn test_into_iterator_preserves_insertion_order() { - let study: Study = Study::new(Direction::Minimize); - - for i in 0..3 { - let trial = study.create_trial(); - study.complete_trial(trial, f64::from(i)); - } - - let ids: Vec = (&study).into_iter().map(|t| t.id).collect(); - assert_eq!(ids, vec![0, 1, 2]); -} - -// ============================================================================= -// Tests: Constraint handling -// ============================================================================= - -#[test] -fn test_is_feasible_all_satisfied() { - let study: Study = Study::new(Direction::Minimize); - let mut trial = study.create_trial(); - trial.set_constraints(vec![-1.0, 0.0, -0.5]); - study.complete_trial(trial, 1.0); - - let completed = study.best_trial().unwrap(); - assert!(completed.is_feasible()); -} - -#[test] -fn test_is_feasible_one_violated() { - let study: Study = Study::new(Direction::Minimize); - let mut trial = study.create_trial(); - trial.set_constraints(vec![-1.0, 0.5, -0.5]); - study.complete_trial(trial, 1.0); - - let completed = study.best_trial().unwrap(); - assert!(!completed.is_feasible()); -} - -#[test] -fn test_is_feasible_empty_constraints() { - let study: Study = Study::new(Direction::Minimize); - let trial = study.create_trial(); - study.complete_trial(trial, 1.0); - - let completed = study.best_trial().unwrap(); - assert!(completed.is_feasible()); -} - -#[test] -fn test_best_trial_prefers_feasible() { - let study: Study = Study::new(Direction::Minimize); - - // Infeasible trial with better objective - let mut trial1 = study.create_trial(); - trial1.set_constraints(vec![1.0]); - study.complete_trial(trial1, 0.1); - - // Feasible trial with worse objective - let mut trial2 = study.create_trial(); - trial2.set_constraints(vec![-1.0]); - study.complete_trial(trial2, 100.0); - - let best = study.best_trial().unwrap(); - assert_eq!(best.id, 1); // feasible trial wins - assert_eq!(best.value, 100.0); -} - -#[test] -fn test_best_trial_feasible_by_objective() { - let study: Study = Study::new(Direction::Minimize); - - // Feasible, worse objective - let mut trial1 = study.create_trial(); - trial1.set_constraints(vec![-1.0]); - study.complete_trial(trial1, 10.0); - - // Feasible, better objective - let mut trial2 = study.create_trial(); - trial2.set_constraints(vec![-0.5]); - study.complete_trial(trial2, 2.0); - - let best = study.best_trial().unwrap(); - assert_eq!(best.id, 1); // lower objective wins among feasible - assert_eq!(best.value, 2.0); -} - -#[test] -fn test_top_trials_ranks_feasible_above_infeasible() { - let study: Study = Study::new(Direction::Minimize); - - // Infeasible, low violation - let mut t0 = study.create_trial(); - t0.set_constraints(vec![0.5]); - study.complete_trial(t0, 1.0); - - // Feasible, worst objective among feasible - let mut t1 = study.create_trial(); - t1.set_constraints(vec![-1.0]); - study.complete_trial(t1, 50.0); - - // Feasible, best objective among feasible - let mut t2 = study.create_trial(); - t2.set_constraints(vec![-0.1]); - study.complete_trial(t2, 5.0); - - // Infeasible, high violation - let mut t3 = study.create_trial(); - t3.set_constraints(vec![3.0]); - study.complete_trial(t3, 0.5); - - let top = study.top_trials(4); - let ids: Vec = top.iter().map(|t| t.id).collect(); - // Feasible sorted by objective first (5.0, 50.0), then infeasible by violation (0.5, 3.0) - assert_eq!(ids, vec![2, 1, 0, 3]); -} - -// ============================================================================= -// Test: StudyBuilder -// ============================================================================= - -#[test] -fn test_builder_defaults() { - let study: Study = Study::builder().build(); - assert_eq!(study.direction(), Direction::Minimize); -} - -#[test] -fn test_builder_maximize() { - let study: Study = Study::builder().maximize().build(); - assert_eq!(study.direction(), Direction::Maximize); -} - -#[test] -fn test_builder_minimize() { - let study: Study = Study::builder().minimize().build(); - assert_eq!(study.direction(), Direction::Minimize); -} - -#[test] -fn test_builder_direction() { - let study: Study = Study::builder().direction(Direction::Maximize).build(); - assert_eq!(study.direction(), Direction::Maximize); -} - -#[test] -fn test_builder_with_sampler() { - let x = FloatParam::new(-5.0, 5.0); - let study: Study = Study::builder().sampler(TpeSampler::new()).build(); - - study - .optimize(10, |trial| { - let val = x.suggest(trial)?; - Ok::<_, Error>(val * val) - }) - .unwrap(); - - assert_eq!(study.trials().len(), 10); -} - -#[test] -fn test_builder_with_pruner() { - use optimizer::pruner::NopPruner; - - let study: Study = Study::builder().pruner(NopPruner).build(); - - assert_eq!(study.direction(), Direction::Minimize); -} - -#[test] -fn test_builder_chaining() { - let study: Study = Study::builder() - .maximize() - .sampler(RandomSampler::with_seed(42)) - .pruner(optimizer::pruner::NopPruner) - .build(); - - assert_eq!(study.direction(), Direction::Maximize); -} - -#[test] -fn test_builder_with_custom_value_type() { - let study: Study = Study::builder().maximize().build(); - assert_eq!(study.direction(), Direction::Maximize); -} - -#[test] -fn test_builder_optimizes_correctly() { - let x = FloatParam::new(-10.0, 10.0); - let study: Study = Study::builder() - .minimize() - .sampler(TpeSampler::builder().seed(42).build().unwrap()) - .build(); - - study - .optimize(100, |trial| { - let val = x.suggest(trial)?; - Ok::<_, Error>((val - 3.0) * (val - 3.0)) - }) - .unwrap(); - - let best = study.best_trial().unwrap(); - assert!( - best.value < 5.0, - "best value should be < 5.0, got {}", - best.value - ); -} diff --git a/tests/pruner/main.rs b/tests/pruner/main.rs new file mode 100644 index 0000000..b9eaa5d --- /dev/null +++ b/tests/pruner/main.rs @@ -0,0 +1,2 @@ +mod median; +mod threshold; diff --git a/tests/median_pruner_tests.rs b/tests/pruner/median.rs similarity index 100% rename from tests/median_pruner_tests.rs rename to tests/pruner/median.rs diff --git a/tests/threshold_pruner_tests.rs b/tests/pruner/threshold.rs similarity index 100% rename from tests/threshold_pruner_tests.rs rename to tests/pruner/threshold.rs diff --git a/tests/bohb_integration.rs b/tests/sampler/bohb.rs similarity index 100% rename from tests/bohb_integration.rs rename to tests/sampler/bohb.rs diff --git a/tests/cma_es_tests.rs b/tests/sampler/cma_es.rs similarity index 99% rename from tests/cma_es_tests.rs rename to tests/sampler/cma_es.rs index 84b3700..2c3f992 100644 --- a/tests/cma_es_tests.rs +++ b/tests/sampler/cma_es.rs @@ -1,5 +1,3 @@ -#![cfg(feature = "cma-es")] - use optimizer::prelude::*; use optimizer::sampler::cma_es::CmaEsSampler; diff --git a/tests/differential_evolution_tests.rs b/tests/sampler/differential_evolution.rs similarity index 100% rename from tests/differential_evolution_tests.rs rename to tests/sampler/differential_evolution.rs diff --git a/tests/gp_tests.rs b/tests/sampler/gp.rs similarity index 99% rename from tests/gp_tests.rs rename to tests/sampler/gp.rs index b2dc23f..6ddf3b5 100644 --- a/tests/gp_tests.rs +++ b/tests/sampler/gp.rs @@ -1,5 +1,3 @@ -#![cfg(feature = "gp")] - use optimizer::prelude::*; use optimizer::sampler::gp::GpSampler; diff --git a/tests/sampler/main.rs b/tests/sampler/main.rs new file mode 100644 index 0000000..9f0f825 --- /dev/null +++ b/tests/sampler/main.rs @@ -0,0 +1,15 @@ +#![allow( + clippy::cast_sign_loss, + clippy::cast_precision_loss, + clippy::cast_possible_truncation +)] + +mod bohb; +#[cfg(feature = "cma-es")] +mod cma_es; +mod differential_evolution; +#[cfg(feature = "gp")] +mod gp; +mod multivariate_tpe; +mod random; +mod tpe; diff --git a/tests/multivariate_tpe_integration.rs b/tests/sampler/multivariate_tpe.rs similarity index 100% rename from tests/multivariate_tpe_integration.rs rename to tests/sampler/multivariate_tpe.rs diff --git a/tests/sampler/random.rs b/tests/sampler/random.rs new file mode 100644 index 0000000..911f192 --- /dev/null +++ b/tests/sampler/random.rs @@ -0,0 +1,142 @@ +use optimizer::parameter::{CategoricalParam, FloatParam, IntParam, Parameter}; +use optimizer::sampler::random::RandomSampler; +use optimizer::{Direction, Error, Study}; + +#[test] +fn test_random_sampler_uniform_float_distribution() { + 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 = x_param.suggest(trial)?; + samples.push(x); + Ok::<_, Error>(x) + }) + .unwrap(); + + // All samples should be in range + for &s in &samples { + assert!((0.0..=1.0).contains(&s), "sample {s} out of range [0, 1]"); + } + + // Check distribution is roughly uniform by looking at quartiles + samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); + + let q1 = samples[n_samples / 4]; + let q2 = samples[n_samples / 2]; + let q3 = samples[3 * n_samples / 4]; + + assert!((q1 - 0.25).abs() < 0.1, "Q1 {q1} should be close to 0.25"); + assert!( + (q2 - 0.5).abs() < 0.1, + "Q2 (median) {q2} should be close to 0.5" + ); + assert!((q3 - 0.75).abs() < 0.1, "Q3 {q3} should be close to 0.75"); +} + +#[test] +fn test_random_sampler_uniform_int_distribution() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(123)); + + 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 = 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(); + + let expected = n_samples as f64 / 10.0; + for (i, &count) in counts.iter().enumerate() { + let diff = (count as f64 - expected).abs() / expected; + assert!( + diff < 0.2, + "value {} appeared {} times, expected ~{}, diff = {:.1}%", + i + 1, + count, + expected, + diff * 100.0 + ); + } +} + +#[test] +fn test_random_sampler_uniform_categorical_distribution() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(456)); + + let n_samples = 2000; + 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 = cat_param.suggest(trial)?; + let idx = choices.iter().position(|&c| c == choice).unwrap(); + counts[idx] += 1; + Ok::<_, Error>(idx as f64) + }) + .unwrap(); + + let expected = n_samples as f64 / 4.0; + for (i, &count) in counts.iter().enumerate() { + let diff = (count as f64 - expected).abs() / expected; + assert!( + diff < 0.15, + "category {} appeared {} times, expected ~{}, diff = {:.1}%", + i, + count, + expected, + diff * 100.0 + ); + } +} + +#[test] +fn test_random_sampler_reproducibility() { + let study1: Study = + Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(999)); + let study2: Study = + Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(999)); + + 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(100, |trial| { + let x = x_param1.suggest(trial)?; + values1.push(x); + Ok::<_, Error>(x) + }) + .unwrap(); + + study2 + .optimize(100, |trial| { + let x = x_param2.suggest(trial)?; + values2.push(x); + Ok::<_, Error>(x) + }) + .unwrap(); + + for (i, (v1, v2)) in values1.iter().zip(values2.iter()).enumerate() { + assert_eq!( + v1, v2, + "values at trial {i} should be identical with same seed: {v1} vs {v2}" + ); + } +} diff --git a/tests/sampler/tpe.rs b/tests/sampler/tpe.rs new file mode 100644 index 0000000..8091077 --- /dev/null +++ b/tests/sampler/tpe.rs @@ -0,0 +1,381 @@ +use optimizer::parameter::{BoolParam, CategoricalParam, FloatParam, IntParam, Parameter}; +use optimizer::sampler::tpe::TpeSampler; +use optimizer::{Direction, Error, Study}; + +#[test] +fn test_tpe_optimizes_quadratic_function() { + // Minimize f(x) = (x - 3)^2 where x in [-10, 10] + // Optimal: x = 3, f(3) = 0 + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(10) + .n_ei_candidates(24) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + let x_param = FloatParam::new(-10.0, 10.0); + + study + .optimize(100, |trial| { + 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 reasonable value over 100 trials + // With random startup + TPE, we expect to get within a few units of optimal + assert!( + best.value < 5.0, + "TPE should find near-optimal: best value {} should be < 5.0", + best.value + ); +} + +#[test] +fn test_tpe_optimizes_multivariate_function() { + // 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) + .n_startup_trials(10) + .build() + .unwrap(); + + 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(100, |trial| { + let x = x_param.suggest(trial)?; + let y = y_param.suggest(trial)?; + Ok::<_, Error>(x * x + y * y) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have at least one trial"); + + // TPE should find a reasonably good solution + assert!( + best.value < 5.0, + "TPE should find near-optimal: best value {} should be < 5.0", + best.value + ); +} + +#[test] +fn test_tpe_maximization() { + // Maximize f(x) = -(x - 2)^2 + 10 where x in [-10, 10] + // Optimal: x = 2, f(2) = 10 + let sampler = TpeSampler::builder() + .seed(456) + .n_startup_trials(5) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Maximize, sampler); + + let x_param = FloatParam::new(-10.0, 10.0); + + study + .optimize(50, |trial| { + 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"); + + assert!( + best.value > 5.0, + "TPE should find reasonably good solution: best value {} should be > 5.0", + best.value + ); +} + +#[test] +fn test_tpe_with_categorical_parameter() { + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(5) + .build() + .unwrap(); + + 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(30, |trial| { + let choice = model_param.suggest(trial)?; + let x = x_param.suggest(trial)?; + + // cubic model is best at x=1 + let value = match choice { + "linear" => x, + "quadratic" => x * x, + "cubic" => -((x - 1.0).powi(2)) + 10.0, // peak at x=1, max value 10 + _ => unreachable!(), + }; + Ok::<_, Error>(value) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have best trial"); + assert!( + best.value > 5.0, + "should find good solution, got {}", + best.value + ); +} + +#[test] +fn test_tpe_with_integer_parameters() { + let sampler = TpeSampler::builder() + .seed(789) + .n_startup_trials(5) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + let n_param = IntParam::new(1, 10); + + // Minimize (n - 7)^2 where n in [1, 10] + study + .optimize(30, |trial| { + 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"); + + assert!( + best.value < 5.0, + "should find n close to 7, best value = {}", + best.value + ); +} + +#[test] +fn test_tpe_with_log_scale_int() { + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(5) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let batch_param = IntParam::new(1, 1024).log_scale(); + + study + .optimize(20, |trial| { + let batch_size = batch_param.suggest(trial)?; + Ok::<_, Error>(((batch_size as f64).log2() - 5.0).powi(2)) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().unwrap(); + assert!(best.value < 10.0, "should find reasonable solution"); +} + +#[test] +fn test_tpe_with_step_distributions() { + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(5) + .build() + .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(20, |trial| { + 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"); + + let best = study.best_trial().unwrap(); + assert!(best.value < 100.0, "should find reasonable solution"); +} + +#[test] +fn test_tpe_with_fixed_kde_bandwidth() { + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(5) + .kde_bandwidth(0.5) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(-5.0, 5.0); + + study + .optimize(20, |trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x * x) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().unwrap(); + assert!(best.value < 10.0, "should find reasonable solution"); +} + +#[test] +fn test_tpe_sampler_invalid_kde_bandwidth() { + let result = TpeSampler::with_config(0.25, 10, 24, Some(-1.0), None); + assert!(matches!(result, Err(Error::InvalidBandwidth(_)))); +} + +#[test] +fn test_tpe_split_trials_with_two_trials() { + let sampler = TpeSampler::builder() + .seed(42) + .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(5, |trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x) + }) + .expect("optimization should succeed with small history"); + + assert_eq!(study.n_trials(), 5); +} + +#[test] +fn test_tpe_empty_good_or_bad_values_fallback() { + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(5) + .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(10, |trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x) + }) + .unwrap(); + + // Now try with a different parameter - TPE won't have history for "y" + study + .optimize(5, |trial| { + let y = y_param.suggest(trial)?; + Ok::<_, Error>(y) + }) + .unwrap(); + + assert_eq!(study.n_trials(), 15); +} + +#[test] +fn test_tpe_sampler_builder_default_trait() { + use optimizer::sampler::tpe::TpeSamplerBuilder; + + let builder = TpeSamplerBuilder::default(); + let sampler = builder.build().unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x_param = FloatParam::new(0.0, 1.0); + + study + .optimize(5, |trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x) + }) + .unwrap(); + + assert_eq!(study.n_trials(), 5); +} + +#[test] +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(5, |trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x) + }) + .unwrap(); + + assert_eq!(study.n_trials(), 5); +} + +#[test] +fn test_suggest_bool_with_tpe() { + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(5) + .build() + .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(20, |trial| { + 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(); + assert!(best.value < 10.0); +} + +#[test] +fn test_params_with_tpe() { + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(5) + .build() + .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(30, |trial| { + 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(); + assert!(best.value < 10.0, "TPE should find good solution"); +} diff --git a/tests/study/ask_tell.rs b/tests/study/ask_tell.rs new file mode 100644 index 0000000..cf7d1d0 --- /dev/null +++ b/tests/study/ask_tell.rs @@ -0,0 +1,98 @@ +use optimizer::parameter::{FloatParam, Parameter}; +use optimizer::sampler::tpe::TpeSampler; +use optimizer::{Direction, Study}; + +#[test] +fn test_ask_and_tell_basic() { + let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(0.0, 10.0); + + for _ in 0..10 { + let mut trial = study.ask(); + let x = x_param.suggest(&mut trial).unwrap(); + let value = x * x; + study.tell(trial, Ok::<_, &str>(value)); + } + + assert_eq!(study.n_trials(), 10); + assert!(study.best_value().unwrap() >= 0.0); +} + +#[test] +fn test_ask_and_tell_with_failures() { + let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(-5.0, 5.0); + + // Alternate success and failure + for i in 0..10 { + let mut trial = study.ask(); + let x = x_param.suggest(&mut trial).unwrap(); + if i % 2 == 0 { + study.tell(trial, Ok::<_, &str>(x * x)); + } else { + study.tell(trial, Err::("simulated failure")); + } + } + + // Only successful trials are counted + assert_eq!(study.n_trials(), 5); +} + +#[test] +fn test_ask_and_tell_with_tpe_sampler() { + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(5) + .build() + .unwrap(); + let study: Study = Study::minimize(sampler); + let x_param = FloatParam::new(-10.0, 10.0); + + for _ in 0..30 { + let mut trial = study.ask(); + let x = x_param.suggest(&mut trial).unwrap(); + study.tell(trial, Ok::<_, &str>((x - 3.0).powi(2))); + } + + assert_eq!(study.n_trials(), 30); + assert!( + study.best_value().unwrap() < 5.0, + "TPE ask-and-tell should find a reasonable value" + ); +} + +#[test] +fn test_ask_and_tell_batch() { + let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(0.0, 10.0); + + // Ask a batch of trials + let batch: Vec<_> = (0..5) + .map(|_| { + let mut t = study.ask(); + let x = x_param.suggest(&mut t).unwrap(); + (t, x) + }) + .collect(); + + // Tell results for the batch + for (trial, x) in batch { + study.tell(trial, Ok::<_, &str>(x * x)); + } + + assert_eq!(study.n_trials(), 5); +} + +#[test] +fn test_ask_and_tell_with_custom_value_type() { + // Ask-and-tell works with non-f64 value types too + let study: Study = Study::new(Direction::Maximize); + + for i in 0..5 { + let trial = study.ask(); + study.tell(trial, Ok::<_, &str>(i * 10)); + } + + assert_eq!(study.n_trials(), 5); + assert_eq!(study.best_value().unwrap(), 40); +} diff --git a/tests/study/builder.rs b/tests/study/builder.rs new file mode 100644 index 0000000..8b421ff --- /dev/null +++ b/tests/study/builder.rs @@ -0,0 +1,92 @@ +use optimizer::parameter::{FloatParam, Parameter}; +use optimizer::sampler::random::RandomSampler; +use optimizer::sampler::tpe::TpeSampler; +use optimizer::{Direction, Error, Study}; + +#[test] +fn test_builder_defaults() { + let study: Study = Study::builder().build(); + assert_eq!(study.direction(), Direction::Minimize); +} + +#[test] +fn test_builder_maximize() { + let study: Study = Study::builder().maximize().build(); + assert_eq!(study.direction(), Direction::Maximize); +} + +#[test] +fn test_builder_minimize() { + let study: Study = Study::builder().minimize().build(); + assert_eq!(study.direction(), Direction::Minimize); +} + +#[test] +fn test_builder_direction() { + let study: Study = Study::builder().direction(Direction::Maximize).build(); + assert_eq!(study.direction(), Direction::Maximize); +} + +#[test] +fn test_builder_with_sampler() { + let x = FloatParam::new(-5.0, 5.0); + let study: Study = Study::builder().sampler(TpeSampler::new()).build(); + + study + .optimize(10, |trial| { + let val = x.suggest(trial)?; + Ok::<_, Error>(val * val) + }) + .unwrap(); + + assert_eq!(study.trials().len(), 10); +} + +#[test] +fn test_builder_with_pruner() { + use optimizer::pruner::NopPruner; + + let study: Study = Study::builder().pruner(NopPruner).build(); + + assert_eq!(study.direction(), Direction::Minimize); +} + +#[test] +fn test_builder_chaining() { + let study: Study = Study::builder() + .maximize() + .sampler(RandomSampler::with_seed(42)) + .pruner(optimizer::pruner::NopPruner) + .build(); + + assert_eq!(study.direction(), Direction::Maximize); +} + +#[test] +fn test_builder_with_custom_value_type() { + let study: Study = Study::builder().maximize().build(); + assert_eq!(study.direction(), Direction::Maximize); +} + +#[test] +fn test_builder_optimizes_correctly() { + let x = FloatParam::new(-10.0, 10.0); + let study: Study = Study::builder() + .minimize() + .sampler(TpeSampler::builder().seed(42).build().unwrap()) + .build(); + + study + .optimize(100, |trial| { + let val = x.suggest(trial)?; + Ok::<_, Error>((val - 3.0) * (val - 3.0)) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert!( + best.value < 5.0, + "best value should be < 5.0, got {}", + best.value + ); +} diff --git a/tests/study/constraints.rs b/tests/study/constraints.rs new file mode 100644 index 0000000..592ac93 --- /dev/null +++ b/tests/study/constraints.rs @@ -0,0 +1,101 @@ +use optimizer::{Direction, Study}; + +#[test] +fn test_is_feasible_all_satisfied() { + let study: Study = Study::new(Direction::Minimize); + let mut trial = study.create_trial(); + trial.set_constraints(vec![-1.0, 0.0, -0.5]); + study.complete_trial(trial, 1.0); + + let completed = study.best_trial().unwrap(); + assert!(completed.is_feasible()); +} + +#[test] +fn test_is_feasible_one_violated() { + let study: Study = Study::new(Direction::Minimize); + let mut trial = study.create_trial(); + trial.set_constraints(vec![-1.0, 0.5, -0.5]); + study.complete_trial(trial, 1.0); + + let completed = study.best_trial().unwrap(); + assert!(!completed.is_feasible()); +} + +#[test] +fn test_is_feasible_empty_constraints() { + let study: Study = Study::new(Direction::Minimize); + let trial = study.create_trial(); + study.complete_trial(trial, 1.0); + + let completed = study.best_trial().unwrap(); + assert!(completed.is_feasible()); +} + +#[test] +fn test_best_trial_prefers_feasible() { + let study: Study = Study::new(Direction::Minimize); + + // Infeasible trial with better objective + let mut trial1 = study.create_trial(); + trial1.set_constraints(vec![1.0]); + study.complete_trial(trial1, 0.1); + + // Feasible trial with worse objective + let mut trial2 = study.create_trial(); + trial2.set_constraints(vec![-1.0]); + study.complete_trial(trial2, 100.0); + + let best = study.best_trial().unwrap(); + assert_eq!(best.id, 1); // feasible trial wins + assert_eq!(best.value, 100.0); +} + +#[test] +fn test_best_trial_feasible_by_objective() { + let study: Study = Study::new(Direction::Minimize); + + // Feasible, worse objective + let mut trial1 = study.create_trial(); + trial1.set_constraints(vec![-1.0]); + study.complete_trial(trial1, 10.0); + + // Feasible, better objective + let mut trial2 = study.create_trial(); + trial2.set_constraints(vec![-0.5]); + study.complete_trial(trial2, 2.0); + + let best = study.best_trial().unwrap(); + assert_eq!(best.id, 1); // lower objective wins among feasible + assert_eq!(best.value, 2.0); +} + +#[test] +fn test_top_trials_ranks_feasible_above_infeasible() { + let study: Study = Study::new(Direction::Minimize); + + // Infeasible, low violation + let mut t0 = study.create_trial(); + t0.set_constraints(vec![0.5]); + study.complete_trial(t0, 1.0); + + // Feasible, worst objective among feasible + let mut t1 = study.create_trial(); + t1.set_constraints(vec![-1.0]); + study.complete_trial(t1, 50.0); + + // Feasible, best objective among feasible + let mut t2 = study.create_trial(); + t2.set_constraints(vec![-0.1]); + study.complete_trial(t2, 5.0); + + // Infeasible, high violation + let mut t3 = study.create_trial(); + t3.set_constraints(vec![3.0]); + study.complete_trial(t3, 0.5); + + let top = study.top_trials(4); + let ids: Vec = top.iter().map(|t| t.id).collect(); + // Feasible sorted by objective first (5.0, 50.0), then infeasible by violation (0.5, 3.0) + assert_eq!(ids, vec![2, 1, 0, 3]); +} diff --git a/tests/study/enqueue.rs b/tests/study/enqueue.rs new file mode 100644 index 0000000..c84e008 --- /dev/null +++ b/tests/study/enqueue.rs @@ -0,0 +1,189 @@ +use std::collections::HashMap; + +use optimizer::parameter::{FloatParam, IntParam, ParamValue, Parameter}; +use optimizer::sampler::random::RandomSampler; +use optimizer::{Direction, Error, Study}; + +#[test] +fn test_enqueue_params_evaluated_first() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 10.0); + let y = IntParam::new(1, 100); + + // Enqueue a specific configuration + study.enqueue(HashMap::from([ + (x.id(), ParamValue::Float(5.0)), + (y.id(), ParamValue::Int(42)), + ])); + + // The first trial should use the enqueued params + let mut trial = study.ask(); + let x_val = x.suggest(&mut trial).unwrap(); + let y_val = y.suggest(&mut trial).unwrap(); + + assert_eq!(x_val, 5.0); + assert_eq!(y_val, 42); +} + +#[test] +fn test_enqueue_fifo_order() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 10.0); + + // Enqueue two configs + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))])); + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))])); + + // First trial gets first enqueued value + let mut trial1 = study.ask(); + assert_eq!(x.suggest(&mut trial1).unwrap(), 1.0); + + // Second trial gets second enqueued value + let mut trial2 = study.ask(); + assert_eq!(x.suggest(&mut trial2).unwrap(), 2.0); +} + +#[test] +fn test_enqueue_then_normal_sampling_resumes() { + let sampler = RandomSampler::with_seed(42); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + let x = FloatParam::new(0.0, 10.0); + + // Enqueue one config + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(5.0))])); + + // First trial uses enqueued value + let mut trial1 = study.ask(); + assert_eq!(x.suggest(&mut trial1).unwrap(), 5.0); + study.tell(trial1, Ok::<_, &str>(25.0)); + + // Second trial uses normal sampling (not 5.0) + let mut trial2 = study.ask(); + let x_val = x.suggest(&mut trial2).unwrap(); + // The sampled value should be in [0, 10] but extremely unlikely to be exactly 5.0 + assert!((0.0..=10.0).contains(&x_val)); +} + +#[test] +fn test_enqueue_with_optimize() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 10.0); + + // Enqueue two specific configs + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))])); + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))])); + + let mut values = Vec::new(); + + study + .optimize(5, |trial| { + let x_val = x.suggest(trial)?; + values.push(x_val); + Ok::<_, Error>(x_val * x_val) + }) + .unwrap(); + + // First two trials should use enqueued values + assert_eq!(values[0], 1.0); + assert_eq!(values[1], 2.0); + // All 5 trials should have completed + assert_eq!(study.n_trials(), 5); +} + +#[test] +fn test_enqueue_partial_params_fall_back_to_sampling() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 10.0); + let y = IntParam::new(1, 100); + + // Enqueue only x, not y + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(3.0))])); + + let mut trial = study.ask(); + let x_val = x.suggest(&mut trial).unwrap(); + let y_val = y.suggest(&mut trial).unwrap(); + + // x should be the enqueued value + assert_eq!(x_val, 3.0); + // y should be sampled (within range) + assert!((1..=100).contains(&y_val)); +} + +#[test] +fn test_enqueue_trials_appear_in_completed_trials() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 10.0); + + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(7.0))])); + + study + .optimize(1, |trial| { + let x_val = x.suggest(trial)?; + Ok::<_, Error>(x_val) + }) + .unwrap(); + + let trials = study.trials(); + assert_eq!(trials.len(), 1); + assert_eq!(trials[0].value, 7.0); + assert_eq!( + *trials[0].params.get(&x.id()).unwrap(), + ParamValue::Float(7.0) + ); +} + +#[test] +fn test_enqueue_with_ask_and_tell() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 10.0); + + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(4.0))])); + + let mut trial = study.ask(); + let x_val = x.suggest(&mut trial).unwrap(); + assert_eq!(x_val, 4.0); + + study.tell(trial, Ok::<_, &str>(x_val * x_val)); + assert_eq!(study.n_trials(), 1); + assert_eq!(study.best_value().unwrap(), 16.0); +} + +#[test] +fn test_n_enqueued() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 10.0); + + assert_eq!(study.n_enqueued(), 0); + + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))])); + assert_eq!(study.n_enqueued(), 1); + + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))])); + assert_eq!(study.n_enqueued(), 2); + + // Creating a trial dequeues one + let _ = study.ask(); + assert_eq!(study.n_enqueued(), 1); + + let _ = study.ask(); + assert_eq!(study.n_enqueued(), 0); +} + +#[test] +fn test_enqueue_counted_in_n_trials() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 10.0); + + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))])); + study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))])); + + study + .optimize(5, |trial| { + let x_val = x.suggest(trial)?; + Ok::<_, Error>(x_val) + }) + .unwrap(); + + // All 5 trials count, including the 2 enqueued ones + assert_eq!(study.n_trials(), 5); +} diff --git a/tests/study/iterator.rs b/tests/study/iterator.rs new file mode 100644 index 0000000..4578a85 --- /dev/null +++ b/tests/study/iterator.rs @@ -0,0 +1,43 @@ +use optimizer::parameter::{FloatParam, Parameter}; +use optimizer::sampler::random::RandomSampler; +use optimizer::{Direction, Study}; + +#[test] +fn test_into_iterator_iterates_all_trials() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); + let x_param = FloatParam::new(0.0, 10.0); + + for _ in 0..5 { + let mut trial = study.create_trial(); + let x = x_param.suggest(&mut trial).unwrap(); + study.complete_trial(trial, x * x); + } + + let mut count = 0; + for trial in &study { + assert_eq!(trial.state, optimizer::TrialState::Complete); + count += 1; + } + assert_eq!(count, 5); +} + +#[test] +fn test_into_iterator_empty_study() { + let study: Study = Study::new(Direction::Minimize); + + let count = (&study).into_iter().count(); + assert_eq!(count, 0); +} + +#[test] +fn test_into_iterator_preserves_insertion_order() { + let study: Study = Study::new(Direction::Minimize); + + for i in 0..3 { + let trial = study.create_trial(); + study.complete_trial(trial, f64::from(i)); + } + + let ids: Vec = (&study).into_iter().map(|t| t.id).collect(); + assert_eq!(ids, vec![0, 1, 2]); +} diff --git a/tests/study/main.rs b/tests/study/main.rs new file mode 100644 index 0000000..d9a788c --- /dev/null +++ b/tests/study/main.rs @@ -0,0 +1,15 @@ +#![allow( + clippy::cast_sign_loss, + clippy::cast_precision_loss, + clippy::cast_possible_truncation +)] + +mod ask_tell; +mod builder; +mod constraints; +mod enqueue; +mod iterator; +mod objective; +mod summary; +mod top_trials; +mod workflow; diff --git a/tests/study/objective.rs b/tests/study/objective.rs new file mode 100644 index 0000000..6bb1944 --- /dev/null +++ b/tests/study/objective.rs @@ -0,0 +1,346 @@ +use optimizer::parameter::{FloatParam, Parameter}; +use optimizer::sampler::random::RandomSampler; +use optimizer::{Direction, Error, Study, Trial}; + +#[test] +fn test_callback_early_stopping() { + use std::ops::ControlFlow; + + use optimizer::Objective; + use optimizer::sampler::CompletedTrial; + + struct EarlyStopAfter5 { + x_param: FloatParam, + } + + impl Objective for EarlyStopAfter5 { + type Error = Error; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial)?; + Ok(x) + } + fn after_trial(&self, study: &Study, _trial: &CompletedTrial) -> ControlFlow<()> { + if study.n_trials() >= 5 { + ControlFlow::Break(()) + } else { + ControlFlow::Continue(()) + } + } + } + + let study: Study = Study::new(Direction::Minimize); + study + .optimize_with( + 100, + EarlyStopAfter5 { + x_param: FloatParam::new(0.0, 10.0), + }, + ) + .expect("optimization should succeed"); + + assert_eq!(study.n_trials(), 5, "should have stopped after 5 trials"); +} + +#[test] +fn test_callback_early_stopping_on_first_trial() { + use std::ops::ControlFlow; + + use optimizer::Objective; + use optimizer::sampler::CompletedTrial; + + struct StopImmediately { + x_param: FloatParam, + } + + impl Objective for StopImmediately { + type Error = Error; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial)?; + Ok(x) + } + fn after_trial( + &self, + _study: &Study, + _trial: &CompletedTrial, + ) -> ControlFlow<()> { + ControlFlow::Break(()) + } + } + + let study: Study = Study::new(Direction::Minimize); + study + .optimize_with( + 100, + StopImmediately { + x_param: FloatParam::new(0.0, 10.0), + }, + ) + .expect("optimization should succeed"); + + assert_eq!(study.n_trials(), 1, "should have stopped after 1 trial"); +} + +#[test] +fn test_callback_sampler_early_stopping() { + use std::ops::ControlFlow; + + use optimizer::Objective; + use optimizer::sampler::CompletedTrial; + + struct StopAfter3 { + x_param: FloatParam, + } + + impl Objective for StopAfter3 { + type Error = Error; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial)?; + Ok(x) + } + fn after_trial(&self, study: &Study, _trial: &CompletedTrial) -> ControlFlow<()> { + if study.n_trials() >= 3 { + ControlFlow::Break(()) + } else { + ControlFlow::Continue(()) + } + } + } + + let sampler = RandomSampler::with_seed(42); + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + study + .optimize_with( + 100, + StopAfter3 { + x_param: FloatParam::new(0.0, 10.0), + }, + ) + .expect("optimization should succeed"); + + assert_eq!(study.n_trials(), 3); +} + +#[test] +fn test_retries_successful_trials_not_retried() { + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; + + use optimizer::Objective; + + struct SuccessObj { + x_param: FloatParam, + call_count: Arc, + } + + impl Objective for SuccessObj { + type Error = Error; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial)?; + self.call_count.fetch_add(1, Ordering::Relaxed); + Ok(x * x) + } + fn max_retries(&self) -> usize { + 3 + } + } + + let study: Study = Study::new(Direction::Minimize); + let call_count = Arc::new(AtomicU32::new(0)); + let obj = SuccessObj { + x_param: FloatParam::new(0.0, 10.0), + call_count: Arc::clone(&call_count), + }; + + study.optimize_with(5, obj).unwrap(); + + // All trials succeed on first try — exactly 5 calls + assert_eq!(call_count.load(Ordering::Relaxed), 5); + assert_eq!(study.n_trials(), 5); +} + +#[test] +fn test_retries_failed_trials_retried_up_to_max() { + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; + + use optimizer::Objective; + + struct AlwaysFailObj { + x_param: FloatParam, + call_count: Arc, + } + + impl Objective for AlwaysFailObj { + type Error = String; + fn evaluate(&self, trial: &mut Trial) -> Result { + let _ = self.x_param.suggest(trial).map_err(|e| e.to_string())?; + self.call_count.fetch_add(1, Ordering::Relaxed); + Err("always fails".to_string()) + } + fn max_retries(&self) -> usize { + 3 + } + } + + let study: Study = Study::new(Direction::Minimize); + let call_count = Arc::new(AtomicU32::new(0)); + let obj = AlwaysFailObj { + x_param: FloatParam::new(0.0, 10.0), + call_count: Arc::clone(&call_count), + }; + + let result = study.optimize_with(1, obj); + + // 1 initial attempt + 3 retries = 4 total calls + assert_eq!(call_count.load(Ordering::Relaxed), 4); + // No trials completed + assert!(matches!(result, Err(Error::NoCompletedTrials))); +} + +#[test] +fn test_retries_permanently_failed_after_exhaustion() { + use optimizer::Objective; + + struct AlwaysFailObj { + x_param: FloatParam, + } + + impl Objective for AlwaysFailObj { + type Error = String; + fn evaluate(&self, trial: &mut Trial) -> Result { + let _ = self.x_param.suggest(trial).map_err(|e| e.to_string())?; + Err("transient error".to_string()) + } + fn max_retries(&self) -> usize { + 2 + } + } + + let study: Study = Study::new(Direction::Minimize); + let obj = AlwaysFailObj { + x_param: FloatParam::new(0.0, 10.0), + }; + + let result = study.optimize_with(3, obj); + + assert!( + matches!(result, Err(Error::NoCompletedTrials)), + "all trials should permanently fail" + ); + assert_eq!( + study.n_trials(), + 0, + "no completed trials should be recorded" + ); +} + +#[test] +fn test_retries_uses_same_parameters() { + use std::sync::atomic::{AtomicU32, Ordering}; + use std::sync::{Arc, Mutex}; + + use optimizer::Objective; + + struct RetryObj { + x_param: FloatParam, + seen_values: Arc>>, + call_count: Arc, + } + + impl Objective for RetryObj { + type Error = String; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial).map_err(|e| e.to_string())?; + self.seen_values.lock().unwrap().push(x); + let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1; + // Fail first two attempts, succeed on third + if count < 3 { + Err("transient".to_string()) + } else { + Ok(x * x) + } + } + fn max_retries(&self) -> usize { + 2 + } + } + + let study: Study = Study::new(Direction::Minimize); + let seen_values = Arc::new(Mutex::new(Vec::new())); + let call_count = Arc::new(AtomicU32::new(0)); + let obj = RetryObj { + x_param: FloatParam::new(0.0, 10.0), + seen_values: Arc::clone(&seen_values), + call_count: Arc::clone(&call_count), + }; + + study.optimize_with(1, obj).unwrap(); + + let values = seen_values.lock().unwrap(); + assert_eq!(values.len(), 3, "should be called 3 times (1 + 2 retries)"); + // All three calls should have gotten the same parameter value + assert_eq!(values[0], values[1]); + assert_eq!(values[1], values[2]); +} + +#[test] +fn test_retries_n_trials_counts_unique_configs() { + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; + + use optimizer::Objective; + + struct FailFirstObj { + x_param: FloatParam, + call_count: Arc, + } + + impl Objective for FailFirstObj { + type Error = String; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial).map_err(|e| e.to_string())?; + let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1; + // Fail first attempt of each config, succeed on retry + if count % 2 == 1 { + Err("transient".to_string()) + } else { + Ok(x * x) + } + } + fn max_retries(&self) -> usize { + 2 + } + } + + let study: Study = Study::new(Direction::Minimize); + let call_count = Arc::new(AtomicU32::new(0)); + let obj = FailFirstObj { + x_param: FloatParam::new(0.0, 10.0), + call_count: Arc::clone(&call_count), + }; + + study.optimize_with(3, obj).unwrap(); + + // 3 unique configs, each needing 2 calls = 6 total calls + assert_eq!(call_count.load(Ordering::Relaxed), 6); + // But only 3 completed trials + assert_eq!(study.n_trials(), 3); +} + +#[test] +fn test_retries_with_zero_max_retries_same_as_optimize() { + let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(0.0, 10.0); + let call_count = std::cell::Cell::new(0u32); + + study + .optimize(5, |trial| { + let x = x_param.suggest(trial)?; + call_count.set(call_count.get() + 1); + Ok::<_, Error>(x * x) + }) + .unwrap(); + + assert_eq!(call_count.get(), 5); + assert_eq!(study.n_trials(), 5); +} diff --git a/tests/study/summary.rs b/tests/study/summary.rs new file mode 100644 index 0000000..8f688af --- /dev/null +++ b/tests/study/summary.rs @@ -0,0 +1,71 @@ +use optimizer::parameter::{FloatParam, Parameter}; +use optimizer::sampler::random::RandomSampler; +use optimizer::{Direction, Error, Study}; + +#[test] +fn test_summary_with_completed_trials() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(1)); + let x = FloatParam::new(0.0, 10.0).name("x"); + + study + .optimize(5, |trial| { + let val = x.suggest(trial)?; + Ok::<_, Error>(val * val) + }) + .unwrap(); + + let summary = study.summary(); + assert!(summary.contains("Minimize")); + assert!(summary.contains("5 trials")); + assert!(summary.contains("Best value:")); + assert!(summary.contains("x = ")); +} + +#[test] +fn test_summary_no_completed_trials() { + let study: Study = Study::new(Direction::Maximize); + let summary = study.summary(); + assert!(summary.contains("Maximize")); + assert!(summary.contains("0 trials")); + assert!(!summary.contains("Best value:")); +} + +#[test] +fn test_summary_with_pruned_trials() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(1)); + let x = FloatParam::new(0.0, 10.0).name("x"); + + // Manually create some complete and pruned trials + for _ in 0..3 { + let mut trial = study.create_trial(); + let val = x.suggest(&mut trial).unwrap(); + study.complete_trial(trial, val); + } + for _ in 0..2 { + let mut trial = study.create_trial(); + let _ = x.suggest(&mut trial).unwrap(); + study.prune_trial(trial); + } + + let summary = study.summary(); + // Should show breakdown when there are pruned trials + if study.n_pruned_trials() > 0 { + assert!(summary.contains("complete")); + assert!(summary.contains("pruned")); + } +} + +#[test] +fn test_display_matches_summary() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(1)); + let x = FloatParam::new(0.0, 10.0).name("x"); + + study + .optimize(3, |trial| { + let val = x.suggest(trial)?; + Ok::<_, Error>(val) + }) + .unwrap(); + + assert_eq!(format!("{study}"), study.summary()); +} diff --git a/tests/study/top_trials.rs b/tests/study/top_trials.rs new file mode 100644 index 0000000..a3844af --- /dev/null +++ b/tests/study/top_trials.rs @@ -0,0 +1,76 @@ +use optimizer::{Direction, Study}; + +#[test] +fn test_top_trials_minimize() { + let study: Study = Study::new(Direction::Minimize); + + // Manually complete trials with known values + for &val in &[5.0, 1.0, 3.0, 2.0, 4.0] { + let trial = study.create_trial(); + study.complete_trial(trial, val); + } + + let top3 = study.top_trials(3); + assert_eq!(top3.len(), 3); + assert_eq!(top3[0].value, 1.0); + assert_eq!(top3[1].value, 2.0); + assert_eq!(top3[2].value, 3.0); +} + +#[test] +fn test_top_trials_maximize() { + let study: Study = Study::new(Direction::Maximize); + + for &val in &[5.0, 1.0, 3.0, 2.0, 4.0] { + let trial = study.create_trial(); + study.complete_trial(trial, val); + } + + let top3 = study.top_trials(3); + assert_eq!(top3.len(), 3); + assert_eq!(top3[0].value, 5.0); + assert_eq!(top3[1].value, 4.0); + assert_eq!(top3[2].value, 3.0); +} + +#[test] +fn test_top_trials_n_greater_than_total() { + let study: Study = Study::new(Direction::Minimize); + + for &val in &[3.0, 1.0] { + let trial = study.create_trial(); + study.complete_trial(trial, val); + } + + let top = study.top_trials(10); + assert_eq!(top.len(), 2); + assert_eq!(top[0].value, 1.0); + assert_eq!(top[1].value, 3.0); +} + +#[test] +fn test_top_trials_empty() { + let study: Study = Study::new(Direction::Minimize); + + let top = study.top_trials(5); + assert!(top.is_empty()); +} + +#[test] +fn test_top_trials_excludes_pruned() { + let study: Study = Study::new(Direction::Minimize); + + // Complete some trials + for &val in &[5.0, 1.0, 3.0] { + let trial = study.create_trial(); + study.complete_trial(trial, val); + } + + // Prune a trial (it gets a default value of 0.0 but should be excluded) + let trial = study.create_trial(); + study.prune_trial(trial); + + let top = study.top_trials(5); + assert_eq!(top.len(), 3, "pruned trial should be excluded"); + assert_eq!(top[0].value, 1.0); +} diff --git a/tests/study/workflow.rs b/tests/study/workflow.rs new file mode 100644 index 0000000..f41bff8 --- /dev/null +++ b/tests/study/workflow.rs @@ -0,0 +1,259 @@ +use optimizer::parameter::{BoolParam, FloatParam, IntParam, Parameter}; +use optimizer::sampler::tpe::TpeSampler; +use optimizer::{Direction, Error, Study}; + +#[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 = x_param.suggest(trial)?; + Ok::<_, Error>(x * x) + }) + .expect("optimization should succeed"); + + assert_eq!(study.n_trials(), 10); + let best = study.best_trial().expect("should have best trial"); + assert!(best.value >= 0.0, "x^2 should be non-negative"); +} + +#[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; + study + .optimize(10, |trial| { + counter += 1; + if counter % 2 == 0 { + return Err::("intentional failure"); + } + let x = x_param.suggest(trial).map_err(|_| "param error")?; + Ok(x * x) + }) + .expect("optimization should succeed with some failures"); + + // Only half the trials should have succeeded + assert_eq!(study.n_trials(), 5, "only 5 trials should have completed"); +} + +#[test] +fn test_no_completed_trials_error() { + let study: Study = Study::new(Direction::Minimize); + + let result = study.best_trial(); + assert!(matches!(result, Err(Error::NoCompletedTrials))); +} + +#[test] +fn test_study_direction() { + let study_min: Study = Study::new(Direction::Minimize); + assert_eq!(study_min.direction(), Direction::Minimize); + + let study_max: Study = Study::new(Direction::Maximize); + assert_eq!(study_max.direction(), Direction::Maximize); +} + +#[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 = x_param.suggest(trial)?; + Ok::<_, Error>(x) + }) + .unwrap(); + + let trials = study.trials(); + assert_eq!(trials.len(), 5); + + for trial in &trials { + assert!( + !trial.params.is_empty(), + "each trial should have parameters" + ); + } +} + +#[test] +fn test_study_set_sampler() { + let mut study: Study = Study::new(Direction::Minimize); + + let tpe = TpeSampler::builder() + .seed(42) + .n_startup_trials(5) + .build() + .unwrap(); + study.set_sampler(tpe); + + let x_param = FloatParam::new(-5.0, 5.0); + + study + .optimize(10, |trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x * x) + }) + .expect("optimization should succeed with new sampler"); + + assert_eq!(study.n_trials(), 10); +} + +#[test] +fn test_study_with_i32_value_type() { + let study: Study = Study::new(Direction::Minimize); + let x_param = IntParam::new(-10, 10); + + study + .optimize(10, |trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x.abs() as i32) + }) + .expect("optimization should succeed"); + + assert_eq!(study.n_trials(), 10); + let best = study.best_trial().expect("should have best trial"); + assert!(best.value >= 0, "absolute value should be non-negative"); +} + +#[test] +fn test_optimize_all_trials_fail() { + let study: Study = Study::new(Direction::Minimize); + + let result = study.optimize(5, |_trial| Err::("always fails")); + + assert!( + matches!(result, Err(Error::NoCompletedTrials)), + "should return NoCompletedTrials when all trials fail" + ); +} + +#[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 = x_param.suggest(trial)?; + Ok::<_, Error>(x) + }) + .unwrap(); + + let best_value = study.best_value().expect("should have best value"); + let best_trial = study.best_trial().expect("should have best trial"); + + assert_eq!( + best_value, best_trial.value, + "best_value should match best_trial.value" + ); +} + +#[test] +fn test_best_trial_with_nan_values() { + let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(0.0, 10.0); + + study + .optimize(5, |trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x) + }) + .unwrap(); + + let best = study.best_trial(); + assert!(best.is_ok()); +} + +#[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 = x_param.suggest(&mut trial).unwrap(); + study.complete_trial(trial, x * x); + + let mut trial2 = study.create_trial(); + let y = x_param.suggest(&mut trial2).unwrap(); + study.complete_trial(trial2, y * y); + + // Manually fail a trial + let trial3 = study.create_trial(); + study.fail_trial(trial3, "test failure"); + + // Only 2 completed trials + assert_eq!(study.n_trials(), 2); +} + +#[test] +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 = x_param.suggest(trial)?; + let n = n_param.suggest(trial)?; + Ok::<_, Error>(x * x + n as f64) + }) + .unwrap(); + + assert_eq!(study.n_trials(), 10); +} + +#[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 = use_feature_param.suggest(trial)?; + let x = x_param.suggest(trial)?; + + let value = if use_feature { x } else { x * 2.0 }; + Ok::<_, Error>(value) + }) + .unwrap(); + + assert_eq!(study.n_trials(), 10); +} + +#[test] +fn test_completed_trial_get() { + let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(-10.0, 10.0).name("x"); + let n_param = IntParam::new(1, 10).name("n"); + + study + .optimize(5, |trial| { + let x = x_param.suggest(trial)?; + let n = n_param.suggest(trial)?; + Ok::<_, Error>(x * x + n as f64) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + let x_val: f64 = best.get(&x_param).unwrap(); + let n_val: i64 = best.get(&n_param).unwrap(); + assert!((-10.0..=10.0).contains(&x_val)); + assert!((1..=10).contains(&n_val)); +} + +#[test] +fn test_single_value_int_range() { + let param = IntParam::new(5, 5); + let mut trial = optimizer::Trial::new(0); + + let n = param.suggest(&mut trial).unwrap(); + assert_eq!(n, 5, "single-value range should return that value"); +}