diff --git a/trieste/acquisition/interface.py b/trieste/acquisition/interface.py index e7c92859d4..817aaa9c57 100644 --- a/trieste/acquisition/interface.py +++ b/trieste/acquisition/interface.py @@ -18,7 +18,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import Callable, Generic, Mapping, Optional +from typing import Any, Callable, Generic, Mapping, Optional from ..data import Dataset from ..models.interfaces import ProbabilisticModelType @@ -57,6 +57,7 @@ def prepare_acquisition_function( self, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: """ Prepare an acquisition function. We assume that this requires at least models, but @@ -64,6 +65,7 @@ def prepare_acquisition_function( :param models: The models for each tag. :param datasets: The data from the observer (optional). + :param metadata: Any metadata to pass to the acquisition function (optional). :return: An acquisition function. """ @@ -72,6 +74,7 @@ def update_acquisition_function( function: AcquisitionFunction, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: """ Update an acquisition function. By default this generates a new acquisition function each @@ -82,6 +85,7 @@ def update_acquisition_function( :param function: The acquisition function to update. :param models: The models for each tag. :param datasets: The data from the observer (optional). + :param metadata: Any metadata to pass to the acquisition function (optional). :return: The updated acquisition function. """ return self.prepare_acquisition_function(models, datasets=datasets) @@ -110,9 +114,12 @@ def prepare_acquisition_function( self, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: return self.single_builder.prepare_acquisition_function( - models[tag], dataset=None if datasets is None else datasets[tag] + models[tag], + dataset=None if datasets is None else datasets[tag], + metadata=metadata, ) def update_acquisition_function( @@ -120,9 +127,13 @@ def update_acquisition_function( function: AcquisitionFunction, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: return self.single_builder.update_acquisition_function( - function, models[tag], dataset=None if datasets is None else datasets[tag] + function, + models[tag], + dataset=None if datasets is None else datasets[tag], + metadata=metadata, ) def __repr__(self) -> str: @@ -135,10 +146,12 @@ def prepare_acquisition_function( self, model: ProbabilisticModelType, dataset: Optional[Dataset] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: """ :param model: The model. :param dataset: The data to use to build the acquisition function (optional). + :param metadata: Any metadata to pass to the acquisition function (optional). :return: An acquisition function. """ @@ -147,14 +160,16 @@ def update_acquisition_function( function: AcquisitionFunction, model: ProbabilisticModelType, dataset: Optional[Dataset] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: """ :param function: The acquisition function to update. :param model: The model. :param dataset: The data from the observer (optional). + :param metadata: Any metadata to pass to the acquisition function (optional). :return: The updated acquisition function. """ - return self.prepare_acquisition_function(model, dataset=dataset) + return self.prepare_acquisition_function(model, dataset=dataset, metadata=metadata) class GreedyAcquisitionFunctionBuilder(Generic[ProbabilisticModelType], ABC): @@ -174,6 +189,7 @@ def prepare_acquisition_function( models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, pending_points: Optional[TensorType] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: """ Generate a new acquisition function. The first time this is called, ``pending_points`` @@ -184,6 +200,7 @@ def prepare_acquisition_function( :param datasets: The data from the observer (optional). :param pending_points: Points already chosen to be in the current batch (of shape [M,D]), where M is the number of pending points and D is the search space dimension. + :param metadata: Any metadata to pass to the acquisition function (optional). :return: An acquisition function. """ @@ -194,6 +211,7 @@ def update_acquisition_function( datasets: Optional[Mapping[Tag, Dataset]] = None, pending_points: Optional[TensorType] = None, new_optimization_step: bool = True, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: """ Update an acquisition function. By default this generates a new acquisition function each @@ -209,6 +227,7 @@ def update_acquisition_function( :param new_optimization_step: Indicates whether this call to update_acquisition_function is to start of a new optimization step, of to continue collecting batch of points for the current step. Defaults to ``True``. + :param metadata: Any metadata to pass to the acquisition function (optional). :return: The updated acquisition function. """ return self.prepare_acquisition_function( @@ -240,11 +259,13 @@ def prepare_acquisition_function( models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, pending_points: Optional[TensorType] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: return self.single_builder.prepare_acquisition_function( models[tag], dataset=None if datasets is None else datasets[tag], pending_points=pending_points, + metadata=metadata, ) def update_acquisition_function( @@ -254,6 +275,7 @@ def update_acquisition_function( datasets: Optional[Mapping[Tag, Dataset]] = None, pending_points: Optional[TensorType] = None, new_optimization_step: bool = True, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: return self.single_builder.update_acquisition_function( function, @@ -261,6 +283,7 @@ def update_acquisition_function( dataset=None if datasets is None else datasets[tag], pending_points=pending_points, new_optimization_step=new_optimization_step, + metadata=metadata, ) def __repr__(self) -> str: @@ -274,12 +297,14 @@ def prepare_acquisition_function( model: ProbabilisticModelType, dataset: Optional[Dataset] = None, pending_points: Optional[TensorType] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: """ :param model: The model. :param dataset: The data from the observer (optional). :param pending_points: Points already chosen to be in the current batch (of shape [M,D]), where M is the number of pending points and D is the search space dimension. + :param metadata: Any metadata to pass to the acquisition function (optional). :return: An acquisition function. """ @@ -290,6 +315,7 @@ def update_acquisition_function( dataset: Optional[Dataset] = None, pending_points: Optional[TensorType] = None, new_optimization_step: bool = True, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: """ :param function: The acquisition function to update. @@ -300,6 +326,7 @@ def update_acquisition_function( :param new_optimization_step: Indicates whether this call to update_acquisition_function is to start of a new optimization step, of to continue collecting batch of points for the current step. Defaults to ``True``. + :param metadata: Any metadata to pass to the acquisition function (optional). :return: The updated acquisition function. """ return self.prepare_acquisition_function( @@ -344,9 +371,12 @@ def prepare_acquisition_function( self, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: return self.single_builder.prepare_acquisition_function( - models[tag], dataset=None if datasets is None else datasets[tag] + models[tag], + dataset=None if datasets is None else datasets[tag], + metadata=metadata, ) def update_acquisition_function( @@ -354,9 +384,13 @@ def update_acquisition_function( function: AcquisitionFunction, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> AcquisitionFunction: return self.single_builder.update_acquisition_function( - function, models[tag], dataset=None if datasets is None else datasets[tag] + function, + models[tag], + dataset=None if datasets is None else datasets[tag], + metadata=metadata, ) def __repr__(self) -> str: diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 5056438e20..9b98818862 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -103,6 +103,7 @@ def acquire( search_space: SearchSpaceType, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> ResultType: """ Return a value of type `T_co`. Typically this will be a set of query points, either on its @@ -117,6 +118,7 @@ def acquire( :param search_space: The local acquisition search space for *this step*. :param models: The model for each tag. :param datasets: The known observer query points and observations for each tag (optional). + :param metadata: Any metadata to use for acquisition (optional). :return: A value of type `T_co`. """ @@ -125,6 +127,7 @@ def acquire_single( search_space: SearchSpaceType, model: ProbabilisticModelType, dataset: Optional[Dataset] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> ResultType: """ A convenience wrapper for :meth:`acquire` that uses only one model, dataset pair. @@ -133,6 +136,7 @@ def acquire_single( is defined. :param model: The model to use. :param dataset: The known observer query points and observations (optional). + :param metadata: Any metadata to use for acquisition (optional). :return: A value of type `T_co`. """ if isinstance(dataset, dict) or isinstance(model, dict): @@ -144,6 +148,7 @@ def acquire_single( search_space, {OBJECTIVE: model}, datasets=None if dataset is None else {OBJECTIVE: dataset}, + metadata=metadata, ) @@ -268,6 +273,7 @@ def acquire( search_space: SearchSpaceType, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> TensorType: """ Return the query point(s) that optimizes the acquisition function produced by ``builder`` @@ -277,18 +283,21 @@ def acquire( :param models: The model for each tag. :param datasets: The known observer query points and observations. Whether this is required depends on the acquisition function used. + :param metadata: Any metadata to pass to the acquisition function (optional). :return: The single (or batch of) points to query. """ if self._acquisition_function is None: self._acquisition_function = self._builder.prepare_acquisition_function( models, datasets=datasets, + metadata=metadata, ) else: self._acquisition_function = self._builder.update_acquisition_function( self._acquisition_function, models, datasets=datasets, + metadata=metadata, ) summary_writer = logging.get_tensorboard_writer() @@ -321,6 +330,7 @@ def acquire( datasets=datasets, pending_points=points, new_optimization_step=False, + metadata=metadata, ) with tf.name_scope(f"EGO.optimizer[{i+1}]"): chosen_point = self._optimizer(search_space, self._acquisition_function) @@ -537,6 +547,7 @@ def acquire( search_space: SearchSpaceType, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> types.State[AsynchronousRuleState | None, TensorType]: """ Constructs a function that, given ``AsynchronousRuleState``, @@ -554,6 +565,7 @@ def acquire( :param search_space: The local acquisition search space for *this step*. :param models: The model of the known data. Uses the single key `OBJECTIVE`. :param datasets: The known observer query points and observations. + :param metadata: Any metadata to pass to the acquisition function (optional). :return: A function that constructs the next acquisition state and the recommended query points from the previous acquisition state. """ @@ -570,12 +582,14 @@ def acquire( self._acquisition_function = self._builder.prepare_acquisition_function( models, datasets=datasets, + metadata=metadata, ) else: self._acquisition_function = self._builder.update_acquisition_function( self._acquisition_function, models, datasets=datasets, + metadata=metadata, ) def state_func( @@ -693,6 +707,7 @@ def acquire( search_space: SearchSpaceType, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> types.State[AsynchronousRuleState | None, TensorType]: """ Constructs a function that, given ``AsynchronousRuleState``, @@ -708,6 +723,7 @@ def acquire( :param search_space: The local acquisition search space for *this step*. :param models: The model of the known data. Uses the single key `OBJECTIVE`. :param datasets: The known observer query points and observations. + :param metadata: Any metadata to pass to the acquisition function (optional). :return: A function that constructs the next acquisition state and the recommended query points from the previous acquisition state. """ @@ -734,6 +750,7 @@ def state_func( models, datasets=datasets, pending_points=state.pending_points, + metadata=metadata, ) else: self._acquisition_function = self._builder.update_acquisition_function( @@ -741,6 +758,7 @@ def state_func( models, datasets=datasets, pending_points=state.pending_points, + metadata=metadata, ) with tf.name_scope("AsynchronousOptimization.optimizer[0]"): @@ -758,6 +776,7 @@ def state_func( datasets=datasets, pending_points=state.pending_points, new_optimization_step=False, + metadata=metadata, ) with tf.name_scope(f"AsynchronousOptimization.optimizer[{i+1}]"): new_point = self._optimizer(search_space, self._acquisition_function) @@ -804,6 +823,7 @@ def acquire( search_space: SearchSpace, models: Mapping[Tag, ProbabilisticModel], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> TensorType: """ Sample ``num_query_points`` (see :meth:`__init__`) points from the @@ -812,6 +832,7 @@ def acquire( :param search_space: The acquisition search space. :param models: Unused. :param datasets: Unused. + :param metadata: Unused. :return: The ``num_query_points`` points to query. """ samples = search_space.sample(self._num_query_points) @@ -905,6 +926,7 @@ def acquire( search_space: SearchSpace, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> TensorType: """ Sample `num_search_space_samples` (see :meth:`__init__`) points from the @@ -914,6 +936,7 @@ def acquire( :param search_space: The local acquisition search space for *this step*. :param models: The model of the known data. Uses the single key `OBJECTIVE`. :param datasets: The known observer query points and observations. + :param metadata: Unused. :return: The ``num_query_points`` points to query. :raise ValueError: If ``models`` do not contain the key `OBJECTIVE`, or it contains any other key. @@ -1019,6 +1042,7 @@ def acquire( search_space: Box, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> types.State[State | None, TensorType]: """ Construct a local search space from ``search_space`` according the trust region algorithm, @@ -1050,6 +1074,7 @@ def acquire( :param models: The model for each tag. :param datasets: The known observer query points and observations. Uses the data for key `OBJECTIVE` to calculate the new trust region. + :param metadata: Any metadata to pass to the subrule (optional). :return: A function that constructs the next acquisition state and the recommended query points from the previous acquisition state. :raise KeyError: If ``datasets`` does not contain the key `OBJECTIVE`. @@ -1095,7 +1120,9 @@ def state_func( tf.reduce_min([global_upper, xmin + eps], axis=0), ) - points = self._rule.acquire(acquisition_space, models, datasets=datasets) + points = self._rule.acquire( + acquisition_space, models, datasets=datasets, metadata=metadata + ) state_ = TrustRegion.State(acquisition_space, eps, y_min, is_global) return state_, points @@ -1231,6 +1258,7 @@ def acquire( search_space: Box, models: Mapping[Tag, TrainableSupportsGetKernel], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> types.State[State | None, TensorType]: """ Construct a local search space from ``search_space`` according the TURBO algorithm, @@ -1256,6 +1284,7 @@ def acquire( :param models: The model for each tag. :param datasets: The known observer query points and observations. Uses the data for key `OBJECTIVE` to calculate the new trust region. + :param metadata: Any metadata to pass to the subrule (optional). :return: A function that constructs the next acquisition state and the recommended query points from the previous acquisition state. :raise KeyError: If ``datasets`` does not contain the key `OBJECTIVE`. @@ -1324,7 +1353,9 @@ def state_func( local_model.optimize(local_dataset) # use local model and local dataset to choose next query point(s) - points = self._rule.acquire_single(acquisition_space, local_model, local_dataset) + points = self._rule.acquire_single( + acquisition_space, local_model, local_dataset, metadata=metadata + ) state_ = TURBO.State(acquisition_space, L, failure_counter, success_counter, y_min) return state_, points @@ -1433,6 +1464,7 @@ def acquire( search_space: SearchSpace, models: Mapping[Tag, ProbabilisticModel], datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, ) -> TensorType: """Acquire a batch of points to observe based on the batch hypervolume Sharpe ratio indicator method. @@ -1443,6 +1475,7 @@ def acquire( :param search_space: The local acquisition search space for *this step*. :param models: The model for each tag. :param datasets: The known observer query points and observations. + :param metadata: Any metadata to pass to the acquisition function (optional). :return: The batch of points to query. """ if models.keys() != {OBJECTIVE}: @@ -1457,13 +1490,14 @@ def acquire( if self._acquisition_function is None: self._acquisition_function = self._builder.prepare_acquisition_function( - models, datasets=datasets + models, datasets=datasets, metadata=metadata ) else: self._acquisition_function = self._builder.update_acquisition_function( self._acquisition_function, models, datasets=datasets, + metadata=metadata, ) # Find non-dominated points