Multi trust regions (#773)

* multiTR with notebook

* notebook

* multi local step TREGO

* Refactor multi trust region implementation

Provide base abstract classes for updateable search space and multi
trust region rule.
Implement concrete classes for box multi trust region rule.

* Fix formatting in notebook & remove orig TR changes

* Remove redundant new assert

* Undo earlier logging change

* Workaround isort & black clash

* Keep old ver of mypy happy

* Fix typo in slicing

* Add new collection space tests

* Address feedback plus other small changes

* Use generic search space type for MultiTrustRegionBox
* Add some private methods to TrustRegionBox
* Add kwargs to reinitialize/update
* Use American spelling
* Move helper function to utils

* Add TrustRegionBox/UpdateableSearchSpace unittests

* Address feedback for space changes

* Some updates to rule/utils from feedback

* Rename reinitialize to initialize
* Replace get_single_model_and_dataset with generic get_value_for_tag;
  plus update unit tests
* Add unit tests for and change get_unique_points_mask implementation as
  previous version erroneously identified some points as duplicates

* Create subspaces outside the rule

Plus revert addition of kwargs

* Fix func to work with old tensorflow

tf.sqrt for array with inf returns nan in older versions.

* Add rule and optimizer unit tests

* Add integ test

* Add shape-checking/docstring for mask func

Plus one other minor tweak.

* Add check_shapes, update docstrings & add tag test

* Remove notebook, to be added in separate PR

* Add deepcopy test

* Check TR copy is deep

* Improve TR box and test slightly

* Address Uri's latest comments

* Move get_value_for_tag to more general utils/misc

* Remove redundant [] checks

* Add fixture for default search space type

* Add test for optim multiple boxes as a batch

* Rename classes and expand some docstrings

* Choose longer class name

* Rename updatable class name

---------

Co-authored-by: Victor Picheny <[email protected]>
Co-authored-by: hstojic <[email protected]>

tests/integration/test_bayesian_optimization.py

@@ -50,8 +50,10 @@
     AsynchronousOptimization,
     AsynchronousRuleState,
     BatchHypervolumeSharpeRatioIndicator,
+    BatchTrustRegionBox,
     DiscreteThompsonSampling,
     EfficientGlobalOptimization,
+    SingleObjectiveTrustRegionBox,
     TrustRegion,
 )
 from trieste.acquisition.sampler import ThompsonSamplerFromTrajectory
@@ -197,6 +199,17 @@ def GPR_OPTIMIZER_PARAMS() -> Tuple[str, List[ParameterSet]]:
                 TURBO(ScaledBranin.search_space, rule=DiscreteThompsonSampling(500, 3)),
                 id="Turbo",
             ),
+            pytest.param(
+                10,
+                BatchTrustRegionBox(
+                    [SingleObjectiveTrustRegionBox(ScaledBranin.search_space) for _ in range(3)],
+                    EfficientGlobalOptimization(
+                        ParallelContinuousThompsonSampling(),
+                        num_query_points=3,
+                    ),
+                ),
+                id="BatchTrustRegionBox",
+            ),
             pytest.param(15, DiscreteThompsonSampling(500, 5), id="DiscreteThompsonSampling"),
             pytest.param(
                 15,

tests/unit/acquisition/test_optimizer.py

@@ -45,6 +45,7 @@
     DiscreteSearchSpace,
     LinearConstraint,
     SearchSpace,
+    TaggedMultiSearchSpace,
     TaggedProductSearchSpace,
 )
 from trieste.types import TensorType
@@ -203,6 +204,40 @@ def test_optimize_continuous_raises_with_invalid_vectorized_batch_size(batch_siz
         generate_continuous_optimizer()(search_space, (acq_fn, batch_size))
 
 
+def test_optimize_continuous_raises_with_mismatch_multi_search_space() -> None:
+    space_A = Box([-1], [2])
+    space_B = Box([3], [4])
+    multi_space = TaggedMultiSearchSpace(spaces=[space_A, space_B])
+    acq_fn = _quadratic_sum([1.0])
+    with pytest.raises(TF_DEBUGGING_ERROR_TYPES, match="The batch shape of initial samples 2 must"):
+        generate_continuous_optimizer()(multi_space, acq_fn)
+
+
+def test_optimize_continuous_finds_points_in_multi_search_space_boxes() -> None:
+    # Test with non-overlapping grid of 2D boxes. Optimize them as a batch and check that each
+    # point is only in the corresponding box.
+    boxes = [Box([x, y], [x + 0.7, y + 0.7]) for x in range(-2, 2) for y in range(-2, 2)]
+    multi_space = TaggedMultiSearchSpace(spaces=boxes)
+    batch_size = len(boxes)
+
+    def target_function(x: TensorType) -> TensorType:  # [N, V, D] -> [N, V]
+        individual_func = [_quadratic_sum([1.0])(x[:, i : i + 1, :]) for i in range(batch_size)]
+        return tf.concat(individual_func, axis=-1)  # vectorize by repeating same function
+
+    optimizer = generate_continuous_optimizer()
+    optimizer = batchify_vectorize(optimizer, batch_size)
+    max_points = optimizer(multi_space, target_function)
+
+    # Ensure order of points is the same as the order of boxes and that each point is only in the
+    # corresponding box.
+    for i, point in enumerate(max_points):
+        for j, box in enumerate(boxes):
+            if i == j:
+                assert point in box
+            else:
+                assert point not in box
+
+
 @pytest.mark.parametrize("num_optimization_runs", [1, 10])
 @pytest.mark.parametrize("num_initial_samples", [1000, 5000])
 def test_optimize_continuous_correctly_uses_init_params(