feat: add BOHB sampler for budget-aware Bayesian optimization

BOHB combines TPE's model-guided sampling with Hyperband's budget-aware evaluation by conditioning the TPE model on trials at specific budget levels, producing better-calibrated parameter proposals.
2026-02-11 18:42:09 +01:00
parent a1dd6d393c
commit aa9734587a
4 changed files with 826 additions and 0 deletions
@@ -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<f64> = 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<f64> = 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<f64> = 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}"
+    );
+}