From 93f44b321e7b9112016288b48325c8863feb173e Mon Sep 17 00:00:00 2001 From: Khurram Ghani <113982802+khurram-ghani@users.noreply.github.com> Date: Fri, 25 Aug 2023 14:02:02 +0100 Subject: [PATCH] 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 Co-authored-by: hstojic --- .../integration/test_bayesian_optimization.py | 13 + tests/unit/acquisition/test_optimizer.py | 35 + tests/unit/acquisition/test_rule.py | 338 ++++++++- tests/unit/acquisition/test_utils.py | 33 + tests/unit/test_space.py | 672 +++++++++++++----- tests/unit/utils/test_misc.py | 29 +- trieste/acquisition/optimizer.py | 110 ++- trieste/acquisition/rule.py | 314 +++++++- trieste/acquisition/utils.py | 36 + trieste/space.py | 303 ++++++-- trieste/utils/misc.py | 24 +- 11 files changed, 1639 insertions(+), 268 deletions(-) diff --git a/tests/integration/test_bayesian_optimization.py b/tests/integration/test_bayesian_optimization.py index 28b0bd3886..ca8981f9a9 100644 --- a/tests/integration/test_bayesian_optimization.py +++ b/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, diff --git a/tests/unit/acquisition/test_optimizer.py b/tests/unit/acquisition/test_optimizer.py index 1c4a090885..e2f7e6a0f3 100644 --- a/tests/unit/acquisition/test_optimizer.py +++ b/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( diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index 13b6acabd2..0e2b575538 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -18,6 +18,7 @@ from typing import Callable, Optional import gpflow +import numpy as np import numpy.testing as npt import pytest import tensorflow as tf @@ -32,6 +33,7 @@ AcquisitionFunction, AcquisitionFunctionBuilder, NegativeLowerConfidenceBound, + ParallelContinuousThompsonSampling, SingleModelAcquisitionBuilder, SingleModelGreedyAcquisitionBuilder, VectorizedAcquisitionFunctionBuilder, @@ -44,9 +46,11 @@ AsynchronousOptimization, AsynchronousRuleState, BatchHypervolumeSharpeRatioIndicator, + BatchTrustRegionBox, DiscreteThompsonSampling, EfficientGlobalOptimization, RandomSampling, + SingleObjectiveTrustRegionBox, TrustRegion, ) from trieste.acquisition.sampler import ( @@ -59,7 +63,7 @@ from trieste.models import ProbabilisticModel from trieste.models.interfaces import TrainableSupportsGetKernel from trieste.observer import OBJECTIVE -from trieste.space import Box +from trieste.space import Box, SearchSpace, TaggedMultiSearchSpace from trieste.types import State, Tag, TensorType @@ -1112,6 +1116,338 @@ def test_turbo_state_deepcopy() -> None: npt.assert_allclose(tr_state_copy.y_min, tr_state.y_min) +# get_local_min raises if dataset is None. +def test_trust_region_box_get_local_min_raises_if_dataset_is_none() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + trb = SingleObjectiveTrustRegionBox(search_space) + with pytest.raises(ValueError, match="dataset must be provided"): + trb.get_local_min(None) + + +# get_local_min picks the minimum x and y values from the dataset. +def test_trust_region_box_get_local_min() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + dataset = Dataset( + tf.constant([[0.1, 0.1], [0.5, 0.5], [0.3, 0.4], [0.8, 0.8], [0.4, 0.4]], dtype=tf.float64), + tf.constant([[0.0], [0.5], [0.2], [0.1], [1.0]], dtype=tf.float64), + ) + trb = SingleObjectiveTrustRegionBox(search_space) + trb._lower = tf.constant([0.2, 0.2], dtype=tf.float64) + trb._upper = tf.constant([0.7, 0.7], dtype=tf.float64) + x_min, y_min = trb.get_local_min(dataset) + npt.assert_array_equal(x_min, tf.constant([0.3, 0.4], dtype=tf.float64)) + npt.assert_array_equal(y_min, tf.constant([0.2], dtype=tf.float64)) + + +# get_local_min returns first x value and inf y value when points in dataset are outside the +# search space. +def test_trust_region_box_get_local_min_outside_search_space() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + dataset = Dataset( + tf.constant([[1.2, 1.3], [-0.4, -0.5]], dtype=tf.float64), + tf.constant([[0.7], [0.9]], dtype=tf.float64), + ) + trb = SingleObjectiveTrustRegionBox(search_space) + x_min, y_min = trb.get_local_min(dataset) + npt.assert_array_equal(x_min, tf.constant([1.2, 1.3], dtype=tf.float64)) + npt.assert_array_equal(y_min, tf.constant([np.inf], dtype=tf.float64)) + + +# Initialize sets the box to a random location, and sets the eps and y_min values. +def test_trust_region_box_initialize() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + datasets = { + OBJECTIVE: Dataset( # Points outside the search space should be ignored. + tf.constant([[1.2, 1.3], [-0.4, -0.5]], dtype=tf.float64), + tf.constant([[0.7], [0.9]], dtype=tf.float64), + ) + } + trb = SingleObjectiveTrustRegionBox(search_space) + trb.initialize(datasets=datasets) + + exp_eps = 0.5 * (search_space.upper - search_space.lower) / 5.0 ** (1.0 / 2.0) + npt.assert_array_equal(trb.eps, exp_eps) + npt.assert_array_compare(np.less_equal, search_space.lower, trb.location) + npt.assert_array_compare(np.less_equal, trb.location, search_space.upper) + npt.assert_array_compare(np.less_equal, search_space.lower, trb.lower) + npt.assert_array_compare(np.less_equal, trb.upper, search_space.upper) + npt.assert_array_compare(np.less_equal, trb.upper - trb.lower, 2 * exp_eps) + npt.assert_array_equal(trb._y_min, tf.constant([np.inf], dtype=tf.float64)) + + +# Update call initializes the box if eps is smaller than min_eps. +def test_trust_region_box_update_initialize() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + datasets = { + OBJECTIVE: Dataset( # Points outside the search space should be ignored. + tf.constant([[1.2, 1.3], [-0.4, -0.5]], dtype=tf.float64), + tf.constant([[0.7], [0.9]], dtype=tf.float64), + ) + } + trb = SingleObjectiveTrustRegionBox(search_space, min_eps=0.5) + trb.initialize(datasets=datasets) + location = trb.location + + trb.update(datasets=datasets) + npt.assert_array_compare(np.not_equal, location, trb.location) + location = trb.location + + trb.update(datasets=datasets) + npt.assert_array_compare(np.not_equal, location, trb.location) + + +# Update call does not initialize the box if eps is larger than min_eps. +def test_trust_region_box_update_no_initialize() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + datasets = { + OBJECTIVE: Dataset( + tf.constant([[0.5, 0.5], [0.0, 0.0], [1.0, 1.0]], dtype=tf.float64), + tf.constant([[0.5], [0.0], [1.0]], dtype=tf.float64), + ) + } + trb = SingleObjectiveTrustRegionBox(search_space, min_eps=0.1) + trb.initialize(datasets=datasets) + trb.location = tf.constant([0.5, 0.5], dtype=tf.float64) + location = trb.location + + trb.update(datasets=datasets) + npt.assert_array_equal(location, trb.location) + + +# Update shrinks/expands box on successful/unsuccessful step. +@pytest.mark.parametrize("success", [True, False]) +def test_trust_region_box_update_size(success: bool) -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + datasets = { + OBJECTIVE: Dataset( + tf.constant([[0.5, 0.5], [0.0, 0.0], [1.0, 1.0]], dtype=tf.float64), + tf.constant([[0.5], [0.3], [1.0]], dtype=tf.float64), + ) + } + trb = SingleObjectiveTrustRegionBox(search_space, min_eps=0.1) + trb.initialize(datasets=datasets) + + # Ensure there is at least one point captured in the box. + orig_point = trb.sample(1) + orig_min = tf.constant([[0.1]], dtype=tf.float64) + datasets[OBJECTIVE] = Dataset( + np.concatenate([datasets[OBJECTIVE].query_points, orig_point], axis=0), + np.concatenate([datasets[OBJECTIVE].observations, orig_min], axis=0), + ) + trb.update(datasets=datasets) + + eps = trb.eps + + if success: + # Sample a point from the box. + new_point = trb.sample(1) + else: + # Pick point outside the box. + new_point = tf.constant([[1.2, 1.3]], dtype=tf.float64) + + # Add a new min point to the dataset. + new_min = tf.constant([[-0.1]], dtype=tf.float64) + datasets[OBJECTIVE] = Dataset( + np.concatenate([datasets[OBJECTIVE].query_points, new_point], axis=0), + np.concatenate([datasets[OBJECTIVE].observations, new_min], axis=0), + ) + # Update the box. + trb.update(datasets=datasets) + + if success: + # Check that the location is the new min point. + new_point = np.squeeze(new_point) + npt.assert_allclose(new_point, trb.location) + npt.assert_allclose(new_min, trb._y_min) + # Check that the box is larger by beta. + npt.assert_allclose(eps / trb._beta, trb.eps) + else: + # Check that the location is the old min point. + orig_point = np.squeeze(orig_point) + npt.assert_allclose(orig_point, trb.location) + npt.assert_allclose(orig_min, trb._y_min) + # Check that the box is smaller by beta. + npt.assert_allclose(eps * trb._beta, trb.eps) + + # Check the new box bounds. + npt.assert_allclose(trb.lower, np.maximum(trb.location - trb.eps, search_space.lower)) + npt.assert_allclose(trb.upper, np.minimum(trb.location + trb.eps, search_space.upper)) + + +# When state is None, acquire returns a multi search space of the correct type. +def test_multi_trust_region_box_acquire_no_state() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + dataset = Dataset( + tf.constant([[0.5, 0.5], [0.0, 0.0], [1.0, 1.0]], dtype=tf.float64), + tf.constant([[0.5], [0.0], [1.0]], dtype=tf.float64), + ) + datasets = {OBJECTIVE: dataset} + model = QuadraticMeanAndRBFKernelWithSamplers( + dataset=dataset, noise_variance=tf.constant(1.0, dtype=tf.float64) + ) + model.kernel = ( + gpflow.kernels.RBF() + ) # need a gpflow kernel object for random feature decompositions + models = {OBJECTIVE: model} + base_rule = EfficientGlobalOptimization( # type: ignore[var-annotated] + builder=ParallelContinuousThompsonSampling(), num_query_points=2 + ) + subspaces = [ + SingleObjectiveTrustRegionBox(search_space, beta=0.1, kappa=1e-3, min_eps=1e-1) + for _ in range(2) + ] + mtb = BatchTrustRegionBox(subspaces, base_rule) + state, points = mtb.acquire(search_space, models, datasets)(None) + + assert state is not None + assert isinstance(state.acquisition_space, TaggedMultiSearchSpace) + assert len(state.acquisition_space.subspace_tags) == 2 + + for index, (tag, point) in enumerate(zip(state.acquisition_space.subspace_tags, points)): + subspace = state.acquisition_space.get_subspace(tag) + assert subspace == subspaces[index] + assert isinstance(subspace, SingleObjectiveTrustRegionBox) + assert subspace.global_search_space == search_space + assert subspace._beta == 0.1 + assert subspace._kappa == 1e-3 + assert subspace._min_eps == 1e-1 + assert point in subspace + + +def test_multi_trust_region_box_raises_on_mismatched_tags() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + dataset = Dataset( + tf.constant([[0.0, 0.0], [1.0, 1.0]], dtype=tf.float64), + tf.constant([[0.0], [1.0]], dtype=tf.float64), + ) + base_rule = EfficientGlobalOptimization( # type: ignore[var-annotated] + builder=ParallelContinuousThompsonSampling(), num_query_points=2 + ) + subspaces = [SingleObjectiveTrustRegionBox(search_space) for _ in range(2)] + mtb = BatchTrustRegionBox(subspaces, base_rule) + + state = BatchTrustRegionBox.State( + acquisition_space=TaggedMultiSearchSpace(subspaces, tags=["a", "b"]) + ) + state_func = mtb.acquire( + search_space, + {OBJECTIVE: QuadraticMeanAndRBFKernelWithSamplers(dataset)}, + {OBJECTIVE: dataset}, + ) + + with pytest.raises(AssertionError, match="The tags of the state acquisition space"): + _, _ = state_func(state) + + +class TestTrustRegionBox(SingleObjectiveTrustRegionBox): + def __init__( + self, + fixed_location: TensorType, + global_search_space: SearchSpace, + beta: float = 0.7, + kappa: float = 1e-4, + min_eps: float = 1e-2, + ): + super().__init__(global_search_space, beta, kappa, min_eps) + self._location = fixed_location + + @property + def location(self) -> TensorType: + return self._location + + @location.setter + def location(self, location: TensorType) -> None: + ... + + def _init_eps(self) -> None: + self.eps = tf.constant(0.07, dtype=tf.float64) + + +# Start with a defined state and dataset. Acquire should return an updated state. +def test_multi_trust_region_box_acquire_with_state() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + init_dataset = Dataset( + tf.constant([[0.25, 0.25], [0.5, 0.5], [0.75, 0.75]], dtype=tf.float64), + tf.constant([[1.0], [1.0], [1.0]], dtype=tf.float64), + ) + model = QuadraticMeanAndRBFKernelWithSamplers( + dataset=init_dataset, noise_variance=tf.constant(1e-6, dtype=tf.float64) + ) + model.kernel = ( + gpflow.kernels.RBF() + ) # need a gpflow kernel object for random feature decompositions + models = {OBJECTIVE: model} + base_rule = EfficientGlobalOptimization( # type: ignore[var-annotated] + builder=ParallelContinuousThompsonSampling(), num_query_points=3 + ) + + # Third region is close to the first. + subspaces = [ + TestTrustRegionBox(tf.constant([0.3, 0.3], dtype=tf.float64), search_space), + TestTrustRegionBox(tf.constant([0.7, 0.7], dtype=tf.float64), search_space), + TestTrustRegionBox(tf.constant([0.3, 0.3], dtype=tf.float64) + 1e-7, search_space), + ] + mtb = BatchTrustRegionBox(subspaces, base_rule) + state = BatchTrustRegionBox.State(acquisition_space=TaggedMultiSearchSpace(subspaces)) + for subspace in subspaces: + subspace.initialize(datasets={OBJECTIVE: init_dataset}) + + dataset = Dataset( + init_dataset.query_points, + tf.constant([[0.1], [0.2], [0.3]], dtype=tf.float64), + ) + state_func = mtb.acquire(search_space, models, {OBJECTIVE: dataset}) + next_state, points = state_func(state) + + assert next_state is not None + assert len(points) == 3 + # The regions correspond to first, third and first points in the dataset. + # First two regions should be updated. + # The third region should be initialized and not updated, as it is too close to the first + # subspace. + for point, subspace, exp_obs, exp_eps in zip( + points, + subspaces, + [dataset.observations[0], dataset.observations[2], dataset.observations[0]], + [0.1, 0.1, 0.07], # First two regions updated, third region initialized. + ): + assert point in subspace + npt.assert_array_equal(subspace._y_min, exp_obs) + # Check the box was updated/initialized correctly. + npt.assert_allclose(subspace.eps, exp_eps) + + +def test_multi_trust_region_box_state_deepcopy() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + dataset = Dataset( + tf.constant([[0.25, 0.25], [0.5, 0.5], [0.75, 0.75]], dtype=tf.float64), + tf.constant([[1.0], [1.0], [1.0]], dtype=tf.float64), + ) + subspaces = [SingleObjectiveTrustRegionBox(search_space, 0.07, 1e-5, 1e-3) for _ in range(3)] + for _subspace in subspaces: + _subspace.initialize(datasets={OBJECTIVE: dataset}) + state = BatchTrustRegionBox.State(acquisition_space=TaggedMultiSearchSpace(subspaces)) + + state_copy = copy.deepcopy(state) + assert state_copy is not state + assert state_copy.acquisition_space is not state.acquisition_space + assert state_copy.acquisition_space._spaces is not state.acquisition_space._spaces + assert state_copy.acquisition_space.subspace_tags == state.acquisition_space.subspace_tags + + for subspace, subspace_copy in zip( + state.acquisition_space._spaces.values(), state_copy.acquisition_space._spaces.values() + ): + assert subspace is not subspace_copy + assert isinstance(subspace, SingleObjectiveTrustRegionBox) + assert isinstance(subspace_copy, SingleObjectiveTrustRegionBox) + assert subspace._beta == subspace_copy._beta + assert subspace._kappa == subspace_copy._kappa + assert subspace._min_eps == subspace_copy._min_eps + npt.assert_array_equal(subspace.eps, subspace_copy.eps) + npt.assert_array_equal(subspace.location, subspace_copy.location) + npt.assert_array_equal(subspace._y_min, subspace_copy._y_min) + + def test_asynchronous_rule_state_pending_points() -> None: pending_points = tf.constant([[1], [2], [3]]) diff --git a/tests/unit/acquisition/test_utils.py b/tests/unit/acquisition/test_utils.py index 13fad5595a..7975cab3a0 100644 --- a/tests/unit/acquisition/test_utils.py +++ b/tests/unit/acquisition/test_utils.py @@ -23,6 +23,7 @@ from trieste.acquisition import AcquisitionFunction from trieste.acquisition.utils import ( get_local_dataset, + get_unique_points_mask, select_nth_output, split_acquisition_function, ) @@ -97,3 +98,35 @@ def test_get_local_dataset_works() -> None: assert tf.shape(get_local_dataset(search_space_1, combined).query_points)[0] == 10 assert tf.shape(get_local_dataset(search_space_2, combined).query_points)[0] == 20 + + +@pytest.mark.parametrize( + "points, tolerance, expected_mask", + [ + ( + tf.constant([[1.0, 1.0], [1.2, 1.1], [2.0, 2.0], [2.2, 2.2], [3.0, 3.0]]), + 0.5, + tf.constant([True, False, True, False, True]), + ), + ( + tf.constant([[1.0, 2.0], [2.0, 3.0], [1.0, 2.1]]), + 0.2, + tf.constant([True, True, False]), + ), + ( + tf.constant([[1.0], [2.0], [1.0], [3.0], [1.71], [1.699999], [3.29], [3.300001]]), + 0.3, + tf.constant([True, True, False, True, False, True, False, True]), + ), + ( + tf.constant([[1.0], [2.0], [1.0], [3.0], [1.699999], [1.71], [3.300001], [3.29]]), + 0.3, + tf.constant([True, True, False, True, True, False, True, False]), + ), + ], +) +def test_get_unique_points_mask( + points: tf.Tensor, tolerance: float, expected_mask: tf.Tensor +) -> None: + mask = get_unique_points_mask(points, tolerance) + np.testing.assert_array_equal(mask, expected_mask) diff --git a/tests/unit/test_space.py b/tests/unit/test_space.py index f94aa67ddc..638093382f 100644 --- a/tests/unit/test_space.py +++ b/tests/unit/test_space.py @@ -17,7 +17,7 @@ import itertools import operator from functools import reduce -from typing import Container, Optional, Sequence +from typing import Any, Container, List, Optional, Sequence, Type import numpy.testing as npt import pytest @@ -27,11 +27,13 @@ from tests.util.misc import TF_DEBUGGING_ERROR_TYPES, ShapeLike, various_shapes from trieste.space import ( Box, + CollectionSearchSpace, Constraint, DiscreteSearchSpace, LinearConstraint, NonlinearConstraint, SearchSpace, + TaggedMultiSearchSpace, TaggedProductSearchSpace, ) from trieste.types import TensorType @@ -681,36 +683,357 @@ def test_box_deepcopy() -> None: npt.assert_allclose(box.upper, box_copy.upper) -def test_product_space_raises_for_non_unqique_subspace_names() -> None: +@pytest.fixture(name="search_space_type", params=[TaggedMultiSearchSpace, TaggedProductSearchSpace]) +def _default_search_space_type_fixture(request: Any) -> Type[CollectionSearchSpace]: + return request.param + + +def test_collection_space_raises_for_non_unqique_subspace_names( + search_space_type: Type[CollectionSearchSpace], +) -> None: space_A = Box([-1, -2], [2, 3]) space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) - with pytest.raises(TF_DEBUGGING_ERROR_TYPES): - TaggedProductSearchSpace(spaces=[space_A, space_B], tags=["A", "A"]) + with pytest.raises(TF_DEBUGGING_ERROR_TYPES, match="Subspace names must be unique"): + search_space_type(spaces=[space_A, space_B], tags=["A", "A"]) -def test_product_space_raises_for_length_mismatch_between_spaces_and_tags() -> None: +def test_collection_space_raises_for_length_mismatch_between_spaces_and_tags( + search_space_type: Type[CollectionSearchSpace], +) -> None: space_A = Box([-1, -2], [2, 3]) space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) - with pytest.raises(TF_DEBUGGING_ERROR_TYPES): - TaggedProductSearchSpace(spaces=[space_A, space_B], tags=["A", "B", "C"]) + with pytest.raises( + TF_DEBUGGING_ERROR_TYPES, match="Number of tags must match number of subspaces" + ): + search_space_type(spaces=[space_A, space_B], tags=["A", "B", "C"]) -def test_product_space_subspace_tags_attribute() -> None: +def test_collection_space_subspace_tags_attribute( + search_space_type: Type[CollectionSearchSpace], +) -> None: decision_space = Box([-1, -2], [2, 3]) context_space = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) - product_space = TaggedProductSearchSpace( + multi_space = search_space_type( spaces=[context_space, decision_space], tags=["context", "decision"] ) - npt.assert_array_equal(product_space.subspace_tags, ["context", "decision"]) + npt.assert_array_equal(multi_space.subspace_tags, ["context", "decision"]) -def test_product_space_subspace_tags_default_behaviour() -> None: +def test_collection_space_subspace_tags_default_behaviour( + search_space_type: Type[CollectionSearchSpace], +) -> None: decision_space = Box([-1, -2], [2, 3]) context_space = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) - product_space = TaggedProductSearchSpace(spaces=[context_space, decision_space]) + multi_space = search_space_type(spaces=[context_space, decision_space]) + + npt.assert_array_equal(multi_space.subspace_tags, ["0", "1"]) + + +def test_collection_space_get_subspace_raises_for_invalid_tag( + search_space_type: Type[CollectionSearchSpace], +) -> None: + space_A = Box([-1, -2], [2, 3]) + space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) + multi_space = search_space_type(spaces=[space_A, space_B], tags=["A", "B"]) + + with pytest.raises( + TF_DEBUGGING_ERROR_TYPES, match="Attempted to access a subspace that does not exist" + ): + multi_space.get_subspace("dummy") + + +def test_collection_space_get_subspace(search_space_type: Type[CollectionSearchSpace]) -> None: + space_A = Box([-1, -2], [2, 3]) + space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) + space_C = Box([-1, -3], [2, 2]) + multi_space = search_space_type(spaces=[space_A, space_B, space_C], tags=["A", "B", "C"]) - npt.assert_array_equal(product_space.subspace_tags, ["0", "1"]) + subspace_A = multi_space.get_subspace("A") + assert isinstance(subspace_A, Box) + npt.assert_array_equal(subspace_A.lower, [-1, -2]) + npt.assert_array_equal(subspace_A.upper, [2, 3]) + + subspace_B = multi_space.get_subspace("B") + assert isinstance(subspace_B, DiscreteSearchSpace) + npt.assert_array_equal(subspace_B.points, tf.constant([[-0.5, 0.5]])) + + subspace_C = multi_space.get_subspace("C") + assert isinstance(subspace_C, Box) + npt.assert_array_equal(subspace_C.lower, [-1, -3]) + npt.assert_array_equal(subspace_C.upper, [2, 2]) + + +@pytest.mark.parametrize( + "search_space_type, point", + [ + (TaggedMultiSearchSpace, tf.constant([-1.0, 0.0], dtype=tf.float64)), + (TaggedMultiSearchSpace, tf.constant([2.0, -2.0], dtype=tf.float64)), + (TaggedMultiSearchSpace, tf.constant([-0.5, 0.5], dtype=tf.float64)), + (TaggedProductSearchSpace, tf.constant([-1.0, 0.0, -0.5, 0.5], dtype=tf.float64)), + (TaggedProductSearchSpace, tf.constant([2.0, -2.0, -0.5, 0.5], dtype=tf.float64)), + ], +) +def test_collection_space_contains_point( + search_space_type: Type[CollectionSearchSpace], + point: tf.Tensor, +) -> None: + space_A = Box([-1.0, -2.0], [2.0, 0.0]) + space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]], dtype=tf.float64)) + collection_space = search_space_type(spaces=[space_A, space_B]) + assert point in collection_space + assert collection_space.contains(point) + + +@pytest.mark.parametrize( + "search_space_type, point", + [ + pytest.param( + TaggedMultiSearchSpace, + tf.constant([-1.1, 0.0], dtype=tf.float64), + id="just outside box", + ), + pytest.param( + TaggedMultiSearchSpace, + tf.constant([-10, 10.0], dtype=tf.float64), + id="well outside box", + ), + pytest.param( + TaggedMultiSearchSpace, + tf.constant([-0.5, 0.51], dtype=tf.float64), + id="just outside discrete", + ), + pytest.param( + TaggedProductSearchSpace, + tf.constant([-1.1, 0.0, -0.5, 0.5], dtype=tf.float64), + id="just outside context space", + ), + pytest.param( + TaggedProductSearchSpace, + tf.constant([-10, 10.0, -0.5, 0.5], dtype=tf.float64), + id="well outside context space", + ), + pytest.param( + TaggedProductSearchSpace, + tf.constant([2.0, 0.0, 2.0, 7.0], dtype=tf.float64), + id="outside decision space", + ), + pytest.param( + TaggedProductSearchSpace, + tf.constant([-10.0, -10.0, -10.0, -10.0], dtype=tf.float64), + id="outside both", + ), + pytest.param( + TaggedProductSearchSpace, + tf.constant([-0.5, 0.5, 1.0, 0.0], dtype=tf.float64), + id="swap order of components", + ), + ], +) +def test_collection_space_does_not_contain_point( + search_space_type: Type[CollectionSearchSpace], + point: tf.Tensor, +) -> None: + space_A = Box([-1.0, -2.0], [2.0, 0.0]) + space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]], dtype=tf.float64)) + collection_space = search_space_type(spaces=[space_A, space_B]) + assert point not in collection_space + assert not collection_space.contains(point) + + +@pytest.mark.parametrize( + "search_space_type, points", + [ + (TaggedMultiSearchSpace, tf.constant([[-1.1, 0.0], [-0.5, 0.5]], dtype=tf.float64)), + ( + TaggedProductSearchSpace, + tf.constant([[-1.1, 0.0, -0.5, 0.5], [-1.0, 0.0, -0.5, 0.5]], dtype=tf.float64), + ), + ], +) +def test_collection_space_contains_broadcasts( + search_space_type: Type[CollectionSearchSpace], + points: tf.Tensor, +) -> None: + space_A = Box([-1.0, -2.0], [2.0, 3.0]) + space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]], dtype=tf.float64)) + collection_space = search_space_type(spaces=[space_A, space_B]) + tf.assert_equal(collection_space.contains(points), [False, True]) + # point in space raises (because python insists on a bool) + with pytest.raises(TF_DEBUGGING_ERROR_TYPES): + _ = points in collection_space + + +@pytest.mark.parametrize( + "spaces", + [ + [DiscreteSearchSpace(tf.constant([[-0.5]]))], + [ + DiscreteSearchSpace(tf.constant([[-0.5, 1.3]])), + DiscreteSearchSpace(tf.constant([[-0.5, -0.3], [1.2, 0.4]])), + ], + [ + Box([-2], [3]), + DiscreteSearchSpace(tf.constant([[-0.5]])), + Box([-1], [2]), + ], + ], +) +def test_collection_space_contains_raises_on_point_of_different_shape( + search_space_type: Type[CollectionSearchSpace], + spaces: Sequence[SearchSpace], +) -> None: + space = search_space_type(spaces=spaces) + dimension = space.dimension + for wrong_input_shape in [dimension - 1, dimension + 1]: + point = tf.zeros([wrong_input_shape]) + with pytest.raises(TF_DEBUGGING_ERROR_TYPES): + _ = point in space + with pytest.raises(TF_DEBUGGING_ERROR_TYPES): + _ = space.contains(point) + + +@pytest.mark.parametrize( + "search_space_type, space_A, exp_shape_tail", + [ + (TaggedMultiSearchSpace, Box([-1, -2], [2, 3]), [2, 2]), + (TaggedProductSearchSpace, Box([-1], [2]), [3]), + ], +) +@pytest.mark.parametrize("num_samples", [0, 1, 10]) +def test_collection_space_sampling_returns_correct_shape( + search_space_type: Type[CollectionSearchSpace], + space_A: SearchSpace, + exp_shape_tail: List[int], + num_samples: int, +) -> None: + space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) + if search_space_type is TaggedMultiSearchSpace: + product: SearchSpace = TaggedMultiSearchSpace([space_A]) * TaggedMultiSearchSpace([space_B]) + else: + product = space_A * space_B + for collection_space in (search_space_type(spaces=[space_A, space_B]), product): + samples = collection_space.sample(num_samples) + npt.assert_array_equal(tf.shape(samples), [num_samples] + exp_shape_tail) + + +@pytest.mark.parametrize( + "search_space_type, space_A", + [ + (TaggedMultiSearchSpace, Box([-1, -2], [2, 3])), + (TaggedProductSearchSpace, Box([-1], [2])), + ], +) +@pytest.mark.parametrize("num_samples", [-1, -10]) +def test_collection_space_sampling_raises_for_invalid_sample_size( + search_space_type: Type[CollectionSearchSpace], + space_A: SearchSpace, + num_samples: int, +) -> None: + space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) + collection_space = search_space_type(spaces=[space_A, space_B]) + with pytest.raises(TF_DEBUGGING_ERROR_TYPES): + collection_space.sample(num_samples) + + +@pytest.mark.parametrize( + "search_space_type, space_A", + [ + (TaggedMultiSearchSpace, Box([-1, -2], [2, 3])), + (TaggedProductSearchSpace, Box([-1], [2])), + ], +) +@pytest.mark.parametrize("num_samples", [0, 1, 10]) +def test_collection_space_discretize_returns_search_space_with_only_points_contained_within_box( + search_space_type: Type[CollectionSearchSpace], + space_A: SearchSpace, + num_samples: int, +) -> None: + space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) + collection_space = search_space_type(spaces=[space_A, space_B]) + + dss = collection_space.discretize(num_samples) + samples = dss.sample(num_samples) + + assert all(sample in collection_space for sample in samples) + + +@pytest.mark.parametrize( + "search_space_type, space_A", + [ + (TaggedMultiSearchSpace, Box([-1, -2], [2, 3])), + (TaggedProductSearchSpace, Box([-1], [2])), + ], +) +@pytest.mark.parametrize("num_samples", [0, 1, 10]) +def test_collection_space_discretize_returns_search_space_with_correct_number_of_points( + search_space_type: Type[CollectionSearchSpace], + space_A: SearchSpace, + num_samples: int, +) -> None: + space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) + collection_space = search_space_type(spaces=[space_A, space_B]) + + dss = collection_space.discretize(num_samples) + samples = dss.sample(num_samples) + + assert len(samples) == num_samples + + +@pytest.mark.parametrize( + "search_space_type, space_A", + [ + (TaggedMultiSearchSpace, Box([-1, -2], [2, 3])), + (TaggedProductSearchSpace, Box([-1], [2])), + ], +) +@pytest.mark.parametrize("seed", [1, 42, 123]) +def test_collection_space_sampling_returns_same_points_for_same_seed( + search_space_type: Type[CollectionSearchSpace], space_A: SearchSpace, seed: int +) -> None: + space_B = DiscreteSearchSpace(tf.random.uniform([100, 2], dtype=tf.float64, seed=42)) + collection_space = search_space_type(spaces=[space_A, space_B]) + random_samples_1 = collection_space.sample(num_samples=100, seed=seed) + random_samples_2 = collection_space.sample(num_samples=100, seed=seed) + npt.assert_allclose(random_samples_1, random_samples_2) + + +@pytest.mark.parametrize( + "search_space_type, space_A", + [ + (TaggedMultiSearchSpace, Box([-1, -2], [2, 3])), + (TaggedProductSearchSpace, Box([-1], [2])), + ], +) +def test_collection_space_sampling_returns_different_points_for_different_call( + search_space_type: Type[CollectionSearchSpace], + space_A: SearchSpace, +) -> None: + space_B = DiscreteSearchSpace(tf.random.uniform([100, 2], dtype=tf.float64, seed=42)) + collection_space = search_space_type(spaces=[space_A, space_B]) + random_samples_1 = collection_space.sample(num_samples=100) + random_samples_2 = collection_space.sample(num_samples=100) + npt.assert_raises(AssertionError, npt.assert_allclose, random_samples_1, random_samples_2) + + +@pytest.mark.parametrize( + "search_space_type, space_A", + [ + (TaggedMultiSearchSpace, Box([-1, -2], [2, 3])), + (TaggedProductSearchSpace, Box([-1], [2])), + ], +) +def test_collection_space_deepcopy( + search_space_type: Type[CollectionSearchSpace], + space_A: SearchSpace, +) -> None: + space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) + collection_space = search_space_type(spaces=[space_A, space_B], tags=["A", "B"]) + + copied_space = copy.deepcopy(collection_space) + npt.assert_allclose(copied_space.get_subspace("A").lower, space_A.lower) + npt.assert_allclose(copied_space.get_subspace("A").upper, space_A.upper) + npt.assert_allclose(copied_space.get_subspace("B").points, space_B.points) # type: ignore @pytest.mark.parametrize( @@ -722,7 +1045,7 @@ def test_product_space_subspace_tags_default_behaviour() -> None: ([Box([-1, -2], [2, 3]), Box([-1, -2], [2, 3]), Box([-1], [2])], 5), ], ) -def test_product_search_space_returns_correct_dimension( +def test_product_space_returns_correct_dimension( spaces: Sequence[SearchSpace], dimension: int ) -> None: for space in (TaggedProductSearchSpace(spaces=spaces), reduce(operator.mul, spaces)): @@ -762,38 +1085,6 @@ def test_product_space_returns_correct_bounds( npt.assert_array_equal(space.upper, upper) -def test_product_space_get_subspace_raises_for_invalid_tag() -> None: - space_A = Box([-1, -2], [2, 3]) - space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) - product_space = TaggedProductSearchSpace(spaces=[space_A, space_B], tags=["A", "B"]) - - with pytest.raises(TF_DEBUGGING_ERROR_TYPES): - product_space.get_subspace("dummy") - - -def test_product_space_get_subspace() -> None: - space_A = Box([-1, -2], [2, 3]) - space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) - space_C = Box([-1], [2]) - product_space = TaggedProductSearchSpace( - spaces=[space_A, space_B, space_C], tags=["A", "B", "C"] - ) - - subspace_A = product_space.get_subspace("A") - assert isinstance(subspace_A, Box) - npt.assert_array_equal(subspace_A.lower, [-1, -2]) - npt.assert_array_equal(subspace_A.upper, [2, 3]) - - subspace_B = product_space.get_subspace("B") - assert isinstance(subspace_B, DiscreteSearchSpace) - npt.assert_array_equal(subspace_B.points, tf.constant([[-0.5, 0.5]])) - - subspace_C = product_space.get_subspace("C") - assert isinstance(subspace_C, Box) - npt.assert_array_equal(subspace_C.lower, [-1]) - npt.assert_array_equal(subspace_C.upper, [2]) - - @pytest.mark.parametrize( "points", [ @@ -879,143 +1170,6 @@ def test_product_space_can_get_subspace_components( npt.assert_array_equal(space.get_subspace_component(tag, points), subspace_points) -@pytest.mark.parametrize( - "point", - [ - tf.constant([-1.0, 0.0, -0.5, 0.5], dtype=tf.float64), - tf.constant([2.0, 3.0, -0.5, 0.5], dtype=tf.float64), - ], -) -def test_product_space_contains_point(point: tf.Tensor) -> None: - space_A = Box([-1.0, -2.0], [2.0, 3.0]) - space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]], dtype=tf.float64)) - product_space = TaggedProductSearchSpace(spaces=[space_A, space_B]) - assert point in product_space - assert product_space.contains(point) - - -@pytest.mark.parametrize( - "point", - [ - tf.constant([-1.1, 0.0, -0.5, 0.5], dtype=tf.float64), # just outside context space - tf.constant([-10, 10.0, -0.5, 0.5], dtype=tf.float64), # well outside context space - tf.constant([2.0, 3.0, 2.0, 7.0], dtype=tf.float64), # outside decision space - tf.constant([-10.0, -10.0, -10.0, -10.0], dtype=tf.float64), # outside both - tf.constant([-0.5, 0.5, 1.0, 2.0], dtype=tf.float64), # swap order of components - ], -) -def test_product_space_does_not_contain_point(point: tf.Tensor) -> None: - space_A = Box([-1.0, -2.0], [2.0, 3.0]) - space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]], dtype=tf.float64)) - product_space = TaggedProductSearchSpace(spaces=[space_A, space_B]) - assert point not in product_space - assert not product_space.contains(point) - - -def test_product_space_contains_broadcasts() -> None: - space_A = Box([-1.0, -2.0], [2.0, 3.0]) - space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]], dtype=tf.float64)) - product_space = TaggedProductSearchSpace(spaces=[space_A, space_B]) - points = tf.constant([[-1.1, 0.0, -0.5, 0.5], [-1.0, 0.0, -0.5, 0.5]], dtype=tf.float64) - tf.assert_equal(product_space.contains(points), [False, True]) - # point in space raises (because python insists on a bool) - with pytest.raises(TF_DEBUGGING_ERROR_TYPES): - _ = points in product_space - - -@pytest.mark.parametrize( - "spaces", - [ - [DiscreteSearchSpace(tf.constant([[-0.5]]))], - [ - DiscreteSearchSpace(tf.constant([[-0.5]])), - DiscreteSearchSpace(tf.constant([[-0.5, -0.3], [1.2, 0.4]])), - ], - [ - Box([-1, -2], [2, 3]), - DiscreteSearchSpace(tf.constant([[-0.5]])), - Box([-1], [2]), - ], - ], -) -def test_product_space_contains_raises_on_point_of_different_shape( - spaces: Sequence[SearchSpace], -) -> None: - space = TaggedProductSearchSpace(spaces=spaces) - dimension = space.dimension - for wrong_input_shape in [dimension - 1, dimension + 1]: - point = tf.zeros([wrong_input_shape]) - with pytest.raises(TF_DEBUGGING_ERROR_TYPES): - _ = point in space - with pytest.raises(TF_DEBUGGING_ERROR_TYPES): - _ = space.contains(point) - - -@pytest.mark.parametrize("num_samples", [0, 1, 10]) -def test_product_space_sampling_returns_correct_shape(num_samples: int) -> None: - space_A = Box([-1], [2]) - space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) - for product_space in (TaggedProductSearchSpace(spaces=[space_A, space_B]), space_A * space_B): - samples = product_space.sample(num_samples) - npt.assert_array_equal(tf.shape(samples), [num_samples, 3]) - - -@pytest.mark.parametrize("num_samples", [-1, -10]) -def test_product_space_sampling_raises_for_invalid_sample_size(num_samples: int) -> None: - space_A = Box([-1], [2]) - space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) - product_space = TaggedProductSearchSpace(spaces=[space_A, space_B]) - with pytest.raises(TF_DEBUGGING_ERROR_TYPES): - product_space.sample(num_samples) - - -@pytest.mark.parametrize("num_samples", [0, 1, 10]) -def test_product_space_discretize_returns_search_space_with_only_points_contained_within_box( - num_samples: int, -) -> None: - space_A = Box([-1], [2]) - space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) - product_space = TaggedProductSearchSpace(spaces=[space_A, space_B]) - - dss = product_space.discretize(num_samples) - samples = dss.sample(num_samples) - - assert all(sample in product_space for sample in samples) - - -@pytest.mark.parametrize("num_samples", [0, 1, 10]) -def test_product_space_discretize_returns_search_space_with_correct_number_of_points( - num_samples: int, -) -> None: - space_A = Box([-1], [2]) - space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) - product_space = TaggedProductSearchSpace(spaces=[space_A, space_B]) - - dss = product_space.discretize(num_samples) - samples = dss.sample(num_samples) - - assert len(samples) == num_samples - - -@pytest.mark.parametrize("seed", [1, 42, 123]) -def test_product_space_sampling_returns_same_points_for_same_seed(seed: int) -> None: - space_A = Box([-1], [2]) - space_B = DiscreteSearchSpace(tf.random.uniform([100, 2], dtype=tf.float64, seed=42)) - product_space = TaggedProductSearchSpace(spaces=[space_A, space_B]) - random_samples_1 = product_space.sample(num_samples=100, seed=seed) - random_samples_2 = product_space.sample(num_samples=100, seed=seed) - npt.assert_allclose(random_samples_1, random_samples_2) - - -def test_product_space_sampling_returns_different_points_for_different_call() -> None: - space_A = Box([-1], [2]) - space_B = DiscreteSearchSpace(tf.random.uniform([100, 2], dtype=tf.float64, seed=42)) - product_space = TaggedProductSearchSpace(spaces=[space_A, space_B]) - random_samples_1 = product_space.sample(num_samples=100) - random_samples_2 = product_space.sample(num_samples=100) - npt.assert_raises(AssertionError, npt.assert_allclose, random_samples_1, random_samples_2) - - def test_product_space___mul___() -> None: space_A = Box([-1], [2]) space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) @@ -1046,17 +1200,6 @@ def test_product_space___mul___() -> None: npt.assert_array_equal(subspace_1_D.points, tf.ones([5, 3], dtype=tf.float64)) -def test_product_search_space_deepcopy() -> None: - space_A = Box([-1], [2]) - space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) - product_space = TaggedProductSearchSpace(spaces=[space_A, space_B], tags=["A", "B"]) - - copied_space = copy.deepcopy(product_space) - npt.assert_allclose(copied_space.get_subspace("A").lower, space_A.lower) - npt.assert_allclose(copied_space.get_subspace("A").upper, space_A.upper) - npt.assert_allclose(copied_space.get_subspace("B").points, space_B.points) # type: ignore - - def test_product_space_handles_empty_spaces() -> None: space_A = Box([-1, -2], [2, 3]) tag_A = TaggedProductSearchSpace(spaces=[space_A], tags=["A"]) @@ -1069,6 +1212,122 @@ def test_product_space_handles_empty_spaces() -> None: npt.assert_array_equal(tag_C.subspace_tags, ["AA", "BB"]) +def test_multi_space_raises_for_empty_collection() -> None: + with pytest.raises(TF_DEBUGGING_ERROR_TYPES, match="At least one subspace is required"): + TaggedMultiSearchSpace([]) + + +@pytest.mark.parametrize( + "spaces", + [ + ([Box([-1], [2]), Box([-1, -2], [2, 3])]), + ([Box([-1, -2], [2, 3]), DiscreteSearchSpace(tf.constant([[-0.5]]))]), + ], +) +def test_multi_space_raises_for_mixed_dimensions(spaces: Sequence[SearchSpace]) -> None: + with pytest.raises( + TF_DEBUGGING_ERROR_TYPES, match="All subspaces must have the same dimension" + ): + TaggedMultiSearchSpace(spaces) + + +@pytest.mark.parametrize( + "spaces, dimension", + [ + ([DiscreteSearchSpace(tf.constant([[-0.5, -0.3], [1.2, 0.4]]))], 2), + ([DiscreteSearchSpace(tf.constant([[-0.5]])), Box([-1], [2])], 1), + ([Box([-1, -2], [2, 3]), DiscreteSearchSpace(tf.constant([[-0.5, -1.5]]))], 2), + ([Box([-1, -2], [2, 3]), Box([-1, -2], [2, 3]), Box([-1, 1], [2, 2])], 2), + ], +) +def test_multi_space_returns_correct_dimension( + spaces: Sequence[SearchSpace], dimension: int +) -> None: + for space in ( + TaggedMultiSearchSpace(spaces=spaces), + reduce(operator.mul, [TaggedMultiSearchSpace([s]) for s in spaces]), + ): + assert space.dimension == dimension + + +@pytest.mark.parametrize( + "spaces, lower, upper", + [ + ( + [DiscreteSearchSpace(tf.constant([[-0.5, 0.4], [1.2, -0.3]]))], + tf.constant([[-0.5, -0.3]]), + tf.constant([[1.2, 0.4]]), + ), + ( + [DiscreteSearchSpace(tf.constant([[-0.5]], dtype=tf.float64)), Box([-1.0], [2.0])], + tf.constant([[-0.5], [-1.0]]), + tf.constant([[-0.5], [2.0]]), + ), + ( + [ + Box([-1, -2], [2, 3]), + DiscreteSearchSpace(tf.constant([[-0.5, 1.5]], dtype=tf.float64)), + ], + tf.constant([[-1.0, -2.0], [-0.5, 1.5]]), + tf.constant([[2.0, 3.0], [-0.5, 1.5]]), + ), + ( + [Box([-1, -2], [2, 3]), Box([-1, -2], [2, 3]), Box([-1, 1], [2, 2])], + tf.constant([[-1.0, -2.0], [-1.0, -2.0], [-1.0, 1.0]]), + tf.constant([[2.0, 3.0], [2.0, 3.0], [2.0, 2.0]]), + ), + ], +) +def test_multi_space_returns_correct_bounds( + spaces: Sequence[SearchSpace], lower: tf.Tensor, upper: tf.Tensor +) -> None: + for space in ( + TaggedMultiSearchSpace(spaces=spaces), + reduce(operator.mul, [TaggedMultiSearchSpace([s]) for s in spaces]), + ): + npt.assert_array_equal(space.lower, lower) + npt.assert_array_equal(space.upper, upper) + + +def test_multi_space___mul___() -> None: + space_A = Box([-1, -2], [2, 3]) + space_B = DiscreteSearchSpace(tf.ones([100, 2], dtype=tf.float64)) + product_space_1 = TaggedMultiSearchSpace(spaces=[space_A, space_B], tags=["A", "B"]) + + space_C = Box([-2, -2], [2, 3]) + space_D = DiscreteSearchSpace(tf.ones([5, 2], dtype=tf.float64)) + product_space_2 = TaggedMultiSearchSpace(spaces=[space_C, space_D], tags=["C", "D"]) + + product_of_product_spaces = product_space_1 * product_space_2 + + subspace_0_A = product_of_product_spaces.get_subspace("0") + assert isinstance(subspace_0_A, Box) + npt.assert_array_equal(subspace_0_A.lower, [-1, -2]) + npt.assert_array_equal(subspace_0_A.upper, [2, 3]) + subspace_0_B = product_of_product_spaces.get_subspace("1") + assert isinstance(subspace_0_B, DiscreteSearchSpace) + npt.assert_array_equal(subspace_0_B.points, tf.ones([100, 2], dtype=tf.float64)) + + subspace_1_C = product_of_product_spaces.get_subspace("2") + assert isinstance(subspace_1_C, Box) + npt.assert_array_equal(subspace_1_C.lower, [-2, -2]) + npt.assert_array_equal(subspace_1_C.upper, [2, 3]) + subspace_1_D = product_of_product_spaces.get_subspace("3") + assert isinstance(subspace_1_D, DiscreteSearchSpace) + npt.assert_array_equal(subspace_1_D.points, tf.ones([5, 2], dtype=tf.float64)) + + +def test_multi_space_handles_empty_spaces() -> None: + tag_A = TaggedProductSearchSpace(spaces=[], tags=[]) + tag_B = TaggedProductSearchSpace(spaces=[], tags=[]) + tag_C = TaggedMultiSearchSpace(spaces=[tag_A, tag_B], tags=["AA", "BB"]) + + assert tag_C.dimension == 0 + npt.assert_array_equal(tag_C.lower, [[], []]) + npt.assert_array_equal(tag_C.upper, [[], []]) + npt.assert_array_equal(tag_C.subspace_tags, ["AA", "BB"]) + + def _nlc_func(x: TensorType) -> TensorType: c0 = x[..., 0] - tf.sin(x[..., 1]) c0 = tf.expand_dims(c0, axis=-1) @@ -1116,6 +1375,31 @@ def _nlc_func(x: TensorType) -> TensorType: TaggedProductSearchSpace([Box([1], [2]), Box([-1], [1])], tags=["B", "A"]), False, ), + ( + TaggedMultiSearchSpace([Box([-1], [1]), Box([1], [2])]), + TaggedMultiSearchSpace([Box([-1], [1]), Box([1], [2])]), + True, + ), + ( + TaggedMultiSearchSpace([Box([-1], [1]), Box([1], [2])]), + TaggedMultiSearchSpace([Box([-1], [1]), Box([3], [4])]), + False, + ), + ( + TaggedMultiSearchSpace([Box([-1], [1]), Box([1], [2])], tags=["A", "B"]), + TaggedMultiSearchSpace([Box([-1], [1]), Box([1], [2])], tags=["B", "A"]), + False, + ), + ( + TaggedMultiSearchSpace([Box([-1], [1]), Box([1], [2])], tags=["A", "B"]), + TaggedMultiSearchSpace([Box([1], [2]), Box([-1], [1])], tags=["B", "A"]), + False, + ), + ( + TaggedProductSearchSpace([Box([-1], [1]), Box([1], [2])]), + TaggedMultiSearchSpace([Box([-1], [1]), Box([1], [2])]), + False, + ), ( Box( [-1], diff --git a/tests/unit/utils/test_misc.py b/tests/unit/utils/test_misc.py index 4c060b3879..ac6dd642eb 100644 --- a/tests/unit/utils/test_misc.py +++ b/tests/unit/utils/test_misc.py @@ -22,8 +22,18 @@ import tensorflow as tf from tests.util.misc import TF_DEBUGGING_ERROR_TYPES, ShapeLike, various_shapes +from trieste.observer import OBJECTIVE from trieste.types import TensorType -from trieste.utils.misc import Err, Ok, Timer, flatten_leading_dims, jit, shapes_equal, to_numpy +from trieste.utils.misc import ( + Err, + Ok, + Timer, + flatten_leading_dims, + get_value_for_tag, + jit, + shapes_equal, + to_numpy, +) @pytest.mark.parametrize("apply", [True, False]) @@ -86,6 +96,23 @@ def test_err() -> None: assert Err(ValueError()).is_err is True +def test_get_value_for_tag_returns_none_if_mapping_is_none() -> None: + assert get_value_for_tag(None) is None + + +def test_get_value_for_tag_raises_if_tag_not_in_mapping() -> None: + with pytest.raises(ValueError, match="tag 'baz' not found in mapping"): + get_value_for_tag({"foo": "bar"}, "baz") + + +def test_get_value_for_tag_returns_value_for_default_tag() -> None: + assert get_value_for_tag({"foo": "bar", OBJECTIVE: "baz"}) == "baz" + + +def test_get_value_for_tag_returns_value_for_specified_tag() -> None: + assert get_value_for_tag({"foo": "bar", OBJECTIVE: "baz"}, "foo") == "bar" + + def test_Timer() -> None: sleep_time = 0.1 with Timer() as timer: diff --git a/trieste/acquisition/optimizer.py b/trieste/acquisition/optimizer.py index 1f9dd77b06..d9b8550d66 100644 --- a/trieste/acquisition/optimizer.py +++ b/trieste/acquisition/optimizer.py @@ -30,10 +30,12 @@ from .. import logging from ..space import ( Box, + CollectionSearchSpace, Constraint, DiscreteSearchSpace, SearchSpace, SearchSpaceType, + TaggedMultiSearchSpace, TaggedProductSearchSpace, ) from ..types import TensorType @@ -100,9 +102,10 @@ def automatic_optimizer_selector( if isinstance(space, DiscreteSearchSpace): return optimize_discrete(space, target_func) - elif isinstance(space, (Box, TaggedProductSearchSpace)): - num_samples = tf.maximum(NUM_SAMPLES_MIN, NUM_SAMPLES_DIM * tf.shape(space.lower)[-1]) - num_runs = NUM_RUNS_DIM * tf.shape(space.lower)[-1] + elif isinstance(space, (Box, CollectionSearchSpace)): + space_dim = space.dimension + num_samples = tf.maximum(NUM_SAMPLES_MIN, NUM_SAMPLES_DIM * space_dim) + num_runs = NUM_RUNS_DIM * space_dim return generate_continuous_optimizer( num_initial_samples=num_samples, num_optimization_runs=num_runs, @@ -170,10 +173,10 @@ def generate_continuous_optimizer( num_optimization_runs: int = 10, num_recovery_runs: int = 10, optimizer_args: Optional[dict[str, Any]] = None, -) -> AcquisitionOptimizer[Box | TaggedProductSearchSpace]: +) -> AcquisitionOptimizer[Box | CollectionSearchSpace]: """ - Generate a gradient-based optimizer for :class:'Box' and :class:'TaggedProductSearchSpace' - spaces and batches of size 1. In the case of a :class:'TaggedProductSearchSpace', We perform + Generate a gradient-based optimizer for :class:'Box' and :class:'CollectionSearchSpace' + spaces and batches of size 1. In the case of a :class:'CollectionSearchSpace', We perform gradient-based optimization across all :class:'Box' subspaces, starting from the best location found across a sample of `num_initial_samples` random points. @@ -216,14 +219,14 @@ def generate_continuous_optimizer( raise ValueError(f"num_recovery_runs must be zero or greater, got {num_recovery_runs}") def optimize_continuous( - space: Box | TaggedProductSearchSpace, + space: Box | CollectionSearchSpace, target_func: Union[AcquisitionFunction, Tuple[AcquisitionFunction, int]], ) -> TensorType: """ A gradient-based :const:`AcquisitionOptimizer` for :class:'Box' - and :class:`TaggedProductSearchSpace' spaces. + and :class:`CollectionSearchSpace' spaces. - For :class:'TaggedProductSearchSpace' we only apply gradient updates to + For :class:'CollectionSearchSpace' we only apply gradient updates to its class:'Box' subspaces. When this functions receives an acquisition-integer tuple as its `target_func`,it @@ -246,8 +249,36 @@ def optimize_continuous( if V < 0: raise ValueError(f"vectorization must be positive, got {V}") - candidates = space.sample(num_initial_samples)[:, None, :] # [num_initial_samples, 1, D] - tiled_candidates = tf.tile(candidates, [1, V, 1]) # [num_initial_samples, V, D] + candidates = space.sample(num_initial_samples) + if tf.rank(candidates) == 3: + # If samples is a tensor of rank 3, then it is a batch of samples. In this case + # the length of the second dimension must be equal to the vectorization of the target + # function. + tf.debugging.assert_equal( + tf.shape(candidates)[1], + V, + message=( + f""" + The batch shape of initial samples {tf.shape(candidates)[1]} must be equal to + the vectorization of the target function {V}. + """ + ), + ) + tiled_candidates = candidates # [num_initial_samples, V, D] + else: + tf.debugging.assert_rank( + candidates, + 2, + message=( + f""" + The initial samples must be a tensor of rank 2, got a tensor of rank + {tf.rank(candidates)}. + """ + ), + ) + tiled_candidates = tf.tile( + candidates[:, None, :], [1, V, 1] + ) # [num_initial_samples, V, D] target_func_values = target_func(tiled_candidates) # [num_samples, V] tf.debugging.assert_shapes( @@ -289,11 +320,38 @@ def optimize_continuous( total_nfev = tf.reduce_max(nfev) # acquisition function is evaluated in parallel recovery_run = False - if ( - num_recovery_runs and not successful_optimization - ): # if all optimizations failed for a function then try again from random starts - random_points = space.sample(num_recovery_runs)[:, None, :] # [num_recovery_runs, 1, D] - tiled_random_points = tf.tile(random_points, [1, V, 1]) # [num_recovery_runs, V, D] + if num_recovery_runs and not successful_optimization: + # if all optimizations failed for a function then try again from random starts + random_points = space.sample(num_recovery_runs) + if tf.rank(random_points) == 3: + # If samples is a tensor of rank 3, then it is a batch of samples. In this case + # the length of the second dimension must be equal to the vectorization of the + # target function. + tf.debugging.assert_equal( + tf.shape(random_points)[1], + V, + message=( + f""" + The batch shape of random samples {tf.shape(random_points)[1]} must be + equal to the vectorization of the target function {V}. + """ + ), + ) + tiled_random_points = random_points # [num_recovery_runs, V, D] + else: + tf.debugging.assert_rank( + random_points, + 2, + message=( + f""" + The random samples must be a tensor of rank 2, got a tensor of rank + {tf.rank(random_points)}. + """ + ), + ) + tiled_random_points = tf.tile( + random_points[:, None, :], [1, V, 1] + ) # [num_recovery_runs, V, D] ( recovery_successes, @@ -383,7 +441,7 @@ def _perform_parallel_continuous_optimization( of our acquisition function (and its gradients) is calculated in parallel (for each optimization step) using Tensorflow. - For :class:'TaggedProductSearchSpace' we only apply gradient updates to + For :class:'CollectionSearchSpace' we only apply gradient updates to its :class:'Box' subspaces, fixing the discrete elements to the best values found across the initial random search. To fix these discrete elements, we optimize over a continuous :class:'Box' relaxation of the discrete subspaces @@ -434,6 +492,24 @@ def _objective_value_and_gradient(x: TensorType) -> Tuple[TensorType, TensorType get_bounds_of_box_relaxation_around_point(space, vectorized_starting_points[i : i + 1]) for i in tf.range(num_optimization_runs) ] + elif isinstance(space, TaggedMultiSearchSpace): + bounds = [ + spo.Bounds(lower, upper) + for lower, upper in zip(space.subspace_lower, space.subspace_upper) + ] + # If bounds is a sequence of tensors, stack them into a single tensor. In this case + # the length of the sequence must be equal to the vectorization of the target function. + tf.debugging.assert_equal( + len(bounds), + V, + message=( + f""" + The length of bounds sequence {len(bounds)} must be equal to the + vectorization of the target function {V}. + """ + ), + ) + bounds = bounds * num_optimization_runs_per_function else: bounds = [spo.Bounds(space.lower, space.upper)] * num_optimization_runs diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 5056438e20..e11977667b 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -22,9 +22,10 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Callable, Generic, Optional, TypeVar, Union, cast, overload +from typing import Any, Callable, Generic, Optional, Sequence, Tuple, TypeVar, Union, cast, overload import numpy as np +from check_shapes import check_shapes, inherit_check_shapes try: import pymoo @@ -47,8 +48,9 @@ TrainableSupportsGetKernel, ) from ..observer import OBJECTIVE -from ..space import Box, SearchSpace +from ..space import Box, SearchSpace, TaggedMultiSearchSpace from ..types import State, Tag, TensorType +from ..utils.misc import get_value_for_tag from .function import ( BatchMonteCarloExpectedImprovement, ExpectedImprovement, @@ -71,7 +73,7 @@ batchify_vectorize, ) from .sampler import ExactThompsonSampler, ThompsonSampler -from .utils import get_local_dataset, select_nth_output +from .utils import get_local_dataset, get_unique_points_mask, select_nth_output ResultType = TypeVar("ResultType", covariant=True) """ Unbound covariant type variable. """ @@ -1103,6 +1105,312 @@ def state_func( return state_func +class UpdatableTrustRegion(SearchSpace): + """A search space that can be updated.""" + + @abstractmethod + def initialize( + self, + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> None: + """ + Initialize the search space using the given models and datasets. + + :param models: The model for each tag. + :param datasets: The dataset for each tag. + """ + ... + + @abstractmethod + def update( + self, + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> None: + """ + Update the search space using the given models and datasets. + + :param models: The model for each tag. + :param datasets: The dataset for each tag. + """ + ... + + +UpdatableTrustRegionType = TypeVar("UpdatableTrustRegionType", bound=UpdatableTrustRegion) +""" A type variable bound to :class:`UpdatableTrustRegion`. """ + + +class BatchTrustRegion( + AcquisitionRule[ + types.State[Optional["BatchTrustRegion.State"], TensorType], + SearchSpace, + ProbabilisticModelType, + ], + Generic[ProbabilisticModelType, UpdatableTrustRegionType], +): + """Abstract class for multi trust region acquisition rules. These are batch algorithms where + each query point is optimized in parallel, with its own separate trust region. + """ + + @dataclass(frozen=True) + class State: + """The acquisition state for the :class:`BatchTrustRegion` acquisition rule.""" + + acquisition_space: TaggedMultiSearchSpace + """ The search space. """ + + def __deepcopy__(self, memo: dict[int, object]) -> BatchTrustRegion.State: + acquisition_space_copy = copy.deepcopy(self.acquisition_space, memo) + return BatchTrustRegion.State(acquisition_space_copy) + + def __init__( + self: "BatchTrustRegion[ProbabilisticModelType, UpdatableTrustRegionType]", + init_subspaces: Sequence[UpdatableTrustRegionType], + rule: AcquisitionRule[TensorType, SearchSpace, ProbabilisticModelType] | None = None, + ): + """ + :param init_subspaces: The initial search spaces for each trust region. + :param rule: The acquisition rule that defines how to search for a new query point in each + subspace. Defaults to :class:`EfficientGlobalOptimization` with default arguments. + """ + if rule is None: + rule = EfficientGlobalOptimization() + + self._init_subspaces = tuple(init_subspaces) + self._tags = tuple([str(index) for index in range(len(init_subspaces))]) + self._rule = rule + + def __repr__(self) -> str: + """""" + return f"""{self.__class__.__name__}({self._init_subspaces!r}, {self._rule!r})""" + + def acquire( + self, + search_space: SearchSpace, + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> types.State[State | None, TensorType]: + def state_func( + state: BatchTrustRegion.State | None, + ) -> Tuple[BatchTrustRegion.State | None, TensorType]: + """ + If state is None, initialize the subspaces by picking new locations. Otherwise, + update the existing subspaces. + + Re-initialize the subspaces if necessary, potentially looking at the entire group. + + Use the rule to acquire points from the acquisition space. + """ + # If state is set, the tags should be the same as the tags of the acquisition space + # in the state. + if state is not None: + assert ( + self._tags == state.acquisition_space.subspace_tags + ), f"""The tags of the state acquisition space + {state.acquisition_space.subspace_tags} should be the same as the tags of the + BatchTrustRegion acquisition rule {self._tags}""" + + subspaces = [] + for tag, init_subspace in zip(self._tags, self._init_subspaces): + if state is None: + subspace = init_subspace + subspace.initialize(models, datasets) + else: + _subspace = state.acquisition_space.get_subspace(tag) + assert isinstance(_subspace, type(init_subspace)) + subspace = _subspace + subspace.update(models, datasets) + + subspaces.append(subspace) + + self.maybe_initialize_subspaces(subspaces, models, datasets) + + if state is None: + acquisition_space = TaggedMultiSearchSpace(subspaces, self._tags) + else: + acquisition_space = state.acquisition_space + + state_ = BatchTrustRegion.State(acquisition_space) + points = self._rule.acquire(acquisition_space, models, datasets=datasets) + + return state_, points + + return state_func + + def maybe_initialize_subspaces( + self, + subspaces: Sequence[UpdatableTrustRegionType], + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> None: + """ + Initialize subspaces if necessary. + Get a mask of subspaces that need to be initialized using an abstract method. + Initialize individual subpaces by calling the method of the UpdatableTrustRegionType class. + + This method can be overridden by subclasses to change this behaviour. + """ + mask = self.get_initialize_subspaces_mask(subspaces, models, datasets) + tf.debugging.assert_equal( + tf.shape(mask), + (len(subspaces),), + message="The mask for initializing subspaces should be of the same length as the " + "number of subspaces", + ) + for ix, subspace in enumerate(subspaces): + if mask[ix]: + subspace.initialize(models, datasets) + + @abstractmethod + @check_shapes("return: [V]") + def get_initialize_subspaces_mask( + self, + subspaces: Sequence[UpdatableTrustRegionType], + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> TensorType: + """ + Return a boolean mask for subspaces that should be initialized. + This method is called during the acquisition step to determine which subspaces should be + initialized and which should be updated. The subspaces corresponding to True values in the + mask will be re-initialized. + + :param subspaces: The sequence of subspaces. + :param models: The model for each tag. + :param datasets: The dataset for each tag. + :return: A boolean mask of length V, where V is the number of subspaces. + """ + ... + + +class SingleObjectiveTrustRegionBox(Box, UpdatableTrustRegion): + """An updatable box search space for use with trust region acquisition rules.""" + + def __init__( + self, + global_search_space: SearchSpace, + beta: float = 0.7, + kappa: float = 1e-4, + min_eps: float = 1e-2, + ): + """ + Calculates the bounds of the box from the location/centre and global bounds. + + :param global_search_space: The global search space this search space lives in. + :param beta: The inverse of the trust region contraction factor. + :param kappa: Scales the threshold for the minimal improvement required for a step to be + considered a success. + :param min_eps: The minimal size of the search space. If the size of the search space is + smaller than this, the search space is reinitialized. + """ + + self._global_search_space = global_search_space + self._beta = beta + self._kappa = kappa + self._min_eps = min_eps + + super().__init__(global_search_space.lower, global_search_space.upper) + + @property + def global_search_space(self) -> SearchSpace: + """The global search space this search space lives in.""" + return self._global_search_space + + def _init_eps(self) -> None: + global_lower = self.global_search_space.lower + global_upper = self.global_search_space.upper + self.eps = 0.5 * (global_upper - global_lower) / (5.0 ** (1.0 / global_lower.shape[-1])) + + def _update_bounds(self) -> None: + self._lower = tf.reduce_max( + [self.global_search_space.lower, self.location - self.eps], axis=0 + ) + self._upper = tf.reduce_min( + [self.global_search_space.upper, self.location + self.eps], axis=0 + ) + + def initialize( + self, + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> None: + """ + Initialize the box by sampling a location from the global search space and setting the + bounds. + """ + dataset = get_value_for_tag(datasets) + + self.location = tf.squeeze(self.global_search_space.sample(1), axis=0) + self._step_is_success = False + self._init_eps() + self._update_bounds() + _, self._y_min = self.get_local_min(dataset) + + def update( + self, + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> None: + """ + Update this box, including centre/location, using the given dataset. If the size of the + box is less than the minimum size, re-initialize the box. + """ + dataset = get_value_for_tag(datasets) + + if tf.reduce_any(self.eps < self._min_eps): + self.initialize(models, datasets) + return + + x_min, y_min = self.get_local_min(dataset) + self.location = x_min + + tr_volume = tf.reduce_prod(self.upper - self.lower) + self._step_is_success = y_min < self._y_min - self._kappa * tr_volume + self.eps = self.eps / self._beta if self._step_is_success else self.eps * self._beta + self._update_bounds() + self._y_min = y_min + + @check_shapes( + "return[0]: [D]", + "return[1]: []", + ) + def get_local_min(self, dataset: Optional[Dataset]) -> Tuple[TensorType, TensorType]: + """Calculate the local minimum of the box using the given dataset.""" + if dataset is None: + raise ValueError("""dataset must be provided""") + + in_tr = self.contains(dataset.query_points) + in_tr_obs = tf.where( + tf.expand_dims(in_tr, axis=-1), + dataset.observations, + tf.constant(np.inf, dtype=dataset.observations.dtype), + ) + ix = tf.argmin(in_tr_obs) + x_min = tf.gather(dataset.query_points, ix) + y_min = tf.gather(in_tr_obs, ix) + + return tf.squeeze(x_min, axis=0), tf.squeeze(y_min) + + +class BatchTrustRegionBox(BatchTrustRegion[ProbabilisticModelType, SingleObjectiveTrustRegionBox]): + """ + Implements the :class:`BatchTrustRegion` *trust region* acquisition algorithm for box regions. + This is intended to be used for single-objective optimization with batching. + """ + + @inherit_check_shapes + def get_initialize_subspaces_mask( + self, + subspaces: Sequence[SingleObjectiveTrustRegionBox], + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> TensorType: + # Initialize the subspaces that have non-unique locations. + centres = tf.stack([subspace.location for subspace in subspaces]) + return tf.logical_not(get_unique_points_mask(centres, tolerance=1e-6)) + + class TURBO( AcquisitionRule[ types.State[Optional["TURBO.State"], TensorType], Box, TrainableSupportsGetKernel diff --git a/trieste/acquisition/utils.py b/trieste/acquisition/utils.py index cb97bd3a2e..590b0bb416 100644 --- a/trieste/acquisition/utils.py +++ b/trieste/acquisition/utils.py @@ -15,6 +15,7 @@ from typing import Tuple, Union import tensorflow as tf +from check_shapes import check_shapes from ..data import Dataset from ..space import SearchSpaceType @@ -135,3 +136,38 @@ def get_local_dataset(local_space: SearchSpaceType, dataset: Dataset) -> Dataset observations=tf.boolean_mask(dataset.observations, is_in_region_mask), ) return local_dataset + + +@check_shapes( + "points: [n_points, ...]", + "return: [n_points]", +) +def get_unique_points_mask(points: TensorType, tolerance: float = 1e-6) -> TensorType: + """Find the boolean mask of unique points in a tensor, within a given tolerance. + + Users can get the actual points with: + + mask = get_unique_points_mask(points, tolerance) + unique_points = tf.boolean_mask(points, mask) + + :param points: A tensor of points, with the first dimension being the number of points. + :param tolerance: The tolerance within which points are considered equal. + :return: A boolean mask for the unique points. + """ + + tolerance = tf.constant(tolerance, dtype=points.dtype) + n_points = tf.shape(points)[0] + mask = tf.zeros(shape=(n_points,), dtype=tf.bool) + + for idx in tf.range(n_points): + # Pairwise distance with previous unique points. + used_points = tf.boolean_mask(points, mask) + distances = tf.norm(points[idx] - used_points, axis=-1) + # Find if there is any point within the tolerance. + min_distance = tf.reduce_min(distances) + + # Update mask. + is_unique_point = min_distance >= tolerance + mask = tf.tensor_scatter_nd_update(mask, [[idx]], [is_unique_point]) + + return mask diff --git a/trieste/space.py b/trieste/space.py index 86aac24b3e..4a3a400246 100644 --- a/trieste/space.py +++ b/trieste/space.py @@ -23,6 +23,7 @@ import scipy.optimize as spo import tensorflow as tf import tensorflow_probability as tfp +from check_shapes import check_shapes from .types import TensorType @@ -869,20 +870,18 @@ def has_constraints(self) -> bool: return len(self._constraints) > 0 -class TaggedProductSearchSpace(SearchSpace): +class CollectionSearchSpace(SearchSpace): r""" - Product :class:`SearchSpace` consisting of a product of - multiple :class:`SearchSpace`. This class provides functionality for - accessing either the resulting combined search space or each individual space. + An abstract :class:`SearchSpace` consisting of a collection of multiple :class:`SearchSpace` + objects, each with a unique tag. This class provides functionality for accessing each individual + space. - Note that this class assumes that individual points in product spaces are - represented with their inputs in the same order as specified when initializing - the space. + Note that the individual spaces are not combined in any way. """ def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] = None): r""" - Build a :class:`TaggedProductSearchSpace` from a list ``spaces`` of other spaces. If + Build a :class:`CollectionSearchSpace` from a list ``spaces`` of other spaces. If ``tags`` are provided then they form the identifiers of the subspaces, otherwise the subspaces are labelled numerically. @@ -914,54 +913,36 @@ def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] ) self._spaces = dict(zip(tags, spaces)) - - subspace_sizes = [space.dimension for space in spaces] - - self._subspace_sizes_by_tag = { - tag: subspace_size for tag, subspace_size in zip(tags, subspace_sizes) - } - - self._subspace_starting_indices = dict(zip(tags, tf.cumsum(subspace_sizes, exclusive=True))) - - self._dimension = tf.cast(tf.reduce_sum(subspace_sizes), dtype=tf.int32) self._tags = tuple(tags) # avoid accidental modification by users def __repr__(self) -> str: - """""" - return f"""TaggedProductSearchSpace(spaces = + return f"""{self.__class__.__name__}(spaces = {[self.get_subspace(tag) for tag in self.subspace_tags]}, tags = {self.subspace_tags}) """ @property - def lower(self) -> TensorType: - """The lowest values taken by each space dimension, concatenated across subspaces.""" - lower_for_each_subspace = [self.get_subspace(tag).lower for tag in self.subspace_tags] - return ( - tf.concat(lower_for_each_subspace, axis=-1) - if lower_for_each_subspace - else tf.constant([], dtype=DEFAULT_DTYPE) - ) + def subspace_lower(self) -> Sequence[TensorType]: + """The lowest values taken by each space dimension, in the same order as specified when + initializing the space.""" + return [self.get_subspace(tag).lower for tag in self.subspace_tags] @property - def upper(self) -> TensorType: - """The highest values taken by each space dimension, concatenated across subspaces.""" - upper_for_each_subspace = [self.get_subspace(tag).upper for tag in self.subspace_tags] - return ( - tf.concat(upper_for_each_subspace, axis=-1) - if upper_for_each_subspace - else tf.constant([], dtype=DEFAULT_DTYPE) - ) + def subspace_upper(self) -> Sequence[TensorType]: + """The highest values taken by each space dimension, in the same order as specified when + initializing the space.""" + return [self.get_subspace(tag).upper for tag in self.subspace_tags] @property def subspace_tags(self) -> tuple[str, ...]: - """Return the names of the subspaces contained in this product space.""" + """Return the names of the subspaces contained in this space.""" return self._tags @property - def dimension(self) -> TensorType: - """The number of inputs in this product search space.""" - return self._dimension + def subspace_dimension(self) -> Sequence[TensorType]: + """The number of inputs in each subspace, in the same order as specified when initializing + the space.""" + return [self.get_subspace(tag).dimension for tag in self.subspace_tags] def get_subspace(self, tag: str) -> SearchSpace: """ @@ -976,10 +957,107 @@ def get_subspace(self, tag: str) -> SearchSpace: message=f""" Attempted to access a subspace that does not exist. This space only contains subspaces with the tags {self.subspace_tags} but received {tag}. - """, + """, ) return self._spaces[tag] + def subspace_sample(self, num_samples: int, seed: Optional[int] = None) -> Sequence[TensorType]: + """ + Sample randomly from the space by sampling from each subspace + and returning the resulting samples in the same order as specified when initializing + the space. + + :param num_samples: The number of points to sample from each subspace. + :param seed: Optional tf.random seed. + :return: ``num_samples`` i.i.d. random points, sampled uniformly, + from each search subspace with shape '[num_samples, D]' , where D is the search space + dimension. + """ + tf.debugging.assert_non_negative(num_samples) + if seed is not None: # ensure reproducibility + tf.random.set_seed(seed) + return [self._spaces[tag].sample(num_samples, seed=seed) for tag in self._tags] + + def __eq__(self, other: object) -> bool: + """ + :param other: A search space. + :return: Whether the search space is identical to this one. + """ + if not isinstance(other, type(self)): + return NotImplemented + return self._tags == other._tags and self._spaces == other._spaces + + +class TaggedProductSearchSpace(CollectionSearchSpace): + r""" + Product :class:`SearchSpace` consisting of a product of + multiple :class:`SearchSpace`. This class provides functionality for + accessing either the resulting combined search space or each individual space. + This class is useful for defining mixed search spaces, for example: + + context_space = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) + decision_space = Box([-1, -2], [2, 3]) + mixed_space = TaggedProductSearchSpace(spaces=[context_space, decision_space]) + + Note that this class assumes that individual points in product spaces are + represented with their inputs in the same order as specified when initializing + the space. + """ + + def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] = None): + r""" + Build a :class:`TaggedProductSearchSpace` from a list ``spaces`` of other spaces. If + ``tags`` are provided then they form the identifiers of the subspaces, otherwise the + subspaces are labelled numerically. + + :param spaces: A sequence of :class:`SearchSpace` objects representing the space's subspaces + :param tags: An optional list of tags giving the unique identifiers of + the space's subspaces. + :raise ValueError (or tf.errors.InvalidArgumentError): If ``spaces`` has a different + length to ``tags`` when ``tags`` is provided or if ``tags`` contains duplicates. + """ + + super().__init__(spaces, tags) + subspace_sizes = self.subspace_dimension + + self._subspace_sizes_by_tag = { + tag: subspace_size for tag, subspace_size in zip(self._tags, subspace_sizes) + } + + self._subspace_starting_indices = dict( + zip(self._tags, tf.cumsum(subspace_sizes, exclusive=True)) + ) + + self._dimension = tf.cast(tf.reduce_sum(subspace_sizes), dtype=tf.int32) + + @property + @check_shapes("return: [D]") + def lower(self) -> TensorType: + """The lowest values taken by each space dimension, concatenated across subspaces.""" + lower_for_each_subspace = self.subspace_lower + return ( + tf.concat(lower_for_each_subspace, axis=-1) + if lower_for_each_subspace + else tf.constant([], dtype=DEFAULT_DTYPE) + ) + + @property + @check_shapes("return: [D]") + def upper(self) -> TensorType: + """The highest values taken by each space dimension, concatenated across subspaces.""" + upper_for_each_subspace = self.subspace_upper + return ( + tf.concat(upper_for_each_subspace, axis=-1) + if upper_for_each_subspace + else tf.constant([], dtype=DEFAULT_DTYPE) + ) + + @property + @check_shapes("return: []") + def dimension(self) -> TensorType: + """The number of inputs in this product search space.""" + return self._dimension + def fix_subspace(self, tag: str, values: TensorType) -> TaggedProductSearchSpace: """ Return a new :class:`TaggedProductSearchSpace` with the specified subspace replaced with @@ -1033,6 +1111,7 @@ def _contains(self, value: TensorType) -> TensorType: ] return tf.reduce_all(in_each_subspace, axis=0) + @check_shapes("return: [num_samples, D]") def sample(self, num_samples: int, seed: Optional[int] = None) -> TensorType: """ Sample randomly from the space by sampling from each subspace @@ -1044,10 +1123,7 @@ def sample(self, num_samples: int, seed: Optional[int] = None) -> TensorType: from this search space with shape '[num_samples, D]' , where D is the search space dimension. """ - tf.debugging.assert_non_negative(num_samples) - if seed is not None: # ensure reproducibility - tf.random.set_seed(seed) - subspace_samples = [self._spaces[tag].sample(num_samples, seed=seed) for tag in self._tags] + subspace_samples = self.subspace_sample(num_samples, seed) return tf.concat(subspace_samples, -1) def product(self, other: TaggedProductSearchSpace) -> TaggedProductSearchSpace: @@ -1060,11 +1136,136 @@ def product(self, other: TaggedProductSearchSpace) -> TaggedProductSearchSpace: """ return TaggedProductSearchSpace(spaces=[self, other]) - def __eq__(self, other: object) -> bool: + +class TaggedMultiSearchSpace(CollectionSearchSpace): + r""" + A :class:`SearchSpace` made up of a collection of multiple :class:`SearchSpace` subspaces, + each with a unique tag. All subspaces must have the same dimensionality. + + Each subspace is treated as an independent space and not combined in any way. This class + provides functionality for accessing all the subspaces at once by using the usual search space + methods, as well as for accessing individual subspaces. + + When accessing all subspaces at once from this class (e.g. `lower()`, `upper()`, `sample()`), + the returned tensors have an extra dimension corresponding to the subspaces. + + This class can be useful to represent a collection of search spaces that do not interact with + each other. For example, it is used to implement batch trust region rules in the + :class:`BatchTrustRegion` class. + """ + + def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] = None): + r""" + Build a :class:`TaggedMultiSearchSpace` from a list ``spaces`` of other spaces. If + ``tags`` are provided then they form the identifiers of the subspaces, otherwise the + subspaces are labelled numerically. + + :param spaces: A sequence of :class:`SearchSpace` objects representing the space's subspaces + :param tags: An optional list of tags giving the unique identifiers of + the space's subspaces. + :raise ValueError (or tf.errors.InvalidArgumentError): If ``spaces`` has a different + length to ``tags`` when ``tags`` is provided or if ``tags`` contains duplicates. + :raise ValueError (or tf.errors.InvalidArgumentError): If ``spaces`` has a different + dimension to each other. """ - :param other: A search space. - :return: Whether the search space is identical to this one. + + # At least one subspace is required. + tf.debugging.assert_greater( + len(spaces), + 0, + message=f""" + At least one subspace is required but received {len(spaces)}. + """, + ) + + tf.debugging.assert_equal( + len(set([int(space.dimension) for space in spaces])), + 1, + message=f""" + All subspaces must have the same dimension but received + {[int(space.dimension) for space in spaces]}. + """, + ) + + super().__init__(spaces, tags) + + @property + @check_shapes("return: [V, D]") + def lower(self) -> TensorType: + """Returns the stacked lower bounds of all the subspaces. + + :return: The lower bounds of shape [V, D], where V is the number of subspaces and D is + the dimensionality of each subspace. """ - if not isinstance(other, TaggedProductSearchSpace): - return NotImplemented - return self._tags == other._tags and self._spaces == other._spaces + return tf.stack(self.subspace_lower, axis=0) + + @property + @check_shapes("return: [V, D]") + def upper(self) -> TensorType: + """Returns the stacked upper bounds of all the subspaces. + + :return: The upper bounds of shape [V, D], where V is the number of subspaces and D is + the dimensionality of each subspace. + """ + return tf.stack(self.subspace_upper, axis=0) + + @property + @check_shapes("return: []") + def dimension(self) -> TensorType: + """The number of inputs in this search space.""" + return self.get_subspace(self.subspace_tags[0]).dimension + + @check_shapes("return: [num_samples, V, D]") + def sample(self, num_samples: int, seed: Optional[int] = None) -> TensorType: + """ + Sample randomly from the space by sampling from each subspace + and returning the resulting samples stacked along the second axis in the same order as + specified when initializing the space. + + :param num_samples: The number of points to sample from each subspace. + :param seed: Optional tf.random seed. + :return: ``num_samples`` i.i.d. random points, sampled uniformly, + from each search subspace with shape '[num_samples, V, D]' , where V is the number of + subspaces and D is the search space dimension. + """ + return tf.stack(self.subspace_sample(num_samples, seed), axis=1) + + def _contains(self, value: TensorType) -> TensorType: + """ + Return `True` if ``value`` is a member of this search space, else `False`. A point + is a member if it is a member of any of the subspaces. + + :param value: A point or points to check for membership of this :class:`SearchSpace`. + :return: A boolean array showing membership for each point in value. + """ + return tf.reduce_any( + [self.get_subspace(tag)._contains(value) for tag in self.subspace_tags], axis=0 + ) + + def product(self, other: TaggedMultiSearchSpace) -> TaggedMultiSearchSpace: + r""" + Return a bigger collection of two :class:`TaggedMultiSearchSpace`\ s, regenerating the + tags. + + :param other: A search space of the same type as this search space. + :return: The product of this search space with the ``other``. + """ + return TaggedMultiSearchSpace( + spaces=tuple(self._spaces.values()) + tuple(other._spaces.values()) + ) + + def discretize(self, num_samples: int) -> DiscreteSearchSpace: + """ + :param num_samples: The number of points in the :class:`DiscreteSearchSpace`. + :return: A discrete search space consisting of ``num_samples`` points sampled uniformly from + this search space. + :raise NotImplementedError: If this :class:`SearchSpace` has constraints. + """ + if self.has_constraints: # Constraints are not supported. + raise NotImplementedError( + "Discretization is currently not supported in the presence of constraints." + ) + samples = self.sample(num_samples) # Sample num_samples points from each subspace. + samples = tf.reshape(samples, [-1, self.dimension]) # Flatten the samples across subspaces. + samples = tf.random.shuffle(samples)[:num_samples] # Randomly pick num_samples points. + return DiscreteSearchSpace(points=samples) diff --git a/trieste/utils/misc.py b/trieste/utils/misc.py index 4116b1f9df..c29b3c4e51 100644 --- a/trieste/utils/misc.py +++ b/trieste/utils/misc.py @@ -23,7 +23,8 @@ from tensorflow.python.util import nest from typing_extensions import Final, final -from ..types import TensorType +from ..observer import OBJECTIVE +from ..types import Tag, TensorType C = TypeVar("C", bound=Callable[..., object]) """ A type variable bound to `typing.Callable`. """ @@ -215,6 +216,27 @@ def map_values(f: Callable[[U], V], mapping: Mapping[K, U]) -> Mapping[K, V]: return {k: f(u) for k, u in mapping.items()} +T = TypeVar("T") +""" An unbound type variable. """ + + +def get_value_for_tag(mapping: Optional[Mapping[Tag, T]], tag: Tag = OBJECTIVE) -> Optional[T]: + """Return the value of a tag in a mapping. + + :param mapping: A mapping from tags to values. + :param tag: A tag. + :return: The value of the tag in the mapping, or None if the mapping is None. + :raises ValueError: If the tag is not in the mapping and the mapping is not None. + """ + + if mapping is None: + return None + elif tag in mapping.keys(): + return mapping[tag] + else: + raise ValueError(f"tag '{tag}' not found in mapping") + + class Timer: """ Functionality for timing chunks of code. For example: