Replace TrustRegion with TREGOBox implementation

secondmind-labs · Sep 7, 2023 · 1c10622 · 1c10622
1 parent 9ee2d21
commit 1c10622
Show file tree

Hide file tree

Showing 7 changed files with 237 additions and 429 deletions.
diff --git a/docs/notebooks/recovering_from_errors.pct.py b/docs/notebooks/recovering_from_errors.pct.py
@@ -1,10 +1,11 @@
 # %% [markdown]
 # # Recovering from errors
 
+import random
+
 # %%
 import numpy as np
 import tensorflow as tf
-import random
 
 np.random.seed(1793)
 tf.random.set_seed(1793)
@@ -41,10 +42,10 @@ def __call__(self, x):
 
 # %% [markdown]
 # ## Set up the problem
-# We'll use the same set up as before, except for the acquisition rule, where we'll use `TrustRegion`. `TrustRegion` is stateful, and we'll need to account for its state to recover, so using this rule gives the reader a more comprehensive overview of how to recover.
+# We'll use the same set up as before, except for the acquisition rule, where we'll use `BatchTrustRegionBox`. `BatchTrustRegionBox` is stateful, and we'll need to account for its state to recover, so using this rule gives the reader a more comprehensive overview of how to recover.
 
 # %%
-from trieste.models.gpflow import build_gpr, GaussianProcessRegression
+from trieste.models.gpflow import GaussianProcessRegression, build_gpr
 
 search_space = trieste.space.Box(
     tf.cast([0.0, 0.0], tf.float64), tf.cast([1.0, 1.0], tf.float64)
@@ -54,7 +55,9 @@ def __call__(self, x):
 gpr = build_gpr(initial_data, search_space)
 model = GaussianProcessRegression(gpr)
 
-acquisition_rule = trieste.acquisition.rule.TrustRegion()
+acquisition_rule = trieste.acquisition.rule.BatchTrustRegionBox(  # type: ignore[var-annotated]
+    trieste.acquisition.rule.TREGOBox(search_space)
+)
 
 # %% [markdown]
 # ## Run the optimization loop
@@ -155,7 +158,10 @@ def __call__(self, x):
     automatic_optimizer_selector, split_size=10_000
 )
 query_rule = EfficientGlobalOptimization(optimizer=optimizer)
-acquisition_rule = trieste.acquisition.rule.TrustRegion(rule=query_rule)
+acquisition_rule = trieste.acquisition.rule.BatchTrustRegionBox(
+    trieste.acquisition.rule.TREGOBox(search_space),
+    rule=query_rule,
+)
 
 # %% [markdown]
 # ## LICENSE

diff --git a/docs/notebooks/trego_compare.pcy.py b/docs/notebooks/trego_compare.pcy.py
diff --git a/tests/integration/test_ask_tell_optimization.py b/tests/integration/test_ask_tell_optimization.py
@@ -28,8 +28,9 @@
     AcquisitionRule,
     AsynchronousGreedy,
     AsynchronousRuleState,
+    BatchTrustRegionBox,
     EfficientGlobalOptimization,
-    TrustRegion,
+    TREGOBox,
 )
 from trieste.ask_tell_optimization import AskTellOptimizer
 from trieste.bayesian_optimizer import OptimizationResult, Record
@@ -57,8 +58,15 @@
             lambda: EfficientGlobalOptimization(),
             id="EfficientGlobalOptimization/reload_state",
         ),
-        pytest.param(15, False, lambda: TrustRegion(), id="TrustRegion"),
-        pytest.param(16, True, lambda: TrustRegion(), id="TrustRegion/reload_state"),
+        pytest.param(
+            15, False, lambda: BatchTrustRegionBox(TREGOBox(ScaledBranin.search_space)), id="TREGO"
+        ),
+        pytest.param(
+            16,
+            True,
+            lambda: BatchTrustRegionBox(TREGOBox(ScaledBranin.search_space)),
+            id="TREGO/reload_state",
+        ),
         pytest.param(
             10,
             False,
@@ -96,7 +104,7 @@ def test_ask_tell_optimizer_finds_minima_of_the_scaled_branin_function(
     | Callable[
         [],
         AcquisitionRule[
-            State[TensorType, AsynchronousRuleState | TrustRegion.State],
+            State[TensorType, AsynchronousRuleState | BatchTrustRegionBox.State],
             Box,
             TrainableProbabilisticModel,
         ],
@@ -116,7 +124,7 @@ def test_ask_tell_optimizer_finds_minima_of_simple_quadratic(
     | Callable[
         [],
         AcquisitionRule[
-            State[TensorType, AsynchronousRuleState | TrustRegion.State],
+            State[TensorType, AsynchronousRuleState | BatchTrustRegionBox.State],
             Box,
             TrainableProbabilisticModel,
         ],
@@ -139,7 +147,7 @@ def _test_ask_tell_optimization_finds_minima(
     | Callable[
         [],
         AcquisitionRule[
-            State[TensorType, AsynchronousRuleState | TrustRegion.State],
+            State[TensorType, AsynchronousRuleState | BatchTrustRegionBox.State],
             Box,
             TrainableProbabilisticModel,
         ],
@@ -173,7 +181,7 @@ def _test_ask_tell_optimization_finds_minima(
 
                 if reload_state:
                     state: Record[
-                        None | State[TensorType, AsynchronousRuleState | TrustRegion.State]
+                        None | State[TensorType, AsynchronousRuleState | BatchTrustRegionBox.State]
                     ] = ask_tell.to_record()
                     written_state = pickle.dumps(state)
 
@@ -188,7 +196,7 @@ def _test_ask_tell_optimization_finds_minima(
                 ask_tell.tell(new_data_point)
 
     result: OptimizationResult[
-        None | State[TensorType, AsynchronousRuleState | TrustRegion.State]
+        None | State[TensorType, AsynchronousRuleState | BatchTrustRegionBox.State]
     ] = ask_tell.to_result()
     dataset = result.try_get_final_dataset()
 

diff --git a/tests/integration/test_bayesian_optimization.py b/tests/integration/test_bayesian_optimization.py
@@ -50,12 +50,12 @@
     AsynchronousOptimization,
     AsynchronousRuleState,
     BatchHypervolumeSharpeRatioIndicator,
+    BatchTrustRegion,
     BatchTrustRegionBox,
     DiscreteThompsonSampling,
     EfficientGlobalOptimization,
     SingleObjectiveTrustRegionBox,
     TREGOBox,
-    TrustRegion,
 )
 from trieste.acquisition.sampler import ThompsonSamplerFromTrajectory
 from trieste.bayesian_optimizer import (
@@ -183,22 +183,22 @@ def GPR_OPTIMIZER_PARAMS() -> Tuple[str, List[ParameterSet]]:
                 ),
                 id="MultipleOptimismNegativeLowerConfidenceBound",
             ),
-            pytest.param(20, TrustRegion(), id="TrustRegion"),
+            pytest.param(
+                20,
+                BatchTrustRegionBox(TREGOBox(ScaledBranin.search_space)),
+                id="TREGO",
+            ),
             pytest.param(
                 15,
-                TrustRegion(
+                BatchTrustRegionBox(
+                    TREGOBox(ScaledBranin.search_space),
                     EfficientGlobalOptimization(
                         MinValueEntropySearch(
                             ScaledBranin.search_space,
                         ).using(OBJECTIVE)
-                    )
+                    ),
                 ),
-                id="TrustRegion/MinValueEntropySearch",
-            ),
-            pytest.param(
-                20,
-                BatchTrustRegionBox([TREGOBox(ScaledBranin.search_space)]),
-                id="TREGO",
+                id="TREGO/MinValueEntropySearch",
             ),
             pytest.param(
                 20,
@@ -211,18 +211,6 @@ def GPR_OPTIMIZER_PARAMS() -> Tuple[str, List[ParameterSet]]:
                 ),
                 id="TREGO/ParallelContinuousThompsonSampling",
             ),
-            pytest.param(
-                15,
-                BatchTrustRegionBox(
-                    [TREGOBox(ScaledBranin.search_space)],
-                    EfficientGlobalOptimization(
-                        MinValueEntropySearch(
-                            ScaledBranin.search_space,
-                        ).using(OBJECTIVE)
-                    ),
-                ),
-                id="TREGO/MinValueEntropySearch",
-            ),
             pytest.param(
                 10,
                 TURBO(ScaledBranin.search_space, rule=DiscreteThompsonSampling(500, 3)),
@@ -281,7 +269,9 @@ def test_bayesian_optimizer_with_gpr_finds_minima_of_scaled_branin(
     num_steps: int,
     acquisition_rule: AcquisitionRule[TensorType, SearchSpace, GaussianProcessRegression]
     | AcquisitionRule[
-        State[TensorType, AsynchronousRuleState | TrustRegion.State], Box, GaussianProcessRegression
+        State[TensorType, AsynchronousRuleState | BatchTrustRegion.State],
+        Box,
+        GaussianProcessRegression,
     ],
 ) -> None:
     _test_optimizer_finds_minimum(
@@ -295,7 +285,9 @@ def test_bayesian_optimizer_with_gpr_finds_minima_of_simple_quadratic(
     num_steps: int,
     acquisition_rule: AcquisitionRule[TensorType, SearchSpace, GaussianProcessRegression]
     | AcquisitionRule[
-        State[TensorType, AsynchronousRuleState | TrustRegion.State], Box, GaussianProcessRegression
+        State[TensorType, AsynchronousRuleState | BatchTrustRegion.State],
+        Box,
+        GaussianProcessRegression,
     ],
 ) -> None:
     # for speed reasons we sometimes test with a simple quadratic defined on the same search space
@@ -566,7 +558,7 @@ def _test_optimizer_finds_minimum(
     num_steps: Optional[int],
     acquisition_rule: AcquisitionRule[TensorType, SearchSpace, TrainableProbabilisticModelType]
     | AcquisitionRule[
-        State[TensorType, AsynchronousRuleState | TrustRegion.State],
+        State[TensorType, AsynchronousRuleState | BatchTrustRegion.State],
         Box,
         TrainableProbabilisticModelType,
     ],