Local models and datasets

docs/notebooks/trust_region.pct.py

@@ -106,7 +106,9 @@ def build_model():
 from trieste.experimental.plotting import plot_bo_points, plot_function_2d
 
 
-def plot_final_result(_dataset: trieste.data.Dataset) -> None:
+def plot_final_result(
+    _dataset: trieste.data.Dataset, num_init_points=num_initial_data_points
+) -> None:
     arg_min_idx = tf.squeeze(tf.argmin(_dataset.observations, axis=0))
     query_points = _dataset.query_points.numpy()
     _, ax = plot_function_2d(
@@ -117,7 +119,7 @@ def plot_final_result(_dataset: trieste.data.Dataset) -> None:
         contour=True,
     )
 
-    plot_bo_points(query_points, ax[0, 0], num_initial_data_points, arg_min_idx)
+    plot_bo_points(query_points, ax[0, 0], num_init_points, arg_min_idx)
 
 
 plot_final_result(dataset)
@@ -144,7 +146,10 @@ def plot_final_result(_dataset: trieste.data.Dataset) -> None:
 )
 
 
-def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None:
+def plot_history(
+    result: trieste.bayesian_optimizer.OptimizationResult,
+    num_init_points=num_initial_data_points,
+) -> None:
     frames = []
     for step, hist in enumerate(
         result.history + [result.final_result.unwrap()]
@@ -154,7 +159,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None:
             search_space.lower,
             search_space.upper,
             hist,
-            num_init=num_initial_data_points,
+            num_init=num_init_points,
         )
 
         if fig is not None:
@@ -240,10 +245,48 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None:
 # We visualize the results as before.
 
 # %%
-plot_final_result(dataset)
+plot_final_result(dataset, num_init_points=0)
 
 # %%
-plot_history(result)
+plot_history(result, num_init_points=0)
+
+# %% [markdown]
+# ## TEST
+
+# %%
+num_query_points = 5
+
+init_subspaces = [
+    trieste.acquisition.rule.SingleObjectiveTrustRegionBox(search_space)
+    for _ in range(num_query_points)
+]
+base_rule = trieste.acquisition.rule.EfficientGlobalOptimization(
+    builder=trieste.acquisition.ParallelContinuousThompsonSampling(),
+    num_query_points=1,
+)
+batch_acq_rule = trieste.acquisition.rule.BatchTrustRegionBox(
+    init_subspaces, base_rule
+)
+
+bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)
+
+num_steps = 5
+result = bo.optimize(
+    num_steps,
+    {trieste.observer.OBJECTIVE: initial_data},
+    trieste.acquisition.utils.copy_to_local_models(
+        build_model(), num_query_points
+    ),
+    batch_acq_rule,
+    track_state=True,
+)
+dataset = result.try_get_final_dataset()
+
+# %%
+plot_final_result(dataset, num_init_points=0)
+
+# %%
+plot_history(result, num_init_points=0)
 
 # %% [markdown]
 # ## Trust region `TurBO` acquisition rule

tests/integration/test_ask_tell_optimization.py

@@ -23,22 +23,24 @@
 import tensorflow as tf
 
 from tests.util.misc import random_seed
-from trieste.acquisition import LocalPenalization
+from trieste.acquisition import LocalPenalization, ParallelContinuousThompsonSampling
 from trieste.acquisition.rule import (
     AcquisitionRule,
     AsynchronousGreedy,
     AsynchronousRuleState,
     BatchTrustRegionBox,
     EfficientGlobalOptimization,
+    SingleObjectiveTrustRegionBox,
     TREGOBox,
 )
+from trieste.acquisition.utils import copy_to_local_models
 from trieste.ask_tell_optimization import AskTellOptimizer
 from trieste.bayesian_optimizer import OptimizationResult, Record
 from trieste.logging import set_step_number, tensorboard_writer
 from trieste.models import TrainableProbabilisticModel
 from trieste.models.gpflow import GaussianProcessRegression, build_gpr
 from trieste.objectives import ScaledBranin, SimpleQuadratic
-from trieste.objectives.utils import mk_observer
+from trieste.objectives.utils import mk_batch_observer, mk_observer
 from trieste.observer import OBJECTIVE
 from trieste.space import Box, SearchSpace
 from trieste.types import State, TensorType
@@ -47,26 +49,58 @@
 # We use a copy of these for a quicker test against a simple quadratic function
 # (copying is necessary as some of the acquisition rules are stateful).
 OPTIMIZER_PARAMS = (
-    "num_steps, reload_state, acquisition_rule_fn",
+    "num_steps, reload_state, acquisition_rule_fn, num_models",
     [
         pytest.param(
-            20, False, lambda: EfficientGlobalOptimization(), id="EfficientGlobalOptimization"
+            20, False, lambda: EfficientGlobalOptimization(), 1, id="EfficientGlobalOptimization"
         ),
         pytest.param(
             20,
             True,
             lambda: EfficientGlobalOptimization(),
+            1,
             id="EfficientGlobalOptimization/reload_state",
         ),
         pytest.param(
-            15, False, lambda: BatchTrustRegionBox(TREGOBox(ScaledBranin.search_space)), id="TREGO"
+            15,
+            False,
+            lambda: BatchTrustRegionBox(TREGOBox(ScaledBranin.search_space)),
+            1,
+            id="TREGO",
         ),
         pytest.param(
             16,
             True,
             lambda: BatchTrustRegionBox(TREGOBox(ScaledBranin.search_space)),
+            1,
             id="TREGO/reload_state",
         ),
+        pytest.param(
+            10,
+            False,
+            lambda: BatchTrustRegionBox(
+                [SingleObjectiveTrustRegionBox(ScaledBranin.search_space) for _ in range(3)],
+                EfficientGlobalOptimization(
+                    ParallelContinuousThompsonSampling(),
+                    num_query_points=3,
+                ),
+            ),
+            1,
+            id="BatchTrustRegionBox",
+        ),
+        pytest.param(
+            10,
+            False,
+            lambda: BatchTrustRegionBox(
+                [SingleObjectiveTrustRegionBox(ScaledBranin.search_space) for _ in range(3)],
+                EfficientGlobalOptimization(
+                    ParallelContinuousThompsonSampling(),
+                    num_query_points=2,
+                ),
+            ),
+            3,
+            id="BatchTrustRegionBox/LocalModels",
+        ),
         pytest.param(
             10,
             False,
@@ -76,6 +110,7 @@
                 ).using(OBJECTIVE),
                 num_query_points=3,
             ),
+            1,
             id="LocalPenalization",
         ),
         pytest.param(
@@ -86,6 +121,7 @@
                     ScaledBranin.search_space,
                 ).using(OBJECTIVE),
             ),
+            1,
             id="LocalPenalization/AsynchronousGreedy",
         ),
     ],
@@ -109,8 +145,11 @@ def test_ask_tell_optimizer_finds_minima_of_the_scaled_branin_function(
             TrainableProbabilisticModel,
         ],
     ],
+    num_models: int,
 ) -> None:
-    _test_ask_tell_optimization_finds_minima(True, num_steps, reload_state, acquisition_rule_fn)
+    _test_ask_tell_optimization_finds_minima(
+        True, num_steps, reload_state, acquisition_rule_fn, num_models
+    )
 
 
 @random_seed
@@ -129,11 +168,12 @@ def test_ask_tell_optimizer_finds_minima_of_simple_quadratic(
             TrainableProbabilisticModel,
         ],
     ],
+    num_models: int,
 ) -> None:
     # for speed reasons we sometimes test with a simple quadratic defined on the same search space
     # branin; currently assume that every rule should be able to solve this in 5 steps
     _test_ask_tell_optimization_finds_minima(
-        False, min(num_steps, 5), reload_state, acquisition_rule_fn
+        False, min(num_steps, 5), reload_state, acquisition_rule_fn, num_models
     )
 
 
@@ -152,6 +192,7 @@ def _test_ask_tell_optimization_finds_minima(
             TrainableProbabilisticModel,
         ],
     ],
+    num_models: int,
 ) -> None:
     # For the case when optimization state is saved and reload on each iteration
     # we need to use new acquisition function object to imitate real life usage
@@ -160,17 +201,22 @@ def _test_ask_tell_optimization_finds_minima(
     search_space = ScaledBranin.search_space
     initial_query_points = search_space.sample(5)
     observer = mk_observer(ScaledBranin.objective if optimize_branin else SimpleQuadratic.objective)
+    batch_observer = mk_batch_observer(observer)
     initial_data = observer(initial_query_points)
 
     model = GaussianProcessRegression(
         build_gpr(initial_data, search_space, likelihood_variance=1e-7)
     )
+    models = copy_to_local_models(model, num_models) if num_models > 1 else {OBJECTIVE: model}
+    initial_dataset = {OBJECTIVE: initial_data}
 
     with tempfile.TemporaryDirectory() as tmpdirname:
         summary_writer = tf.summary.create_file_writer(tmpdirname)
         with tensorboard_writer(summary_writer):
             set_step_number(0)
-            ask_tell = AskTellOptimizer(search_space, initial_data, model, acquisition_rule_fn())
+            ask_tell = AskTellOptimizer(
+                search_space, initial_dataset, models, acquisition_rule_fn()
+            )
 
             for i in range(1, num_steps + 1):
                 # two scenarios are tested here, depending on `reload_state` parameter
@@ -185,7 +231,11 @@ def _test_ask_tell_optimization_finds_minima(
                     ] = ask_tell.to_record()
                     written_state = pickle.dumps(state)
 
-                new_data_point = observer(new_point)
+                # If query points are rank 3, then use a batched observer.
+                if tf.rank(new_point) == 3:
+                    new_data_point = batch_observer(new_point)
+                else:
+                    new_data_point = observer(new_point)
 
                 if reload_state:
                     state = pickle.loads(written_state)