diff --git a/docs/notebooks/mixed_search_spaces.pct.py b/docs/notebooks/mixed_search_spaces.pct.py index 60127116dc..b6d803e081 100644 --- a/docs/notebooks/mixed_search_spaces.pct.py +++ b/docs/notebooks/mixed_search_spaces.pct.py @@ -67,6 +67,8 @@ # We create our mixed search space by instantiating this class with a list containing the discrete # and continuous spaces, without any explicit tags (hence using default tags). # This can be easily extended to more than two search spaces by adding more elements to the list. +# +# Note: the dtype of all the component search spaces must be the same. # %% from trieste.space import Box, DiscreteSearchSpace, TaggedProductSearchSpace @@ -217,6 +219,168 @@ alpha=0.6, ) +# %% [markdown] +# ## Trust region with mixed search spaces +# +# In this section, we demonstrate the use of trust region acquisition rules with mixed search +# spaces. We use the same mixed search space and observer as before, and the same initial data. +# See [trust region Bayesian optimization notebook](trust_region.ipynb) for an introduction to +# trust region acquisition rules. +# +# First we build a Gaussian process model of the objective function using the initial data. + +# %% +gpflow_model = build_gpr( + initial_data, mixed_search_space, likelihood_variance=1e-7 +) +model = GaussianProcessRegression(gpflow_model) + +# %% [markdown] +# We create a trust region acquisition rule that uses the efficient global optimization (EGO) +# acquisition rule as the base rule. The trust region acquisition rule is initialized with a set of +# trust regions; 5 in this example. Each trust regions is defined as a product of a discrete and a +# continuous trust sub-region. The base rule is then used to optimize the acquisition function +# within each trust region. This setup is similar to the one used in the "Batch trust region rule" +# section of the [trust region Bayesian optimization notebook](trust_region.ipynb). +# +# Analogous to a `TaggedProductSearchSpace`, each trust region is a product of a discrete and a +# continuous trust (sub-)region. The discrete sub-region is defined by +# `FixedPointTrustRegionDiscrete` that selects a random point from the discrete space, which is then +# fixed for the duration of the optimization. The continuous sub-region is defined by +# `SingleObjectiveTrustRegionBox`, just as in the trust region notebook, where the region is updated +# at each step of the optimization. + +# %% +from trieste.acquisition import ParallelContinuousThompsonSampling +from trieste.acquisition.rule import ( + BatchTrustRegionProduct, + EfficientGlobalOptimization, + FixedPointTrustRegionDiscrete, + SingleObjectiveTrustRegionBox, + UpdatableTrustRegionProduct, +) + +num_query_points = 5 +init_regions = [ + UpdatableTrustRegionProduct( + [ + FixedPointTrustRegionDiscrete(discrete_space), + SingleObjectiveTrustRegionBox(continuous_space), + ] + ) + for _ in range(num_query_points) +] +base_rule = EfficientGlobalOptimization( # type: ignore[var-annotated] + builder=ParallelContinuousThompsonSampling(), + num_query_points=num_query_points, +) +tr_acq_rule = BatchTrustRegionProduct(init_regions, base_rule) + +# %% [markdown] +# We run the optimization loop for 15 steps using the trust region acquisition rule. + +# %% +bo = BayesianOptimizer(observer, mixed_search_space) + +num_steps = 15 +tr_result = bo.optimize( + num_steps, initial_data, model, tr_acq_rule, track_state=True +) +dataset = tr_result.try_get_final_dataset() + +# %% [markdown] +# The best point found by the optimizer is obtained as before. + +# %% +query_point, observation, arg_min_idx = tr_result.try_get_optimal_point() + +print(f"query point: {query_point}") +print(f"observation: {observation}") + +# %% [markdown] +# Plot of the optimization loop over the mixed search space, similar to the previous plot. + +# %% +query_points = dataset.query_points.numpy() +observations = dataset.observations.numpy() + +_, ax = plot_function_2d( + scaled_branin, + ScaledBranin.search_space.lower, + ScaledBranin.search_space.upper, + contour=True, +) +plot_bo_points(query_points, ax[0, 0], num_initial_points, arg_min_idx) +ax[0, 0].set_xlabel(r"$x_1$") +ax[0, 0].set_ylabel(r"$x_2$") + +for point in points: + ax[0, 0].vlines( + point, + mixed_search_space.lower[1], + mixed_search_space.upper[1], + colors="b", + linestyles="dashed", + alpha=0.6, + ) + +# %% [markdown] +# Finally, we visualize the optimization progress by plotting the 5 (product) trust regions at each +# step. The trust regions are shown as translucent boxes, with each box in a different color. The +# new query point for earch region is plotted in matching color. +# +# Note that since the discrete dimension is on the x-axis, the trust regions appear as vertical +# lines with zero width. + +# %% +import base64 +from typing import Optional + +import IPython +import matplotlib.pyplot as plt + +from trieste.bayesian_optimizer import OptimizationResult +from trieste.experimental.plotting import ( + convert_figure_to_frame, + convert_frames_to_gif, + plot_trust_region_history_2d, +) + + +def plot_history( + result: OptimizationResult, + num_query_points: Optional[int] = None, +) -> None: + frames = [] + for step, hist in enumerate( + result.history + [result.final_result.unwrap()] + ): + fig, _ = plot_trust_region_history_2d( + scaled_branin, + ScaledBranin.search_space.lower, + ScaledBranin.search_space.upper, + hist, + num_query_points=num_query_points, + num_init=num_initial_points, + alpha=1.0, + ) + + if fig is not None: + fig.suptitle(f"step number {step}") + frames.append(convert_figure_to_frame(fig)) + plt.close(fig) + + gif_file = convert_frames_to_gif(frames) + gif = IPython.display.HTML( + ''.format( + base64.b64encode(gif_file.getvalue()).decode() + ) + ) + IPython.display.display(gif) + + +plot_history(tr_result) + # %% [markdown] # ## LICENSE # diff --git a/tests/integration/test_mixed_space_bayesian_optimization.py b/tests/integration/test_mixed_space_bayesian_optimization.py index 5cb112b76a..b87e036ad2 100644 --- a/tests/integration/test_mixed_space_bayesian_optimization.py +++ b/tests/integration/test_mixed_space_bayesian_optimization.py @@ -13,6 +13,9 @@ # limitations under the License. from __future__ import annotations +from typing import cast + +import numpy as np import numpy.testing as npt import pytest import tensorflow as tf @@ -22,8 +25,16 @@ AcquisitionFunctionClass, BatchMonteCarloExpectedImprovement, LocalPenalization, + ParallelContinuousThompsonSampling, +) +from trieste.acquisition.rule import ( + AcquisitionRule, + BatchTrustRegionProduct, + EfficientGlobalOptimization, + FixedPointTrustRegionDiscrete, + SingleObjectiveTrustRegionBox, + UpdatableTrustRegionProduct, ) -from trieste.acquisition.rule import AcquisitionRule, EfficientGlobalOptimization from trieste.bayesian_optimizer import BayesianOptimizer from trieste.models import TrainableProbabilisticModel from trieste.models.gpflow import GaussianProcessRegression, build_gpr @@ -34,6 +45,34 @@ from trieste.types import TensorType +def _get_mixed_search_space() -> TaggedProductSearchSpace: + # The discrete space is defined by a set of 10 points that are equally spaced, ensuring that + # the three Branin minimizers (of dimension 0) are included in this set. The continuous + # dimension is defined by the interval [0, 1]. + # We observe that the first and third minimizers are equidistant from the middle minimizer, so + # we choose the discretization points to be equally spaced around the middle minimizer. + minimizers0 = ScaledBranin.minimizers[:, 0] + step = (minimizers0[1] - minimizers0[0]) / 4 + points = np.concatenate( + [ + # Equally spaced points to the left of the middle minimizer. Skip the last point as it + # is the same as the first point in the next array. + np.flip(np.arange(minimizers0[1], 0.0, -step))[:-1], + # Equally spaced points to the right of the middle minimizer. + np.arange(minimizers0[1], 1.0, step), + ] + ) + discrete_space = DiscreteSearchSpace(points[:, None]) + continuous_space = Box([0], [1]) + return TaggedProductSearchSpace( + spaces=[discrete_space, continuous_space], + tags=["discrete", "continuous"], + ) + + +mixed_search_space = _get_mixed_search_space() + + @random_seed @pytest.mark.parametrize( "num_steps, acquisition_rule", @@ -57,6 +96,35 @@ ), id="LocalPenalization", ), + pytest.param( + 8, + BatchTrustRegionProduct( + [ + UpdatableTrustRegionProduct( + [ + FixedPointTrustRegionDiscrete( + cast( + DiscreteSearchSpace, mixed_search_space.get_subspace("discrete") + ) + ), + SingleObjectiveTrustRegionBox( + mixed_search_space.get_subspace("continuous") + ), + ], + tags=mixed_search_space.subspace_tags, + ) + for _ in range(10) + ], + EfficientGlobalOptimization( + ParallelContinuousThompsonSampling(), + # Use a large batch to ensure discrete init finds a good point. + # We are using a fixed point trust region for the discrete space, so + # the init point is randomly chosen and then never updated. + num_query_points=10, + ), + ), + id="TrustRegionSingleObjectiveFixed", + ), ], ) def test_optimizer_finds_minima_of_the_scaled_branin_function( @@ -65,20 +133,15 @@ def test_optimizer_finds_minima_of_the_scaled_branin_function( TensorType, TaggedProductSearchSpace, TrainableProbabilisticModel ], ) -> None: - search_space = TaggedProductSearchSpace( - spaces=[Box([0], [1]), DiscreteSearchSpace(tf.linspace(0, 1, 15)[:, None])], - tags=["continuous", "discrete"], - ) - - initial_query_points = search_space.sample(5) + initial_query_points = mixed_search_space.sample(5) observer = mk_observer(ScaledBranin.objective) initial_data = observer(initial_query_points) model = GaussianProcessRegression( - build_gpr(initial_data, search_space, likelihood_variance=1e-8) + build_gpr(initial_data, mixed_search_space, likelihood_variance=1e-8) ) dataset = ( - BayesianOptimizer(observer, search_space) + BayesianOptimizer(observer, mixed_search_space) .optimize(num_steps, initial_data, model, acquisition_rule) .try_get_final_dataset() ) diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index 1cab4f6a75..32777ff8f8 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -15,7 +15,7 @@ import copy from collections.abc import Mapping -from typing import Callable, Optional +from typing import Any, Callable, List, Optional, Union from unittest.mock import ANY, MagicMock import gpflow @@ -49,12 +49,16 @@ AsynchronousRuleState, BatchHypervolumeSharpeRatioIndicator, BatchTrustRegionBox, + BatchTrustRegionProduct, DiscreteThompsonSampling, EfficientGlobalOptimization, + FixedPointTrustRegionDiscrete, RandomSampling, SingleObjectiveTrustRegionBox, TREGOBox, TURBOBox, + UpdatableTrustRegion, + UpdatableTrustRegionProduct, ) from trieste.acquisition.sampler import ( ExactThompsonSampler, @@ -68,7 +72,13 @@ from trieste.models.interfaces import TrainableSupportsGetKernel from trieste.objectives.utils import mk_batch_observer from trieste.observer import OBJECTIVE -from trieste.space import Box, SearchSpace, TaggedMultiSearchSpace +from trieste.space import ( + Box, + DiscreteSearchSpace, + SearchSpace, + TaggedMultiSearchSpace, + TaggedProductSearchSpace, +) from trieste.types import State, Tag, TensorType from trieste.utils.misc import LocalizedTag, get_value_for_tag @@ -589,11 +599,12 @@ def test_trego_for_default_state( rule: AcquisitionRule[TensorType, SearchSpace, ProbabilisticModel], expected_query_point: TensorType, ) -> None: - dataset = Dataset(tf.constant([[0.1, 0.2]]), tf.constant([[0.012]])) + dataset = Dataset(tf.constant([[0.0, 0.1, 0.3, 0.2]]), tf.constant([[0.012]])) lower_bound = tf.constant([-2.2, -1.0]) upper_bound = tf.constant([1.3, 3.3]) search_space = Box(lower_bound, upper_bound) - subspace = TREGOBox(search_space) + # Includes a quick test of input_active_dims. The irrelevant input dimension should be ignored. + subspace = TREGOBox(search_space, input_active_dims=[1, 3]) tr = BatchTrustRegionBox(subspace, rule) model = QuadraticMeanAndRBFKernel() @@ -952,21 +963,35 @@ def test_turbo_acquire_returns_correct_shape(num_query_points: int) -> None: @random_seed -def test_turbo_for_default_state() -> None: +@pytest.mark.parametrize( + "lengthscales, exp_upper", + [ + (4.0, [0.8, 0.8]), + ( + [4.0, 0.1, 0.1, 1.0], # Unused lengthscales should be ignored due to input_active_dims. + [0.8, 0.2], + ), + ], +) +def test_turbo_for_default_state( + lengthscales: Union[float, List[float]], exp_upper: List[float] +) -> None: dataset = Dataset( - tf.constant([[0.0, 0.0]], dtype=tf.float64), tf.constant([[0.012]], dtype=tf.float64) + tf.constant([[0.0, 0.3, 0.2, 0.0]], dtype=tf.float64), + tf.constant([[0.012]], dtype=tf.float64), ) lower_bound = tf.constant([0.0, 0.0], dtype=tf.float64) upper_bound = tf.constant([1.0, 1.0], dtype=tf.float64) search_space = Box(lower_bound, upper_bound) - orig_region = TURBOBox(search_space) + # Includes a quick test of input_active_dims. The irrelevant input dimension should be ignored. + orig_region = TURBOBox(search_space, input_active_dims=[0, 3]) region = copy.deepcopy(orig_region) tr = BatchTrustRegionBox(region, rule=DiscreteThompsonSampling(100, 1)) model = QuadraticMeanAndRBFKernelWithSamplers( dataset, noise_variance=tf.constant(1e-5, dtype=tf.float64) ) model.kernel = gpflow.kernels.RBF( - lengthscales=tf.constant([4.0, 1.0], dtype=tf.float64), variance=1e-5 + lengthscales=tf.constant(lengthscales, dtype=tf.float64), variance=1e-5 ) # need a gpflow kernel for TURBOBox state, query_point = tr.acquire_single(search_space, model, dataset=dataset)(None) tr.filter_datasets({OBJECTIVE: model}, {OBJECTIVE: dataset}) @@ -977,7 +1002,7 @@ def test_turbo_for_default_state() -> None: npt.assert_array_almost_equal(state_region.lower, orig_region.lower) npt.assert_array_almost_equal(state_region.upper, orig_region.upper) npt.assert_array_almost_equal(region.lower, lower_bound) - npt.assert_array_almost_equal(region.upper, tf.constant([0.8, 0.2], dtype=tf.float64)) + npt.assert_array_almost_equal(region.upper, tf.constant(exp_upper, dtype=tf.float64)) npt.assert_array_almost_equal(region.y_min, [0.012]) npt.assert_array_almost_equal(region.L, tf.cast(0.8, dtype=tf.float64)) assert region.success_counter == 0 @@ -1242,6 +1267,79 @@ def test_turbo_state_deepcopy() -> None: npt.assert_allclose(tr_subspace_copy.y_min, tr_subspace.y_min) +@pytest.mark.parametrize( + "active_dims, in_values, exp_values", + [ + ([0], tf.constant([[0.1, 0.2], [-0.1, -0.2]]), tf.constant([[0.1], [-0.1]])), + ([1], tf.constant([[0.1, 0.2], [-0.1, -0.2]]), tf.constant([[0.2], [-0.2]])), + ([0, 1], tf.constant([[0.1, 0.2], [-0.1, -0.2]]), tf.constant([[0.1, 0.2], [-0.1, -0.2]])), + ( + [1, 3], + tf.constant([[0.1, 0.2, 0.3, 0.4], [-0.1, -0.2, -0.3, -0.4]]), + tf.constant([[0.2, 0.4], [-0.2, -0.4]]), + ), + ( + None, + tf.constant([[0.1, 0.2, 0.3, 0.4], [-0.1, -0.2, -0.3, -0.4]]), + tf.constant([[0.1, 0.2, 0.3, 0.4], [-0.1, -0.2, -0.3, -0.4]]), + ), + ( + slice(1, 3), + tf.constant([[0.1, 0.2, 0.3], [-0.1, -0.2, -0.3]]), + tf.constant([[0.2, 0.3], [-0.2, -0.3]]), + ), + ( + slice(0, 4, 2), + tf.constant([[0.1, 0.2, 0.3, 0.4], [-0.1, -0.2, -0.3, -0.4]]), + tf.constant([[0.1, 0.3], [-0.1, -0.3]]), + ), + ( + [0], + Dataset(tf.constant([[0.1, 0.2], [-0.1, -0.2]]), tf.constant([[0.4], [0.5]])), + Dataset(tf.constant([[0.1], [-0.1]]), tf.constant([[0.4], [0.5]])), + ), + ( + [2], + Dataset( + tf.constant([[0.1, 0.2, 0.3], [-0.1, -0.2, -0.3]]), tf.constant([[0.4], [0.5]]) + ), + Dataset(tf.constant([[0.3], [-0.3]]), tf.constant([[0.4], [0.5]])), + ), + ( + [0, 1], + Dataset(tf.constant([[0.1, 0.2], [-0.1, -0.2]]), tf.constant([[0.4], [0.5]])), + Dataset(tf.constant([[0.1, 0.2], [-0.1, -0.2]]), tf.constant([[0.4], [0.5]])), + ), + ( + slice(1, 3), + Dataset( + tf.constant([[0.1, 0.2, 0.3], [-0.1, -0.2, -0.3]]), tf.constant([[0.4], [0.5]]) + ), + Dataset(tf.constant([[0.2, 0.3], [-0.2, -0.3]]), tf.constant([[0.4], [0.5]])), + ), + ([1], QuadraticMeanAndRBFKernel(), None), # exp_values in unused. + ([7, 10], QuadraticMeanAndRBFKernel(x_shift=0.3), None), # exp_values in unused. + ], +) +def test_trust_region_with_input_active_dims( + active_dims: Optional[Union[slice, List[int]]], + in_values: Union[TensorType, Dataset, ProbabilisticModel], + exp_values: Union[TensorType, Dataset, ProbabilisticModel], +) -> None: + dummy_search_space = Box([0.0], [1.0]) + tr = SingleObjectiveTrustRegionBox(dummy_search_space, input_active_dims=active_dims) + out_values = tr.with_input_active_dims(in_values) + if isinstance(in_values, Dataset): + assert isinstance(exp_values, Dataset) + npt.assert_allclose(out_values.query_points, exp_values.query_points) + npt.assert_allclose(out_values.observations, exp_values.observations) + elif isinstance(in_values, ProbabilisticModel): + assert out_values is in_values + else: + assert isinstance(exp_values, tf.Tensor) + npt.assert_allclose(out_values, exp_values) + + @pytest.mark.parametrize( "datasets", [ @@ -1293,11 +1391,12 @@ 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.5, 1.2, 1.3, 0.6], [0.4, -0.4, -0.5, 0.3]], dtype=tf.float64), tf.constant([[0.7], [0.9]], dtype=tf.float64), ) } - trb = SingleObjectiveTrustRegionBox(search_space) + # Includes a quick test of input_active_dims. The irrelevant input dimension should be ignored. + trb = SingleObjectiveTrustRegionBox(search_space, input_active_dims=[1, 2]) trb.initialize(datasets=datasets) exp_eps = 0.5 * (search_space.upper - search_space.lower) / 5.0 ** (1.0 / 2.0) @@ -1336,11 +1435,15 @@ 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.5, 0.5, 0.6], [0.0, -0.4, 0.0, 0.3], [1.0, 0.9, 1.0, 0.1]], + dtype=tf.float64, + ), tf.constant([[0.5], [0.0], [1.0]], dtype=tf.float64), ) } - trb = SingleObjectiveTrustRegionBox(search_space, min_eps=0.1) + # Includes a quick test of input_active_dims. The irrelevant input dimension should be ignored. + trb = SingleObjectiveTrustRegionBox(search_space, min_eps=0.1, input_active_dims=[0, 2]) trb.initialize(datasets=datasets) trb.location = tf.constant([0.5, 0.5], dtype=tf.float64) location = trb.location @@ -1860,6 +1963,286 @@ def test_multi_trust_region_box_state_deepcopy() -> None: npt.assert_array_equal(subspace._y_min, subspace_copy._y_min) +@pytest.fixture +def discrete_search_space() -> DiscreteSearchSpace: + return DiscreteSearchSpace(np.arange(10, dtype=np.float64)[:, None]) + + +@pytest.fixture +def continuous_search_space() -> Box: + return Box([0.0], [1.0]) + + +@pytest.mark.parametrize("with_initialize", [True, False]) +def test_fixed_trust_region_discrete_initialize( + discrete_search_space: DiscreteSearchSpace, with_initialize: bool +) -> None: + # Check that FixedTrustRegionDiscrete inits correctly by picking a single point from the global + # search space. + tr = FixedPointTrustRegionDiscrete(discrete_search_space) + if with_initialize: + tr.initialize() + assert tr.location.shape == (1,) + assert tr.location in discrete_search_space + + +def test_fixed_trust_region_discrete_update( + discrete_search_space: DiscreteSearchSpace, +) -> None: + # Update call should not change the location of the region. + tr = FixedPointTrustRegionDiscrete(discrete_search_space) + tr.initialize() + orig_location = tr.location.numpy() + tr.update() + npt.assert_equal(orig_location, tr.location) + + +def test_updatable_tr_product_raises_on_no_regions() -> None: + with pytest.raises(AssertionError, match="at least one region should be provided"): + UpdatableTrustRegionProduct([]) + + +def test_updatable_tr_product_raises_on_missing_index( + discrete_search_space: DiscreteSearchSpace, continuous_search_space: Box +) -> None: + region1 = FixedPointTrustRegionDiscrete(discrete_search_space, region_index=0) + region2 = SingleObjectiveTrustRegionBox(continuous_search_space, region_index=1) + with pytest.raises(AssertionError, match="regions can only have a region_index"): + UpdatableTrustRegionProduct([region1, region2]) + + +def test_updatable_tr_product_raises_on_mismatch_index( + discrete_search_space: DiscreteSearchSpace, continuous_search_space: Box +) -> None: + region1 = FixedPointTrustRegionDiscrete(discrete_search_space, region_index=0) + region2 = SingleObjectiveTrustRegionBox(continuous_search_space, region_index=1) + with pytest.raises(AssertionError, match="all regions should have the same index"): + UpdatableTrustRegionProduct([region1, region2], region_index=0) + + +def test_updatable_tr_product_raises_on_active_dims_set( + discrete_search_space: DiscreteSearchSpace, continuous_search_space: Box +) -> None: + region1 = FixedPointTrustRegionDiscrete(discrete_search_space, input_active_dims=[0]) + region2 = SingleObjectiveTrustRegionBox(continuous_search_space) + with pytest.raises(AssertionError, match="input_active_dims ..0.. should not be set"): + UpdatableTrustRegionProduct([region1, region2]) + + +def test_updatable_tr_product_sets_all_region_indices( + discrete_search_space: DiscreteSearchSpace, continuous_search_space: Box +) -> None: + region1 = FixedPointTrustRegionDiscrete(discrete_search_space, region_index=None) + region2 = SingleObjectiveTrustRegionBox(continuous_search_space, region_index=1) + tr = UpdatableTrustRegionProduct([region1, region2], region_index=1) + + assert tuple(tr.regions.keys()) == tr.subspace_tags + assert list(tr.regions.values()) == [region1, region2] + + assert next(iter(tr.regions.values())).region_index == 1 + assert len(set([region.region_index for region in tr.regions.values()])) == 1 + tr.region_index = 10 + assert next(iter(tr.regions.values())).region_index == 10 + assert len(set([region.region_index for region in tr.regions.values()])) == 1 + + +def test_updatable_tr_product_location( + discrete_search_space: DiscreteSearchSpace, continuous_search_space: Box +) -> None: + # Check the combined locations of the subregions. + region1 = FixedPointTrustRegionDiscrete(discrete_search_space) + region2 = SingleObjectiveTrustRegionBox(continuous_search_space) + tr = UpdatableTrustRegionProduct([region1, region2]) + + assert tr.location.dtype == tf.float64 + npt.assert_array_equal( + tr.location, np.concatenate([region1.location, region2.location], axis=-1) + ) + + +@pytest.mark.parametrize( + "datasets_only_arg, method", + [ + (False, lambda tr: tr.initialize), + (False, lambda tr: tr.update), + (True, lambda tr: tr.get_datasets_filter_mask), + ], +) +@pytest.mark.parametrize( + "datasets", + [ + None, + { + OBJECTIVE: Dataset( + tf.constant([[3.0, 0.5], [1.0, 0.0], [2.0, 1.0]], dtype=tf.float64), + tf.constant([[0.5], [0.0], [1.0]], dtype=tf.float64), + ) + }, + ], +) +def test_updatable_tr_product_method_calls_subregions( + datasets_only_arg: bool, + method: Callable[ + [ + UpdatableTrustRegion, + ], + Callable[..., Any], # We can have different signatures for the methods. + ], + datasets: Optional[Mapping[Tag, Dataset]], +) -> None: + # Check that calling initialize/update/* should call the initialize/update/* method of all + # subregions with the correct arguments. + region1 = MagicMock( + spec=FixedPointTrustRegionDiscrete, region_index=None, input_active_dims=None, dimension=1 + ) + region2 = MagicMock( + spec=SingleObjectiveTrustRegionBox, region_index=None, input_active_dims=None, dimension=1 + ) + tr = UpdatableTrustRegionProduct([region1, region2], region_index=2) + + models = {OBJECTIVE: QuadraticMeanAndRBFKernel()} + + if datasets_only_arg: + method(tr)(datasets) + else: + method(tr)(models, datasets, "dummy_arg", dummy_kwarg="dummy_kwarg_value") + + for region in [region1, region2]: + # Can't use region1.*.assert_called_once_with() directly as bool comparison + # doesn't work with datasets. So we check the call_args instead. + mock = method(region) + mock.assert_called_once() # type: ignore[attr-defined] + call_args = mock.call_args # type: ignore[attr-defined] + if datasets_only_arg: + call_dataset = call_args[0][0] + else: + print(call_args[1]) + assert call_args[1] == {"dummy_kwarg": "dummy_kwarg_value"} + assert call_args[0][0] == models + call_dataset = call_args[0][1] + assert call_args[0][2] == "dummy_arg" + + if datasets is None: + assert call_dataset is None + else: + assert datasets.keys() == call_dataset.keys() + for key in datasets: + npt.assert_array_equal(datasets[key].query_points, call_dataset[key].query_points) + npt.assert_array_equal( + datasets[key].observations, + call_dataset[key].observations, + ) + + +def test_updatable_tr_product_datasets_filter_mask_raises_on_missing_index() -> None: + region1 = MagicMock( + spec=FixedPointTrustRegionDiscrete, region_index=None, input_active_dims=None, dimension=1 + ) + region2 = MagicMock( + spec=SingleObjectiveTrustRegionBox, region_index=None, input_active_dims=None, dimension=1 + ) + tr = UpdatableTrustRegionProduct([region1, region2], region_index=None) + + datasets = {OBJECTIVE: empty_dataset([2], [1])} + with pytest.raises(AssertionError, match="the region_index should be set for filtering"): + tr.get_datasets_filter_mask(datasets) + + +def test_updatable_tr_product_datasets_filter_mask_value() -> None: + # Calling get_datasets_filter_mask on the product region returns a boolean AND of the masks + # returned by the subregions. + region1 = MagicMock( + spec=FixedPointTrustRegionDiscrete, region_index=None, input_active_dims=None, dimension=1 + ) + region1.get_datasets_filter_mask.return_value = { + "tag1": tf.constant([True, False, True], dtype=tf.bool), + "tag2": tf.constant([True, True, False], dtype=tf.bool), + } + region2 = MagicMock( + spec=SingleObjectiveTrustRegionBox, region_index=None, input_active_dims=None, dimension=1 + ) + region2.get_datasets_filter_mask.return_value = { + "tag1": tf.constant([True, False, False], dtype=tf.bool), + "tag2": tf.constant([True, True, True], dtype=tf.bool), + } + tr = UpdatableTrustRegionProduct([region1, region2], region_index=3) + + datasets = {OBJECTIVE: empty_dataset([2], [1])} + mask = tr.get_datasets_filter_mask(datasets) + assert mask is not None + assert mask.keys() == {"tag1", "tag2"} + npt.assert_array_equal(mask["tag1"], [True, False, False]) + npt.assert_array_equal(mask["tag2"], [True, True, False]) + + +@pytest.mark.parametrize( + "rule, exp_num_subspaces", + [ + (EfficientGlobalOptimization(), 1), + (EfficientGlobalOptimization(ParallelContinuousThompsonSampling(), num_query_points=2), 2), + (RandomSampling(num_query_points=2), 1), + ], +) +def test_batch_trust_region_product_no_subspace( + discrete_search_space: DiscreteSearchSpace, + continuous_search_space: Box, + rule: AcquisitionRule[TensorType, SearchSpace, ProbabilisticModel], + exp_num_subspaces: int, +) -> None: + # Check batch trust region creates default subspaces when none are provided at init. + search_space = TaggedProductSearchSpace( + [discrete_search_space, continuous_search_space, discrete_search_space] + ) + tr_rule = BatchTrustRegionProduct(rule=rule) + tr_rule.acquire(search_space, {}) + + assert tr_rule._tags is not None + assert tr_rule._subspaces is not None + assert len(tr_rule._subspaces) == exp_num_subspaces + for i, (subspace, tag) in enumerate(zip(tr_rule._subspaces, tr_rule._tags)): + assert isinstance(subspace, UpdatableTrustRegionProduct) + assert subspace.global_search_space == search_space + assert tag == f"{i}" + + subregions = subspace.regions + assert len(subregions) == 3 + assert subregions.keys() == {"0", "1", "2"} + for r, t, g in zip( + subregions.values(), + [ + FixedPointTrustRegionDiscrete, + SingleObjectiveTrustRegionBox, + FixedPointTrustRegionDiscrete, + ], + [discrete_search_space, continuous_search_space, discrete_search_space], + ): + assert isinstance(r, t) + assert r.global_search_space == g + + +def test_batch_trust_region_product_raises_for_wrong_search_space() -> None: + search_space = Box([0.0], [1.0]) + tr_rule = BatchTrustRegionProduct() # type: ignore[var-annotated] + with pytest.raises(AssertionError, match="search_space should be a TaggedProductSearchSpace"): + tr_rule.acquire(search_space, {}) + + +def test_batch_trust_region_product_raises_for_mismatched_search_space( + discrete_search_space: DiscreteSearchSpace, continuous_search_space: Box +) -> None: + search_space = TaggedProductSearchSpace( + [discrete_search_space, continuous_search_space, discrete_search_space] + ) + tr_rule = BatchTrustRegionProduct() # type: ignore[var-annotated] + tr_rule.acquire(search_space, {}) + + different_search_space = TaggedProductSearchSpace( + [discrete_search_space, continuous_search_space] + ) + with pytest.raises(AssertionError, match="The global search space of the subspaces should"): + tr_rule.acquire(different_search_space, {}) + + def test_asynchronous_rule_state_pending_points() -> None: pending_points = tf.constant([[1], [2], [3]]) diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index ae1a1f468f..774cb6c719 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -26,6 +26,7 @@ Any, Callable, Generic, + List, Optional, Sequence, Set, @@ -60,7 +61,13 @@ TrainableSupportsGetKernel, ) from ..observer import OBJECTIVE -from ..space import Box, SearchSpace, TaggedMultiSearchSpace +from ..space import ( + Box, + DiscreteSearchSpace, + SearchSpace, + TaggedMultiSearchSpace, + TaggedProductSearchSpace, +) from ..types import State, Tag, TensorType from ..utils.misc import LocalizedTag from .function import ( @@ -987,13 +994,24 @@ def acquire( class UpdatableTrustRegion(SearchSpace): """A search space that can be updated.""" - def __init__(self, region_index: Optional[int] = None) -> None: + def __init__( + self, + region_index: Optional[int] = None, + input_active_dims: Optional[Union[slice, Sequence[int]]] = None, + ) -> None: """ :param region_index: The index of the region in a multi-region search space. This is used to identify the local models and datasets to use for acquisition. If `None`, the global models and datasets are used. + :param input_active_dims: The active dimensions of the input space, either a slice or list + of indices into the columns of the space. If `None`, all dimensions are active. + + When this region is part of a product search-space (via `UpdatableTrustRegionProduct`), + this is used to select the active dimensions of the full input space that belong to this + region. """ self.region_index = region_index + self.input_active_dims = input_active_dims @abstractmethod def initialize( @@ -1038,7 +1056,75 @@ def _get_tags(self, tags: Set[Tag]) -> Tuple[Set[Tag], Set[Tag]]: return local_gtags, global_tags - def select_in_region(self, mapping: Optional[Mapping[Tag, T]]) -> Optional[Mapping[Tag, T]]: + @overload + def with_input_active_dims(self, value: TensorType) -> TensorType: + ... + + @overload + def with_input_active_dims(self, value: Dataset) -> Dataset: + ... + + @overload + def with_input_active_dims(self, value: ProbabilisticModel) -> ProbabilisticModel: + ... + + def with_input_active_dims( + self, value: Union[TensorType, Dataset, ProbabilisticModel] + ) -> Union[TensorType, Dataset, ProbabilisticModel]: + """ + Select and return active components from the input dimensions of the given value, using + `input_active_dims` of this search space. If `input_active_dims` is `None`, all dimensions + are returned. + + For datasets, the active selection is applied to the query points. For models, no + selection is applied; they are returned as is. + + :param value: The value to select the active input dimensions for. + :return: The value with the active input dimensions selected. + """ + + # No selection for models. + # Nothing to do if active dimensions are not set. + if isinstance(value, ProbabilisticModel) or self.input_active_dims is None: + return value + + # Select components of query points for datasets. + if isinstance(value, Dataset): + input = value.query_points + else: + input = value + + if isinstance(self.input_active_dims, slice): + selected_input = input[..., self.input_active_dims] + elif self.input_active_dims is not None: + selected_input = tf.gather(input, self.input_active_dims, axis=-1) + + if isinstance(value, Dataset): + return Dataset(selected_input, value.observations) + else: + return selected_input + + @overload + def select_in_region(self, mapping: None) -> None: + ... + + @overload + def select_in_region(self, mapping: Mapping[Tag, TensorType]) -> Mapping[Tag, TensorType]: + ... + + @overload + def select_in_region(self, mapping: Mapping[Tag, Dataset]) -> Mapping[Tag, Dataset]: + ... + + @overload + def select_in_region( + self, mapping: Mapping[Tag, ProbabilisticModel] + ) -> Mapping[Tag, ProbabilisticModel]: + ... + + def select_in_region( + self, mapping: Optional[Mapping[Tag, Union[TensorType, Dataset, ProbabilisticModel]]] + ) -> Optional[Mapping[Tag, Union[TensorType, Dataset, ProbabilisticModel]]]: """ Select items belonging to this region for acquisition. @@ -1050,7 +1136,7 @@ def select_in_region(self, mapping: Optional[Mapping[Tag, T]]) -> Optional[Mappi elif self.region_index is None: # If no index, then return the global items. _mapping = { - tag: item + tag: self.with_input_active_dims(item) for tag, item in mapping.items() if not LocalizedTag.from_tag(tag).is_local } @@ -1061,9 +1147,9 @@ def select_in_region(self, mapping: Optional[Mapping[Tag, T]]) -> Optional[Mappi _mapping = {} for tag in local_gtags: ltag = LocalizedTag(tag, self.region_index) - _mapping[ltag] = mapping[ltag] + _mapping[ltag] = self.with_input_active_dims(mapping[ltag]) for tag in global_tags: - _mapping[tag] = mapping[tag] + _mapping[tag] = self.with_input_active_dims(mapping[tag]) return _mapping if _mapping else None @@ -1088,7 +1174,7 @@ def get_datasets_filter_mask( else: # Only keep points that are in the region. return { - tag: self.contains(dataset.query_points) + tag: self.contains(self.with_input_active_dims(dataset.query_points)) for tag, dataset in datasets.items() if LocalizedTag.from_tag(tag).local_index == self.region_index } @@ -1394,15 +1480,18 @@ def __init__( self, global_search_space: SearchSpace, region_index: Optional[int] = None, + input_active_dims: Optional[Union[slice, Sequence[int]]] = None, ): """ :param global_search_space: The global search space this search space lives in. :param region_index: The index of the region in a multi-region search space. This is used to identify the local models and datasets to use for acquisition. If `None`, the global models and datasets are used. + :param input_active_dims: The active dimensions of the input space, either a slice or list + of indices into the columns of the space. If `None`, all dimensions are active. """ Box.__init__(self, global_search_space.lower, global_search_space.upper) - UpdatableTrustRegion.__init__(self, region_index) + UpdatableTrustRegion.__init__(self, region_index, input_active_dims) self._global_search_space = global_search_space # Random initial location in the global search space. self.location = tf.squeeze(global_search_space.sample(1), axis=0) @@ -1432,6 +1521,7 @@ def __init__( kappa: float = 1e-4, min_eps: float = 1e-2, region_index: Optional[int] = None, + input_active_dims: Optional[Union[slice, Sequence[int]]] = None, ): """ Calculates the bounds of the box from the location/centre and global bounds. @@ -1445,8 +1535,10 @@ def __init__( :param region_index: The index of the region in a multi-region search space. This is used to identify the local models and datasets to use for acquisition. If `None`, the global models and datasets are used. + :param input_active_dims: The active dimensions of the input space, either a slice or list + of indices into the columns of the space. If `None`, all dimensions are active. """ - super().__init__(global_search_space, region_index) + super().__init__(global_search_space, region_index, input_active_dims) self._beta = beta self._kappa = kappa self._min_eps = min_eps @@ -1631,9 +1723,17 @@ def __init__( kappa: float = 1e-4, min_eps: float = 1e-2, region_index: Optional[int] = None, + input_active_dims: Optional[Union[slice, Sequence[int]]] = None, ): self._is_global = False - super().__init__(global_search_space, beta, kappa, min_eps, region_index) + super().__init__( + global_search_space, + beta, + kappa, + min_eps, + region_index=region_index, + input_active_dims=input_active_dims, + ) @property def eps(self) -> TensorType: @@ -1684,7 +1784,9 @@ def get_datasets_filter_mask( else: # Don't filter out any points from the dataset. Always keep the entire dataset. return { - tag: tf.ones(tf.shape(dataset.query_points)[:-1], dtype=tf.bool) + tag: tf.ones( + tf.shape(self.with_input_active_dims(dataset.query_points))[:-1], dtype=tf.bool + ) for tag, dataset in datasets.items() if LocalizedTag.from_tag(tag).local_index == self.region_index } @@ -1722,6 +1824,7 @@ def __init__( success_tolerance: int = 3, failure_tolerance: Optional[int] = None, region_index: Optional[int] = None, + input_active_dims: Optional[Union[slice, Sequence[int]]] = None, ): """ Note that the optional parameters are set by a heuristic if not given by the user. @@ -1735,8 +1838,10 @@ def __init__( :param region_index: The index of the region in a multi-region search space. This is used to identify the local models and datasets to use for acquisition. If `None`, the global models and datasets are used. + :param input_active_dims: The active dimensions of the input space, either a slice or list + of indices into the columns of the space. If `None`, all dimensions are active. """ - super().__init__(global_search_space, region_index) + super().__init__(global_search_space, region_index, input_active_dims) search_space_max_width = tf.reduce_max( global_search_space.upper - global_search_space.lower @@ -1798,6 +1903,11 @@ def _set_tr_width(self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = lengthscales = ( model.get_kernel().lengthscales ) # stretch region according to model lengthscales + + # Select the input lengthscales that are active for this region. + if tf.size(lengthscales) > 1: + lengthscales = self.with_input_active_dims(lengthscales) + self.tr_width = ( lengthscales * self.L @@ -1889,6 +1999,306 @@ def get_dataset_min( return tf.squeeze(x_min, axis=0), tf.squeeze(y_min) +class UpdatableTrustRegionDiscrete(DiscreteSearchSpace, UpdatableTrustRegion): + """ + An updatable discrete search space with an associated global search space. + """ + + def __init__( + self, + global_search_space: DiscreteSearchSpace, + region_index: Optional[int] = None, + input_active_dims: Optional[Union[slice, Sequence[int]]] = None, + ): + """ + :param global_search_space: The global search space this search space lives in. + :param region_index: The index of the region in a multi-region search space. This is used to + identify the local models and datasets to use for acquisition. If `None`, the + global models and datasets are used. + :param input_active_dims: The active dimensions of the input space, either a slice or list + of indices into the columns of the space. If `None`, all dimensions are active. + """ + # Ensure global_points is a copied tensor, in case a variable is passed in. + DiscreteSearchSpace.__init__(self, tf.constant(global_search_space.points)) + UpdatableTrustRegion.__init__(self, region_index, input_active_dims) + self._global_search_space = global_search_space + + @property + @abstractmethod + def location(self) -> TensorType: + """Center of the region.""" + + @property + def global_search_space(self) -> DiscreteSearchSpace: + """The global space this search space lives in.""" + return self._global_search_space + + +class FixedPointTrustRegionDiscrete(UpdatableTrustRegionDiscrete): + """ + A discrete trust region with a fixed point location that does not change across active learning + steps. The fixed point is selected at random from the global (discrete) search space at + initialization time. + """ + + def __init__( + self, + global_search_space: DiscreteSearchSpace, + region_index: Optional[int] = None, + input_active_dims: Optional[Union[slice, Sequence[int]]] = None, + ): + super().__init__(global_search_space, region_index, input_active_dims) + # Random initial point from the global search space. + self._points = self.global_search_space.sample(1) + + @property + def location(self) -> TensorType: + """The location point of the region.""" + return tf.squeeze(self._points, axis=0) # Only one point. + + def initialize( + self, + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> None: + # Pick a random point from the global search space. + self._points = self.global_search_space.sample(1) + + def update( + self, + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> None: + # Keep the point fixed, no updates needed. + pass + + +UpdatableTrustRegionWithGlobalSearchSpace = Union[ + UpdatableTrustRegionBox, UpdatableTrustRegionDiscrete +] +"""A type alias for updatable trust regions with a global search space.""" + + +class UpdatableTrustRegionProduct(TaggedProductSearchSpace, UpdatableTrustRegion): + """ + An updatable mixed search space that is the product of multiple updatable trust sub-regions. + + This is useful for combining different types of search spaces, such as continuous and discrete, + to form a mixed search space for trust region acquisition rules. + + Note: the dtype of all the component search spaces must be the same. + """ + + def __init__( + self, + regions: Sequence[UpdatableTrustRegionWithGlobalSearchSpace], + tags: Optional[Sequence[str]] = None, + region_index: Optional[int] = None, + ): + """ + :param regions: The trust sub-regions to be combined to create a product trust region. + :param tags: An optional list of tags giving the unique identifiers of the region's + sub-regions. + :param region_index: The index of the region in a multi-region search space. This is used to + identify the local models and datasets to use for acquisition. If `None`, the + global models and datasets are used. + """ + assert len(regions) > 0, "at least one region should be provided" + + # If set, assert all regions have the same index and matching the product index. + if region_index is not None: + assert all( + region.region_index == region_index + for region in regions + if region.region_index is not None + ), ( + "all regions should have the same index, if set, as the " + f"product region ({region_index})" + ) + else: + assert all(region.region_index is None for region in regions), ( + f"regions can only have a region_index if the product region ({region_index}) " + "has one" + ) + + self._global_search_space = TaggedProductSearchSpace( + [region.global_search_space for region in regions], tags + ) + + TaggedProductSearchSpace.__init__(self, regions) + # When UpdatableTrustRegion sets the region_index, it will also set the region_index for + # each region. + # Setting of input active dims is not supported for product regions. All input dims + # are always active. + UpdatableTrustRegion.__init__(self, region_index, input_active_dims=None) + + # Set active dimensions for each sub-region. + dim_ix = 0 + for region, dims in zip(regions, self.subspace_dimension): + # Check the region's input active dims are not already set. + assert region.input_active_dims is None, ( + f"input_active_dims ({region.input_active_dims}) should not be set for sub-regions " + f" ({region}) of a product region" + ) + region.input_active_dims = slice(dim_ix, dim_ix + dims) + dim_ix += dims + + @property + def region_index(self) -> Optional[int]: + """The index of the region in a multi-region search space.""" + return self._region_index + + @region_index.setter + def region_index(self, region_index: Optional[int]) -> None: + """Set the index of the region in a multi-region search space, including all sub-regions.""" + self._region_index = region_index + # Override the region index for each sub-region. These would either already be set to the + # same value (assert in __init__), or None. + for region in self.regions.values(): + region.region_index = region_index + + @property + def regions(self) -> Mapping[str, UpdatableTrustRegionWithGlobalSearchSpace]: + """The sub-regions of the product trust region.""" + _regions = {} + for tag, region in self._spaces.items(): + assert isinstance(region, (UpdatableTrustRegionBox, UpdatableTrustRegionDiscrete)) + _regions[tag] = region + return _regions + + @property + def location(self) -> TensorType: + """ + The location of the product trust region, concatenated from the locations of the + sub-regions. + """ + return tf.concat([region.location for region in self.regions.values()], axis=-1) + + @property + def global_search_space(self) -> TaggedProductSearchSpace: + """The global search space this search space lives in.""" + return self._global_search_space + + def initialize( + self, + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, + datasets: Optional[Mapping[Tag, Dataset]] = None, + *args: Any, + **kwargs: Any, + ) -> None: + for region in self.regions.values(): + region.initialize(models, datasets, *args, **kwargs) + + def update( + self, + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, + datasets: Optional[Mapping[Tag, Dataset]] = None, + *args: Any, + **kwargs: Any, + ) -> None: + for region in self.regions.values(): + region.update(models, datasets, *args, **kwargs) + + def get_datasets_filter_mask( + self, datasets: Optional[Mapping[Tag, Dataset]] + ) -> Optional[Mapping[Tag, tf.Tensor]]: + # Return a boolean AND of the masks of each sub-region. + + # Only select the region datasets for filtering. Don't directly filter the global dataset. + assert ( + self.region_index is not None + ), "the region_index should be set for filtering local datasets" + + # Mask for each sub-region. + masks = [region.get_datasets_filter_mask(datasets) for region in self.regions.values()] + + if masks[0] is not None: # There is always at least one region. + assert all( + set(mask.keys()) == set(masks[0].keys()) for mask in masks if mask is not None + ), "all region masks should have the same keys" + + return { + tag: tf.reduce_all([mask[tag] for mask in masks if mask is not None], axis=0) + for tag in masks[0].keys() + } + else: + return None + + +class BatchTrustRegionProduct( + BatchTrustRegion[ProbabilisticModelType, UpdatableTrustRegionProduct] +): + """ + Implements the :class:`BatchTrustRegion` *trust region* acquisition rule for mixed search + spaces. This is intended to be used for single-objective optimization with batching. + """ + + def acquire( + self, + search_space: SearchSpace, + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> types.State[BatchTrustRegion.State | None, TensorType]: + if self._subspaces is None: + # If no initial subspaces were provided, create N default subspaces, where N is the + # number of query points in the base-rule. + # Currently the detection for N is only implemented for EGO. + # Note: the reason we don't create the default subspaces in `__init__` is because we + # don't have the global search space at that point. + if isinstance(self._rule, EfficientGlobalOptimization): + num_query_points = self._rule._num_query_points + else: + num_query_points = 1 + + def create_subregions() -> Sequence[UpdatableTrustRegionWithGlobalSearchSpace]: + # Take a global product search space and convert each of its subspaces to an + # updatable trust sub-region. These sub-regions are then used to create a + # trust region product. + assert isinstance( + search_space, TaggedProductSearchSpace + ), "search_space should be a TaggedProductSearchSpace" + + subregions: List[UpdatableTrustRegionWithGlobalSearchSpace] = [] + for tag in search_space.subspace_tags: + subspace = search_space.get_subspace(tag) + if isinstance(subspace, DiscreteSearchSpace): + subregions.append(FixedPointTrustRegionDiscrete(subspace)) + else: + subregions.append(SingleObjectiveTrustRegionBox(subspace)) + return subregions + + init_subspaces: Tuple[UpdatableTrustRegionProduct, ...] = tuple( + UpdatableTrustRegionProduct(create_subregions()) for _ in range(num_query_points) + ) + self._subspaces = init_subspaces + for index, subspace in enumerate(self._subspaces): + subspace.region_index = index # Override the index. + self._tags = tuple(str(index) for index in range(self.num_local_datasets)) + + # Ensure passed in global search space is always the same as the search space passed to + # the subspaces. + for subspace in self._subspaces: + assert subspace.global_search_space == search_space, ( + "The global search space of the subspaces should be the same as the " + "search space passed to the BatchTrustRegionProduct acquisition rule. " + "If you want to change the global search space, you should recreate the rule. " + "Note: all subspaces should be initialized with the same global search space." + ) + + return super().acquire(search_space, models, datasets) + + @inherit_check_shapes + def get_initialize_subspaces_mask( + self, + subspaces: Sequence[UpdatableTrustRegionProduct], + 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 BatchHypervolumeSharpeRatioIndicator( AcquisitionRule[TensorType, SearchSpace, ProbabilisticModel] ): diff --git a/trieste/experimental/plotting/plotting.py b/trieste/experimental/plotting/plotting.py index 521b34f2c0..48200912a0 100644 --- a/trieste/experimental/plotting/plotting.py +++ b/trieste/experimental/plotting/plotting.py @@ -555,6 +555,7 @@ def plot_trust_region_history_2d( history: Record[StateType, ProbabilisticModel] | FrozenRecord[StateType, ProbabilisticModel], num_query_points: Optional[int] = None, num_init: Optional[int] = None, + alpha: float = 0.3, ) -> tuple[Optional[Figure], Optional[Axes]]: """ Plot the contour of the objective function, query points and the trust regions for a particular @@ -566,6 +567,7 @@ def plot_trust_region_history_2d( :param history: the optimization history for a particular step of the optimization process :param num_query_points: total number of query points in this step :param num_init: initial number of BO points + :param alpha: transparency for the trust regions :return: figure and axes """ @@ -645,7 +647,7 @@ def plot_trust_region_history_2d( ub[1] - lb[1], facecolor=space_colors[i], edgecolor=space_colors[i], - alpha=0.3, + alpha=alpha, ) ) diff --git a/trieste/space.py b/trieste/space.py index ae28253215..3c750eb583 100644 --- a/trieste/space.py +++ b/trieste/space.py @@ -999,6 +999,8 @@ class TaggedProductSearchSpace(CollectionSearchSpace): decision_space = Box([-1, -2], [2, 3]) mixed_space = TaggedProductSearchSpace(spaces=[context_space, decision_space]) + Note: the dtype of all the component search spaces must be the same. + 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.