From 905dc149cae2acaf70fd59567b5fdf32e6985bb0 Mon Sep 17 00:00:00 2001 From: Victor Picheny Date: Thu, 18 May 2023 09:13:56 +0100 Subject: [PATCH 01/33] multiTR with notebook --- trieste/acquisition/optimizer.py | 61 ++++++--- trieste/acquisition/rule.py | 209 ++++++++++++++++++++++++++++++- trieste/space.py | 9 ++ 3 files changed, 262 insertions(+), 17 deletions(-) diff --git a/trieste/acquisition/optimizer.py b/trieste/acquisition/optimizer.py index 1f9dd77b06..cb9481e2ea 100644 --- a/trieste/acquisition/optimizer.py +++ b/trieste/acquisition/optimizer.py @@ -32,6 +32,7 @@ Box, Constraint, DiscreteSearchSpace, + MultiBoxSearchSpace, SearchSpace, SearchSpaceType, TaggedProductSearchSpace, @@ -100,9 +101,14 @@ 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, TaggedProductSearchSpace, MultiBoxSearchSpace)): + if isinstance(space, MultiBoxSearchSpace): + space_dim = tf.shape(space.get_subspace(space.subspace_tags[0]).lower)[-1] + else: + space_dim = tf.shape(space.lower)[-1] + + 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,7 +176,7 @@ 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 | TaggedProductSearchSpace | MultiBoxSearchSpace]: """ 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 @@ -216,7 +222,7 @@ 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 | TaggedProductSearchSpace | MultiBoxSearchSpace, target_func: Union[AcquisitionFunction, Tuple[AcquisitionFunction, int]], ) -> TensorType: """ @@ -246,8 +252,17 @@ 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] + if isinstance(space, MultiBoxSearchSpace): + candidates = [ + space.get_subspace(tag).sample(num_initial_samples)[:, None, :] + for tag in space.subspace_tags + ] # list of [num_initial_samples, 1, D] + tiled_candidates = tf.concat(candidates, axis=1) # [num_initial_samples, V, D] + else: + 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] target_func_values = target_func(tiled_candidates) # [num_samples, V] tf.debugging.assert_shapes( @@ -289,12 +304,19 @@ 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 + if isinstance(space, MultiBoxSearchSpace): + candidates = [ + space.get_subspace(tag).sample(num_initial_samples)[:, None, :] + for tag in space.subspace_tags + ] # list of [num_initial_samples, 1, D] + tiled_random_points = tf.concat(candidates, axis=1) # [num_initial_samples, V, D] + else: + 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] ( recovery_successes, recovery_fun_values, @@ -350,7 +372,9 @@ def improvements() -> tf.Tensor: if V == 1: logging.scalar("spo_improvement_on_initial_samples", improvements) else: - logging.histogram("spo_improvements_on_initial_samples", improvements) + logging.histogram( + "spo_improvement_on_initial_samples_across_subspaces", improvements + ) best_run_ids = tf.math.argmax(fun_values, axis=0) # [V] chosen_points = tf.gather( @@ -427,7 +451,14 @@ def _objective_value(vectorized_x: TensorType) -> TensorType: # [N, D] -> [N, 1 def _objective_value_and_gradient(x: TensorType) -> Tuple[TensorType, TensorType]: return tfp.math.value_and_gradient(_objective_value, x) # [len(x), 1], [len(x), D] - if isinstance( + if isinstance(space, MultiBoxSearchSpace): + bounds = [ + spo.Bounds(space.get_subspace(tag).lower, space.get_subspace(tag).upper) + for tag in space.subspace_tags + ] + # TODO: check that this creates the right set of bou + bounds = bounds * num_optimization_runs_per_function + elif isinstance( space, TaggedProductSearchSpace ): # build continuous relaxation of discrete subspaces bounds = [ diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index d731ec59d4..82f2ae3b2e 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -21,7 +21,7 @@ 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, TypeVar, Union, cast, overload import numpy as np @@ -41,7 +41,7 @@ from ..models import ProbabilisticModel from ..models.interfaces import HasReparamSampler, ModelStack, ProbabilisticModelType from ..observer import OBJECTIVE -from ..space import Box, SearchSpace +from ..space import Box, MultiBoxSearchSpace, SearchSpace from ..types import State, Tag, TensorType from .function import ( BatchMonteCarloExpectedImprovement, @@ -1097,6 +1097,211 @@ def state_func( return state_func +class MultiTrustRegion( + AcquisitionRule[ + types.State[Optional["MultiTrustRegion.State"], TensorType], Box, ProbabilisticModelType + ] +): + """Implements the *trust region* acquisition algorithm.""" + + @dataclass(frozen=True) + class State: + """The acquisition state for the :class:`TrustRegion` acquisition rule.""" + + acquisition_space: MultiBoxSearchSpace + """ The search space. """ + + eps: Sequence[TensorType] + """ + The (maximum) vector from the current best point to each bound of the acquisition space. + """ + + y_min: Sequence[TensorType] + """ The minimum observed value. """ + + # is_global: Sequence[bool] | Sequence[TensorType] + # """ + # `True` if the search space was global, else `False` if it was local. + # May be a scalar boolean `TensorType` instead of a `bool`. + # """ + + def __deepcopy__(self, memo: dict[int, object]) -> MultiTrustRegion.State: + box_copy = copy.deepcopy(self.acquisition_space, memo) + return MultiTrustRegion.State(box_copy, self.eps, self.y_min) # , self.is_global) + + @overload + def __init__( + self: "MultiTrustRegion[ProbabilisticModel]", + rule: None = None, + beta: float = 0.7, + kappa: float = 1e-4, + number_of_tr: int = 1, + ): + ... + + @overload + def __init__( + self: "MultiTrustRegion[ProbabilisticModelType]", + rule: AcquisitionRule[TensorType, Box, ProbabilisticModelType], + beta: float = 0.7, + kappa: float = 1e-4, + number_of_tr: int = 1, + ): + ... + + def __init__( + self, + rule: AcquisitionRule[TensorType, Box, ProbabilisticModelType] | None = None, + beta: float = 0.7, + kappa: float = 1e-4, + number_of_tr: int = 1, + ): + if rule is None: + rule = EfficientGlobalOptimization() + + self._rule = rule + self._beta = beta + self._kappa = kappa + self._min_eps = 1e-2 + self._tags = [str(index) for index in range(number_of_tr)] + + def __repr__(self) -> str: + """""" + return f"MultiTrustRegion({self._rule!r}, {self._beta!r}, {self._kappa!r})" + + def acquire( + self, + search_space: Box, + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> types.State[State | None, TensorType]: + if datasets is None or OBJECTIVE not in datasets.keys(): + raise ValueError(f"""datasets must be provided and contain the key {OBJECTIVE}""") + + dataset = datasets[OBJECTIVE] + + global_lower = search_space.lower + global_upper = search_space.upper + + # global_y_min = tf.reduce_min(dataset.observations, axis=0) + def state_func( + state: MultiTrustRegion.State | None, + ) -> tuple[MultiTrustRegion.State | None, TensorType]: + if state is None: + eps = { + tag: 0.5 + * (global_upper - global_lower) + / (5.0 ** (1.0 / global_lower.shape[-1])) + for tag in self._tags + } + + x_min = {tag: search_space.sample(1) for tag in self._tags} + + aspace = [ + Box( + tf.reduce_max([global_lower, (x_min[tag] - eps[tag])[0, :]], axis=0), + tf.reduce_min([global_upper, (x_min[tag] + eps[tag])[0, :]], axis=0), + ) + for tag in self._tags + ] + + acquisition_space = MultiBoxSearchSpace(aspace, self._tags) + + y_min = { + tag: get_local_y_min(dataset, acquisition_space.get_subspace(tag)) + for tag in self._tags + } + + else: + y_min = { + tag: get_local_y_min(dataset, state.acquisition_space.get_subspace(tag)) + for tag in state.acquisition_space.subspace_tags + } + + x_min = { + tag: get_local_x_min(dataset, state.acquisition_space.get_subspace(tag)) + for tag in state.acquisition_space.subspace_tags + } + eps = {} + for tag in state.acquisition_space.subspace_tags: + tr_volume = tf.reduce_prod( + state.acquisition_space.get_subspace(tag).upper + - state.acquisition_space.get_subspace(tag).lower + ) + step_is_success = y_min[tag] < state.y_min[tag] - self._kappa * tr_volume + eps[tag] = ( + state.eps[tag] / self._beta + if step_is_success + else state.eps[tag] * self._beta + ) + + x_min, y_min, eps = reinitialise_useless_trs( + x_min, y_min, eps, search_space, dataset, self._min_eps + ) + + aspace = [ + Box( + tf.reduce_max([global_lower, (x_min[tag] - eps[tag])[0, :]], axis=0), + tf.reduce_min([global_upper, (x_min[tag] + eps[tag])[0, :]], axis=0), + ) + for tag in state.acquisition_space.subspace_tags + ] + + acquisition_space = MultiBoxSearchSpace( + aspace, state.acquisition_space.subspace_tags + ) + + points = self._rule.acquire(acquisition_space, models, datasets=datasets) + state_ = MultiTrustRegion.State(acquisition_space, eps, y_min) + + return state_, points + + return state_func + + +def get_local_y_min(dataset: Dataset, search_space: SearchSpace): + in_tr = search_space.contains(dataset.query_points) + ones = tf.ones_like(dataset.observations) * 1e6 + in_tr_obs = tf.where(in_tr[:, None], dataset.observations, ones) + return tf.reduce_min(in_tr_obs, axis=0) + + +def get_local_x_min(dataset: Dataset, search_space: SearchSpace): + in_tr = search_space.contains(dataset.query_points) + ones = tf.ones_like(dataset.observations) * 1e6 + in_tr_obs = tf.where(in_tr[:, None], dataset.observations, ones) + return dataset.query_points[tf.argmin(in_tr_obs)[0], :][None, :] + + +def reinitialise_useless_trs( + x_min: Sequence[TensorType], + y_min: Sequence[TensorType], + eps: Sequence[TensorType], + search_space: Box, + dataset: Dataset, + min_eps: float, +) -> tuple[Sequence[TensorType], Sequence[TensorType], Sequence[TensorType]]: + xx = tf.zeros([0, dataset.query_points.shape[1]], dtype=dataset.query_points.dtype) + + global_lower = search_space.lower + global_upper = search_space.upper + + for tag, x in x_min.items(): + duplicated = tf.less_equal(tf.reduce_min(tf.reduce_sum(tf.abs(x - xx), axis=1)), 1e-6) + + if duplicated or (np.min(eps[tag].numpy()) < min_eps): + eps[tag] = 0.1 * (global_upper - global_lower) / (5.0 ** (1.0 / global_lower.shape[-1])) + x_min[tag] = search_space.sample(1) + temp_space = Box( + tf.reduce_max([global_lower, (x_min[tag] - eps[tag])[0, :]], axis=0), + tf.reduce_min([global_upper, (x_min[tag] + eps[tag])[0, :]], axis=0), + ) + y_min[tag] = get_local_y_min(dataset, temp_space) + xx = tf.concat([xx, x_min[tag]], axis=0) + + return x_min, y_min, eps + + class BatchHypervolumeSharpeRatioIndicator( AcquisitionRule[TensorType, SearchSpace, ProbabilisticModel] ): diff --git a/trieste/space.py b/trieste/space.py index b99818f6e4..1bfb819003 100644 --- a/trieste/space.py +++ b/trieste/space.py @@ -1077,3 +1077,12 @@ def __eq__(self, other: object) -> bool: def __deepcopy__(self, memo: dict[int, object]) -> TaggedProductSearchSpace: return self + + +class MultiBoxSearchSpace(TaggedProductSearchSpace): + def __repr__(self) -> str: + """""" + return f"""MultiBoxSearchSpace(spaces = + {[self.get_subspace(tag) for tag in self.subspace_tags]}, + tags = {self.subspace_tags}) + """ From 5c27203a65d8fd6ef63b059e33cb9e1236561616 Mon Sep 17 00:00:00 2001 From: Victor Picheny Date: Thu, 18 May 2023 12:06:23 +0100 Subject: [PATCH 02/33] notebook --- docs/notebooks/multi_trust_region.pct.py | 177 +++++++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 docs/notebooks/multi_trust_region.pct.py diff --git a/docs/notebooks/multi_trust_region.pct.py b/docs/notebooks/multi_trust_region.pct.py new file mode 100644 index 0000000000..a10d11fb52 --- /dev/null +++ b/docs/notebooks/multi_trust_region.pct.py @@ -0,0 +1,177 @@ +import numpy as np +import tensorflow as tf +from trieste.acquisition.rule import EfficientGlobalOptimization +from trieste.acquisition.rule import MultiTrustRegion, get_local_x_min +from trieste.acquisition import ParallelContinuousThompsonSampling +from trieste.acquisition.optimizer import automatic_optimizer_selector +from trieste.acquisition.utils import split_acquisition_function_calls +from trieste.ask_tell_optimization import AskTellOptimizer +from trieste.experimental.plotting import plot_regret +from matplotlib import pyplot as plt +from trieste.models.gpflow import ( + SparseVariational, + build_svgp, + build_gpr, + GaussianProcessRegression, + ConditionalImprovementReduction, +) +from trieste.models.optimizer import BatchOptimizer +import trieste +from trieste.objectives import ScaledBranin, Hartmann6 +from trieste.types import TensorType +from trieste.logging import pyplot +from matplotlib.pyplot import cm +# from aim.ext.tensorboard_tracker import Run +from datetime import datetime +from trieste.experimental.plotting.plotting import create_grid +from matplotlib.patches import Rectangle + + +np.random.seed(179) +tf.random.set_seed(179) + +# CONFIG +tensorboard_dir_1 = f'./results/{datetime.now()}/tensorboard' + +summary_writer = tf.summary.create_file_writer(tensorboard_dir_1) +trieste.logging.set_tensorboard_writer(summary_writer) + +obj = ScaledBranin.objective +search_space = ScaledBranin.search_space + + +def obj_fun( + x: TensorType, +) -> TensorType: # contaminate observations with Gaussian noise + return obj(x) # + tf.random.normal([len(x), 1], 0, .1, tf.float64) + + +num_initial_data_points = 6 +num_query_points = 3 +num_steps = 10 + +# aim_repo_dir = './results' +# +# run_1 = Run( +# experiment='TR', +# sync_tensorboard_log_dir=tensorboard_dir_1, +# repo=aim_repo_dir, +# system_tracking_interval=None, # Tracks CPU, GPU, memory usage periodically +# log_system_params=True, # Logs environment variables (DANGER), git hash, packages installed etc. +# capture_terminal_logs=True, # Stores the stdout (VERY USEFUL TO SEE WHY EXPERIMENTS BROKE) +# ) +# This can be any arbitrary tags that you want to associate with your run, I often at least +# add a semi-unique name that summarises the main components of the run, so its easy to see +# run_1.add_tag('gpr') + +# These are parameters associated wit the run. Ideally this would be *everything* you are likely +# to configure between different runs. If you use Hydra, you can store the dictionary representation +# of your config here, so no manual entry. +# run_1["hparams"] = { +# 'num_initial_data_points': num_initial_data_points, +# 'num_query_points': num_query_points, +# 'num_steps': num_steps, +# } + +initial_query_points = search_space.sample(num_initial_data_points) +observer = trieste.objectives.utils.mk_observer(obj_fun) +initial_data = observer(initial_query_points) + + +# gpflow_model = build_svgp( +# initial_data, search_space, likelihood_variance=0.001, num_inducing_points=50 +# ) +# +# inducing_point_selector = ConditionalImprovementReduction() +# +# model = SparseVariational( +# gpflow_model, +# num_rff_features=1000, +# inducing_point_selector=inducing_point_selector, +# optimizer=BatchOptimizer( +# tf.optimizers.Adam(0.05), max_iter=100, batch_size=50, compile=True +# ), +# ) +gpflow_model = build_gpr(initial_data, search_space, likelihood_variance=1e-4, trainable_likelihood=False) +model = GaussianProcessRegression(gpflow_model) + +base_rule = EfficientGlobalOptimization( + builder=ParallelContinuousThompsonSampling(), + num_query_points=num_query_points, + optimizer=split_acquisition_function_calls( + automatic_optimizer_selector, split_size=100_000), +) + +acq_rule = MultiTrustRegion(base_rule, number_of_tr=num_query_points) + +ask_tell = AskTellOptimizer(search_space, initial_data, model, fit_model=True, acquisition_rule=acq_rule) + +color = cm.rainbow(np.linspace(0, 1, num_query_points)) + +Xplot, xx, yy = create_grid(mins=search_space.lower, maxs=search_space.upper, grid_density=90) +ff = obj_fun(Xplot).numpy() + +for step in range(num_steps): + print(f"step number {step}") + trieste.logging.set_step_number(step) + + new_points = ask_tell.ask() + new_data = observer(new_points) + # monitor models after each tell + if summary_writer: + models = ask_tell._models # pylint: disable=protected-access + trieste.logging.set_step_number(step) + + with summary_writer.as_default(step=step): + for tag, model in models.items(): + with tf.name_scope(f"{tag}.model"): + model.log() + + fig, ax = plt.subplots(1, 2, squeeze=False, figsize=(15, 5)) + fig.suptitle(f"step number {step}") + ax[0, 0].scatter(ask_tell.dataset.query_points[:, 0].numpy(), ask_tell.dataset.query_points[:, 1].numpy(), color="blue") + ax[0, 0].scatter(new_points[:, 0].numpy(), new_points[:, 1].numpy(), color="red") + + state = ask_tell.acquisition_state + + xmin = {tag: get_local_x_min(ask_tell.dataset, state.acquisition_space.get_subspace(tag)) for tag in + state.acquisition_space.subspace_tags} + i = 0 + + ax[0, 1].contour(xx, yy, ff.reshape(*xx.shape), 80, alpha=0.5) + + for tag in state.acquisition_space.subspace_tags: + ax[0, 1].scatter(xmin[tag].numpy()[0, 0], xmin[tag].numpy()[0, 1], color=color[i], marker="x", alpha=0.5) + lb = state.acquisition_space.get_subspace(tag).lower + ub = state.acquisition_space.get_subspace(tag).upper + ax[0, 1].add_patch(Rectangle((lb[0], lb[1]), ub[0] - lb[0], ub[1] - lb[1], + facecolor=color[i], + edgecolor=color[i], + alpha=0.3)) + ax[0, 1].scatter(new_points[i, 0].numpy(), new_points[i, 1].numpy(), color=color[i], alpha=0.5) + ax[0, 1].scatter(ask_tell.dataset.query_points[:, 0].numpy(), + ask_tell.dataset.query_points[:, 1].numpy(), color="black", alpha=0.2) + i = i + 1 + + pyplot("Query points", fig) + plt.close(fig) + + ask_tell.tell(new_data) + +dataset = ask_tell.dataset + +ground_truth_regret = obj(dataset.query_points) - Hartmann6.minimum +best_found_truth_idx = tf.squeeze(tf.argmin(ground_truth_regret, axis=0)) + +fig, ax = plt.subplots() +plot_regret( + ground_truth_regret.numpy(), ax, num_init=10, idx_best=best_found_truth_idx +) + +ax.set_yscale("log") +ax.set_ylabel("Regret") +ax.set_xlabel("# evaluations") + +fig, ax = plt.subplots() +ax.scatter(dataset.query_points[:, 0], dataset.query_points[:, 1]) + From e46150d7d5146b736db7ca0c013d693012337572 Mon Sep 17 00:00:00 2001 From: Victor Picheny Date: Fri, 2 Jun 2023 09:44:29 +0100 Subject: [PATCH 03/33] multi local step TREGO --- trieste/acquisition/rule.py | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 82f2ae3b2e..d8623b098b 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -961,9 +961,15 @@ class State: `TensorType` instead of a `bool`. """ + num_local_steps: int | TensorType + """ + `True` if the search space was global, else `False` if it was local. May be a scalar boolean + `TensorType` instead of a `bool`. + """ + def __deepcopy__(self, memo: dict[int, object]) -> TrustRegion.State: box_copy = copy.deepcopy(self.acquisition_space, memo) - return TrustRegion.State(box_copy, self.eps, self.y_min, self.is_global) + return TrustRegion.State(box_copy, self.eps, self.y_min, self.is_global, self.num_local_steps) @overload def __init__( @@ -971,6 +977,7 @@ def __init__( rule: None = None, beta: float = 0.7, kappa: float = 1e-4, + max_num_local_steps: int = 1, ): ... @@ -980,6 +987,7 @@ def __init__( rule: AcquisitionRule[TensorType, Box, ProbabilisticModelType], beta: float = 0.7, kappa: float = 1e-4, + max_num_local_steps: int = 1, ): ... @@ -988,6 +996,7 @@ def __init__( rule: AcquisitionRule[TensorType, Box, ProbabilisticModelType] | None = None, beta: float = 0.7, kappa: float = 1e-4, + max_num_local_steps: int = 1, ): """ :param rule: The acquisition rule that defines how to search for a new query point in a @@ -1003,6 +1012,7 @@ def __init__( self._rule = rule self._beta = beta self._kappa = kappa + self._max_num_local_steps = max_num_local_steps def __repr__(self) -> str: """""" @@ -1078,19 +1088,26 @@ def state_func( else state.eps * self._beta ) - is_global = step_is_success or not state.is_global + # is_global = step_is_success or not state.is_global + + if state.is_global: + is_global = step_is_success + else: + is_global = (state.num_local_steps >= self._max_num_local_steps) if is_global: acquisition_space = search_space + num_local_steps = 0 else: xmin = dataset.query_points[tf.argmin(dataset.observations)[0], :] acquisition_space = Box( tf.reduce_max([global_lower, xmin - eps], axis=0), tf.reduce_min([global_upper, xmin + eps], axis=0), ) + num_local_steps = state.num_local_steps + 1 points = self._rule.acquire(acquisition_space, models, datasets=datasets) - state_ = TrustRegion.State(acquisition_space, eps, y_min, is_global) + state_ = TrustRegion.State(acquisition_space, eps, y_min, is_global, num_local_steps) return state_, points From ebe27f41e189a66a576b048e59740d4509124973 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Fri, 2 Jun 2023 13:53:53 +0100 Subject: [PATCH 04/33] 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. --- docs/notebooks/multi_trust_region.pct.py | 16 +- trieste/acquisition/optimizer.py | 111 +++-- trieste/acquisition/rule.py | 495 ++++++++++++++++------- trieste/space.py | 270 ++++++++++--- 4 files changed, 632 insertions(+), 260 deletions(-) diff --git a/docs/notebooks/multi_trust_region.pct.py b/docs/notebooks/multi_trust_region.pct.py index a10d11fb52..17669f1b60 100644 --- a/docs/notebooks/multi_trust_region.pct.py +++ b/docs/notebooks/multi_trust_region.pct.py @@ -1,7 +1,8 @@ +# %% import numpy as np import tensorflow as tf from trieste.acquisition.rule import EfficientGlobalOptimization -from trieste.acquisition.rule import MultiTrustRegion, get_local_x_min +from trieste.acquisition.rule import MultiTrustRegionBox from trieste.acquisition import ParallelContinuousThompsonSampling from trieste.acquisition.optimizer import automatic_optimizer_selector from trieste.acquisition.utils import split_acquisition_function_calls @@ -27,6 +28,7 @@ from matplotlib.patches import Rectangle +# %% np.random.seed(179) tf.random.set_seed(179) @@ -36,6 +38,7 @@ summary_writer = tf.summary.create_file_writer(tensorboard_dir_1) trieste.logging.set_tensorboard_writer(summary_writer) +# %% obj = ScaledBranin.objective search_space = ScaledBranin.search_space @@ -50,6 +53,7 @@ def obj_fun( num_query_points = 3 num_steps = 10 +# %% [markdown] # aim_repo_dir = './results' # # run_1 = Run( @@ -73,6 +77,7 @@ def obj_fun( # 'num_steps': num_steps, # } +# %% initial_query_points = search_space.sample(num_initial_data_points) observer = trieste.objectives.utils.mk_observer(obj_fun) initial_data = observer(initial_query_points) @@ -102,10 +107,11 @@ def obj_fun( automatic_optimizer_selector, split_size=100_000), ) -acq_rule = MultiTrustRegion(base_rule, number_of_tr=num_query_points) +acq_rule = MultiTrustRegionBox(base_rule, number_of_tr=num_query_points) ask_tell = AskTellOptimizer(search_space, initial_data, model, fit_model=True, acquisition_rule=acq_rule) +# %% color = cm.rainbow(np.linspace(0, 1, num_query_points)) Xplot, xx, yy = create_grid(mins=search_space.lower, maxs=search_space.upper, grid_density=90) @@ -134,14 +140,14 @@ def obj_fun( state = ask_tell.acquisition_state - xmin = {tag: get_local_x_min(ask_tell.dataset, state.acquisition_space.get_subspace(tag)) for tag in + xmin = {tag: state.acquisition_space.get_subspace(tag).get_local_min(ask_tell.dataset)[0] for tag in state.acquisition_space.subspace_tags} i = 0 ax[0, 1].contour(xx, yy, ff.reshape(*xx.shape), 80, alpha=0.5) for tag in state.acquisition_space.subspace_tags: - ax[0, 1].scatter(xmin[tag].numpy()[0, 0], xmin[tag].numpy()[0, 1], color=color[i], marker="x", alpha=0.5) + ax[0, 1].scatter(xmin[tag].numpy()[0], xmin[tag].numpy()[1], color=color[i], marker="x", alpha=0.5) lb = state.acquisition_space.get_subspace(tag).lower ub = state.acquisition_space.get_subspace(tag).upper ax[0, 1].add_patch(Rectangle((lb[0], lb[1]), ub[0] - lb[0], ub[1] - lb[1], @@ -158,6 +164,7 @@ def obj_fun( ask_tell.tell(new_data) +# %% dataset = ask_tell.dataset ground_truth_regret = obj(dataset.query_points) - Hartmann6.minimum @@ -172,6 +179,7 @@ def obj_fun( ax.set_ylabel("Regret") ax.set_xlabel("# evaluations") +# %% fig, ax = plt.subplots() ax.scatter(dataset.query_points[:, 0], dataset.query_points[:, 1]) diff --git a/trieste/acquisition/optimizer.py b/trieste/acquisition/optimizer.py index cb9481e2ea..5a8337296a 100644 --- a/trieste/acquisition/optimizer.py +++ b/trieste/acquisition/optimizer.py @@ -30,11 +30,12 @@ from .. import logging from ..space import ( Box, + CollectionSearchSpace, Constraint, DiscreteSearchSpace, - MultiBoxSearchSpace, SearchSpace, SearchSpaceType, + TaggedMultiSearchSpace, TaggedProductSearchSpace, ) from ..types import TensorType @@ -101,12 +102,8 @@ def automatic_optimizer_selector( if isinstance(space, DiscreteSearchSpace): return optimize_discrete(space, target_func) - elif isinstance(space, (Box, TaggedProductSearchSpace, MultiBoxSearchSpace)): - if isinstance(space, MultiBoxSearchSpace): - space_dim = tf.shape(space.get_subspace(space.subspace_tags[0]).lower)[-1] - else: - space_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( @@ -176,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 | MultiBoxSearchSpace]: +) -> 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. @@ -222,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 | MultiBoxSearchSpace, + 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 @@ -252,17 +249,26 @@ def optimize_continuous( if V < 0: raise ValueError(f"vectorization must be positive, got {V}") - if isinstance(space, MultiBoxSearchSpace): - candidates = [ - space.get_subspace(tag).sample(num_initial_samples)[:, None, :] - for tag in space.subspace_tags - ] # list of [num_initial_samples, 1, D] - tiled_candidates = tf.concat(candidates, axis=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: - 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] + 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( @@ -306,17 +312,27 @@ def optimize_continuous( recovery_run = False if num_recovery_runs and not successful_optimization: # if all optimizations failed for a function then try again from random starts - if isinstance(space, MultiBoxSearchSpace): - candidates = [ - space.get_subspace(tag).sample(num_initial_samples)[:, None, :] - for tag in space.subspace_tags - ] # list of [num_initial_samples, 1, D] - tiled_random_points = tf.concat(candidates, axis=1) # [num_initial_samples, V, D] + 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: - 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] + tiled_random_points = tf.tile( + random_points[:, None, :], [1, V, 1] + ) # [num_recovery_runs, V, D] + ( recovery_successes, recovery_fun_values, @@ -407,7 +423,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 @@ -451,20 +467,31 @@ def _objective_value(vectorized_x: TensorType) -> TensorType: # [N, D] -> [N, 1 def _objective_value_and_gradient(x: TensorType) -> Tuple[TensorType, TensorType]: return tfp.math.value_and_gradient(_objective_value, x) # [len(x), 1], [len(x), D] - if isinstance(space, MultiBoxSearchSpace): - bounds = [ - spo.Bounds(space.get_subspace(tag).lower, space.get_subspace(tag).upper) - for tag in space.subspace_tags - ] - # TODO: check that this creates the right set of bou - bounds = bounds * num_optimization_runs_per_function - elif isinstance( + if isinstance( space, TaggedProductSearchSpace ): # build continuous relaxation of discrete subspaces bounds = [ 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 609f8e1473..2f97320841 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -22,7 +22,19 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Callable, Generic, Optional, Sequence, TypeVar, Union, cast, overload +from typing import ( + Any, + Callable, + Generic, + Optional, + Sequence, + Tuple, + Type, + TypeVar, + Union, + cast, + overload, +) import numpy as np @@ -47,7 +59,7 @@ TrainableSupportsGetKernel, ) from ..observer import OBJECTIVE -from ..space import Box, MultiBoxSearchSpace, SearchSpace +from ..space import Box, SearchSpace, TaggedMultiSearchSpace from ..types import State, Tag, TensorType from .function import ( BatchMonteCarloExpectedImprovement, @@ -975,7 +987,9 @@ class State: def __deepcopy__(self, memo: dict[int, object]) -> TrustRegion.State: box_copy = copy.deepcopy(self.acquisition_space, memo) - return TrustRegion.State(box_copy, self.eps, self.y_min, self.is_global, self.num_local_steps) + return TrustRegion.State( + box_copy, self.eps, self.y_min, self.is_global, self.num_local_steps + ) @overload def __init__( @@ -1099,7 +1113,7 @@ def state_func( if state.is_global: is_global = step_is_success else: - is_global = (state.num_local_steps >= self._max_num_local_steps) + is_global = state.num_local_steps >= self._max_num_local_steps if is_global: acquisition_space = search_space @@ -1110,6 +1124,7 @@ def state_func( tf.reduce_max([global_lower, xmin - eps], axis=0), tf.reduce_min([global_upper, xmin + eps], axis=0), ) + assert state is not None num_local_steps = state.num_local_steps + 1 points = self._rule.acquire(acquisition_space, models, datasets=datasets) @@ -1120,209 +1135,378 @@ def state_func( return state_func +class UpdateableSearchSpace(SearchSpace): + """A search space that can be updated.""" + + @abstractmethod + def reinitialise( + self, + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> None: + """Reinitialise the search space using the given models and datasets.""" + ... + + @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.""" + ... + + def get_single_model_and_dataset( + self, + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> Tuple[Optional[ProbabilisticModelType], Optional[Dataset]]: + """Get a single model and dataset from the given arguments.""" + if models is None: + model = None + elif OBJECTIVE in models.keys(): + model = models[OBJECTIVE] + else: + raise ValueError(f"""if models is provided, it must contain the key {OBJECTIVE}""") + + if datasets is None: + dataset = None + elif OBJECTIVE in datasets.keys(): + dataset = datasets[OBJECTIVE] + else: + raise ValueError(f"""if datasets is provided, it must contain the key {OBJECTIVE}""") + + return model, dataset + + +UpdateableSearchSpaceType = TypeVar("UpdateableSearchSpaceType", bound=UpdateableSearchSpace) +""" A type variable bound to :class:`UpdateableSearchSpace`. """ + + class MultiTrustRegion( AcquisitionRule[ - types.State[Optional["MultiTrustRegion.State"], TensorType], Box, ProbabilisticModelType - ] + types.State[Optional["MultiTrustRegion.State"], TensorType], + SearchSpace, + ProbabilisticModelType, + ], + Generic[ProbabilisticModelType, UpdateableSearchSpaceType], ): - """Implements the *trust region* acquisition algorithm.""" + """Abstrct class for multi trust region acquisition rules.""" @dataclass(frozen=True) class State: - """The acquisition state for the :class:`TrustRegion` acquisition rule.""" + """The acquisition state for the :class:`MultiTrustRegion` acquisition rule.""" - acquisition_space: MultiBoxSearchSpace + acquisition_space: TaggedMultiSearchSpace """ The search space. """ - eps: Sequence[TensorType] + def __deepcopy__(self, memo: dict[int, object]) -> MultiTrustRegion.State: + acquisition_space_copy = copy.deepcopy(self.acquisition_space, memo) + return MultiTrustRegion.State(acquisition_space_copy) + + def __init__( + self: "MultiTrustRegion[ProbabilisticModelType, UpdateableSearchSpaceType]", + subspace_type: Type[UpdateableSearchSpaceType], + rule: AcquisitionRule[TensorType, SearchSpace, ProbabilisticModelType] | None = None, + number_of_tr: int = 1, + ): """ - The (maximum) vector from the current best point to each bound of the acquisition space. + :param subspace_type: The type of the subspace to use. + :param rule: The acquisition rule that defines how to search for a new query point in a + given search space. Defaults to :class:`EfficientGlobalOptimization` with default + arguments. + :param number_of_tr: The number of trust regions. """ + if rule is None: + rule = EfficientGlobalOptimization() - y_min: Sequence[TensorType] - """ The minimum observed value. """ + self._subspace_type = subspace_type + self._rule = rule + self._number_of_tr = number_of_tr + self._tags = tuple([str(index) for index in range(number_of_tr)]) - # is_global: Sequence[bool] | Sequence[TensorType] - # """ - # `True` if the search space was global, else `False` if it was local. - # May be a scalar boolean `TensorType` instead of a `bool`. - # """ + def __repr__(self) -> str: + """""" + return f"""MultiTrustRegion( + {self._subspace_type!r}, + {self._rule!r}, + {self._number_of_tr!r})""" - def __deepcopy__(self, memo: dict[int, object]) -> MultiTrustRegion.State: - box_copy = copy.deepcopy(self.acquisition_space, memo) - return MultiTrustRegion.State(box_copy, self.eps, self.y_min) # , self.is_global) + 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: MultiTrustRegion.State | None, + ) -> Tuple[MultiTrustRegion.State | None, TensorType]: + """If state is None, initialise the subspaces by picking new locations. Otherwise, + update the existing subspaces. + Reinitialise 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 + MultiTrustRegion acquisition rule {self._tags}""" + + subspaces = [] + for tag in self._tags: + if state is None: + subspace = self.create_subspace(search_space) + subspace.reinitialise(models, datasets) + else: + _subspace = state.acquisition_space.get_subspace(tag) + assert isinstance(_subspace, self._subspace_type) + subspace = _subspace + subspace.update(models, datasets) - @overload - def __init__( - self: "MultiTrustRegion[ProbabilisticModel]", - rule: None = None, - beta: float = 0.7, - kappa: float = 1e-4, - number_of_tr: int = 1, - ): - ... + subspaces.append(subspace) - @overload - def __init__( - self: "MultiTrustRegion[ProbabilisticModelType]", - rule: AcquisitionRule[TensorType, Box, ProbabilisticModelType], - beta: float = 0.7, - kappa: float = 1e-4, - number_of_tr: int = 1, - ): + self.maybe_reinitialise_subspaces(subspaces, models, datasets) + + if state is None: + acquisition_space = TaggedMultiSearchSpace(subspaces, self._tags) + else: + acquisition_space = state.acquisition_space + + state_ = MultiTrustRegion.State(acquisition_space) + points = self._rule.acquire(acquisition_space, models, datasets=datasets) + + return state_, points + + return state_func + + def create_subspace( + self, + search_space: SearchSpace, + ) -> UpdateableSearchSpaceType: + """Create a subspace from the given search space. + This is the default implementation. Can be overridden by subclasses.""" + return self._subspace_type() + + def maybe_reinitialise_subspaces( + self, + subspaces: Sequence[UpdateableSearchSpaceType], + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> None: + """Reinitialise subspaces if necessary. + Get a mask of subspaces that need to be reinitialised using an abstract method. + Reinitialise individual subpaces by calling the method of the UpdateableSearchSpace class. + """ + mask = self.get_reinitialise_subspaces_mask(subspaces, models, datasets) + for ix, subspace in enumerate(subspaces): + if mask[ix]: + subspace.reinitialise(models, datasets) + + @abstractmethod + def get_reinitialise_subspaces_mask( + self, + subspaces: Sequence[UpdateableSearchSpaceType], + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> TensorType: + """Get mask for subspaces that need to be reinitialised.""" ... + +class TrustRegionBox(Box, UpdateableSearchSpace): + """An updateable box search space for use with trust region acquisition rules.""" + def __init__( self, - rule: AcquisitionRule[TensorType, Box, ProbabilisticModelType] | None = None, + global_search_space: SearchSpace, beta: float = 0.7, kappa: float = 1e-4, - number_of_tr: int = 1, + min_eps: float = 1e-2, ): - if rule is None: - rule = EfficientGlobalOptimization() + """ + Calculates the bounds of the box from the location/centre and global bounds. - self._rule = rule + :param global_search_space: The global search space this search space lives in. + """ + + self._global_search_space = global_search_space self._beta = beta self._kappa = kappa - self._min_eps = 1e-2 - self._tags = [str(index) for index in range(number_of_tr)] + self._min_eps = min_eps + self._initialised = False - def __repr__(self) -> str: - """""" - return f"MultiTrustRegion({self._rule!r}, {self._beta!r}, {self._kappa!r})" + super().__init__(global_search_space.lower, global_search_space.upper) - def acquire( + @property + def global_search_space(self) -> SearchSpace: + """The global search space this search space lives in.""" + return self._global_search_space + + @property + def location(self) -> TensorType: + """The location of the search space.""" + return self._location + + @location.setter + def location(self, location: TensorType) -> None: + """Set the location of the search space.""" + self._location = location + + def reinitialise( self, - search_space: Box, - models: Mapping[Tag, ProbabilisticModelType], + models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, - ) -> types.State[State | None, TensorType]: - if datasets is None or OBJECTIVE not in datasets.keys(): - raise ValueError(f"""datasets must be provided and contain the key {OBJECTIVE}""") + ) -> None: + """Reinitialise the box.""" + _, dataset = self.get_single_model_and_dataset(models, datasets) - dataset = datasets[OBJECTIVE] + self.location = tf.squeeze(self.global_search_space.sample(1), axis=0) - global_lower = search_space.lower - global_upper = search_space.upper + global_lower = self.global_search_space.lower + global_upper = self.global_search_space.upper - # global_y_min = tf.reduce_min(dataset.observations, axis=0) - def state_func( - state: MultiTrustRegion.State | None, - ) -> tuple[MultiTrustRegion.State | None, TensorType]: - if state is None: - eps = { - tag: 0.5 - * (global_upper - global_lower) - / (5.0 ** (1.0 / global_lower.shape[-1])) - for tag in self._tags - } - - x_min = {tag: search_space.sample(1) for tag in self._tags} - - aspace = [ - Box( - tf.reduce_max([global_lower, (x_min[tag] - eps[tag])[0, :]], axis=0), - tf.reduce_min([global_upper, (x_min[tag] + eps[tag])[0, :]], axis=0), - ) - for tag in self._tags - ] + if self._initialised: + scale = 0.1 + else: + scale = 0.5 + self._initialised = True - acquisition_space = MultiBoxSearchSpace(aspace, self._tags) + self._eps = scale * (global_upper - global_lower) / (5.0 ** (1.0 / global_lower.shape[-1])) + self._lower = tf.reduce_max([global_lower, self.location - self._eps], axis=0) + self._upper = tf.reduce_min([global_upper, self.location + self._eps], axis=0) - y_min = { - tag: get_local_y_min(dataset, acquisition_space.get_subspace(tag)) - for tag in self._tags - } + _, self.y_min = self.get_local_min(dataset) - else: - y_min = { - tag: get_local_y_min(dataset, state.acquisition_space.get_subspace(tag)) - for tag in state.acquisition_space.subspace_tags - } - - x_min = { - tag: get_local_x_min(dataset, state.acquisition_space.get_subspace(tag)) - for tag in state.acquisition_space.subspace_tags - } - eps = {} - for tag in state.acquisition_space.subspace_tags: - tr_volume = tf.reduce_prod( - state.acquisition_space.get_subspace(tag).upper - - state.acquisition_space.get_subspace(tag).lower - ) - step_is_success = y_min[tag] < state.y_min[tag] - self._kappa * tr_volume - eps[tag] = ( - state.eps[tag] / self._beta - if step_is_success - else state.eps[tag] * self._beta - ) + 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, reinitialise the box.""" + _, dataset = self.get_single_model_and_dataset(models, datasets) - x_min, y_min, eps = reinitialise_useless_trs( - x_min, y_min, eps, search_space, dataset, self._min_eps - ) + if tf.reduce_any(self._eps < self._min_eps): + self.reinitialise(models, datasets) + return - aspace = [ - Box( - tf.reduce_max([global_lower, (x_min[tag] - eps[tag])[0, :]], axis=0), - tf.reduce_min([global_upper, (x_min[tag] + eps[tag])[0, :]], axis=0), - ) - for tag in state.acquisition_space.subspace_tags - ] + _, y_min = self.get_local_min(dataset) - acquisition_space = MultiBoxSearchSpace( - aspace, state.acquisition_space.subspace_tags - ) + tr_volume = tf.reduce_prod(self.upper - self.lower) + step_is_success = y_min < self.y_min - self._kappa * tr_volume + self._eps = self._eps / self._beta if step_is_success else self._eps * self._beta - points = self._rule.acquire(acquisition_space, models, datasets=datasets) - state_ = MultiTrustRegion.State(acquisition_space, eps, y_min) + 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 + ) - return state_, points + _, self.y_min = self.get_local_min(dataset) - return state_func + 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) -def get_local_y_min(dataset: Dataset, search_space: SearchSpace): - in_tr = search_space.contains(dataset.query_points) - ones = tf.ones_like(dataset.observations) * 1e6 - in_tr_obs = tf.where(in_tr[:, None], dataset.observations, ones) - return tf.reduce_min(in_tr_obs, axis=0) + return tf.squeeze(x_min, axis=0), tf.squeeze(y_min) -def get_local_x_min(dataset: Dataset, search_space: SearchSpace): - in_tr = search_space.contains(dataset.query_points) - ones = tf.ones_like(dataset.observations) * 1e6 - in_tr_obs = tf.where(in_tr[:, None], dataset.observations, ones) - return dataset.query_points[tf.argmin(in_tr_obs)[0], :][None, :] +def get_unique_points_mask(points: TensorType, tolerance: float = 1e-6) -> TensorType: + """Find unique points in a list, within a given tolerance.""" + # Calculate the pairwise distances between points. + distances = tf.norm(tf.expand_dims(points, 1) - tf.expand_dims(points, 0), axis=-1) -def reinitialise_useless_trs( - x_min: Sequence[TensorType], - y_min: Sequence[TensorType], - eps: Sequence[TensorType], - search_space: Box, - dataset: Dataset, - min_eps: float, -) -> tuple[Sequence[TensorType], Sequence[TensorType], Sequence[TensorType]]: - xx = tf.zeros([0, dataset.query_points.shape[1]], dtype=dataset.query_points.dtype) + # Create a mask that is True for the lower triangle of the distance matrix, not including + # the diagonal. + mask_tri = tf.linalg.band_part(tf.ones_like(distances, dtype=tf.bool), -1, 0) - global_lower = search_space.lower - global_upper = search_space.upper + # Use the mask to ignore the upper triangle and diagonal of the distance matrix. + distances_masked = tf.where(mask_tri, distances, tf.constant(np.inf, dtype=distances.dtype)) - for tag, x in x_min.items(): - duplicated = tf.less_equal(tf.reduce_min(tf.reduce_sum(tf.abs(x - xx), axis=1)), 1e-6) + tolerance = tf.constant(tolerance, dtype=distances.dtype) - if duplicated or (np.min(eps[tag].numpy()) < min_eps): - eps[tag] = 0.1 * (global_upper - global_lower) / (5.0 ** (1.0 / global_lower.shape[-1])) - x_min[tag] = search_space.sample(1) - temp_space = Box( - tf.reduce_max([global_lower, (x_min[tag] - eps[tag])[0, :]], axis=0), - tf.reduce_min([global_upper, (x_min[tag] + eps[tag])[0, :]], axis=0), - ) - y_min[tag] = get_local_y_min(dataset, temp_space) - xx = tf.concat([xx, x_min[tag]], axis=0) + # Create a boolean mask for each point. + mask = tf.reduce_all( + tf.logical_or( + distances_masked > tolerance, tf.eye(distances_masked.shape[0], dtype=tf.bool) + ), + axis=1, + ) - return x_min, y_min, eps + # Return the mask of unique points; can get the actual points with: + # tf.boolean_mask(points, mask) + return mask + + +class MultiTrustRegionBox(MultiTrustRegion[ProbabilisticModelType, TrustRegionBox]): + """Implements the *trust region* acquisition algorithm for a box.""" + + def __init__( + self, + rule: AcquisitionRule[TensorType, SearchSpace, ProbabilisticModelType] | None = None, + number_of_tr: int = 1, + beta: float = 0.7, + kappa: float = 1e-4, + min_eps: float = 1e-2, + ): + """Concrete implementation of :class:`MultiTrustRegion` for a box. + + :param rule: The acquisition rule that defines how to search for a new query point in a + given search space. Defaults to :class:`EfficientGlobalOptimization` with default + arguments. + :param number_of_tr: The number of trust regions. + :param beta: The inverse of the trust region contraction factor. + :param kappa: The trust region volume scaling factor. + :param min_eps: The minimum size of the trust region. + """ + super().__init__(TrustRegionBox, rule, number_of_tr) + self._beta = beta + self._kappa = kappa + self._min_eps = min_eps + + def __repr__(self) -> str: + """""" + return f"""MultiTrustRegionBox( + {self._rule!r}, + {self._number_of_tr!r}, + {self._beta!r}, + {self._kappa!r}, + {self._min_eps!r})""" + + def create_subspace( + self, + search_space: SearchSpace, + ) -> TrustRegionBox: + """Create a subspace from the given global search space.""" + return TrustRegionBox(search_space, self._beta, self._kappa, self._min_eps) + + def get_reinitialise_subspaces_mask( + self, + subspaces: Sequence[TrustRegionBox], + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> TensorType: + """Get mask for subspaces that need to be reinitialised. + # Reinitalise 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( @@ -1448,7 +1632,6 @@ def __repr__(self) -> str: """""" return f"TURBO({self._num_trust_regions!r}, {self._rule})" - def acquire( self, search_space: Box, diff --git a/trieste/space.py b/trieste/space.py index 1bfb819003..04bbe78c30 100644 --- a/trieste/space.py +++ b/trieste/space.py @@ -875,20 +875,17 @@ 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. + A :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. @@ -920,17 +917,217 @@ def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] ) self._spaces = dict(zip(tags, spaces)) + self._tags = tuple(tags) # avoid accidental modification by users + + def __repr__(self) -> str: + """""" + return f"""CollectionSearchSpace(spaces = + {[self.get_subspace(tag) for tag in self.subspace_tags]}, + tags = {self.subspace_tags}) + """ + + @property + 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 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 space.""" + return self._tags + + @property + 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] - subspace_sizes = [space.dimension for space in spaces] + def get_subspace(self, tag: str) -> SearchSpace: + """ + Return the domain of a particular subspace. + + :param tag: The tag specifying the target subspace. + :return: Target subspace. + """ + tf.debugging.assert_equal( + tag in self.subspace_tags, + True, + 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 TaggedMultiSearchSpace(CollectionSearchSpace): + r""" + A :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 the individual spaces are not combined in any way, however all subspaces should have + the same dimension. + """ + + 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. + :raise ValueError (or tf.errors.InvalidArgumentError): If ``spaces`` has a different + dimension to each other. + """ + + tf.debugging.assert_equal( + len(set([int(space.dimension) for space in spaces])), + 1, + message=f""" + All spaces must have the same dimension but received + {[space.dimension for space in spaces]}. + """, + ) + + super().__init__(spaces, tags) + + def __repr__(self) -> str: + """""" + return f"""TaggedMultiSearchSpace(spaces = + {[self.get_subspace(tag) for tag in self.subspace_tags]}, + tags = {self.subspace_tags}) + """ + + @property + def lower(self) -> TensorType: + lower = self.subspace_lower + return tf.stack(lower, axis=0) if lower else tf.constant([], dtype=DEFAULT_DTYPE) + + @property + def upper(self) -> TensorType: + upper = self.subspace_upper + return tf.stack(upper, axis=0) if upper else tf.constant([], dtype=DEFAULT_DTYPE) + + @property + def dimension(self) -> TensorType: + """The number of inputs in this search space.""" + return self.get_subspace(self.subspace_tags[0]).dimension + + 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 __deepcopy__(self, memo: dict[int, object]) -> TaggedMultiSearchSpace: + return self + + +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. + + 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(tags, subspace_sizes) + tag: subspace_size for tag, subspace_size in zip(self._tags, subspace_sizes) } - self._subspace_starting_indices = dict(zip(tags, tf.cumsum(subspace_sizes, exclusive=True))) + 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) - self._tags = tuple(tags) # avoid accidental modification by users def __repr__(self) -> str: """""" @@ -942,7 +1139,7 @@ def __repr__(self) -> str: @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] + lower_for_each_subspace = self.subspace_lower return ( tf.concat(lower_for_each_subspace, axis=-1) if lower_for_each_subspace @@ -952,40 +1149,18 @@ def lower(self) -> TensorType: @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] + 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 - def subspace_tags(self) -> tuple[str, ...]: - """Return the names of the subspaces contained in this product space.""" - return self._tags - @property def dimension(self) -> TensorType: """The number of inputs in this product search space.""" return self._dimension - def get_subspace(self, tag: str) -> SearchSpace: - """ - Return the domain of a particular subspace. - - :param tag: The tag specifying the target subspace. - :return: Target subspace. - """ - tf.debugging.assert_equal( - tag in self.subspace_tags, - True, - 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 fix_subspace(self, tag: str, values: TensorType) -> TaggedProductSearchSpace: """ Return a new :class:`TaggedProductSearchSpace` with the specified subspace replaced with @@ -1050,10 +1225,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: @@ -1066,23 +1238,5 @@ def product(self, other: TaggedProductSearchSpace) -> TaggedProductSearchSpace: """ return TaggedProductSearchSpace(spaces=[self, other]) - 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, TaggedProductSearchSpace): - return NotImplemented - return self._tags == other._tags and self._spaces == other._spaces - def __deepcopy__(self, memo: dict[int, object]) -> TaggedProductSearchSpace: return self - - -class MultiBoxSearchSpace(TaggedProductSearchSpace): - def __repr__(self) -> str: - """""" - return f"""MultiBoxSearchSpace(spaces = - {[self.get_subspace(tag) for tag in self.subspace_tags]}, - tags = {self.subspace_tags}) - """ From 683eacc25df9c7e5398f046bb78d8229b7d0652f Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 26 Jul 2023 16:18:34 +0100 Subject: [PATCH 05/33] Fix formatting in notebook & remove orig TR changes --- docs/notebooks/multi_trust_region.pct.py | 137 ++++++++++++----------- trieste/acquisition/rule.py | 25 +---- 2 files changed, 75 insertions(+), 87 deletions(-) diff --git a/docs/notebooks/multi_trust_region.pct.py b/docs/notebooks/multi_trust_region.pct.py index 17669f1b60..8450443041 100644 --- a/docs/notebooks/multi_trust_region.pct.py +++ b/docs/notebooks/multi_trust_region.pct.py @@ -1,39 +1,32 @@ # %% +# from aim.ext.tensorboard_tracker import Run +from datetime import datetime + import numpy as np import tensorflow as tf -from trieste.acquisition.rule import EfficientGlobalOptimization -from trieste.acquisition.rule import MultiTrustRegionBox +from matplotlib import pyplot as plt +from matplotlib.patches import Rectangle +from matplotlib.pyplot import cm + +import trieste from trieste.acquisition import ParallelContinuousThompsonSampling from trieste.acquisition.optimizer import automatic_optimizer_selector +from trieste.acquisition.rule import EfficientGlobalOptimization, MultiTrustRegionBox from trieste.acquisition.utils import split_acquisition_function_calls from trieste.ask_tell_optimization import AskTellOptimizer from trieste.experimental.plotting import plot_regret -from matplotlib import pyplot as plt -from trieste.models.gpflow import ( - SparseVariational, - build_svgp, - build_gpr, - GaussianProcessRegression, - ConditionalImprovementReduction, -) -from trieste.models.optimizer import BatchOptimizer -import trieste -from trieste.objectives import ScaledBranin, Hartmann6 -from trieste.types import TensorType -from trieste.logging import pyplot -from matplotlib.pyplot import cm -# from aim.ext.tensorboard_tracker import Run -from datetime import datetime from trieste.experimental.plotting.plotting import create_grid -from matplotlib.patches import Rectangle - +from trieste.logging import pyplot +from trieste.models.gpflow import GaussianProcessRegression, build_gpr +from trieste.objectives import Hartmann6, ScaledBranin +from trieste.types import TensorType # %% np.random.seed(179) tf.random.set_seed(179) # CONFIG -tensorboard_dir_1 = f'./results/{datetime.now()}/tensorboard' +tensorboard_dir_1 = f"./results/{datetime.now()}/tensorboard" summary_writer = tf.summary.create_file_writer(tensorboard_dir_1) trieste.logging.set_tensorboard_writer(summary_writer) @@ -53,30 +46,6 @@ def obj_fun( num_query_points = 3 num_steps = 10 -# %% [markdown] -# aim_repo_dir = './results' -# -# run_1 = Run( -# experiment='TR', -# sync_tensorboard_log_dir=tensorboard_dir_1, -# repo=aim_repo_dir, -# system_tracking_interval=None, # Tracks CPU, GPU, memory usage periodically -# log_system_params=True, # Logs environment variables (DANGER), git hash, packages installed etc. -# capture_terminal_logs=True, # Stores the stdout (VERY USEFUL TO SEE WHY EXPERIMENTS BROKE) -# ) -# This can be any arbitrary tags that you want to associate with your run, I often at least -# add a semi-unique name that summarises the main components of the run, so its easy to see -# run_1.add_tag('gpr') - -# These are parameters associated wit the run. Ideally this would be *everything* you are likely -# to configure between different runs. If you use Hydra, you can store the dictionary representation -# of your config here, so no manual entry. -# run_1["hparams"] = { -# 'num_initial_data_points': num_initial_data_points, -# 'num_query_points': num_query_points, -# 'num_steps': num_steps, -# } - # %% initial_query_points = search_space.sample(num_initial_data_points) observer = trieste.objectives.utils.mk_observer(obj_fun) @@ -97,19 +66,25 @@ def obj_fun( # tf.optimizers.Adam(0.05), max_iter=100, batch_size=50, compile=True # ), # ) -gpflow_model = build_gpr(initial_data, search_space, likelihood_variance=1e-4, trainable_likelihood=False) +gpflow_model = build_gpr( + initial_data, + search_space, + likelihood_variance=1e-4, + trainable_likelihood=False, +) model = GaussianProcessRegression(gpflow_model) base_rule = EfficientGlobalOptimization( builder=ParallelContinuousThompsonSampling(), num_query_points=num_query_points, - optimizer=split_acquisition_function_calls( - automatic_optimizer_selector, split_size=100_000), + optimizer=split_acquisition_function_calls(automatic_optimizer_selector, split_size=100_000), ) acq_rule = MultiTrustRegionBox(base_rule, number_of_tr=num_query_points) -ask_tell = AskTellOptimizer(search_space, initial_data, model, fit_model=True, acquisition_rule=acq_rule) +ask_tell = AskTellOptimizer( + search_space, initial_data, model, fit_model=True, acquisition_rule=acq_rule +) # %% color = cm.rainbow(np.linspace(0, 1, num_query_points)) @@ -135,28 +110,63 @@ def obj_fun( fig, ax = plt.subplots(1, 2, squeeze=False, figsize=(15, 5)) fig.suptitle(f"step number {step}") - ax[0, 0].scatter(ask_tell.dataset.query_points[:, 0].numpy(), ask_tell.dataset.query_points[:, 1].numpy(), color="blue") + ax[0, 0].scatter( + ask_tell.dataset.query_points[:, 0].numpy(), + ask_tell.dataset.query_points[:, 1].numpy(), + color="blue", + ) ax[0, 0].scatter(new_points[:, 0].numpy(), new_points[:, 1].numpy(), color="red") state = ask_tell.acquisition_state - - xmin = {tag: state.acquisition_space.get_subspace(tag).get_local_min(ask_tell.dataset)[0] for tag in - state.acquisition_space.subspace_tags} + assert state is not None + assert isinstance(state, MultiTrustRegionBox.State) + + xmin = { + tag: state.acquisition_space.get_subspace( + tag + ).get_local_min( # type: ignore[attr-defined] + ask_tell.dataset + )[ + 0 + ] + for tag in state.acquisition_space.subspace_tags + } i = 0 ax[0, 1].contour(xx, yy, ff.reshape(*xx.shape), 80, alpha=0.5) for tag in state.acquisition_space.subspace_tags: - ax[0, 1].scatter(xmin[tag].numpy()[0], xmin[tag].numpy()[1], color=color[i], marker="x", alpha=0.5) + ax[0, 1].scatter( + xmin[tag].numpy()[0], + xmin[tag].numpy()[1], + color=color[i], + marker="x", + alpha=0.5, + ) lb = state.acquisition_space.get_subspace(tag).lower ub = state.acquisition_space.get_subspace(tag).upper - ax[0, 1].add_patch(Rectangle((lb[0], lb[1]), ub[0] - lb[0], ub[1] - lb[1], - facecolor=color[i], - edgecolor=color[i], - alpha=0.3)) - ax[0, 1].scatter(new_points[i, 0].numpy(), new_points[i, 1].numpy(), color=color[i], alpha=0.5) - ax[0, 1].scatter(ask_tell.dataset.query_points[:, 0].numpy(), - ask_tell.dataset.query_points[:, 1].numpy(), color="black", alpha=0.2) + ax[0, 1].add_patch( + Rectangle( + (lb[0], lb[1]), + ub[0] - lb[0], + ub[1] - lb[1], + facecolor=color[i], + edgecolor=color[i], + alpha=0.3, + ) + ) + ax[0, 1].scatter( + new_points[i, 0].numpy(), + new_points[i, 1].numpy(), + color=color[i], + alpha=0.5, + ) + ax[0, 1].scatter( + ask_tell.dataset.query_points[:, 0].numpy(), + ask_tell.dataset.query_points[:, 1].numpy(), + color="black", + alpha=0.2, + ) i = i + 1 pyplot("Query points", fig) @@ -171,9 +181,7 @@ def obj_fun( best_found_truth_idx = tf.squeeze(tf.argmin(ground_truth_regret, axis=0)) fig, ax = plt.subplots() -plot_regret( - ground_truth_regret.numpy(), ax, num_init=10, idx_best=best_found_truth_idx -) +plot_regret(ground_truth_regret.numpy(), ax, num_init=10, idx_best=best_found_truth_idx) ax.set_yscale("log") ax.set_ylabel("Regret") @@ -182,4 +190,3 @@ def obj_fun( # %% fig, ax = plt.subplots() ax.scatter(dataset.query_points[:, 0], dataset.query_points[:, 1]) - diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 2f97320841..a583c84f2a 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -979,17 +979,9 @@ class State: `TensorType` instead of a `bool`. """ - num_local_steps: int | TensorType - """ - `True` if the search space was global, else `False` if it was local. May be a scalar boolean - `TensorType` instead of a `bool`. - """ - def __deepcopy__(self, memo: dict[int, object]) -> TrustRegion.State: box_copy = copy.deepcopy(self.acquisition_space, memo) - return TrustRegion.State( - box_copy, self.eps, self.y_min, self.is_global, self.num_local_steps - ) + return TrustRegion.State(box_copy, self.eps, self.y_min, self.is_global) @overload def __init__( @@ -997,7 +989,6 @@ def __init__( rule: None = None, beta: float = 0.7, kappa: float = 1e-4, - max_num_local_steps: int = 1, ): ... @@ -1007,7 +998,6 @@ def __init__( rule: AcquisitionRule[TensorType, Box, ProbabilisticModelType], beta: float = 0.7, kappa: float = 1e-4, - max_num_local_steps: int = 1, ): ... @@ -1016,7 +1006,6 @@ def __init__( rule: AcquisitionRule[TensorType, Box, ProbabilisticModelType] | None = None, beta: float = 0.7, kappa: float = 1e-4, - max_num_local_steps: int = 1, ): """ :param rule: The acquisition rule that defines how to search for a new query point in a @@ -1032,7 +1021,6 @@ def __init__( self._rule = rule self._beta = beta self._kappa = kappa - self._max_num_local_steps = max_num_local_steps def __repr__(self) -> str: """""" @@ -1108,16 +1096,10 @@ def state_func( else state.eps * self._beta ) - # is_global = step_is_success or not state.is_global - - if state.is_global: - is_global = step_is_success - else: - is_global = state.num_local_steps >= self._max_num_local_steps + is_global = step_is_success or not state.is_global if is_global: acquisition_space = search_space - num_local_steps = 0 else: xmin = dataset.query_points[tf.argmin(dataset.observations)[0], :] acquisition_space = Box( @@ -1125,10 +1107,9 @@ def state_func( tf.reduce_min([global_upper, xmin + eps], axis=0), ) assert state is not None - num_local_steps = state.num_local_steps + 1 points = self._rule.acquire(acquisition_space, models, datasets=datasets) - state_ = TrustRegion.State(acquisition_space, eps, y_min, is_global, num_local_steps) + state_ = TrustRegion.State(acquisition_space, eps, y_min, is_global) return state_, points From 809153c0aff6beeb391afe3812b25497d66a9b74 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 26 Jul 2023 16:24:54 +0100 Subject: [PATCH 06/33] Remove redundant new assert --- trieste/acquisition/rule.py | 1 - 1 file changed, 1 deletion(-) diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index a583c84f2a..1d2d2deaa4 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -1106,7 +1106,6 @@ def state_func( tf.reduce_max([global_lower, xmin - eps], axis=0), tf.reduce_min([global_upper, xmin + eps], axis=0), ) - assert state is not None points = self._rule.acquire(acquisition_space, models, datasets=datasets) state_ = TrustRegion.State(acquisition_space, eps, y_min, is_global) From b6ee392bd2d4ab1e63f3840a81a00647e249764e Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 26 Jul 2023 16:31:58 +0100 Subject: [PATCH 07/33] Undo earlier logging change --- trieste/acquisition/optimizer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/trieste/acquisition/optimizer.py b/trieste/acquisition/optimizer.py index 5a8337296a..7eacec596e 100644 --- a/trieste/acquisition/optimizer.py +++ b/trieste/acquisition/optimizer.py @@ -388,9 +388,7 @@ def improvements() -> tf.Tensor: if V == 1: logging.scalar("spo_improvement_on_initial_samples", improvements) else: - logging.histogram( - "spo_improvement_on_initial_samples_across_subspaces", improvements - ) + logging.histogram("spo_improvements_on_initial_samples", improvements) best_run_ids = tf.math.argmax(fun_values, axis=0) # [V] chosen_points = tf.gather( From fcc16e200b49867ca495dbefd6889a58885c4d10 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 26 Jul 2023 16:40:27 +0100 Subject: [PATCH 08/33] Workaround isort & black clash --- docs/notebooks/multi_trust_region.pct.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/docs/notebooks/multi_trust_region.pct.py b/docs/notebooks/multi_trust_region.pct.py index 8450443041..296ce8f81c 100644 --- a/docs/notebooks/multi_trust_region.pct.py +++ b/docs/notebooks/multi_trust_region.pct.py @@ -11,7 +11,7 @@ import trieste from trieste.acquisition import ParallelContinuousThompsonSampling from trieste.acquisition.optimizer import automatic_optimizer_selector -from trieste.acquisition.rule import EfficientGlobalOptimization, MultiTrustRegionBox +from trieste.acquisition.rule import MultiTrustRegionBox from trieste.acquisition.utils import split_acquisition_function_calls from trieste.ask_tell_optimization import AskTellOptimizer from trieste.experimental.plotting import plot_regret @@ -74,10 +74,12 @@ def obj_fun( ) model = GaussianProcessRegression(gpflow_model) -base_rule = EfficientGlobalOptimization( +base_rule = trieste.acquisition.rule.EfficientGlobalOptimization( builder=ParallelContinuousThompsonSampling(), num_query_points=num_query_points, - optimizer=split_acquisition_function_calls(automatic_optimizer_selector, split_size=100_000), + optimizer=split_acquisition_function_calls( + automatic_optimizer_selector, split_size=100_000 + ), ) acq_rule = MultiTrustRegionBox(base_rule, number_of_tr=num_query_points) @@ -89,7 +91,9 @@ def obj_fun( # %% color = cm.rainbow(np.linspace(0, 1, num_query_points)) -Xplot, xx, yy = create_grid(mins=search_space.lower, maxs=search_space.upper, grid_density=90) +Xplot, xx, yy = create_grid( + mins=search_space.lower, maxs=search_space.upper, grid_density=90 +) ff = obj_fun(Xplot).numpy() for step in range(num_steps): @@ -115,7 +119,9 @@ def obj_fun( ask_tell.dataset.query_points[:, 1].numpy(), color="blue", ) - ax[0, 0].scatter(new_points[:, 0].numpy(), new_points[:, 1].numpy(), color="red") + ax[0, 0].scatter( + new_points[:, 0].numpy(), new_points[:, 1].numpy(), color="red" + ) state = ask_tell.acquisition_state assert state is not None @@ -181,7 +187,9 @@ def obj_fun( best_found_truth_idx = tf.squeeze(tf.argmin(ground_truth_regret, axis=0)) fig, ax = plt.subplots() -plot_regret(ground_truth_regret.numpy(), ax, num_init=10, idx_best=best_found_truth_idx) +plot_regret( + ground_truth_regret.numpy(), ax, num_init=10, idx_best=best_found_truth_idx +) ax.set_yscale("log") ax.set_ylabel("Regret") From 9f2152734ea9d3d02d78a9d423cdc4b6439864ca Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 26 Jul 2023 16:48:57 +0100 Subject: [PATCH 09/33] Keep old ver of mypy happy --- docs/notebooks/multi_trust_region.pct.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/notebooks/multi_trust_region.pct.py b/docs/notebooks/multi_trust_region.pct.py index 296ce8f81c..9a60660d12 100644 --- a/docs/notebooks/multi_trust_region.pct.py +++ b/docs/notebooks/multi_trust_region.pct.py @@ -128,9 +128,9 @@ def obj_fun( assert isinstance(state, MultiTrustRegionBox.State) xmin = { - tag: state.acquisition_space.get_subspace( + tag: state.acquisition_space.get_subspace( # type: ignore[attr-defined] tag - ).get_local_min( # type: ignore[attr-defined] + ).get_local_min( ask_tell.dataset )[ 0 From 1762d82775915fab1dc4cd1fee385a962ef8e84e Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 26 Jul 2023 17:00:59 +0100 Subject: [PATCH 10/33] Fix typo in slicing --- trieste/acquisition/optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trieste/acquisition/optimizer.py b/trieste/acquisition/optimizer.py index 7eacec596e..b5fc55591e 100644 --- a/trieste/acquisition/optimizer.py +++ b/trieste/acquisition/optimizer.py @@ -267,7 +267,7 @@ def optimize_continuous( tiled_candidates = candidates # [num_initial_samples, V, D] else: tiled_candidates = tf.tile( - candidates[:None, :], [1, V, 1] + candidates[:, None, :], [1, V, 1] ) # [num_initial_samples, V, D] target_func_values = target_func(tiled_candidates) # [num_samples, V] From 535dbb1901015ce86c01139567ff02fc03890eba Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Thu, 27 Jul 2023 15:23:06 +0100 Subject: [PATCH 11/33] Add new collection space tests --- tests/unit/test_space.py | 671 +++++++++++++++++++++++++----------- trieste/acquisition/rule.py | 2 +- trieste/space.py | 29 +- 3 files changed, 507 insertions(+), 195 deletions(-) diff --git a/tests/unit/test_space.py b/tests/unit/test_space.py index 5320d6b457..cb4705fe07 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 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 @@ -662,38 +664,362 @@ def test_box_deepcopy() -> None: npt.assert_allclose(box.upper, box_copy.upper) -def test_product_space_raises_for_non_unqique_subspace_names() -> None: +# Parameterize for both TaggedMultiSearchSpace and TaggedProductSearchSpace. +@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) +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: +@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) +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: +@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) +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( + product_space = search_space_type( spaces=[context_space, decision_space], tags=["context", "decision"] ) npt.assert_array_equal(product_space.subspace_tags, ["context", "decision"]) -def test_product_space_subspace_tags_default_behaviour() -> None: +@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) +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]) + product_space = search_space_type(spaces=[context_space, decision_space]) npt.assert_array_equal(product_space.subspace_tags, ["0", "1"]) +@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) +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]])) + product_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" + ): + product_space.get_subspace("dummy") + + +@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) +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]) + product_space = search_space_type(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, -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("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) +@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( "spaces, dimension", [ @@ -703,7 +1029,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)): @@ -743,38 +1069,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", [ @@ -860,143 +1154,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)) @@ -1027,17 +1184,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"]) @@ -1050,6 +1196,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) @@ -1097,6 +1359,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/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 1d2d2deaa4..c682af6d70 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -1171,7 +1171,7 @@ class MultiTrustRegion( ], Generic[ProbabilisticModelType, UpdateableSearchSpaceType], ): - """Abstrct class for multi trust region acquisition rules.""" + """Abstract class for multi trust region acquisition rules.""" @dataclass(frozen=True) class State: diff --git a/trieste/space.py b/trieste/space.py index 04bbe78c30..bfc5048437 100644 --- a/trieste/space.py +++ b/trieste/space.py @@ -1017,12 +1017,21 @@ def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] dimension to each other. """ + # 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 spaces must have the same dimension but received - {[space.dimension for space in spaces]}. + All subspaces must have the same dimension but received + {[int(space.dimension) for space in spaces]}. """, ) @@ -1088,6 +1097,22 @@ def product(self, other: TaggedMultiSearchSpace) -> 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) + def __deepcopy__(self, memo: dict[int, object]) -> TaggedMultiSearchSpace: return self From c4ce7ab78e48bc1a8d5efe88adbd7c11533c1b0f Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 2 Aug 2023 16:14:55 +0100 Subject: [PATCH 12/33] 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 --- docs/notebooks/multi_trust_region.pct.py | 6 +- tests/unit/test_space.py | 1 - trieste/acquisition/rule.py | 129 ++++++++++------------- trieste/acquisition/utils.py | 34 ++++++ 4 files changed, 93 insertions(+), 77 deletions(-) diff --git a/docs/notebooks/multi_trust_region.pct.py b/docs/notebooks/multi_trust_region.pct.py index 9a60660d12..df50a5aad1 100644 --- a/docs/notebooks/multi_trust_region.pct.py +++ b/docs/notebooks/multi_trust_region.pct.py @@ -11,7 +11,7 @@ import trieste from trieste.acquisition import ParallelContinuousThompsonSampling from trieste.acquisition.optimizer import automatic_optimizer_selector -from trieste.acquisition.rule import MultiTrustRegionBox +from trieste.acquisition.rule import MultiTrustRegionBox, TrustRegionBox from trieste.acquisition.utils import split_acquisition_function_calls from trieste.ask_tell_optimization import AskTellOptimizer from trieste.experimental.plotting import plot_regret @@ -82,7 +82,9 @@ def obj_fun( ), ) -acq_rule = MultiTrustRegionBox(base_rule, number_of_tr=num_query_points) +acq_rule = MultiTrustRegionBox( + TrustRegionBox, base_rule, number_of_tr=num_query_points +) ask_tell = AskTellOptimizer( search_space, initial_data, model, fit_model=True, acquisition_rule=acq_rule diff --git a/tests/unit/test_space.py b/tests/unit/test_space.py index cb4705fe07..9564f4840e 100644 --- a/tests/unit/test_space.py +++ b/tests/unit/test_space.py @@ -664,7 +664,6 @@ def test_box_deepcopy() -> None: npt.assert_allclose(box.upper, box_copy.upper) -# Parameterize for both TaggedMultiSearchSpace and TaggedProductSearchSpace. @pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) def test_collection_space_raises_for_non_unqique_subspace_names( search_space_type: Type[CollectionSearchSpace], diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index c682af6d70..d2d4e6717f 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -83,7 +83,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. """ @@ -1119,12 +1119,13 @@ class UpdateableSearchSpace(SearchSpace): """A search space that can be updated.""" @abstractmethod - def reinitialise( + def reinitialize( self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, + **kwargs: Any, ) -> None: - """Reinitialise the search space using the given models and datasets.""" + """Reinitialize the search space using the given models and datasets.""" ... @abstractmethod @@ -1132,6 +1133,7 @@ def update( self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, + **kwargs: Any, ) -> None: """Update the search space using the given models and datasets.""" ... @@ -1171,7 +1173,9 @@ class MultiTrustRegion( ], Generic[ProbabilisticModelType, UpdateableSearchSpaceType], ): - """Abstract class for multi trust region acquisition rules.""" + """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: @@ -1223,7 +1227,7 @@ def state_func( ) -> Tuple[MultiTrustRegion.State | None, TensorType]: """If state is None, initialise the subspaces by picking new locations. Otherwise, update the existing subspaces. - Reinitialise the subspaces if necessary, potentially looking at the entire group. + Reinitialize 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 @@ -1239,7 +1243,7 @@ def state_func( for tag in self._tags: if state is None: subspace = self.create_subspace(search_space) - subspace.reinitialise(models, datasets) + subspace.reinitialize(models, datasets) else: _subspace = state.acquisition_space.get_subspace(tag) assert isinstance(_subspace, self._subspace_type) @@ -1248,7 +1252,7 @@ def state_func( subspaces.append(subspace) - self.maybe_reinitialise_subspaces(subspaces, models, datasets) + self.maybe_reinitialize_subspaces(subspaces, models, datasets) if state is None: acquisition_space = TaggedMultiSearchSpace(subspaces, self._tags) @@ -1270,29 +1274,29 @@ def create_subspace( This is the default implementation. Can be overridden by subclasses.""" return self._subspace_type() - def maybe_reinitialise_subspaces( + def maybe_reinitialize_subspaces( self, subspaces: Sequence[UpdateableSearchSpaceType], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> None: - """Reinitialise subspaces if necessary. - Get a mask of subspaces that need to be reinitialised using an abstract method. - Reinitialise individual subpaces by calling the method of the UpdateableSearchSpace class. + """Reinitialize subspaces if necessary. + Get a mask of subspaces that need to be reinitialized using an abstract method. + Reinitialize individual subpaces by calling the method of the UpdateableSearchSpace class. """ - mask = self.get_reinitialise_subspaces_mask(subspaces, models, datasets) + mask = self.get_reinitialize_subspaces_mask(subspaces, models, datasets) for ix, subspace in enumerate(subspaces): if mask[ix]: - subspace.reinitialise(models, datasets) + subspace.reinitialize(models, datasets) @abstractmethod - def get_reinitialise_subspaces_mask( + def get_reinitialize_subspaces_mask( self, subspaces: Sequence[UpdateableSearchSpaceType], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> TensorType: - """Get mask for subspaces that need to be reinitialised.""" + """Get mask for subspaces that need to be reinitialized.""" ... @@ -1316,7 +1320,6 @@ def __init__( self._beta = beta self._kappa = kappa self._min_eps = min_eps - self._initialised = False super().__init__(global_search_space.lower, global_search_space.upper) @@ -1335,28 +1338,31 @@ def location(self, location: TensorType) -> None: """Set the location of the search space.""" self._location = location - def reinitialise( + 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 reinitialize( self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, + **kwargs: Any, ) -> None: - """Reinitialise the box.""" + """Reinitialize the box.""" _, dataset = self.get_single_model_and_dataset(models, datasets) self.location = tf.squeeze(self.global_search_space.sample(1), axis=0) - - global_lower = self.global_search_space.lower - global_upper = self.global_search_space.upper - - if self._initialised: - scale = 0.1 - else: - scale = 0.5 - self._initialised = True - - self._eps = scale * (global_upper - global_lower) / (5.0 ** (1.0 / global_lower.shape[-1])) - self._lower = tf.reduce_max([global_lower, self.location - self._eps], axis=0) - self._upper = tf.reduce_min([global_upper, self.location + self._eps], axis=0) + self._init_eps() + self._update_bounds() _, self.y_min = self.get_local_min(dataset) @@ -1364,27 +1370,23 @@ def update( self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, + **kwargs: Any, ) -> None: """Update this box, including centre/location, using the given dataset. If the size of the - box is less than the minimum size, reinitialise the box.""" + box is less than the minimum size, reinitialize the box.""" _, dataset = self.get_single_model_and_dataset(models, datasets) if tf.reduce_any(self._eps < self._min_eps): - self.reinitialise(models, datasets) + self.reinitialize(models, datasets, **kwargs) return - _, y_min = self.get_local_min(dataset) + x_min, y_min = self.get_local_min(dataset) + self.location = x_min tr_volume = tf.reduce_prod(self.upper - self.lower) step_is_success = y_min < self.y_min - self._kappa * tr_volume self._eps = self._eps / self._beta if step_is_success else self._eps * self._beta - - 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 - ) + self._update_bounds() _, self.y_min = self.get_local_min(dataset) @@ -1406,39 +1408,16 @@ def get_local_min(self, dataset: Optional[Dataset]) -> Tuple[TensorType, TensorT return tf.squeeze(x_min, axis=0), tf.squeeze(y_min) -def get_unique_points_mask(points: TensorType, tolerance: float = 1e-6) -> TensorType: - """Find unique points in a list, within a given tolerance.""" - - # Calculate the pairwise distances between points. - distances = tf.norm(tf.expand_dims(points, 1) - tf.expand_dims(points, 0), axis=-1) +TrustRegionBoxType = TypeVar("TrustRegionBoxType", bound=TrustRegionBox) +""" A type variable bound to :class:`TrustRegionBox`. """ - # Create a mask that is True for the lower triangle of the distance matrix, not including - # the diagonal. - mask_tri = tf.linalg.band_part(tf.ones_like(distances, dtype=tf.bool), -1, 0) - # Use the mask to ignore the upper triangle and diagonal of the distance matrix. - distances_masked = tf.where(mask_tri, distances, tf.constant(np.inf, dtype=distances.dtype)) - - tolerance = tf.constant(tolerance, dtype=distances.dtype) - - # Create a boolean mask for each point. - mask = tf.reduce_all( - tf.logical_or( - distances_masked > tolerance, tf.eye(distances_masked.shape[0], dtype=tf.bool) - ), - axis=1, - ) - - # Return the mask of unique points; can get the actual points with: - # tf.boolean_mask(points, mask) - return mask - - -class MultiTrustRegionBox(MultiTrustRegion[ProbabilisticModelType, TrustRegionBox]): +class MultiTrustRegionBox(MultiTrustRegion[ProbabilisticModelType, TrustRegionBoxType]): """Implements the *trust region* acquisition algorithm for a box.""" def __init__( self, + subspace_type: Type[TrustRegionBoxType], rule: AcquisitionRule[TensorType, SearchSpace, ProbabilisticModelType] | None = None, number_of_tr: int = 1, beta: float = 0.7, @@ -1447,6 +1426,7 @@ def __init__( ): """Concrete implementation of :class:`MultiTrustRegion` for a box. + :param subspace_type: The type of the subspace to use. :param rule: The acquisition rule that defines how to search for a new query point in a given search space. Defaults to :class:`EfficientGlobalOptimization` with default arguments. @@ -1455,7 +1435,7 @@ def __init__( :param kappa: The trust region volume scaling factor. :param min_eps: The minimum size of the trust region. """ - super().__init__(TrustRegionBox, rule, number_of_tr) + super().__init__(subspace_type, rule, number_of_tr) self._beta = beta self._kappa = kappa self._min_eps = min_eps @@ -1463,6 +1443,7 @@ def __init__( def __repr__(self) -> str: """""" return f"""MultiTrustRegionBox( + {self._subspace_type!r}, {self._rule!r}, {self._number_of_tr!r}, {self._beta!r}, @@ -1472,17 +1453,17 @@ def __repr__(self) -> str: def create_subspace( self, search_space: SearchSpace, - ) -> TrustRegionBox: + ) -> TrustRegionBoxType: """Create a subspace from the given global search space.""" - return TrustRegionBox(search_space, self._beta, self._kappa, self._min_eps) + return self._subspace_type(search_space, self._beta, self._kappa, self._min_eps) - def get_reinitialise_subspaces_mask( + def get_reinitialize_subspaces_mask( self, - subspaces: Sequence[TrustRegionBox], + subspaces: Sequence[TrustRegionBoxType], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> TensorType: - """Get mask for subspaces that need to be reinitialised. + """Get mask for subspaces that need to be reinitialized. # Reinitalise the subspaces that have non-unique locations. """ centres = tf.stack([subspace.location for subspace in subspaces]) diff --git a/trieste/acquisition/utils.py b/trieste/acquisition/utils.py index cb97bd3a2e..b13c6cb35b 100644 --- a/trieste/acquisition/utils.py +++ b/trieste/acquisition/utils.py @@ -14,6 +14,7 @@ import functools from typing import Tuple, Union +import numpy as np import tensorflow as tf from ..data import Dataset @@ -135,3 +136,36 @@ def get_local_dataset(local_space: SearchSpaceType, dataset: Dataset) -> Dataset observations=tf.boolean_mask(dataset.observations, is_in_region_mask), ) return local_dataset + + +def get_unique_points_mask(points: TensorType, tolerance: float = 1e-6) -> TensorType: + """Find unique points in a list, within a given tolerance. + + :param points: A tensor of points. + :param tolerance: The tolerance within which points are considered equal. + :return: A boolean mask for the unique points. + """ + + # Calculate the pairwise distances between points. + distances = tf.norm(tf.expand_dims(points, 1) - tf.expand_dims(points, 0), axis=-1) + + # Create a mask that is True for the lower triangle of the distance matrix, not including + # the diagonal. + mask_tri = tf.linalg.band_part(tf.ones_like(distances, dtype=tf.bool), -1, 0) + + # Use the mask to ignore the upper triangle and diagonal of the distance matrix. + distances_masked = tf.where(mask_tri, distances, tf.constant(np.inf, dtype=distances.dtype)) + + tolerance = tf.constant(tolerance, dtype=distances.dtype) + + # Create a boolean mask for each point. + mask = tf.reduce_all( + tf.logical_or( + distances_masked > tolerance, tf.eye(distances_masked.shape[0], dtype=tf.bool) + ), + axis=1, + ) + + # Return the mask of unique points; can get the actual points with: + # tf.boolean_mask(points, mask) + return mask From 0c38efb9cab676c3930e49fca33ff2e9557a368a Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 2 Aug 2023 16:57:32 +0100 Subject: [PATCH 13/33] Add TrustRegionBox/UpdateableSearchSpace unittests --- tests/unit/acquisition/test_rule.py | 164 ++++++++++++++++++++++++++++ 1 file changed, 164 insertions(+) diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index 13b6acabd2..5e53ea87cc 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 @@ -48,6 +49,7 @@ EfficientGlobalOptimization, RandomSampling, TrustRegion, + TrustRegionBox, ) from trieste.acquisition.sampler import ( ExactThompsonSampler, @@ -1112,6 +1114,168 @@ 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 = TrustRegionBox(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 = TrustRegionBox(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 = TrustRegionBox(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)) + + +# get_single_model_and_dataset returns model and dataset with the OBJECTIVE tag. +def test_trust_region_box_get_single_model_dataset() -> None: + search_space = Box([0.0, 0.0], [1.0, 1.0]) + dataset = Dataset(tf.zeros([1, 2], dtype=tf.float64), tf.zeros([1, 1], dtype=tf.float64)) + models = { + "foo": QuadraticMeanAndRBFKernel(), + OBJECTIVE: QuadraticMeanAndRBFKernelWithSamplers(dataset), + } + datasets = {"foo": empty_dataset([2], [1]), OBJECTIVE: dataset} + trb = TrustRegionBox(search_space) + _model, _dataset = trb.get_single_model_and_dataset(models, datasets) + assert isinstance(_model, QuadraticMeanAndRBFKernelWithSamplers) + assert _dataset is dataset + + +# Reinitialize sets the box to a random location, and sets the eps and y_min values. +def test_trust_region_box_reinitialize() -> 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 = TrustRegionBox(search_space) + trb.reinitialize(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 reintializes the box if eps is smaller than min_eps. +def test_trust_region_box_update_reinitialize() -> 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 = TrustRegionBox(search_space, min_eps=0.5) + trb.reinitialize(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 reintialize the box if eps is larger than min_eps. +def test_trust_region_box_update_no_reinitialize() -> 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 = TrustRegionBox(search_space, min_eps=0.1) + trb.reinitialize(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.0], [1.0]], dtype=tf.float64), + ) + } + trb = TrustRegionBox(search_space, min_eps=0.1) + trb.reinitialize(datasets=datasets) + eps = trb._eps + + if success: + # Sample a point from the box. + point = trb.sample(1) + else: + # Pick point outside the box. + point = tf.constant([[1.2, 1.3]], dtype=tf.float64) + + # Add a new min point to the dataset. + datasets[OBJECTIVE] = Dataset( + np.concatenate([datasets[OBJECTIVE].query_points, point], axis=0), + np.concatenate([datasets[OBJECTIVE].observations, [[-0.1]]], axis=0), + ) + # Update the box. + trb.update(datasets=datasets) + + if success: + # Check that the location is the new min point. + point = np.squeeze(point) + npt.assert_allclose(point, trb.location) + npt.assert_allclose(tf.constant([-0.1], dtype=tf.float64), trb.y_min) + # Check that the box is smaller by beta. + npt.assert_allclose(eps / trb._beta, trb._eps) + else: + # Check that the location is the old min point. + point, y_min = trb.get_local_min(datasets[OBJECTIVE]) + npt.assert_allclose(point, trb.location) + npt.assert_allclose(y_min, trb.y_min) + # Check that the box is larger by beta. + npt.assert_allclose(eps * trb._beta, trb._eps) + + # Check the new box bounds. + npt.assert_allclose(trb.lower, np.maximum(point - trb._eps, search_space.lower)) + npt.assert_allclose(trb.upper, np.minimum(point + trb._eps, search_space.upper)) + + def test_asynchronous_rule_state_pending_points() -> None: pending_points = tf.constant([[1], [2], [3]]) From 309ef09dafd170d195b273a2a3fb2a46626947fb Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Thu, 3 Aug 2023 15:00:25 +0100 Subject: [PATCH 14/33] Address feedback for space changes --- trieste/acquisition/optimizer.py | 20 +++ trieste/acquisition/utils.py | 16 +- trieste/space.py | 272 ++++++++++++++++--------------- 3 files changed, 169 insertions(+), 139 deletions(-) diff --git a/trieste/acquisition/optimizer.py b/trieste/acquisition/optimizer.py index b5fc55591e..d9b8550d66 100644 --- a/trieste/acquisition/optimizer.py +++ b/trieste/acquisition/optimizer.py @@ -266,6 +266,16 @@ def optimize_continuous( ) 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] @@ -329,6 +339,16 @@ def optimize_continuous( ) 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] diff --git a/trieste/acquisition/utils.py b/trieste/acquisition/utils.py index b13c6cb35b..2716ac72f7 100644 --- a/trieste/acquisition/utils.py +++ b/trieste/acquisition/utils.py @@ -16,6 +16,7 @@ import numpy as np import tensorflow as tf +from check_shapes import check_shapes from ..data import Dataset from ..space import SearchSpaceType @@ -138,10 +139,19 @@ def get_local_dataset(local_space: SearchSpaceType, dataset: Dataset) -> Dataset return local_dataset +@check_shapes( + "points: [n_points, ...]", + "return: [n_points]", +) def get_unique_points_mask(points: TensorType, tolerance: float = 1e-6) -> TensorType: - """Find unique points in a list, within a given tolerance. + """Find the boolean mask of unique points in a tensor, within a given tolerance. - :param points: A tensor of points. + 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. """ @@ -166,6 +176,4 @@ def get_unique_points_mask(points: TensorType, tolerance: float = 1e-6) -> Tenso axis=1, ) - # Return the mask of unique points; can get the actual points with: - # tf.boolean_mask(points, mask) return mask diff --git a/trieste/space.py b/trieste/space.py index bfc5048437..28f84f651e 100644 --- a/trieste/space.py +++ b/trieste/space.py @@ -877,8 +877,9 @@ def has_constraints(self) -> bool: class CollectionSearchSpace(SearchSpace): r""" - A :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. + 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 the individual spaces are not combined in any way. """ @@ -920,8 +921,7 @@ def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] self._tags = tuple(tags) # avoid accidental modification by users def __repr__(self) -> str: - """""" - return f"""CollectionSearchSpace(spaces = + return f"""{self.__class__.__name__}(spaces = {[self.get_subspace(tag) for tag in self.subspace_tags]}, tags = {self.subspace_tags}) """ @@ -993,130 +993,6 @@ def __eq__(self, other: object) -> bool: return self._tags == other._tags and self._spaces == other._spaces -class TaggedMultiSearchSpace(CollectionSearchSpace): - r""" - A :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 the individual spaces are not combined in any way, however all subspaces should have - the same dimension. - """ - - 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. - :raise ValueError (or tf.errors.InvalidArgumentError): If ``spaces`` has a different - dimension to each other. - """ - - # 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) - - def __repr__(self) -> str: - """""" - return f"""TaggedMultiSearchSpace(spaces = - {[self.get_subspace(tag) for tag in self.subspace_tags]}, - tags = {self.subspace_tags}) - """ - - @property - def lower(self) -> TensorType: - lower = self.subspace_lower - return tf.stack(lower, axis=0) if lower else tf.constant([], dtype=DEFAULT_DTYPE) - - @property - def upper(self) -> TensorType: - upper = self.subspace_upper - return tf.stack(upper, axis=0) if upper else tf.constant([], dtype=DEFAULT_DTYPE) - - @property - def dimension(self) -> TensorType: - """The number of inputs in this search space.""" - return self.get_subspace(self.subspace_tags[0]).dimension - - 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) - - def __deepcopy__(self, memo: dict[int, object]) -> TaggedMultiSearchSpace: - return self - - class TaggedProductSearchSpace(CollectionSearchSpace): r""" Product :class:`SearchSpace` consisting of a product of @@ -1154,13 +1030,6 @@ def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] self._dimension = tf.cast(tf.reduce_sum(subspace_sizes), dtype=tf.int32) - def __repr__(self) -> str: - """""" - return f"""TaggedProductSearchSpace(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.""" @@ -1265,3 +1134,136 @@ def product(self, other: TaggedProductSearchSpace) -> TaggedProductSearchSpace: def __deepcopy__(self, memo: dict[int, object]) -> TaggedProductSearchSpace: return self + + +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. + """ + + 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. + :raise ValueError (or tf.errors.InvalidArgumentError): If ``spaces`` has a different + dimension to each other. + """ + + # 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 + 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. + """ + + lower = self.subspace_lower + return tf.stack(lower, axis=0) if lower else tf.constant([], dtype=DEFAULT_DTYPE) + + @property + 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. + """ + + upper = self.subspace_upper + return tf.stack(upper, axis=0) if upper else tf.constant([], dtype=DEFAULT_DTYPE) + + @property + def dimension(self) -> TensorType: + """The number of inputs in this search space.""" + return self.get_subspace(self.subspace_tags[0]).dimension + + 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) + + def __deepcopy__(self, memo: dict[int, object]) -> TaggedMultiSearchSpace: + return self From d58306ecdab7f38e6a60a70d23911c8d1b5617ac Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Fri, 4 Aug 2023 16:40:35 +0100 Subject: [PATCH 15/33] 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 --- tests/unit/acquisition/test_rule.py | 35 ++++---------- tests/unit/acquisition/test_utils.py | 52 +++++++++++++++++++++ trieste/acquisition/rule.py | 68 ++++++++++------------------ trieste/acquisition/utils.py | 57 +++++++++++++++-------- 4 files changed, 124 insertions(+), 88 deletions(-) diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index 5e53ea87cc..11d1bf507c 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -1151,23 +1151,8 @@ def test_trust_region_box_get_local_min_outside_search_space() -> None: npt.assert_array_equal(y_min, tf.constant([np.inf], dtype=tf.float64)) -# get_single_model_and_dataset returns model and dataset with the OBJECTIVE tag. -def test_trust_region_box_get_single_model_dataset() -> None: - search_space = Box([0.0, 0.0], [1.0, 1.0]) - dataset = Dataset(tf.zeros([1, 2], dtype=tf.float64), tf.zeros([1, 1], dtype=tf.float64)) - models = { - "foo": QuadraticMeanAndRBFKernel(), - OBJECTIVE: QuadraticMeanAndRBFKernelWithSamplers(dataset), - } - datasets = {"foo": empty_dataset([2], [1]), OBJECTIVE: dataset} - trb = TrustRegionBox(search_space) - _model, _dataset = trb.get_single_model_and_dataset(models, datasets) - assert isinstance(_model, QuadraticMeanAndRBFKernelWithSamplers) - assert _dataset is dataset - - -# Reinitialize sets the box to a random location, and sets the eps and y_min values. -def test_trust_region_box_reinitialize() -> None: +# 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. @@ -1176,7 +1161,7 @@ def test_trust_region_box_reinitialize() -> None: ) } trb = TrustRegionBox(search_space) - trb.reinitialize(datasets=datasets) + 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) @@ -1188,8 +1173,8 @@ def test_trust_region_box_reinitialize() -> None: npt.assert_array_equal(trb.y_min, tf.constant([np.inf], dtype=tf.float64)) -# Update call reintializes the box if eps is smaller than min_eps. -def test_trust_region_box_update_reinitialize() -> None: +# 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. @@ -1198,7 +1183,7 @@ def test_trust_region_box_update_reinitialize() -> None: ) } trb = TrustRegionBox(search_space, min_eps=0.5) - trb.reinitialize(datasets=datasets) + trb.initialize(datasets=datasets) location = trb.location trb.update(datasets=datasets) @@ -1209,8 +1194,8 @@ def test_trust_region_box_update_reinitialize() -> None: npt.assert_array_compare(np.not_equal, location, trb.location) -# Update call does not reintialize the box if eps is larger than min_eps. -def test_trust_region_box_update_no_reinitialize() -> None: +# 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( @@ -1219,7 +1204,7 @@ def test_trust_region_box_update_no_reinitialize() -> None: ) } trb = TrustRegionBox(search_space, min_eps=0.1) - trb.reinitialize(datasets=datasets) + trb.initialize(datasets=datasets) trb.location = tf.constant([0.5, 0.5], dtype=tf.float64) location = trb.location @@ -1238,7 +1223,7 @@ def test_trust_region_box_update_size(success: bool) -> None: ) } trb = TrustRegionBox(search_space, min_eps=0.1) - trb.reinitialize(datasets=datasets) + trb.initialize(datasets=datasets) eps = trb._eps if success: diff --git a/tests/unit/acquisition/test_utils.py b/tests/unit/acquisition/test_utils.py index 13fad5595a..e84ab6a42b 100644 --- a/tests/unit/acquisition/test_utils.py +++ b/tests/unit/acquisition/test_utils.py @@ -23,10 +23,13 @@ from trieste.acquisition import AcquisitionFunction from trieste.acquisition.utils import ( get_local_dataset, + get_unique_points_mask, + get_value_for_tag, select_nth_output, split_acquisition_function, ) from trieste.data import Dataset +from trieste.observer import OBJECTIVE from trieste.space import Box, SearchSpaceType @@ -97,3 +100,52 @@ 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) + + +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" diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index d2d4e6717f..590d830db8 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -83,7 +83,7 @@ batchify_vectorize, ) from .sampler import ExactThompsonSampler, ThompsonSampler -from .utils import get_local_dataset, get_unique_points_mask, select_nth_output +from .utils import get_local_dataset, get_unique_points_mask, get_value_for_tag, select_nth_output ResultType = TypeVar("ResultType", covariant=True) """ Unbound covariant type variable. """ @@ -1119,13 +1119,13 @@ class UpdateableSearchSpace(SearchSpace): """A search space that can be updated.""" @abstractmethod - def reinitialize( + def initialize( self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, **kwargs: Any, ) -> None: - """Reinitialize the search space using the given models and datasets.""" + """Initialize the search space using the given models and datasets.""" ... @abstractmethod @@ -1138,28 +1138,6 @@ def update( """Update the search space using the given models and datasets.""" ... - def get_single_model_and_dataset( - self, - models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, - datasets: Optional[Mapping[Tag, Dataset]] = None, - ) -> Tuple[Optional[ProbabilisticModelType], Optional[Dataset]]: - """Get a single model and dataset from the given arguments.""" - if models is None: - model = None - elif OBJECTIVE in models.keys(): - model = models[OBJECTIVE] - else: - raise ValueError(f"""if models is provided, it must contain the key {OBJECTIVE}""") - - if datasets is None: - dataset = None - elif OBJECTIVE in datasets.keys(): - dataset = datasets[OBJECTIVE] - else: - raise ValueError(f"""if datasets is provided, it must contain the key {OBJECTIVE}""") - - return model, dataset - UpdateableSearchSpaceType = TypeVar("UpdateableSearchSpaceType", bound=UpdateableSearchSpace) """ A type variable bound to :class:`UpdateableSearchSpace`. """ @@ -1227,7 +1205,7 @@ def state_func( ) -> Tuple[MultiTrustRegion.State | None, TensorType]: """If state is None, initialise the subspaces by picking new locations. Otherwise, update the existing subspaces. - Reinitialize the subspaces if necessary, potentially looking at the entire group. + 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 @@ -1243,7 +1221,7 @@ def state_func( for tag in self._tags: if state is None: subspace = self.create_subspace(search_space) - subspace.reinitialize(models, datasets) + subspace.initialize(models, datasets) else: _subspace = state.acquisition_space.get_subspace(tag) assert isinstance(_subspace, self._subspace_type) @@ -1252,7 +1230,7 @@ def state_func( subspaces.append(subspace) - self.maybe_reinitialize_subspaces(subspaces, models, datasets) + self.maybe_initialize_subspaces(subspaces, models, datasets) if state is None: acquisition_space = TaggedMultiSearchSpace(subspaces, self._tags) @@ -1274,29 +1252,29 @@ def create_subspace( This is the default implementation. Can be overridden by subclasses.""" return self._subspace_type() - def maybe_reinitialize_subspaces( + def maybe_initialize_subspaces( self, subspaces: Sequence[UpdateableSearchSpaceType], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> None: - """Reinitialize subspaces if necessary. - Get a mask of subspaces that need to be reinitialized using an abstract method. - Reinitialize individual subpaces by calling the method of the UpdateableSearchSpace class. + """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 UpdateableSearchSpace class. """ - mask = self.get_reinitialize_subspaces_mask(subspaces, models, datasets) + mask = self.get_initialize_subspaces_mask(subspaces, models, datasets) for ix, subspace in enumerate(subspaces): if mask[ix]: - subspace.reinitialize(models, datasets) + subspace.initialize(models, datasets) @abstractmethod - def get_reinitialize_subspaces_mask( + def get_initialize_subspaces_mask( self, subspaces: Sequence[UpdateableSearchSpaceType], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> TensorType: - """Get mask for subspaces that need to be reinitialized.""" + """Get mask for subspaces that need to be initialized.""" ... @@ -1351,14 +1329,14 @@ def _update_bounds(self) -> None: [self.global_search_space.upper, self.location + self._eps], axis=0 ) - def reinitialize( + def initialize( self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, **kwargs: Any, ) -> None: - """Reinitialize the box.""" - _, dataset = self.get_single_model_and_dataset(models, datasets) + """Initialize the box.""" + dataset = get_value_for_tag(datasets) self.location = tf.squeeze(self.global_search_space.sample(1), axis=0) self._init_eps() @@ -1373,11 +1351,11 @@ def update( **kwargs: Any, ) -> None: """Update this box, including centre/location, using the given dataset. If the size of the - box is less than the minimum size, reinitialize the box.""" - _, dataset = self.get_single_model_and_dataset(models, datasets) + box is less than the minimum size, initialize the box.""" + dataset = get_value_for_tag(datasets) if tf.reduce_any(self._eps < self._min_eps): - self.reinitialize(models, datasets, **kwargs) + self.initialize(models, datasets, **kwargs) return x_min, y_min = self.get_local_min(dataset) @@ -1457,14 +1435,14 @@ def create_subspace( """Create a subspace from the given global search space.""" return self._subspace_type(search_space, self._beta, self._kappa, self._min_eps) - def get_reinitialize_subspaces_mask( + def get_initialize_subspaces_mask( self, subspaces: Sequence[TrustRegionBoxType], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> TensorType: - """Get mask for subspaces that need to be reinitialized. - # Reinitalise the subspaces that have non-unique locations. + """Get mask for subspaces that need to be initialized. + # 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)) diff --git a/trieste/acquisition/utils.py b/trieste/acquisition/utils.py index 2716ac72f7..51faf57e61 100644 --- a/trieste/acquisition/utils.py +++ b/trieste/acquisition/utils.py @@ -12,15 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Tuple, Union +from typing import Mapping, Optional, Tuple, TypeVar, Union import numpy as np import tensorflow as tf from check_shapes import check_shapes from ..data import Dataset +from ..observer import OBJECTIVE from ..space import SearchSpaceType -from ..types import TensorType +from ..types import Tag, TensorType from .interface import AcquisitionFunction from .optimizer import AcquisitionOptimizer @@ -156,24 +157,44 @@ def get_unique_points_mask(points: TensorType, tolerance: float = 1e-6) -> Tenso :return: A boolean mask for the unique points. """ - # Calculate the pairwise distances between points. - distances = tf.norm(tf.expand_dims(points, 1) - tf.expand_dims(points, 0), axis=-1) + tolerance = tf.constant(tolerance, dtype=points.dtype) + n_points = tf.shape(points)[0] + used_points = tf.fill(value=tf.constant(np.inf, dtype=points.dtype), dims=tf.shape(points)) + mask = tf.zeros(shape=(n_points,), dtype=tf.bool) - # Create a mask that is True for the lower triangle of the distance matrix, not including - # the diagonal. - mask_tri = tf.linalg.band_part(tf.ones_like(distances, dtype=tf.bool), -1, 0) + for idx in tf.range(n_points): + # Pairwise distance with remaining points. + distances = tf.norm(points[idx] - used_points, axis=-1) + # Find if there is any point within the tolerance. + min_distance = tf.reduce_min(distances) + is_unique_point = min_distance >= tolerance - # Use the mask to ignore the upper triangle and diagonal of the distance matrix. - distances_masked = tf.where(mask_tri, distances, tf.constant(np.inf, dtype=distances.dtype)) + # Update mask. + mask = tf.tensor_scatter_nd_update(mask, [[idx]], [is_unique_point]) - tolerance = tf.constant(tolerance, dtype=distances.dtype) - - # Create a boolean mask for each point. - mask = tf.reduce_all( - tf.logical_or( - distances_masked > tolerance, tf.eye(distances_masked.shape[0], dtype=tf.bool) - ), - axis=1, - ) + if is_unique_point: + # Save this point to 'used_points' for future iterations. + used_points = tf.tensor_scatter_nd_update(used_points, [[idx]], [points[idx]]) return mask + + +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") From c2abf11906704f38267526d50d352308ba3dc07c Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 7 Aug 2023 12:24:47 +0100 Subject: [PATCH 16/33] Create subspaces outside the rule Plus revert addition of kwargs --- docs/notebooks/multi_trust_region.pct.py | 5 +- trieste/acquisition/rule.py | 106 ++++------------------- 2 files changed, 17 insertions(+), 94 deletions(-) diff --git a/docs/notebooks/multi_trust_region.pct.py b/docs/notebooks/multi_trust_region.pct.py index df50a5aad1..dce4ada0d5 100644 --- a/docs/notebooks/multi_trust_region.pct.py +++ b/docs/notebooks/multi_trust_region.pct.py @@ -82,9 +82,8 @@ def obj_fun( ), ) -acq_rule = MultiTrustRegionBox( - TrustRegionBox, base_rule, number_of_tr=num_query_points -) +subspaces = [TrustRegionBox(search_space) for _ in range(num_query_points)] +acq_rule = MultiTrustRegionBox(subspaces, base_rule) ask_tell = AskTellOptimizer( search_space, initial_data, model, fit_model=True, acquisition_rule=acq_rule diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 590d830db8..211211eb8e 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -22,19 +22,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import dataclass -from typing import ( - Any, - Callable, - Generic, - Optional, - Sequence, - Tuple, - Type, - TypeVar, - Union, - cast, - overload, -) +from typing import Any, Callable, Generic, Optional, Sequence, Tuple, TypeVar, Union, cast, overload import numpy as np @@ -1123,7 +1111,6 @@ def initialize( self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, - **kwargs: Any, ) -> None: """Initialize the search space using the given models and datasets.""" ... @@ -1133,7 +1120,6 @@ def update( self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, - **kwargs: Any, ) -> None: """Update the search space using the given models and datasets.""" ... @@ -1168,31 +1154,25 @@ def __deepcopy__(self, memo: dict[int, object]) -> MultiTrustRegion.State: def __init__( self: "MultiTrustRegion[ProbabilisticModelType, UpdateableSearchSpaceType]", - subspace_type: Type[UpdateableSearchSpaceType], + init_subspaces: Sequence[UpdateableSearchSpaceType], rule: AcquisitionRule[TensorType, SearchSpace, ProbabilisticModelType] | None = None, - number_of_tr: int = 1, ): """ - :param subspace_type: The type of the subspace to use. + :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 a given search space. Defaults to :class:`EfficientGlobalOptimization` with default arguments. - :param number_of_tr: The number of trust regions. """ if rule is None: rule = EfficientGlobalOptimization() - self._subspace_type = subspace_type + self._init_subspaces = tuple(init_subspaces) + self._tags = tuple([str(index) for index in range(len(init_subspaces))]) self._rule = rule - self._number_of_tr = number_of_tr - self._tags = tuple([str(index) for index in range(number_of_tr)]) def __repr__(self) -> str: """""" - return f"""MultiTrustRegion( - {self._subspace_type!r}, - {self._rule!r}, - {self._number_of_tr!r})""" + return f"""{self.__class__.__name__}({self._init_subspaces!r}, {self._rule!r})""" def acquire( self, @@ -1218,13 +1198,13 @@ def state_func( MultiTrustRegion acquisition rule {self._tags}""" subspaces = [] - for tag in self._tags: + for index, tag in enumerate(self._tags): if state is None: - subspace = self.create_subspace(search_space) + subspace = self._init_subspaces[index] subspace.initialize(models, datasets) else: _subspace = state.acquisition_space.get_subspace(tag) - assert isinstance(_subspace, self._subspace_type) + assert isinstance(_subspace, type(self._init_subspaces[index])) subspace = _subspace subspace.update(models, datasets) @@ -1244,14 +1224,6 @@ def state_func( return state_func - def create_subspace( - self, - search_space: SearchSpace, - ) -> UpdateableSearchSpaceType: - """Create a subspace from the given search space. - This is the default implementation. Can be overridden by subclasses.""" - return self._subspace_type() - def maybe_initialize_subspaces( self, subspaces: Sequence[UpdateableSearchSpaceType], @@ -1260,7 +1232,7 @@ def maybe_initialize_subspaces( ) -> 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 UpdateableSearchSpace class. + Initialize individual subpaces by calling the method of the UpdateableSearchSpaceType class. """ mask = self.get_initialize_subspaces_mask(subspaces, models, datasets) for ix, subspace in enumerate(subspaces): @@ -1333,7 +1305,6 @@ def initialize( self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, - **kwargs: Any, ) -> None: """Initialize the box.""" dataset = get_value_for_tag(datasets) @@ -1348,14 +1319,13 @@ def update( self, models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, - **kwargs: Any, ) -> None: """Update this box, including centre/location, using the given dataset. If the size of the box is less than the minimum size, initialize the box.""" dataset = get_value_for_tag(datasets) if tf.reduce_any(self._eps < self._min_eps): - self.initialize(models, datasets, **kwargs) + self.initialize(models, datasets) return x_min, y_min = self.get_local_min(dataset) @@ -1386,63 +1356,17 @@ def get_local_min(self, dataset: Optional[Dataset]) -> Tuple[TensorType, TensorT return tf.squeeze(x_min, axis=0), tf.squeeze(y_min) -TrustRegionBoxType = TypeVar("TrustRegionBoxType", bound=TrustRegionBox) -""" A type variable bound to :class:`TrustRegionBox`. """ - - -class MultiTrustRegionBox(MultiTrustRegion[ProbabilisticModelType, TrustRegionBoxType]): - """Implements the *trust region* acquisition algorithm for a box.""" - - def __init__( - self, - subspace_type: Type[TrustRegionBoxType], - rule: AcquisitionRule[TensorType, SearchSpace, ProbabilisticModelType] | None = None, - number_of_tr: int = 1, - beta: float = 0.7, - kappa: float = 1e-4, - min_eps: float = 1e-2, - ): - """Concrete implementation of :class:`MultiTrustRegion` for a box. - - :param subspace_type: The type of the subspace to use. - :param rule: The acquisition rule that defines how to search for a new query point in a - given search space. Defaults to :class:`EfficientGlobalOptimization` with default - arguments. - :param number_of_tr: The number of trust regions. - :param beta: The inverse of the trust region contraction factor. - :param kappa: The trust region volume scaling factor. - :param min_eps: The minimum size of the trust region. - """ - super().__init__(subspace_type, rule, number_of_tr) - self._beta = beta - self._kappa = kappa - self._min_eps = min_eps - - def __repr__(self) -> str: - """""" - return f"""MultiTrustRegionBox( - {self._subspace_type!r}, - {self._rule!r}, - {self._number_of_tr!r}, - {self._beta!r}, - {self._kappa!r}, - {self._min_eps!r})""" - - def create_subspace( - self, - search_space: SearchSpace, - ) -> TrustRegionBoxType: - """Create a subspace from the given global search space.""" - return self._subspace_type(search_space, self._beta, self._kappa, self._min_eps) +class MultiTrustRegionBox(MultiTrustRegion[ProbabilisticModelType, TrustRegionBox]): + """Implements the :class:`MultiTrustRegion` *trust region* acquisition algorithm for a box.""" def get_initialize_subspaces_mask( self, - subspaces: Sequence[TrustRegionBoxType], + subspaces: Sequence[TrustRegionBox], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> TensorType: """Get mask for subspaces that need to be initialized. - # Initialize the subspaces that have non-unique locations. + 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)) From 12c8f1f99ae361edeeac0484dd2be929a08459ba Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 7 Aug 2023 14:20:31 +0100 Subject: [PATCH 17/33] Fix func to work with old tensorflow tf.sqrt for array with inf returns nan in older versions. --- trieste/acquisition/utils.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/trieste/acquisition/utils.py b/trieste/acquisition/utils.py index 51faf57e61..04f9b1dd94 100644 --- a/trieste/acquisition/utils.py +++ b/trieste/acquisition/utils.py @@ -14,7 +14,6 @@ import functools from typing import Mapping, Optional, Tuple, TypeVar, Union -import numpy as np import tensorflow as tf from check_shapes import check_shapes @@ -159,23 +158,19 @@ def get_unique_points_mask(points: TensorType, tolerance: float = 1e-6) -> Tenso tolerance = tf.constant(tolerance, dtype=points.dtype) n_points = tf.shape(points)[0] - used_points = tf.fill(value=tf.constant(np.inf, dtype=points.dtype), dims=tf.shape(points)) mask = tf.zeros(shape=(n_points,), dtype=tf.bool) for idx in tf.range(n_points): - # Pairwise distance with remaining 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) - is_unique_point = min_distance >= tolerance # Update mask. + is_unique_point = min_distance >= tolerance mask = tf.tensor_scatter_nd_update(mask, [[idx]], [is_unique_point]) - if is_unique_point: - # Save this point to 'used_points' for future iterations. - used_points = tf.tensor_scatter_nd_update(used_points, [[idx]], [points[idx]]) - return mask From 3c725bd9e8ee5d9933ae4f684ccc7ba6dbc3fa2c Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 7 Aug 2023 17:48:01 +0100 Subject: [PATCH 18/33] Add rule and optimizer unit tests --- tests/unit/acquisition/test_optimizer.py | 10 ++ tests/unit/acquisition/test_rule.py | 120 ++++++++++++++++++++++- tests/unit/test_space.py | 20 ++-- 3 files changed, 139 insertions(+), 11 deletions(-) diff --git a/tests/unit/acquisition/test_optimizer.py b/tests/unit/acquisition/test_optimizer.py index 1c4a090885..650ad50068 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,15 @@ 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) + + @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 11d1bf507c..3d65f391d1 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -33,6 +33,7 @@ AcquisitionFunction, AcquisitionFunctionBuilder, NegativeLowerConfidenceBound, + ParallelContinuousThompsonSampling, SingleModelAcquisitionBuilder, SingleModelGreedyAcquisitionBuilder, VectorizedAcquisitionFunctionBuilder, @@ -47,6 +48,7 @@ BatchHypervolumeSharpeRatioIndicator, DiscreteThompsonSampling, EfficientGlobalOptimization, + MultiTrustRegionBox, RandomSampling, TrustRegion, TrustRegionBox, @@ -61,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 @@ -1261,6 +1263,122 @@ def test_trust_region_box_update_size(success: bool) -> None: npt.assert_allclose(trb.upper, np.minimum(point + 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} + # models = {OBJECTIVE: svgp_model(dataset.query_points, dataset.observations)} + 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 = [TrustRegionBox(search_space, beta=0.1, kappa=1e-3, min_eps=1e-1) for _ in range(2)] + mtb = MultiTrustRegionBox(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, TrustRegionBox) + 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 + + +class TestTrustRegionBox(TrustRegionBox): + 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 = MultiTrustRegionBox(subspaces, base_rule) + state = MultiTrustRegionBox.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_asynchronous_rule_state_pending_points() -> None: pending_points = tf.constant([[1], [2], [3]]) diff --git a/tests/unit/test_space.py b/tests/unit/test_space.py index 9564f4840e..dde53430cf 100644 --- a/tests/unit/test_space.py +++ b/tests/unit/test_space.py @@ -692,11 +692,11 @@ def test_collection_space_subspace_tags_attribute( ) -> None: decision_space = Box([-1, -2], [2, 3]) context_space = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) - product_space = search_space_type( + 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"]) @pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) @@ -705,9 +705,9 @@ def test_collection_space_subspace_tags_default_behaviour( ) -> None: decision_space = Box([-1, -2], [2, 3]) context_space = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) - product_space = search_space_type(spaces=[context_space, decision_space]) + multi_space = search_space_type(spaces=[context_space, decision_space]) - npt.assert_array_equal(product_space.subspace_tags, ["0", "1"]) + npt.assert_array_equal(multi_space.subspace_tags, ["0", "1"]) @pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) @@ -716,12 +716,12 @@ def test_collection_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 = search_space_type(spaces=[space_A, space_B], tags=["A", "B"]) + 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" ): - product_space.get_subspace("dummy") + multi_space.get_subspace("dummy") @pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) @@ -729,18 +729,18 @@ def test_collection_space_get_subspace(search_space_type: Type[CollectionSearchS space_A = Box([-1, -2], [2, 3]) space_B = DiscreteSearchSpace(tf.constant([[-0.5, 0.5]])) space_C = Box([-1, -3], [2, 2]) - product_space = search_space_type(spaces=[space_A, space_B, space_C], tags=["A", "B", "C"]) + multi_space = search_space_type(spaces=[space_A, space_B, space_C], tags=["A", "B", "C"]) - subspace_A = product_space.get_subspace("A") + 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 = product_space.get_subspace("B") + 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 = product_space.get_subspace("C") + 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]) From be06c63d1371b034b135f36b037a75ef189d2cff Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Tue, 8 Aug 2023 10:05:52 +0100 Subject: [PATCH 19/33] Add integ test --- tests/integration/test_bayesian_optimization.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/integration/test_bayesian_optimization.py b/tests/integration/test_bayesian_optimization.py index 28b0bd3886..a64f04ed54 100644 --- a/tests/integration/test_bayesian_optimization.py +++ b/tests/integration/test_bayesian_optimization.py @@ -52,7 +52,9 @@ BatchHypervolumeSharpeRatioIndicator, DiscreteThompsonSampling, EfficientGlobalOptimization, + MultiTrustRegionBox, TrustRegion, + TrustRegionBox, ) from trieste.acquisition.sampler import ThompsonSamplerFromTrajectory from trieste.bayesian_optimizer import ( @@ -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, + MultiTrustRegionBox( + [TrustRegionBox(ScaledBranin.search_space) for _ in range(3)], + EfficientGlobalOptimization( + ParallelContinuousThompsonSampling(), + num_query_points=3, + ), + ), + id="MultiTrustRegionBox", + ), pytest.param(15, DiscreteThompsonSampling(500, 5), id="DiscreteThompsonSampling"), pytest.param( 15, From 07066569485d87836242ba5626504c856ba11c27 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Tue, 8 Aug 2023 14:45:09 +0100 Subject: [PATCH 20/33] Add shape-checking/docstring for mask func Plus one other minor tweak. --- trieste/acquisition/rule.py | 31 ++++++++++++++++++++++++------- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 211211eb8e..677b0db5ab 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -25,6 +25,7 @@ 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 @@ -1198,13 +1199,13 @@ def state_func( MultiTrustRegion acquisition rule {self._tags}""" subspaces = [] - for index, tag in enumerate(self._tags): + for tag, init_subspace in zip(self._tags, self._init_subspaces): if state is None: - subspace = self._init_subspaces[index] + subspace = init_subspace subspace.initialize(models, datasets) else: _subspace = state.acquisition_space.get_subspace(tag) - assert isinstance(_subspace, type(self._init_subspaces[index])) + assert isinstance(_subspace, type(init_subspace)) subspace = _subspace subspace.update(models, datasets) @@ -1235,18 +1236,35 @@ def maybe_initialize_subspaces( Initialize individual subpaces by calling the method of the UpdateableSearchSpaceType class. """ 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[UpdateableSearchSpaceType], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> TensorType: - """Get mask for subspaces that need to be initialized.""" + """ + 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. + """ ... @@ -1359,15 +1377,14 @@ def get_local_min(self, dataset: Optional[Dataset]) -> Tuple[TensorType, TensorT class MultiTrustRegionBox(MultiTrustRegion[ProbabilisticModelType, TrustRegionBox]): """Implements the :class:`MultiTrustRegion` *trust region* acquisition algorithm for a box.""" + @inherit_check_shapes def get_initialize_subspaces_mask( self, subspaces: Sequence[TrustRegionBox], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> TensorType: - """Get mask for subspaces that need to be initialized. - Initialize the subspaces that have non-unique locations. - """ + # 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)) From 9def5622f29fc96001fa7cfd3b09351455af52a4 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Tue, 8 Aug 2023 17:47:37 +0100 Subject: [PATCH 21/33] Add check_shapes, update docstrings & add tag test --- tests/unit/acquisition/test_rule.py | 26 +++++++++++++++- trieste/acquisition/rule.py | 46 ++++++++++++++++++++++------- trieste/space.py | 11 ++++++- 3 files changed, 70 insertions(+), 13 deletions(-) diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index 3d65f391d1..98b38338fe 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -1271,7 +1271,6 @@ def test_multi_trust_region_box_acquire_no_state() -> None: tf.constant([[0.5], [0.0], [1.0]], dtype=tf.float64), ) datasets = {OBJECTIVE: dataset} - # models = {OBJECTIVE: svgp_model(dataset.query_points, dataset.observations)} model = QuadraticMeanAndRBFKernelWithSamplers( dataset=dataset, noise_variance=tf.constant(1.0, dtype=tf.float64) ) @@ -1301,6 +1300,31 @@ def test_multi_trust_region_box_acquire_no_state() -> None: 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 = [TrustRegionBox(search_space) for _ in range(2)] + mtb = MultiTrustRegionBox(subspaces, base_rule) + + state = MultiTrustRegionBox.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(TrustRegionBox): def __init__( self, diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 677b0db5ab..160d7e491e 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -1113,7 +1113,12 @@ def initialize( models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> None: - """Initialize the search space using the given models and datasets.""" + """ + 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 @@ -1122,7 +1127,12 @@ def update( models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> None: - """Update the search space using the given models and datasets.""" + """ + Update the search space using the given models and datasets. + + :param models: The model for each tag. + :param datasets: The dataset for each tag. + """ ... @@ -1160,9 +1170,8 @@ def __init__( ): """ :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 a - given search space. Defaults to :class:`EfficientGlobalOptimization` with default - arguments. + :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() @@ -1184,9 +1193,12 @@ def acquire( def state_func( state: MultiTrustRegion.State | None, ) -> Tuple[MultiTrustRegion.State | None, TensorType]: - """If state is None, initialise the subspaces by picking new locations. Otherwise, + """ + If state is None, initialize the subspaces by picking new locations. Otherwise, update the existing subspaces. - Initialize the subspaces if necessary, potentially looking at the entire group. + + 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 @@ -1231,9 +1243,12 @@ def maybe_initialize_subspaces( models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> None: - """Initialize subspaces if necessary. + """ + 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 UpdateableSearchSpaceType 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( @@ -1324,7 +1339,10 @@ def initialize( models: Optional[Mapping[Tag, ProbabilisticModelType]] = None, datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> None: - """Initialize the box.""" + """ + 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) @@ -1338,8 +1356,10 @@ def update( 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, initialize the box.""" + """ + 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): @@ -1356,6 +1376,10 @@ def update( _, self.y_min = self.get_local_min(dataset) + @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: diff --git a/trieste/space.py b/trieste/space.py index 28f84f651e..a11ecf4c05 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 @@ -1031,6 +1032,7 @@ def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] 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 @@ -1041,6 +1043,7 @@ def lower(self) -> TensorType: ) @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 @@ -1051,6 +1054,7 @@ def upper(self) -> TensorType: ) @property + @check_shapes("return: []") def dimension(self) -> TensorType: """The number of inputs in this product search space.""" return self._dimension @@ -1108,6 +1112,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 @@ -1151,7 +1156,7 @@ class TaggedMultiSearchSpace(CollectionSearchSpace): 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:`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. @@ -1185,6 +1190,7 @@ def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] super().__init__(spaces, tags) @property + @check_shapes("return: [V, D]") def lower(self) -> TensorType: """Returns the stacked lower bounds of all the subspaces. @@ -1196,6 +1202,7 @@ def lower(self) -> TensorType: return tf.stack(lower, axis=0) if lower else tf.constant([], dtype=DEFAULT_DTYPE) @property + @check_shapes("return: [V, D]") def upper(self) -> TensorType: """Returns the stacked upper bounds of all the subspaces. @@ -1207,10 +1214,12 @@ def upper(self) -> TensorType: return tf.stack(upper, axis=0) if upper else tf.constant([], dtype=DEFAULT_DTYPE) @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 From 0973630084bccabaac282472360c02bb14c56751 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Tue, 8 Aug 2023 17:50:12 +0100 Subject: [PATCH 22/33] Remove notebook, to be added in separate PR --- docs/notebooks/multi_trust_region.pct.py | 201 ----------------------- 1 file changed, 201 deletions(-) delete mode 100644 docs/notebooks/multi_trust_region.pct.py diff --git a/docs/notebooks/multi_trust_region.pct.py b/docs/notebooks/multi_trust_region.pct.py deleted file mode 100644 index dce4ada0d5..0000000000 --- a/docs/notebooks/multi_trust_region.pct.py +++ /dev/null @@ -1,201 +0,0 @@ -# %% -# from aim.ext.tensorboard_tracker import Run -from datetime import datetime - -import numpy as np -import tensorflow as tf -from matplotlib import pyplot as plt -from matplotlib.patches import Rectangle -from matplotlib.pyplot import cm - -import trieste -from trieste.acquisition import ParallelContinuousThompsonSampling -from trieste.acquisition.optimizer import automatic_optimizer_selector -from trieste.acquisition.rule import MultiTrustRegionBox, TrustRegionBox -from trieste.acquisition.utils import split_acquisition_function_calls -from trieste.ask_tell_optimization import AskTellOptimizer -from trieste.experimental.plotting import plot_regret -from trieste.experimental.plotting.plotting import create_grid -from trieste.logging import pyplot -from trieste.models.gpflow import GaussianProcessRegression, build_gpr -from trieste.objectives import Hartmann6, ScaledBranin -from trieste.types import TensorType - -# %% -np.random.seed(179) -tf.random.set_seed(179) - -# CONFIG -tensorboard_dir_1 = f"./results/{datetime.now()}/tensorboard" - -summary_writer = tf.summary.create_file_writer(tensorboard_dir_1) -trieste.logging.set_tensorboard_writer(summary_writer) - -# %% -obj = ScaledBranin.objective -search_space = ScaledBranin.search_space - - -def obj_fun( - x: TensorType, -) -> TensorType: # contaminate observations with Gaussian noise - return obj(x) # + tf.random.normal([len(x), 1], 0, .1, tf.float64) - - -num_initial_data_points = 6 -num_query_points = 3 -num_steps = 10 - -# %% -initial_query_points = search_space.sample(num_initial_data_points) -observer = trieste.objectives.utils.mk_observer(obj_fun) -initial_data = observer(initial_query_points) - - -# gpflow_model = build_svgp( -# initial_data, search_space, likelihood_variance=0.001, num_inducing_points=50 -# ) -# -# inducing_point_selector = ConditionalImprovementReduction() -# -# model = SparseVariational( -# gpflow_model, -# num_rff_features=1000, -# inducing_point_selector=inducing_point_selector, -# optimizer=BatchOptimizer( -# tf.optimizers.Adam(0.05), max_iter=100, batch_size=50, compile=True -# ), -# ) -gpflow_model = build_gpr( - initial_data, - search_space, - likelihood_variance=1e-4, - trainable_likelihood=False, -) -model = GaussianProcessRegression(gpflow_model) - -base_rule = trieste.acquisition.rule.EfficientGlobalOptimization( - builder=ParallelContinuousThompsonSampling(), - num_query_points=num_query_points, - optimizer=split_acquisition_function_calls( - automatic_optimizer_selector, split_size=100_000 - ), -) - -subspaces = [TrustRegionBox(search_space) for _ in range(num_query_points)] -acq_rule = MultiTrustRegionBox(subspaces, base_rule) - -ask_tell = AskTellOptimizer( - search_space, initial_data, model, fit_model=True, acquisition_rule=acq_rule -) - -# %% -color = cm.rainbow(np.linspace(0, 1, num_query_points)) - -Xplot, xx, yy = create_grid( - mins=search_space.lower, maxs=search_space.upper, grid_density=90 -) -ff = obj_fun(Xplot).numpy() - -for step in range(num_steps): - print(f"step number {step}") - trieste.logging.set_step_number(step) - - new_points = ask_tell.ask() - new_data = observer(new_points) - # monitor models after each tell - if summary_writer: - models = ask_tell._models # pylint: disable=protected-access - trieste.logging.set_step_number(step) - - with summary_writer.as_default(step=step): - for tag, model in models.items(): - with tf.name_scope(f"{tag}.model"): - model.log() - - fig, ax = plt.subplots(1, 2, squeeze=False, figsize=(15, 5)) - fig.suptitle(f"step number {step}") - ax[0, 0].scatter( - ask_tell.dataset.query_points[:, 0].numpy(), - ask_tell.dataset.query_points[:, 1].numpy(), - color="blue", - ) - ax[0, 0].scatter( - new_points[:, 0].numpy(), new_points[:, 1].numpy(), color="red" - ) - - state = ask_tell.acquisition_state - assert state is not None - assert isinstance(state, MultiTrustRegionBox.State) - - xmin = { - tag: state.acquisition_space.get_subspace( # type: ignore[attr-defined] - tag - ).get_local_min( - ask_tell.dataset - )[ - 0 - ] - for tag in state.acquisition_space.subspace_tags - } - i = 0 - - ax[0, 1].contour(xx, yy, ff.reshape(*xx.shape), 80, alpha=0.5) - - for tag in state.acquisition_space.subspace_tags: - ax[0, 1].scatter( - xmin[tag].numpy()[0], - xmin[tag].numpy()[1], - color=color[i], - marker="x", - alpha=0.5, - ) - lb = state.acquisition_space.get_subspace(tag).lower - ub = state.acquisition_space.get_subspace(tag).upper - ax[0, 1].add_patch( - Rectangle( - (lb[0], lb[1]), - ub[0] - lb[0], - ub[1] - lb[1], - facecolor=color[i], - edgecolor=color[i], - alpha=0.3, - ) - ) - ax[0, 1].scatter( - new_points[i, 0].numpy(), - new_points[i, 1].numpy(), - color=color[i], - alpha=0.5, - ) - ax[0, 1].scatter( - ask_tell.dataset.query_points[:, 0].numpy(), - ask_tell.dataset.query_points[:, 1].numpy(), - color="black", - alpha=0.2, - ) - i = i + 1 - - pyplot("Query points", fig) - plt.close(fig) - - ask_tell.tell(new_data) - -# %% -dataset = ask_tell.dataset - -ground_truth_regret = obj(dataset.query_points) - Hartmann6.minimum -best_found_truth_idx = tf.squeeze(tf.argmin(ground_truth_regret, axis=0)) - -fig, ax = plt.subplots() -plot_regret( - ground_truth_regret.numpy(), ax, num_init=10, idx_best=best_found_truth_idx -) - -ax.set_yscale("log") -ax.set_ylabel("Regret") -ax.set_xlabel("# evaluations") - -# %% -fig, ax = plt.subplots() -ax.scatter(dataset.query_points[:, 0], dataset.query_points[:, 1]) From 01f370b97137fbc27ca2eec6710df2f3419c9d8d Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Tue, 8 Aug 2023 18:32:19 +0100 Subject: [PATCH 23/33] Add deepcopy test --- tests/unit/acquisition/test_rule.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index 98b38338fe..caac4ab316 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -1403,6 +1403,34 @@ def test_multi_trust_region_box_acquire_with_state() -> None: 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 = [TrustRegionBox(search_space, 0.07, 1e-5, 1e-3) for _ in range(3)] + for _subspace in subspaces: + _subspace.initialize(datasets={OBJECTIVE: dataset}) + state = MultiTrustRegionBox.State(acquisition_space=TaggedMultiSearchSpace(subspaces)) + + state_copy = copy.deepcopy(state) + assert state_copy is not state + 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 isinstance(subspace, TrustRegionBox) + assert isinstance(subspace_copy, TrustRegionBox) + 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]]) From 66b4aa453c339f58ce0bb36da857eed169a48ebe Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 9 Aug 2023 12:41:27 +0100 Subject: [PATCH 24/33] Check TR copy is deep --- tests/unit/acquisition/test_rule.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index caac4ab316..4433a42173 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -1416,11 +1416,14 @@ def test_multi_trust_region_box_state_deepcopy() -> None: 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, TrustRegionBox) assert isinstance(subspace_copy, TrustRegionBox) assert subspace._beta == subspace_copy._beta From 9bb6a8070bfc428aed05a81231e741198de5f0ee Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Thu, 10 Aug 2023 13:15:39 +0100 Subject: [PATCH 25/33] Improve TR box and test slightly --- tests/unit/acquisition/test_rule.py | 47 ++++++++++++++++++----------- trieste/acquisition/rule.py | 27 +++++++++++------ 2 files changed, 47 insertions(+), 27 deletions(-) diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index 4433a42173..32122f569f 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -1172,7 +1172,7 @@ def test_trust_region_box_initialize() -> None: 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)) + 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. @@ -1221,46 +1221,57 @@ def test_trust_region_box_update_size(success: bool) -> None: 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), + tf.constant([[0.5], [0.3], [1.0]], dtype=tf.float64), ) } trb = TrustRegionBox(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. - point = trb.sample(1) + new_point = trb.sample(1) else: # Pick point outside the box. - point = tf.constant([[1.2, 1.3]], dtype=tf.float64) + 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, point], axis=0), - np.concatenate([datasets[OBJECTIVE].observations, [[-0.1]]], axis=0), + 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. - point = np.squeeze(point) - npt.assert_allclose(point, trb.location) - npt.assert_allclose(tf.constant([-0.1], dtype=tf.float64), trb.y_min) - # Check that the box is smaller by beta. + 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. - point, y_min = trb.get_local_min(datasets[OBJECTIVE]) - npt.assert_allclose(point, trb.location) - npt.assert_allclose(y_min, trb.y_min) - # Check that the box is larger by beta. + 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(point - trb._eps, search_space.lower)) - npt.assert_allclose(trb.upper, np.minimum(point + trb._eps, search_space.upper)) + 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. @@ -1398,7 +1409,7 @@ def test_multi_trust_region_box_acquire_with_state() -> None: [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) + npt.assert_array_equal(subspace._y_min, exp_obs) # Check the box was updated/initialized correctly. npt.assert_allclose(subspace._eps, exp_eps) @@ -1431,7 +1442,7 @@ def test_multi_trust_region_box_state_deepcopy() -> None: 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) + npt.assert_array_equal(subspace._y_min, subspace_copy._y_min) def test_asynchronous_rule_state_pending_points() -> None: diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 160d7e491e..7a73477588 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -1321,6 +1321,16 @@ def location(self, location: TensorType) -> None: """Set the location of the search space.""" self._location = location + @property + def eps(self) -> TensorType: + """The size of the search space.""" + return self._eps + + @eps.setter + def eps(self, eps: TensorType) -> None: + """Set the size of the search space.""" + self._eps = eps + def _init_eps(self) -> None: global_lower = self.global_search_space.lower global_upper = self.global_search_space.upper @@ -1328,10 +1338,10 @@ def _init_eps(self) -> None: def _update_bounds(self) -> None: self._lower = tf.reduce_max( - [self.global_search_space.lower, self.location - self._eps], axis=0 + [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 + [self.global_search_space.upper, self.location + self.eps], axis=0 ) def initialize( @@ -1346,10 +1356,10 @@ def initialize( 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) + _, self._y_min = self.get_local_min(dataset) def update( self, @@ -1362,7 +1372,7 @@ def update( """ dataset = get_value_for_tag(datasets) - if tf.reduce_any(self._eps < self._min_eps): + if tf.reduce_any(self.eps < self._min_eps): self.initialize(models, datasets) return @@ -1370,11 +1380,10 @@ def update( self.location = x_min tr_volume = tf.reduce_prod(self.upper - self.lower) - step_is_success = y_min < self.y_min - self._kappa * tr_volume - self._eps = self._eps / self._beta if step_is_success else self._eps * self._beta + 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 = self.get_local_min(dataset) + self._y_min = y_min @check_shapes( "return[0]: [D]", From 492455722fa5099a0a6068bf51e9e215d71a092b Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 21 Aug 2023 13:47:31 +0100 Subject: [PATCH 26/33] Address Uri's latest comments --- tests/unit/acquisition/test_rule.py | 20 ++++++++++---------- trieste/acquisition/rule.py | 27 ++++++--------------------- 2 files changed, 16 insertions(+), 31 deletions(-) diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index 32122f569f..a85b30d4b3 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -1166,7 +1166,7 @@ def test_trust_region_box_initialize() -> None: 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_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) @@ -1236,7 +1236,7 @@ def test_trust_region_box_update_size(success: bool) -> None: ) trb.update(datasets=datasets) - eps = trb._eps + eps = trb.eps if success: # Sample a point from the box. @@ -1260,18 +1260,18 @@ def test_trust_region_box_update_size(success: bool) -> None: 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) + 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) + 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)) + 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. @@ -1357,7 +1357,7 @@ def location(self, location: TensorType) -> None: ... def _init_eps(self) -> None: - self._eps = tf.constant(0.07, dtype=tf.float64) + self.eps = tf.constant(0.07, dtype=tf.float64) # Start with a defined state and dataset. Acquire should return an updated state. @@ -1411,7 +1411,7 @@ def test_multi_trust_region_box_acquire_with_state() -> None: 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) + npt.assert_allclose(subspace.eps, exp_eps) def test_multi_trust_region_box_state_deepcopy() -> None: @@ -1440,8 +1440,8 @@ def test_multi_trust_region_box_state_deepcopy() -> None: 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.eps, subspace_copy.eps) + npt.assert_array_equal(subspace.location, subspace_copy.location) npt.assert_array_equal(subspace._y_min, subspace_copy._y_min) diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 7a73477588..81daeb09fb 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -1297,6 +1297,11 @@ def __init__( 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 @@ -1311,30 +1316,10 @@ def global_search_space(self) -> SearchSpace: """The global search space this search space lives in.""" return self._global_search_space - @property - def location(self) -> TensorType: - """The location of the search space.""" - return self._location - - @location.setter - def location(self, location: TensorType) -> None: - """Set the location of the search space.""" - self._location = location - - @property - def eps(self) -> TensorType: - """The size of the search space.""" - return self._eps - - @eps.setter - def eps(self, eps: TensorType) -> None: - """Set the size of the search space.""" - self._eps = eps - 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])) + 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( From f0e125d4646f769f30a1b0404318a6dc9f9f4c82 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 21 Aug 2023 15:39:13 +0100 Subject: [PATCH 27/33] Move get_value_for_tag to more general utils/misc --- tests/unit/acquisition/test_utils.py | 19 ------------------ tests/unit/utils/test_misc.py | 29 +++++++++++++++++++++++++++- trieste/acquisition/rule.py | 3 ++- trieste/acquisition/utils.py | 26 ++----------------------- trieste/utils/misc.py | 24 ++++++++++++++++++++++- 5 files changed, 55 insertions(+), 46 deletions(-) diff --git a/tests/unit/acquisition/test_utils.py b/tests/unit/acquisition/test_utils.py index e84ab6a42b..7975cab3a0 100644 --- a/tests/unit/acquisition/test_utils.py +++ b/tests/unit/acquisition/test_utils.py @@ -24,12 +24,10 @@ from trieste.acquisition.utils import ( get_local_dataset, get_unique_points_mask, - get_value_for_tag, select_nth_output, split_acquisition_function, ) from trieste.data import Dataset -from trieste.observer import OBJECTIVE from trieste.space import Box, SearchSpaceType @@ -132,20 +130,3 @@ def test_get_unique_points_mask( ) -> None: mask = get_unique_points_mask(points, tolerance) np.testing.assert_array_equal(mask, expected_mask) - - -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" 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/rule.py b/trieste/acquisition/rule.py index 81daeb09fb..ef561ba70f 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -50,6 +50,7 @@ from ..observer import OBJECTIVE 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, @@ -72,7 +73,7 @@ batchify_vectorize, ) from .sampler import ExactThompsonSampler, ThompsonSampler -from .utils import get_local_dataset, get_unique_points_mask, get_value_for_tag, 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. """ diff --git a/trieste/acquisition/utils.py b/trieste/acquisition/utils.py index 04f9b1dd94..590b0bb416 100644 --- a/trieste/acquisition/utils.py +++ b/trieste/acquisition/utils.py @@ -12,15 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Mapping, Optional, Tuple, TypeVar, Union +from typing import Tuple, Union import tensorflow as tf from check_shapes import check_shapes from ..data import Dataset -from ..observer import OBJECTIVE from ..space import SearchSpaceType -from ..types import Tag, TensorType +from ..types import TensorType from .interface import AcquisitionFunction from .optimizer import AcquisitionOptimizer @@ -172,24 +171,3 @@ def get_unique_points_mask(points: TensorType, tolerance: float = 1e-6) -> Tenso mask = tf.tensor_scatter_nd_update(mask, [[idx]], [is_unique_point]) return mask - - -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") 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: From 3daa81f6abbf3fb2a5f085c0b7fbd6d600b170a7 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 21 Aug 2023 16:04:36 +0100 Subject: [PATCH 28/33] Remove redundant [] checks --- trieste/space.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/trieste/space.py b/trieste/space.py index d37a7d7194..23f50786e0 100644 --- a/trieste/space.py +++ b/trieste/space.py @@ -1188,9 +1188,7 @@ def lower(self) -> TensorType: :return: The lower bounds of shape [V, D], where V is the number of subspaces and D is the dimensionality of each subspace. """ - - lower = self.subspace_lower - return tf.stack(lower, axis=0) if lower else tf.constant([], dtype=DEFAULT_DTYPE) + return tf.stack(self.subspace_lower, axis=0) @property @check_shapes("return: [V, D]") @@ -1200,9 +1198,7 @@ def upper(self) -> TensorType: :return: The upper bounds of shape [V, D], where V is the number of subspaces and D is the dimensionality of each subspace. """ - - upper = self.subspace_upper - return tf.stack(upper, axis=0) if upper else tf.constant([], dtype=DEFAULT_DTYPE) + return tf.stack(self.subspace_upper, axis=0) @property @check_shapes("return: []") From 0b19246438cfa60ecc2f3202533d9aec2de0f0f5 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 21 Aug 2023 16:25:58 +0100 Subject: [PATCH 29/33] Add fixture for default search space type --- tests/unit/test_space.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_space.py b/tests/unit/test_space.py index dde53430cf..f170d8a139 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, List, Optional, Sequence, Type +from typing import Any, Container, List, Optional, Sequence, Type import numpy.testing as npt import pytest @@ -664,7 +664,11 @@ def test_box_deepcopy() -> None: npt.assert_allclose(box.upper, box_copy.upper) -@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) +@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: @@ -674,7 +678,6 @@ def test_collection_space_raises_for_non_unqique_subspace_names( search_space_type(spaces=[space_A, space_B], tags=["A", "A"]) -@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) def test_collection_space_raises_for_length_mismatch_between_spaces_and_tags( search_space_type: Type[CollectionSearchSpace], ) -> None: @@ -686,7 +689,6 @@ def test_collection_space_raises_for_length_mismatch_between_spaces_and_tags( search_space_type(spaces=[space_A, space_B], tags=["A", "B", "C"]) -@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) def test_collection_space_subspace_tags_attribute( search_space_type: Type[CollectionSearchSpace], ) -> None: @@ -699,7 +701,6 @@ def test_collection_space_subspace_tags_attribute( npt.assert_array_equal(multi_space.subspace_tags, ["context", "decision"]) -@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) def test_collection_space_subspace_tags_default_behaviour( search_space_type: Type[CollectionSearchSpace], ) -> None: @@ -710,7 +711,6 @@ def test_collection_space_subspace_tags_default_behaviour( npt.assert_array_equal(multi_space.subspace_tags, ["0", "1"]) -@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) def test_collection_space_get_subspace_raises_for_invalid_tag( search_space_type: Type[CollectionSearchSpace], ) -> None: @@ -724,7 +724,6 @@ def test_collection_space_get_subspace_raises_for_invalid_tag( multi_space.get_subspace("dummy") -@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) 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]])) @@ -846,7 +845,6 @@ def test_collection_space_contains_broadcasts( _ = points in collection_space -@pytest.mark.parametrize("search_space_type", [TaggedMultiSearchSpace, TaggedProductSearchSpace]) @pytest.mark.parametrize( "spaces", [ From 044f07cc52101bf5f5bb83acfe8c08af71ab00cc Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Tue, 22 Aug 2023 12:08:32 +0100 Subject: [PATCH 30/33] Add test for optim multiple boxes as a batch --- tests/unit/acquisition/test_optimizer.py | 25 ++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/unit/acquisition/test_optimizer.py b/tests/unit/acquisition/test_optimizer.py index 650ad50068..e2f7e6a0f3 100644 --- a/tests/unit/acquisition/test_optimizer.py +++ b/tests/unit/acquisition/test_optimizer.py @@ -213,6 +213,31 @@ def test_optimize_continuous_raises_with_mismatch_multi_search_space() -> None: 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( From 8ca39e96b578bcb7b7bd6e337d113180054b1258 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Tue, 22 Aug 2023 14:18:39 +0100 Subject: [PATCH 31/33] Rename classes and expand some docstrings --- .../integration/test_bayesian_optimization.py | 10 ++-- tests/unit/acquisition/test_rule.py | 46 ++++++++++--------- trieste/acquisition/rule.py | 31 +++++++------ trieste/space.py | 9 ++++ 4 files changed, 55 insertions(+), 41 deletions(-) diff --git a/tests/integration/test_bayesian_optimization.py b/tests/integration/test_bayesian_optimization.py index a64f04ed54..66e769fcfb 100644 --- a/tests/integration/test_bayesian_optimization.py +++ b/tests/integration/test_bayesian_optimization.py @@ -50,11 +50,11 @@ AsynchronousOptimization, AsynchronousRuleState, BatchHypervolumeSharpeRatioIndicator, + BatchTrustRegionBox, DiscreteThompsonSampling, EfficientGlobalOptimization, - MultiTrustRegionBox, + SingleObjectiveTRBox, TrustRegion, - TrustRegionBox, ) from trieste.acquisition.sampler import ThompsonSamplerFromTrajectory from trieste.bayesian_optimizer import ( @@ -201,14 +201,14 @@ def GPR_OPTIMIZER_PARAMS() -> Tuple[str, List[ParameterSet]]: ), pytest.param( 10, - MultiTrustRegionBox( - [TrustRegionBox(ScaledBranin.search_space) for _ in range(3)], + BatchTrustRegionBox( + [SingleObjectiveTRBox(ScaledBranin.search_space) for _ in range(3)], EfficientGlobalOptimization( ParallelContinuousThompsonSampling(), num_query_points=3, ), ), - id="MultiTrustRegionBox", + id="BatchTrustRegionBox", ), pytest.param(15, DiscreteThompsonSampling(500, 5), id="DiscreteThompsonSampling"), pytest.param( diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index a85b30d4b3..94b157e136 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -46,12 +46,12 @@ AsynchronousOptimization, AsynchronousRuleState, BatchHypervolumeSharpeRatioIndicator, + BatchTrustRegionBox, DiscreteThompsonSampling, EfficientGlobalOptimization, - MultiTrustRegionBox, RandomSampling, + SingleObjectiveTRBox, TrustRegion, - TrustRegionBox, ) from trieste.acquisition.sampler import ( ExactThompsonSampler, @@ -1119,7 +1119,7 @@ def test_turbo_state_deepcopy() -> None: # 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 = TrustRegionBox(search_space) + trb = SingleObjectiveTRBox(search_space) with pytest.raises(ValueError, match="dataset must be provided"): trb.get_local_min(None) @@ -1131,7 +1131,7 @@ def test_trust_region_box_get_local_min() -> None: 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 = TrustRegionBox(search_space) + trb = SingleObjectiveTRBox(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) @@ -1147,7 +1147,7 @@ def test_trust_region_box_get_local_min_outside_search_space() -> None: tf.constant([[1.2, 1.3], [-0.4, -0.5]], dtype=tf.float64), tf.constant([[0.7], [0.9]], dtype=tf.float64), ) - trb = TrustRegionBox(search_space) + trb = SingleObjectiveTRBox(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)) @@ -1162,7 +1162,7 @@ def test_trust_region_box_initialize() -> None: tf.constant([[0.7], [0.9]], dtype=tf.float64), ) } - trb = TrustRegionBox(search_space) + trb = SingleObjectiveTRBox(search_space) trb.initialize(datasets=datasets) exp_eps = 0.5 * (search_space.upper - search_space.lower) / 5.0 ** (1.0 / 2.0) @@ -1184,7 +1184,7 @@ def test_trust_region_box_update_initialize() -> None: tf.constant([[0.7], [0.9]], dtype=tf.float64), ) } - trb = TrustRegionBox(search_space, min_eps=0.5) + trb = SingleObjectiveTRBox(search_space, min_eps=0.5) trb.initialize(datasets=datasets) location = trb.location @@ -1205,7 +1205,7 @@ def test_trust_region_box_update_no_initialize() -> None: tf.constant([[0.5], [0.0], [1.0]], dtype=tf.float64), ) } - trb = TrustRegionBox(search_space, min_eps=0.1) + trb = SingleObjectiveTRBox(search_space, min_eps=0.1) trb.initialize(datasets=datasets) trb.location = tf.constant([0.5, 0.5], dtype=tf.float64) location = trb.location @@ -1224,7 +1224,7 @@ def test_trust_region_box_update_size(success: bool) -> None: tf.constant([[0.5], [0.3], [1.0]], dtype=tf.float64), ) } - trb = TrustRegionBox(search_space, min_eps=0.1) + trb = SingleObjectiveTRBox(search_space, min_eps=0.1) trb.initialize(datasets=datasets) # Ensure there is at least one point captured in the box. @@ -1292,8 +1292,10 @@ def test_multi_trust_region_box_acquire_no_state() -> None: base_rule = EfficientGlobalOptimization( # type: ignore[var-annotated] builder=ParallelContinuousThompsonSampling(), num_query_points=2 ) - subspaces = [TrustRegionBox(search_space, beta=0.1, kappa=1e-3, min_eps=1e-1) for _ in range(2)] - mtb = MultiTrustRegionBox(subspaces, base_rule) + subspaces = [ + SingleObjectiveTRBox(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 @@ -1303,7 +1305,7 @@ def test_multi_trust_region_box_acquire_no_state() -> None: 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, TrustRegionBox) + assert isinstance(subspace, SingleObjectiveTRBox) assert subspace.global_search_space == search_space assert subspace._beta == 0.1 assert subspace._kappa == 1e-3 @@ -1320,10 +1322,10 @@ def test_multi_trust_region_box_raises_on_mismatched_tags() -> None: base_rule = EfficientGlobalOptimization( # type: ignore[var-annotated] builder=ParallelContinuousThompsonSampling(), num_query_points=2 ) - subspaces = [TrustRegionBox(search_space) for _ in range(2)] - mtb = MultiTrustRegionBox(subspaces, base_rule) + subspaces = [SingleObjectiveTRBox(search_space) for _ in range(2)] + mtb = BatchTrustRegionBox(subspaces, base_rule) - state = MultiTrustRegionBox.State( + state = BatchTrustRegionBox.State( acquisition_space=TaggedMultiSearchSpace(subspaces, tags=["a", "b"]) ) state_func = mtb.acquire( @@ -1336,7 +1338,7 @@ def test_multi_trust_region_box_raises_on_mismatched_tags() -> None: _, _ = state_func(state) -class TestTrustRegionBox(TrustRegionBox): +class TestTrustRegionBox(SingleObjectiveTRBox): def __init__( self, fixed_location: TensorType, @@ -1384,8 +1386,8 @@ def test_multi_trust_region_box_acquire_with_state() -> None: 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 = MultiTrustRegionBox(subspaces, base_rule) - state = MultiTrustRegionBox.State(acquisition_space=TaggedMultiSearchSpace(subspaces)) + mtb = BatchTrustRegionBox(subspaces, base_rule) + state = BatchTrustRegionBox.State(acquisition_space=TaggedMultiSearchSpace(subspaces)) for subspace in subspaces: subspace.initialize(datasets={OBJECTIVE: init_dataset}) @@ -1420,10 +1422,10 @@ def test_multi_trust_region_box_state_deepcopy() -> None: 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 = [TrustRegionBox(search_space, 0.07, 1e-5, 1e-3) for _ in range(3)] + subspaces = [SingleObjectiveTRBox(search_space, 0.07, 1e-5, 1e-3) for _ in range(3)] for _subspace in subspaces: _subspace.initialize(datasets={OBJECTIVE: dataset}) - state = MultiTrustRegionBox.State(acquisition_space=TaggedMultiSearchSpace(subspaces)) + state = BatchTrustRegionBox.State(acquisition_space=TaggedMultiSearchSpace(subspaces)) state_copy = copy.deepcopy(state) assert state_copy is not state @@ -1435,8 +1437,8 @@ def test_multi_trust_region_box_state_deepcopy() -> None: state.acquisition_space._spaces.values(), state_copy.acquisition_space._spaces.values() ): assert subspace is not subspace_copy - assert isinstance(subspace, TrustRegionBox) - assert isinstance(subspace_copy, TrustRegionBox) + assert isinstance(subspace, SingleObjectiveTRBox) + assert isinstance(subspace_copy, SingleObjectiveTRBox) assert subspace._beta == subspace_copy._beta assert subspace._kappa == subspace_copy._kappa assert subspace._min_eps == subspace_copy._min_eps diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index ef561ba70f..680ed87c25 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -1141,9 +1141,9 @@ def update( """ A type variable bound to :class:`UpdateableSearchSpace`. """ -class MultiTrustRegion( +class BatchTrustRegion( AcquisitionRule[ - types.State[Optional["MultiTrustRegion.State"], TensorType], + types.State[Optional["BatchTrustRegion.State"], TensorType], SearchSpace, ProbabilisticModelType, ], @@ -1155,17 +1155,17 @@ class MultiTrustRegion( @dataclass(frozen=True) class State: - """The acquisition state for the :class:`MultiTrustRegion` acquisition rule.""" + """The acquisition state for the :class:`BatchTrustRegion` acquisition rule.""" acquisition_space: TaggedMultiSearchSpace """ The search space. """ - def __deepcopy__(self, memo: dict[int, object]) -> MultiTrustRegion.State: + def __deepcopy__(self, memo: dict[int, object]) -> BatchTrustRegion.State: acquisition_space_copy = copy.deepcopy(self.acquisition_space, memo) - return MultiTrustRegion.State(acquisition_space_copy) + return BatchTrustRegion.State(acquisition_space_copy) def __init__( - self: "MultiTrustRegion[ProbabilisticModelType, UpdateableSearchSpaceType]", + self: "BatchTrustRegion[ProbabilisticModelType, UpdateableSearchSpaceType]", init_subspaces: Sequence[UpdateableSearchSpaceType], rule: AcquisitionRule[TensorType, SearchSpace, ProbabilisticModelType] | None = None, ): @@ -1192,8 +1192,8 @@ def acquire( datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> types.State[State | None, TensorType]: def state_func( - state: MultiTrustRegion.State | None, - ) -> Tuple[MultiTrustRegion.State | None, TensorType]: + 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. @@ -1209,7 +1209,7 @@ def state_func( 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 - MultiTrustRegion acquisition rule {self._tags}""" + BatchTrustRegion acquisition rule {self._tags}""" subspaces = [] for tag, init_subspace in zip(self._tags, self._init_subspaces): @@ -1231,7 +1231,7 @@ def state_func( else: acquisition_space = state.acquisition_space - state_ = MultiTrustRegion.State(acquisition_space) + state_ = BatchTrustRegion.State(acquisition_space) points = self._rule.acquire(acquisition_space, models, datasets=datasets) return state_, points @@ -1284,7 +1284,7 @@ def get_initialize_subspaces_mask( ... -class TrustRegionBox(Box, UpdateableSearchSpace): +class SingleObjectiveTRBox(Box, UpdateableSearchSpace): """An updateable box search space for use with trust region acquisition rules.""" def __init__( @@ -1393,13 +1393,16 @@ def get_local_min(self, dataset: Optional[Dataset]) -> Tuple[TensorType, TensorT return tf.squeeze(x_min, axis=0), tf.squeeze(y_min) -class MultiTrustRegionBox(MultiTrustRegion[ProbabilisticModelType, TrustRegionBox]): - """Implements the :class:`MultiTrustRegion` *trust region* acquisition algorithm for a box.""" +class BatchTrustRegionBox(BatchTrustRegion[ProbabilisticModelType, SingleObjectiveTRBox]): + """ + 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[TrustRegionBox], + subspaces: Sequence[SingleObjectiveTRBox], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> TensorType: diff --git a/trieste/space.py b/trieste/space.py index 23f50786e0..222d64df89 100644 --- a/trieste/space.py +++ b/trieste/space.py @@ -993,6 +993,11 @@ class TaggedProductSearchSpace(CollectionSearchSpace): 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 @@ -1143,6 +1148,10 @@ class TaggedMultiSearchSpace(CollectionSearchSpace): 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): From ef62eebcf68a52052a25a8eb73efce10361a15f8 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Thu, 24 Aug 2023 10:02:00 +0100 Subject: [PATCH 32/33] Choose longer class name --- .../integration/test_bayesian_optimization.py | 4 +-- tests/unit/acquisition/test_rule.py | 31 ++++++++++--------- trieste/acquisition/rule.py | 6 ++-- 3 files changed, 21 insertions(+), 20 deletions(-) diff --git a/tests/integration/test_bayesian_optimization.py b/tests/integration/test_bayesian_optimization.py index 66e769fcfb..ca8981f9a9 100644 --- a/tests/integration/test_bayesian_optimization.py +++ b/tests/integration/test_bayesian_optimization.py @@ -53,7 +53,7 @@ BatchTrustRegionBox, DiscreteThompsonSampling, EfficientGlobalOptimization, - SingleObjectiveTRBox, + SingleObjectiveTrustRegionBox, TrustRegion, ) from trieste.acquisition.sampler import ThompsonSamplerFromTrajectory @@ -202,7 +202,7 @@ def GPR_OPTIMIZER_PARAMS() -> Tuple[str, List[ParameterSet]]: pytest.param( 10, BatchTrustRegionBox( - [SingleObjectiveTRBox(ScaledBranin.search_space) for _ in range(3)], + [SingleObjectiveTrustRegionBox(ScaledBranin.search_space) for _ in range(3)], EfficientGlobalOptimization( ParallelContinuousThompsonSampling(), num_query_points=3, diff --git a/tests/unit/acquisition/test_rule.py b/tests/unit/acquisition/test_rule.py index 94b157e136..0e2b575538 100644 --- a/tests/unit/acquisition/test_rule.py +++ b/tests/unit/acquisition/test_rule.py @@ -50,7 +50,7 @@ DiscreteThompsonSampling, EfficientGlobalOptimization, RandomSampling, - SingleObjectiveTRBox, + SingleObjectiveTrustRegionBox, TrustRegion, ) from trieste.acquisition.sampler import ( @@ -1119,7 +1119,7 @@ def test_turbo_state_deepcopy() -> None: # 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 = SingleObjectiveTRBox(search_space) + trb = SingleObjectiveTrustRegionBox(search_space) with pytest.raises(ValueError, match="dataset must be provided"): trb.get_local_min(None) @@ -1131,7 +1131,7 @@ def test_trust_region_box_get_local_min() -> None: 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 = SingleObjectiveTRBox(search_space) + 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) @@ -1147,7 +1147,7 @@ def test_trust_region_box_get_local_min_outside_search_space() -> None: tf.constant([[1.2, 1.3], [-0.4, -0.5]], dtype=tf.float64), tf.constant([[0.7], [0.9]], dtype=tf.float64), ) - trb = SingleObjectiveTRBox(search_space) + 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)) @@ -1162,7 +1162,7 @@ def test_trust_region_box_initialize() -> None: tf.constant([[0.7], [0.9]], dtype=tf.float64), ) } - trb = SingleObjectiveTRBox(search_space) + trb = SingleObjectiveTrustRegionBox(search_space) trb.initialize(datasets=datasets) exp_eps = 0.5 * (search_space.upper - search_space.lower) / 5.0 ** (1.0 / 2.0) @@ -1184,7 +1184,7 @@ def test_trust_region_box_update_initialize() -> None: tf.constant([[0.7], [0.9]], dtype=tf.float64), ) } - trb = SingleObjectiveTRBox(search_space, min_eps=0.5) + trb = SingleObjectiveTrustRegionBox(search_space, min_eps=0.5) trb.initialize(datasets=datasets) location = trb.location @@ -1205,7 +1205,7 @@ def test_trust_region_box_update_no_initialize() -> None: tf.constant([[0.5], [0.0], [1.0]], dtype=tf.float64), ) } - trb = SingleObjectiveTRBox(search_space, min_eps=0.1) + 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 @@ -1224,7 +1224,7 @@ def test_trust_region_box_update_size(success: bool) -> None: tf.constant([[0.5], [0.3], [1.0]], dtype=tf.float64), ) } - trb = SingleObjectiveTRBox(search_space, min_eps=0.1) + trb = SingleObjectiveTrustRegionBox(search_space, min_eps=0.1) trb.initialize(datasets=datasets) # Ensure there is at least one point captured in the box. @@ -1293,7 +1293,8 @@ def test_multi_trust_region_box_acquire_no_state() -> None: builder=ParallelContinuousThompsonSampling(), num_query_points=2 ) subspaces = [ - SingleObjectiveTRBox(search_space, beta=0.1, kappa=1e-3, min_eps=1e-1) for _ in range(2) + 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) @@ -1305,7 +1306,7 @@ def test_multi_trust_region_box_acquire_no_state() -> None: 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, SingleObjectiveTRBox) + assert isinstance(subspace, SingleObjectiveTrustRegionBox) assert subspace.global_search_space == search_space assert subspace._beta == 0.1 assert subspace._kappa == 1e-3 @@ -1322,7 +1323,7 @@ def test_multi_trust_region_box_raises_on_mismatched_tags() -> None: base_rule = EfficientGlobalOptimization( # type: ignore[var-annotated] builder=ParallelContinuousThompsonSampling(), num_query_points=2 ) - subspaces = [SingleObjectiveTRBox(search_space) for _ in range(2)] + subspaces = [SingleObjectiveTrustRegionBox(search_space) for _ in range(2)] mtb = BatchTrustRegionBox(subspaces, base_rule) state = BatchTrustRegionBox.State( @@ -1338,7 +1339,7 @@ def test_multi_trust_region_box_raises_on_mismatched_tags() -> None: _, _ = state_func(state) -class TestTrustRegionBox(SingleObjectiveTRBox): +class TestTrustRegionBox(SingleObjectiveTrustRegionBox): def __init__( self, fixed_location: TensorType, @@ -1422,7 +1423,7 @@ def test_multi_trust_region_box_state_deepcopy() -> None: 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 = [SingleObjectiveTRBox(search_space, 0.07, 1e-5, 1e-3) for _ in range(3)] + 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)) @@ -1437,8 +1438,8 @@ def test_multi_trust_region_box_state_deepcopy() -> None: state.acquisition_space._spaces.values(), state_copy.acquisition_space._spaces.values() ): assert subspace is not subspace_copy - assert isinstance(subspace, SingleObjectiveTRBox) - assert isinstance(subspace_copy, SingleObjectiveTRBox) + 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 diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 680ed87c25..47536f741f 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -1284,7 +1284,7 @@ def get_initialize_subspaces_mask( ... -class SingleObjectiveTRBox(Box, UpdateableSearchSpace): +class SingleObjectiveTrustRegionBox(Box, UpdateableSearchSpace): """An updateable box search space for use with trust region acquisition rules.""" def __init__( @@ -1393,7 +1393,7 @@ def get_local_min(self, dataset: Optional[Dataset]) -> Tuple[TensorType, TensorT return tf.squeeze(x_min, axis=0), tf.squeeze(y_min) -class BatchTrustRegionBox(BatchTrustRegion[ProbabilisticModelType, SingleObjectiveTRBox]): +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. @@ -1402,7 +1402,7 @@ class BatchTrustRegionBox(BatchTrustRegion[ProbabilisticModelType, SingleObjecti @inherit_check_shapes def get_initialize_subspaces_mask( self, - subspaces: Sequence[SingleObjectiveTRBox], + subspaces: Sequence[SingleObjectiveTrustRegionBox], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> TensorType: From 8f178395ad0dc191f0cb2cfb509ee4899f2e1358 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Fri, 25 Aug 2023 11:04:14 +0100 Subject: [PATCH 33/33] Rename updatable class name --- trieste/acquisition/rule.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 47536f741f..e11977667b 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -1105,7 +1105,7 @@ def state_func( return state_func -class UpdateableSearchSpace(SearchSpace): +class UpdatableTrustRegion(SearchSpace): """A search space that can be updated.""" @abstractmethod @@ -1137,8 +1137,8 @@ def update( ... -UpdateableSearchSpaceType = TypeVar("UpdateableSearchSpaceType", bound=UpdateableSearchSpace) -""" A type variable bound to :class:`UpdateableSearchSpace`. """ +UpdatableTrustRegionType = TypeVar("UpdatableTrustRegionType", bound=UpdatableTrustRegion) +""" A type variable bound to :class:`UpdatableTrustRegion`. """ class BatchTrustRegion( @@ -1147,7 +1147,7 @@ class BatchTrustRegion( SearchSpace, ProbabilisticModelType, ], - Generic[ProbabilisticModelType, UpdateableSearchSpaceType], + 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. @@ -1165,8 +1165,8 @@ def __deepcopy__(self, memo: dict[int, object]) -> BatchTrustRegion.State: return BatchTrustRegion.State(acquisition_space_copy) def __init__( - self: "BatchTrustRegion[ProbabilisticModelType, UpdateableSearchSpaceType]", - init_subspaces: Sequence[UpdateableSearchSpaceType], + self: "BatchTrustRegion[ProbabilisticModelType, UpdatableTrustRegionType]", + init_subspaces: Sequence[UpdatableTrustRegionType], rule: AcquisitionRule[TensorType, SearchSpace, ProbabilisticModelType] | None = None, ): """ @@ -1240,14 +1240,14 @@ def state_func( def maybe_initialize_subspaces( self, - subspaces: Sequence[UpdateableSearchSpaceType], + 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 UpdateableSearchSpaceType class. + Initialize individual subpaces by calling the method of the UpdatableTrustRegionType class. This method can be overridden by subclasses to change this behaviour. """ @@ -1266,7 +1266,7 @@ def maybe_initialize_subspaces( @check_shapes("return: [V]") def get_initialize_subspaces_mask( self, - subspaces: Sequence[UpdateableSearchSpaceType], + subspaces: Sequence[UpdatableTrustRegionType], models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> TensorType: @@ -1284,8 +1284,8 @@ def get_initialize_subspaces_mask( ... -class SingleObjectiveTrustRegionBox(Box, UpdateableSearchSpace): - """An updateable box search space for use with trust region acquisition rules.""" +class SingleObjectiveTrustRegionBox(Box, UpdatableTrustRegion): + """An updatable box search space for use with trust region acquisition rules.""" def __init__( self,