diff --git a/src/lib.rs b/src/lib.rs index 6afd342..2283d5d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,6 +19,7 @@ //! - **Grid Search** - Exhaustive search over a specified parameter grid //! - **Sobol (QMC)** - Quasi-random sampling for better space coverage (requires `sobol` feature) //! - **CMA-ES** - Covariance Matrix Adaptation Evolution Strategy for continuous optimization (requires `cma-es` feature) +//! - **BOHB** - Bayesian Optimization + `HyperBand` for budget-aware TPE sampling //! //! Additional features include: //! @@ -236,6 +237,7 @@ pub use pruner::{ SuccessiveHalvingPruner, ThresholdPruner, WilcoxonPruner, }; pub use sampler::CompletedTrial; +pub use sampler::bohb::BohbSampler; #[cfg(feature = "cma-es")] pub use sampler::cma_es::CmaEsSampler; pub use sampler::grid::GridSearchSampler; @@ -268,6 +270,7 @@ pub mod prelude { SuccessiveHalvingPruner, ThresholdPruner, }; pub use crate::sampler::CompletedTrial; + pub use crate::sampler::bohb::BohbSampler; #[cfg(feature = "cma-es")] pub use crate::sampler::cma_es::CmaEsSampler; pub use crate::sampler::grid::GridSearchSampler; diff --git a/src/sampler/bohb.rs b/src/sampler/bohb.rs new file mode 100644 index 0000000..a7121c1 --- /dev/null +++ b/src/sampler/bohb.rs @@ -0,0 +1,689 @@ +//! BOHB (Bayesian Optimization + `HyperBand`) sampler. +//! +//! BOHB combines TPE's model-guided sampling with Hyperband's budget-aware +//! evaluation. Instead of building one global TPE model, BOHB conditions +//! its TPE model on trials evaluated at a specific budget level, giving +//! better-calibrated proposals for each rung of the Hyperband schedule. +//! +//! # How it works +//! +//! 1. Compute all Hyperband rung steps (budget levels) from the config. +//! 2. On each `sample()` call, scan the history's `intermediate_values` +//! to find the **largest budget level** with enough observations +//! (`>= min_points_in_model`). +//! 3. Build a filtered history where each trial's `value` is replaced +//! with its intermediate value at that budget level. +//! 4. Delegate to an internal [`TpeSampler`] for the actual sampling. +//! 5. Fall back to random sampling if no budget level has enough data. +//! +//! # Examples +//! +//! ``` +//! use optimizer::sampler::bohb::BohbSampler; +//! use optimizer::{Direction, Study}; +//! +//! let bohb = BohbSampler::new(); +//! let pruner = bohb.matching_pruner(Direction::Minimize); +//! let study: Study = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner); +//! ``` +//! +//! Using the builder for custom configuration: +//! +//! ``` +//! use optimizer::sampler::bohb::BohbSampler; +//! +//! let bohb = BohbSampler::builder() +//! .min_resource(1) +//! .max_resource(81) +//! .reduction_factor(3) +//! .min_points_in_model(10) +//! .seed(42) +//! .build() +//! .unwrap(); +//! ``` + +use crate::distribution::Distribution; +use crate::error::Result; +use crate::param::ParamValue; +use crate::pruner::HyperbandPruner; +use crate::sampler::tpe::TpeSampler; +use crate::sampler::{CompletedTrial, Sampler}; +use crate::types::Direction; + +/// A BOHB sampler that combines TPE with Hyperband budget awareness. +/// +/// BOHB filters trial history by budget level before delegating to TPE, +/// so the surrogate model is conditioned on trials evaluated at the same +/// resource level. This produces better-calibrated parameter proposals +/// than using a single global model across all budgets. +/// +/// Use [`BohbSampler::matching_pruner`] to create a [`HyperbandPruner`] +/// with matching parameters. +pub struct BohbSampler { + min_resource: u64, + max_resource: u64, + reduction_factor: u64, + min_points_in_model: usize, + tpe: TpeSampler, +} + +impl BohbSampler { + /// Creates a new BOHB sampler with default settings. + /// + /// Defaults: + /// - `min_resource`: 1 + /// - `max_resource`: 81 + /// - `reduction_factor`: 3 + /// - `min_points_in_model`: 10 + /// - TPE: default settings + #[must_use] + pub fn new() -> Self { + Self { + min_resource: 1, + max_resource: 81, + reduction_factor: 3, + min_points_in_model: 10, + tpe: TpeSampler::new(), + } + } + + /// Creates a builder for configuring a BOHB sampler. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::bohb::BohbSampler; + /// + /// let sampler = BohbSampler::builder() + /// .min_resource(1) + /// .max_resource(27) + /// .reduction_factor(3) + /// .min_points_in_model(5) + /// .seed(42) + /// .build() + /// .unwrap(); + /// ``` + #[must_use] + pub fn builder() -> BohbSamplerBuilder { + BohbSamplerBuilder::new() + } + + /// Creates a [`HyperbandPruner`] with matching Hyperband parameters. + /// + /// This ensures the pruner's budget schedule is consistent with the + /// budget levels used by BOHB for model conditioning. + #[must_use] + pub fn matching_pruner(&self, direction: Direction) -> HyperbandPruner { + HyperbandPruner::new() + .min_resource(self.min_resource) + .max_resource(self.max_resource) + .reduction_factor(self.reduction_factor) + .direction(direction) + } + + /// Compute all unique budget levels (rung steps) across all Hyperband brackets. + /// + /// Returns sorted ascending. + #[allow( + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::cast_sign_loss + )] + fn all_budget_levels(&self) -> Vec { + let eta = self.reduction_factor as f64; + let ratio = self.max_resource as f64 / self.min_resource as f64; + let s_max = (ratio.ln() / eta.ln()).floor() as u64; + + let mut levels = Vec::new(); + for bracket in 0..=s_max { + let exponent = s_max.saturating_sub(bracket); + let min_resource_bracket = + (self.max_resource as f64 / eta.powi(exponent as i32)).ceil() as u64; + + let mut rung: u32 = 0; + while let Some(power) = self.reduction_factor.checked_pow(rung) { + let step = min_resource_bracket.saturating_mul(power); + if step > self.max_resource { + break; + } + levels.push(step); + rung += 1; + } + } + + levels.sort_unstable(); + levels.dedup(); + levels + } + + /// Build a filtered history for a specific budget level. + /// + /// For each trial that has an intermediate value at the given budget step, + /// creates a new `CompletedTrial` with `value` replaced by the intermediate + /// value at that step. + fn filter_history_for_budget(history: &[CompletedTrial], budget: u64) -> Vec { + history + .iter() + .filter_map(|trial| { + trial + .intermediate_values + .iter() + .find(|(step, _)| *step == budget) + .map(|(_, iv)| CompletedTrial { + id: trial.id, + params: trial.params.clone(), + distributions: trial.distributions.clone(), + param_labels: trial.param_labels.clone(), + value: *iv, + intermediate_values: trial.intermediate_values.clone(), + state: trial.state, + user_attrs: trial.user_attrs.clone(), + }) + }) + .collect() + } +} + +impl Default for BohbSampler { + fn default() -> Self { + Self::new() + } +} + +impl Sampler for BohbSampler { + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + history: &[CompletedTrial], + ) -> ParamValue { + // Find the largest budget level with enough observations + let levels = self.all_budget_levels(); + + for &budget in levels.iter().rev() { + let count = history + .iter() + .filter(|t| { + t.intermediate_values + .iter() + .any(|(step, _)| *step == budget) + }) + .count(); + + if count >= self.min_points_in_model { + let filtered = Self::filter_history_for_budget(history, budget); + return self.tpe.sample(distribution, trial_id, &filtered); + } + } + + // Not enough data at any budget level: delegate to TPE with empty history + // which triggers its uniform-random startup behavior. + self.tpe.sample(distribution, trial_id, &[]) + } +} + +/// Builder for configuring a [`BohbSampler`]. +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::bohb::BohbSamplerBuilder; +/// +/// let sampler = BohbSamplerBuilder::new() +/// .min_resource(1) +/// .max_resource(81) +/// .reduction_factor(3) +/// .gamma(0.15) +/// .seed(42) +/// .build() +/// .unwrap(); +/// ``` +pub struct BohbSamplerBuilder { + min_resource: u64, + max_resource: u64, + reduction_factor: u64, + min_points_in_model: usize, + tpe_builder: crate::sampler::tpe::TpeSamplerBuilder, +} + +impl BohbSamplerBuilder { + /// Creates a new builder with default settings. + #[must_use] + pub fn new() -> Self { + Self { + min_resource: 1, + max_resource: 81, + reduction_factor: 3, + min_points_in_model: 10, + tpe_builder: crate::sampler::tpe::TpeSamplerBuilder::new(), + } + } + + /// Sets the minimum resource (budget) per trial. + /// + /// # Panics + /// + /// Panics if `r` is 0. + #[must_use] + pub fn min_resource(mut self, r: u64) -> Self { + assert!(r > 0, "min_resource must be > 0, got {r}"); + self.min_resource = r; + self + } + + /// Sets the maximum resource (budget) per trial. + /// + /// # Panics + /// + /// Panics if `r` is 0. + #[must_use] + pub fn max_resource(mut self, r: u64) -> Self { + assert!(r > 0, "max_resource must be > 0, got {r}"); + self.max_resource = r; + self + } + + /// Sets the reduction factor (eta). + /// + /// # Panics + /// + /// Panics if `eta` is less than 2. + #[must_use] + pub fn reduction_factor(mut self, eta: u64) -> Self { + assert!(eta >= 2, "reduction_factor must be >= 2, got {eta}"); + self.reduction_factor = eta; + self + } + + /// Sets the minimum number of observations at a budget level before + /// BOHB uses TPE instead of random sampling. + #[must_use] + pub fn min_points_in_model(mut self, n: usize) -> Self { + self.min_points_in_model = n; + self + } + + /// Sets a fixed gamma value for the internal TPE sampler. + #[must_use] + pub fn gamma(mut self, gamma: f64) -> Self { + self.tpe_builder = self.tpe_builder.gamma(gamma); + self + } + + /// Sets a custom gamma strategy for the internal TPE sampler. + #[must_use] + pub fn gamma_strategy( + mut self, + strategy: G, + ) -> Self { + self.tpe_builder = self.tpe_builder.gamma_strategy(strategy); + self + } + + /// Sets the number of EI candidates for the internal TPE sampler. + #[must_use] + pub fn n_ei_candidates(mut self, n: usize) -> Self { + self.tpe_builder = self.tpe_builder.n_ei_candidates(n); + self + } + + /// Sets a fixed KDE bandwidth for the internal TPE sampler. + #[must_use] + pub fn kde_bandwidth(mut self, bandwidth: f64) -> Self { + self.tpe_builder = self.tpe_builder.kde_bandwidth(bandwidth); + self + } + + /// Sets a seed for reproducible sampling. + #[must_use] + pub fn seed(mut self, seed: u64) -> Self { + self.tpe_builder = self.tpe_builder.seed(seed); + self + } + + /// Builds the configured [`BohbSampler`]. + /// + /// # Errors + /// + /// Returns an error if the TPE configuration is invalid (e.g. gamma + /// not in (0, 1) or bandwidth not positive). + pub fn build(self) -> Result { + let tpe = self.tpe_builder.build()?; + Ok(BohbSampler { + min_resource: self.min_resource, + max_resource: self.max_resource, + reduction_factor: self.reduction_factor, + min_points_in_model: self.min_points_in_model, + tpe, + }) + } +} + +impl Default for BohbSamplerBuilder { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +#[allow(clippy::cast_precision_loss)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::distribution::{FloatDistribution, IntDistribution}; + use crate::parameter::ParamId; + use crate::types::TrialState; + + fn make_trial_with_intermediates( + id: u64, + value: f64, + params: Vec<(ParamId, ParamValue, Distribution)>, + intermediate_values: Vec<(u64, f64)>, + ) -> CompletedTrial { + let mut param_map = HashMap::new(); + let mut dist_map = HashMap::new(); + for (param_id, pv, dist) in params { + param_map.insert(param_id, pv); + dist_map.insert(param_id, dist); + } + CompletedTrial { + id, + params: param_map, + distributions: dist_map, + param_labels: HashMap::new(), + value, + intermediate_values, + state: TrialState::Complete, + user_attrs: HashMap::new(), + } + } + + #[test] + fn budget_levels_default() { + let bohb = BohbSampler::new(); + let levels = bohb.all_budget_levels(); + // With min=1, max=81, eta=3: + // bracket 0: 1, 3, 9, 27, 81 + // bracket 1: 3, 9, 27, 81 + // bracket 2: 9, 27, 81 + // bracket 3: 27, 81 + // bracket 4: 81 + // Unique sorted: [1, 3, 9, 27, 81] + assert_eq!(levels, vec![1, 3, 9, 27, 81]); + } + + #[test] + fn budget_levels_eta2() { + let bohb = BohbSampler::builder() + .min_resource(1) + .max_resource(16) + .reduction_factor(2) + .build() + .unwrap(); + let levels = bohb.all_budget_levels(); + // s_max = floor(ln(16)/ln(2)) = 4 + // bracket 0: 1, 2, 4, 8, 16 + // bracket 1: 2, 4, 8, 16 + // bracket 2: 4, 8, 16 + // bracket 3: 8, 16 + // bracket 4: 16 + // Unique sorted: [1, 2, 4, 8, 16] + assert_eq!(levels, vec![1, 2, 4, 8, 16]); + } + + #[test] + fn filter_history_selects_correct_budget() { + let x_id = ParamId::new(); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + let history = vec![ + make_trial_with_intermediates( + 0, + 0.5, + vec![(x_id, ParamValue::Float(0.3), dist.clone())], + vec![(1, 0.9), (3, 0.7), (9, 0.5)], + ), + make_trial_with_intermediates( + 1, + 0.4, + vec![(x_id, ParamValue::Float(0.6), dist.clone())], + vec![(1, 0.8), (3, 0.4)], + ), + make_trial_with_intermediates( + 2, + 0.3, + vec![(x_id, ParamValue::Float(0.1), dist.clone())], + vec![(1, 0.7)], + ), + ]; + + // Budget 3: trials 0 and 1 have intermediate values at step 3 + let filtered = BohbSampler::filter_history_for_budget(&history, 3); + assert_eq!(filtered.len(), 2); + assert!((filtered[0].value - 0.7).abs() < f64::EPSILON); + assert!((filtered[1].value - 0.4).abs() < f64::EPSILON); + + // Budget 9: only trial 0 + let filtered = BohbSampler::filter_history_for_budget(&history, 9); + assert_eq!(filtered.len(), 1); + assert!((filtered[0].value - 0.5).abs() < f64::EPSILON); + + // Budget 27: nobody + let filtered = BohbSampler::filter_history_for_budget(&history, 27); + assert!(filtered.is_empty()); + } + + #[test] + fn matching_pruner_has_same_params() { + let bohb = BohbSampler::builder() + .min_resource(2) + .max_resource(64) + .reduction_factor(4) + .build() + .unwrap(); + let pruner = bohb.matching_pruner(Direction::Minimize); + + // We can't directly inspect HyperbandPruner fields, but we can + // verify it was created without panicking with the same params. + // The pruner's rung steps should match BOHB's budget levels. + // Just verify it doesn't panic. + drop(pruner); + } + + #[test] + fn fallback_to_random_when_insufficient_data() { + let bohb = BohbSampler::builder() + .min_points_in_model(10) + .seed(42) + .build() + .unwrap(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Only 3 trials with intermediate values (< min_points_in_model=10) + let x_id = ParamId::new(); + let history: Vec = (0..3) + .map(|i| { + make_trial_with_intermediates( + i, + i as f64, + vec![(x_id, ParamValue::Float(i as f64 / 3.0), dist.clone())], + vec![(1, i as f64)], + ) + }) + .collect(); + + // Should not panic, should sample within bounds + for trial_id in 0..20 { + let val = bohb.sample(&dist, trial_id, &history); + if let ParamValue::Float(v) = val { + assert!((0.0..=1.0).contains(&v)); + } else { + panic!("Expected Float"); + } + } + } + + #[test] + fn uses_budget_level_when_enough_data() { + let bohb = BohbSampler::builder() + .min_points_in_model(5) + .seed(42) + .build() + .unwrap(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 10.0, + log_scale: false, + step: None, + }); + + // Create 20 trials with intermediate values at budget 1. + // Good trials have x near 2.0, bad trials have x far from 2.0. + let x_id = ParamId::new(); + let history: Vec = (0..20) + .map(|i| { + let x = i as f64 / 2.0; + let iv_at_1 = (x - 2.0).powi(2); + make_trial_with_intermediates( + i, + iv_at_1, // final value same as intermediate for simplicity + vec![(x_id, ParamValue::Float(x), dist.clone())], + vec![(1, iv_at_1)], + ) + }) + .collect(); + + // Should use TPE on filtered history at budget 1 + let val = bohb.sample(&dist, 100, &history); + if let ParamValue::Float(v) = val { + assert!((0.0..=10.0).contains(&v), "Value {v} out of bounds"); + } else { + panic!("Expected Float"); + } + } + + #[test] + fn prefers_largest_budget_level() { + let bohb = BohbSampler::builder() + .min_resource(1) + .max_resource(9) + .reduction_factor(3) + .min_points_in_model(3) + .seed(42) + .build() + .unwrap(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 10.0, + log_scale: false, + step: None, + }); + + // Budget levels: [1, 3, 9] + assert_eq!(bohb.all_budget_levels(), vec![1, 3, 9]); + + // Create 5 trials with intermediates at budget 1 and 3 + let x_id = ParamId::new(); + let history: Vec = (0..5) + .map(|i| { + let x = i as f64; + make_trial_with_intermediates( + i, + x, + vec![(x_id, ParamValue::Float(x), dist.clone())], + vec![(1, x * 2.0), (3, x)], + ) + }) + .collect(); + + // Budget 3 has 5 observations (>= 3), budget 9 has 0. + // BOHB should pick budget 3 (largest with enough data). + // The filtered history at budget 3 has values [0, 1, 2, 3, 4]. + let filtered_3 = BohbSampler::filter_history_for_budget(&history, 3); + assert_eq!(filtered_3.len(), 5); + let filtered_9 = BohbSampler::filter_history_for_budget(&history, 9); + assert_eq!(filtered_9.len(), 0); + + // Should sample successfully + let val = bohb.sample(&dist, 100, &history); + assert!(matches!(val, ParamValue::Float(_))); + } + + #[test] + fn builder_validates_tpe_params() { + // Invalid gamma + let result = BohbSampler::builder().gamma(1.5).build(); + assert!(result.is_err()); + + // Invalid bandwidth + let result = BohbSampler::builder().kde_bandwidth(-1.0).build(); + assert!(result.is_err()); + } + + #[test] + #[should_panic(expected = "min_resource must be > 0")] + fn builder_rejects_zero_min_resource() { + let _ = BohbSampler::builder().min_resource(0); + } + + #[test] + #[should_panic(expected = "max_resource must be > 0")] + fn builder_rejects_zero_max_resource() { + let _ = BohbSampler::builder().max_resource(0); + } + + #[test] + #[should_panic(expected = "reduction_factor must be >= 2")] + fn builder_rejects_small_reduction_factor() { + let _ = BohbSampler::builder().reduction_factor(1); + } + + #[test] + fn int_distribution_works() { + let bohb = BohbSampler::builder() + .min_points_in_model(3) + .seed(42) + .build() + .unwrap(); + + let dist = Distribution::Int(IntDistribution { + low: 0, + high: 100, + log_scale: false, + step: None, + }); + + let x_id = ParamId::new(); + let history: Vec = (0..10) + .map(|i| { + make_trial_with_intermediates( + i, + i as f64, + vec![(x_id, ParamValue::Int(i.cast_signed() * 10), dist.clone())], + vec![(1, i as f64)], + ) + }) + .collect(); + + let val = bohb.sample(&dist, 100, &history); + if let ParamValue::Int(v) = val { + assert!((0..=100).contains(&v)); + } else { + panic!("Expected Int"); + } + } +} diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index 8495c3b..a7ca85c 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -1,5 +1,6 @@ //! Sampler trait and implementations for parameter sampling. +pub mod bohb; #[cfg(feature = "cma-es")] pub mod cma_es; pub mod grid; diff --git a/tests/bohb_integration.rs b/tests/bohb_integration.rs new file mode 100644 index 0000000..980de7c --- /dev/null +++ b/tests/bohb_integration.rs @@ -0,0 +1,133 @@ +//! Integration tests for the BOHB sampler. + +#![allow( + clippy::cast_sign_loss, + clippy::cast_precision_loss, + clippy::cast_possible_truncation +)] + +use optimizer::parameter::{FloatParam, Parameter}; +use optimizer::sampler::bohb::BohbSampler; +use optimizer::{Direction, Error, Study, TrialPruned}; + +#[test] +fn bohb_converges_on_quadratic() { + let bohb = BohbSampler::builder() + .min_resource(1) + .max_resource(9) + .reduction_factor(3) + .min_points_in_model(5) + .seed(42) + .build() + .unwrap(); + + let pruner = bohb.matching_pruner(Direction::Minimize); + let study: Study = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner); + + let x_param = FloatParam::new(-10.0, 10.0); + + study + .optimize(60, |trial| { + let x = x_param.suggest(trial)?; + + // Report intermediate values at budget steps 1, 3, 9 + let obj = (x - 3.0).powi(2); + // Simulate budget-based evaluation with noise decreasing at higher budgets + trial.report(1, obj + 5.0); + trial.report(3, obj + 1.0); + trial.report(9, obj); + + Ok::<_, Error>(obj) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have trials"); + assert!( + best.value < 10.0, + "BOHB should find a reasonable solution, got {}", + best.value + ); +} + +#[test] +fn bohb_with_pruning() { + let bohb = BohbSampler::builder() + .min_resource(1) + .max_resource(27) + .reduction_factor(3) + .min_points_in_model(3) + .seed(123) + .build() + .unwrap(); + + let pruner = bohb.matching_pruner(Direction::Minimize); + let study: Study = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner); + + let x_param = FloatParam::new(-5.0, 5.0); + + study + .optimize(40, |trial| { + let x = x_param.suggest(trial)?; + let obj = x * x; + + // Report at each rung step and check for pruning + for &step in &[1u64, 3, 9, 27] { + let noisy_obj = obj + 10.0 / step as f64; + trial.report(step, noisy_obj); + + if trial.should_prune() { + return Err(TrialPruned.into()); + } + } + + Ok::<_, Error>(obj) + }) + .expect("optimization should succeed"); + + // Verify we have completed trials + let best = study.best_trial().expect("should have at least one trial"); + assert!( + best.value < 25.0, + "best value {} should be reasonable", + best.value + ); +} + +#[test] +fn bohb_uses_budget_conditioned_history() { + // Verify that BOHB conditions on budget level by testing that samples + // are influenced by intermediate values, not just final values. + let bohb = BohbSampler::builder() + .min_resource(1) + .max_resource(9) + .reduction_factor(3) + .min_points_in_model(3) + .seed(42) + .build() + .unwrap(); + + let pruner = bohb.matching_pruner(Direction::Minimize); + let study: Study = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner); + + let x_param = FloatParam::new(0.0, 10.0); + + study + .optimize(30, |trial| { + let x = x_param.suggest(trial)?; + // Intermediate values that guide optimization toward x=2 + trial.report(1, (x - 2.0).powi(2) + 1.0); + trial.report(3, (x - 2.0).powi(2) + 0.5); + trial.report(9, (x - 2.0).powi(2)); + + Ok::<_, Error>((x - 2.0).powi(2)) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().unwrap(); + let best_x: f64 = best.get(&x_param).unwrap(); + // Should find x reasonably close to 2.0 + assert!( + (best_x - 2.0).abs() < 5.0, + "BOHB should explore near x=2, got x={best_x}" + ); +}