diff --git a/trieste/acquisition/interface.py b/trieste/acquisition/interface.py index e7c92859d4..3828eb53fd 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 @@ -67,6 +67,23 @@ def prepare_acquisition_function( :return: An acquisition function. """ + def prepare_acquisition_function_with_metadata( + self, + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, + ) -> AcquisitionFunction: + """ + Prepare an acquisition function using additional metadata. By default, this is just + dropped, but you can override this method to use the metadata during acquisition. + + :param models: The models for each tag. + :param datasets: The data from the observer (optional). + :param metadata: Metadata from the observer (optional). + :return: An acquisition function. + """ + return self.prepare_acquisition_function(models, datasets=datasets) + def update_acquisition_function( self, function: AcquisitionFunction, @@ -74,7 +91,7 @@ def update_acquisition_function( datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> AcquisitionFunction: """ - Update an acquisition function. By default this generates a new acquisition function each + Update an acquisition function. By default, this generates a new acquisition function each time. However, if the function is decorated with `@tf.function`, then you can override this method to update its variables instead and avoid retracing the acquisition function on every optimization loop. @@ -86,6 +103,25 @@ def update_acquisition_function( """ return self.prepare_acquisition_function(models, datasets=datasets) + def update_acquisition_function_with_metadata( + self, + 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 is just + dropped, but you can override this method to use the metadata during acquisition. + + :param function: The acquisition function to update. + :param models: The models for each tag. + :param datasets: The data from the observer (optional). + :param metadata: Metadata from the observer (optional). + :return: The updated acquisition function. + """ + return self.update_acquisition_function(function, models, datasets=datasets) + class SingleModelAcquisitionBuilder(Generic[ProbabilisticModelType], ABC): """ @@ -115,6 +151,18 @@ def prepare_acquisition_function( models[tag], dataset=None if datasets is None else datasets[tag] ) + def prepare_acquisition_function_with_metadata( + 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_with_metadata( + models[tag], + dataset=None if datasets is None else datasets[tag], + metadata=metadata, + ) + def update_acquisition_function( self, function: AcquisitionFunction, @@ -125,6 +173,20 @@ def update_acquisition_function( function, models[tag], dataset=None if datasets is None else datasets[tag] ) + def update_acquisition_function_with_metadata( + self, + 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_with_metadata( + function, + models[tag], + dataset=None if datasets is None else datasets[tag], + metadata=metadata, + ) + def __repr__(self) -> str: return f"{self.single_builder!r} using tag {tag!r}" @@ -142,6 +204,20 @@ def prepare_acquisition_function( :return: An acquisition function. """ + def prepare_acquisition_function_with_metadata( + 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: Metadata from the observer (optional). + :return: An acquisition function. + """ + return self.prepare_acquisition_function(model, dataset=dataset) + def update_acquisition_function( self, function: AcquisitionFunction, @@ -156,6 +232,21 @@ def update_acquisition_function( """ return self.prepare_acquisition_function(model, dataset=dataset) + def update_acquisition_function_with_metadata( + self, + 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). + :return: The updated acquisition function. + """ + return self.update_acquisition_function(function, model, dataset=dataset) + class GreedyAcquisitionFunctionBuilder(Generic[ProbabilisticModelType], ABC): """ @@ -187,6 +278,20 @@ def prepare_acquisition_function( :return: An acquisition function. """ + def prepare_acquisition_function_with_metadata( + self, + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + pending_points: Optional[TensorType] = None, + metadata: Optional[Mapping[str, Any]] = None, + ) -> AcquisitionFunction: + """ + Same as prepare_acquisition_function but accepts additional metadata argument. + """ + return self.prepare_acquisition_function( + models, datasets=datasets, pending_points=pending_points + ) + def update_acquisition_function( self, function: AcquisitionFunction, @@ -215,6 +320,26 @@ def update_acquisition_function( models, datasets=datasets, pending_points=pending_points ) + def update_acquisition_function_with_metadata( + self, + function: AcquisitionFunction, + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + pending_points: Optional[TensorType] = None, + new_optimization_step: bool = True, + metadata: Optional[Mapping[str, Any]] = None, + ) -> AcquisitionFunction: + """ + Same as update_acquisition_function but accepts additional metadata argument. + """ + return self.update_acquisition_function( + function, + models, + datasets=datasets, + pending_points=pending_points, + new_optimization_step=new_optimization_step, + ) + class SingleModelGreedyAcquisitionBuilder(Generic[ProbabilisticModelType], ABC): """ @@ -247,6 +372,20 @@ def prepare_acquisition_function( pending_points=pending_points, ) + def prepare_acquisition_function_with_metadata( + self, + 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_with_metadata( + models[tag], + dataset=None if datasets is None else datasets[tag], + pending_points=pending_points, + metadata=metadata, + ) + def update_acquisition_function( self, function: AcquisitionFunction, @@ -263,6 +402,24 @@ def update_acquisition_function( new_optimization_step=new_optimization_step, ) + def update_acquisition_function_with_metadata( + self, + function: AcquisitionFunction, + models: Mapping[Tag, ProbabilisticModelType], + 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_with_metadata( + function, + models[tag], + 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: return f"{self.single_builder!r} using tag {tag!r}" @@ -283,6 +440,20 @@ def prepare_acquisition_function( :return: An acquisition function. """ + def prepare_acquisition_function_with_metadata( + self, + model: ProbabilisticModelType, + dataset: Optional[Dataset] = None, + pending_points: Optional[TensorType] = None, + metadata: Optional[Mapping[str, Any]] = None, + ) -> AcquisitionFunction: + """ + Same as prepare_acquisition_function but accepts additional metadata argument. + """ + return self.prepare_acquisition_function( + model, dataset=dataset, pending_points=pending_points + ) + def update_acquisition_function( self, function: AcquisitionFunction, @@ -308,6 +479,26 @@ def update_acquisition_function( pending_points=pending_points, ) + def update_acquisition_function_with_metadata( + self, + function: AcquisitionFunction, + model: ProbabilisticModelType, + dataset: Optional[Dataset] = None, + pending_points: Optional[TensorType] = None, + new_optimization_step: bool = True, + metadata: Optional[Mapping[str, Any]] = None, + ) -> AcquisitionFunction: + """ + Same as prepare_acquisition_function but accepts additional metadata argument. + """ + return self.update_acquisition_function( + function, + model, + dataset=dataset, + pending_points=pending_points, + new_optimization_step=new_optimization_step, + ) + class VectorizedAcquisitionFunctionBuilder(AcquisitionFunctionBuilder[ProbabilisticModelType]): """ @@ -349,6 +540,18 @@ def prepare_acquisition_function( models[tag], dataset=None if datasets is None else datasets[tag] ) + def prepare_acquisition_function_with_metadata( + 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_with_metadata( + models[tag], + dataset=None if datasets is None else datasets[tag], + metadata=metadata, + ) + def update_acquisition_function( self, function: AcquisitionFunction, @@ -359,6 +562,20 @@ def update_acquisition_function( function, models[tag], dataset=None if datasets is None else datasets[tag] ) + def update_acquisition_function_with_metadata( + self, + 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_with_metadata( + function, + models[tag], + dataset=None if datasets is None else datasets[tag], + metadata=metadata, + ) + def __repr__(self) -> str: return f"{self.single_builder!r} using tag {tag!r}" diff --git a/trieste/acquisition/rule.py b/trieste/acquisition/rule.py index 5056438e20..16651c1fed 100644 --- a/trieste/acquisition/rule.py +++ b/trieste/acquisition/rule.py @@ -60,7 +60,6 @@ GreedyAcquisitionFunctionBuilder, SingleModelAcquisitionBuilder, SingleModelGreedyAcquisitionBuilder, - SingleModelVectorizedAcquisitionBuilder, VectorizedAcquisitionFunctionBuilder, ) from .multi_objective import Pareto @@ -105,7 +104,7 @@ def acquire( datasets: Optional[Mapping[Tag, Dataset]] = None, ) -> ResultType: """ - Return a value of type `T_co`. Typically this will be a set of query points, either on its + Return a value of type `T_co`. Typically, this will be a set of query points, either on its own as a `TensorType` (see e.g. :class:`EfficientGlobalOptimization`), or within some context (see e.g. :class:`TrustRegion`). We assume that this requires at least models, but it may sometimes also need data. @@ -120,11 +119,25 @@ def acquire( :return: A value of type `T_co`. """ + def acquire_with_metadata( + self, + search_space: SearchSpaceType, + models: Mapping[Tag, ProbabilisticModelType], + datasets: Optional[Mapping[Tag, Dataset]] = None, + metadata: Optional[Mapping[str, Any]] = None, + ) -> ResultType: + """ + Same as acquire, but accepts an additional metadata argument. By default, this is just + dropped, but you can override this method to use the metadata during acquisition. + """ + return self.acquire(search_space, models, datasets=datasets) + def acquire_single( self, 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 +146,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 additional acquisition metadata (optional). :return: A value of type `T_co`. """ if isinstance(dataset, dict) or isinstance(model, dict): @@ -140,10 +154,11 @@ def acquire_single( "AcquisitionRule.acquire_single method does not support multiple datasets " "or models: use acquire instead" ) - return self.acquire( + return self.acquire_with_metadata( search_space, {OBJECTIVE: model}, datasets=None if dataset is None else {OBJECTIVE: dataset}, + metadata=metadata, ) @@ -182,10 +197,8 @@ def __init__( builder: Optional[ AcquisitionFunctionBuilder[ProbabilisticModelType] | GreedyAcquisitionFunctionBuilder[ProbabilisticModelType] - | VectorizedAcquisitionFunctionBuilder[ProbabilisticModelType] | SingleModelAcquisitionBuilder[ProbabilisticModelType] | SingleModelGreedyAcquisitionBuilder[ProbabilisticModelType] - | SingleModelVectorizedAcquisitionBuilder[ProbabilisticModelType] ] = None, optimizer: AcquisitionOptimizer[SearchSpaceType] | None = None, num_query_points: int = 1, @@ -226,7 +239,6 @@ def __init__( ( SingleModelAcquisitionBuilder, SingleModelGreedyAcquisitionBuilder, - SingleModelVectorizedAcquisitionBuilder, ), ): builder = builder.using(OBJECTIVE) @@ -245,7 +257,6 @@ def __init__( self._builder: Union[ AcquisitionFunctionBuilder[ProbabilisticModelType], GreedyAcquisitionFunctionBuilder[ProbabilisticModelType], - VectorizedAcquisitionFunctionBuilder[ProbabilisticModelType], ] = builder self._optimizer = optimizer self._num_query_points = num_query_points @@ -268,6 +279,15 @@ def acquire( search_space: SearchSpaceType, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> TensorType: + return self.acquire_with_metadata(search_space, models, datasets=datasets) + + def acquire_with_metadata( + self, + 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 +297,16 @@ 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 additional acquisition metadata. (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, + self._acquisition_function = self._builder.prepare_acquisition_function_with_metadata( + models, datasets=datasets, metadata=metadata ) else: - self._acquisition_function = self._builder.update_acquisition_function( - self._acquisition_function, - models, - datasets=datasets, + self._acquisition_function = self._builder.update_acquisition_function_with_metadata( + self._acquisition_function, models, datasets=datasets, metadata=metadata ) summary_writer = logging.get_tensorboard_writer() @@ -315,12 +333,15 @@ def acquire( for i in range( self._num_query_points - 1 ): # greedily allocate remaining batch elements - self._acquisition_function = self._builder.update_acquisition_function( - self._acquisition_function, - models, - datasets=datasets, - pending_points=points, - new_optimization_step=False, + self._acquisition_function = ( + self._builder.update_acquisition_function_with_metadata( + self._acquisition_function, + models, + 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 +558,15 @@ def acquire( search_space: SearchSpaceType, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> types.State[AsynchronousRuleState | None, TensorType]: + return self.acquire_with_metadata(search_space, models, datasets=datasets) + + def acquire_with_metadata( + self, + 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 +584,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 additional acquisition metadata. (optional) :return: A function that constructs the next acquisition state and the recommended query points from the previous acquisition state. """ @@ -567,15 +598,12 @@ def acquire( ) if self._acquisition_function is None: - self._acquisition_function = self._builder.prepare_acquisition_function( - models, - datasets=datasets, + self._acquisition_function = self._builder.prepare_acquisition_function_with_metadata( + models, datasets=datasets, metadata=metadata ) else: - self._acquisition_function = self._builder.update_acquisition_function( - self._acquisition_function, - models, - datasets=datasets, + self._acquisition_function = self._builder.update_acquisition_function_with_metadata( + self._acquisition_function, models, datasets=datasets, metadata=metadata ) def state_func( @@ -1019,6 +1047,15 @@ def acquire( search_space: Box, models: Mapping[Tag, ProbabilisticModelType], datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> types.State[State | None, TensorType]: + return self.acquire_with_metadata(search_space, models, datasets=datasets) + + def acquire_with_metadata( + self, + 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 +1087,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 additional acquisition metadata (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 +1133,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_with_metadata( + acquisition_space, models, datasets=datasets, metadata=metadata + ) state_ = TrustRegion.State(acquisition_space, eps, y_min, is_global) return state_, points @@ -1231,6 +1271,15 @@ def acquire( search_space: Box, models: Mapping[Tag, TrainableSupportsGetKernel], datasets: Optional[Mapping[Tag, Dataset]] = None, + ) -> types.State[State | None, TensorType]: + return self.acquire_with_metadata(search_space, models, datasets=datasets) + + def acquire_with_metadata( + self, + 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 +1305,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 additional acquisition metadata (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 +1374,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 diff --git a/trieste/ask_tell_optimization.py b/trieste/ask_tell_optimization.py index 3b5c973963..87b8db53d8 100644 --- a/trieste/ask_tell_optimization.py +++ b/trieste/ask_tell_optimization.py @@ -21,7 +21,7 @@ from __future__ import annotations from copy import deepcopy -from typing import Dict, Generic, Mapping, TypeVar, cast, overload +from typing import Any, Dict, Generic, Mapping, Optional, TypeVar, cast, overload try: import pandas as pd @@ -375,10 +375,11 @@ def to_result(self, copy: bool = True) -> OptimizationResult[StateType]: record: Record[StateType] = self.to_record(copy=copy) return OptimizationResult(Ok(record), []) - def ask(self) -> TensorType: + def ask(self, metadata: Optional[Mapping[str, Any]] = None) -> TensorType: """Suggests a point (or points in batch mode) to observe by optimizing the acquisition function. If the acquisition is stateful, its state is saved. + :param metadata: Any acquisition metadata (optional). :return: A :class:`TensorType` instance representing suggested point(s). """ # This trick deserves a comment to explain what's going on @@ -389,8 +390,8 @@ def ask(self) -> TensorType: # so code below is needed to cater for both cases with Timer() as query_point_generation_timer: - points_or_stateful = self._acquisition_rule.acquire( - self._search_space, self._models, datasets=self._datasets + points_or_stateful = self._acquisition_rule.acquire_with_metadata( + self._search_space, self._models, datasets=self._datasets, metadata=metadata ) if callable(points_or_stateful):