trieste.acquisition.function.greedy_batch#

This module contains local penalization-based acquisition function builders.

Module Contents#

class LocalPenalization(search_space: trieste.space.SearchSpace, num_samples: int = 500, penalizer: Callable[[trieste.models.ProbabilisticModel, trieste.types.TensorType, trieste.types.TensorType, trieste.types.TensorType], trieste.acquisition.interface.PenalizationFunction | trieste.acquisition.interface.UpdatablePenalizationFunction] | None = None, base_acquisition_function_builder: trieste.acquisition.function.function.ExpectedImprovement | trieste.acquisition.function.entropy.MinValueEntropySearch[trieste.models.ProbabilisticModel] | trieste.acquisition.function.function.MakePositive[trieste.models.ProbabilisticModel] | None = None)[source]#

Bases: trieste.acquisition.interface.SingleModelGreedyAcquisitionBuilder[trieste.models.ProbabilisticModel]

Builder of the acquisition function maker for greedily collecting batches by local penalization. The resulting AcquisitionFunctionMaker takes in a set of pending points and returns a base acquisition function penalized around those points. An estimate of the objective function’s Lipschitz constant is used to control the size of penalization.

Local penalization allows us to perform batch Bayesian optimization with a standard (non-batch) acquisition function. All that we require is that the acquisition function takes strictly positive values. By iteratively building a batch of points though sequentially maximizing this acquisition function but down-weighted around locations close to the already chosen (pending) points, local penalization provides diverse batches of candidate points.

Local penalization is applied to the acquisition function multiplicatively. However, to improve numerical stability, we perform additive penalization in a log space.

The Lipschitz constant and additional penalization parameters are estimated once when first preparing the acquisition function with no pending points. These estimates are reused for all subsequent function calls.

Parameters:

search_space – The global search space over which the optimisation is defined.
num_samples – Size of the random sample over which the Lipschitz constant is estimated. We recommend scaling this with search space dimension.
penalizer – The chosen penalization method (defaults to soft penalization). This should be a function that accepts a model, pending points, lipschitz constant and eta and returns a PenalizationFunction.
base_acquisition_function_builder – Base acquisition function to be penalized (defaults to expected improvement). Local penalization only supports strictly positive acquisition functions.

Raises:

tf.errors.InvalidArgumentError – If num_samples is not positive.

prepare_acquisition_function(model: trieste.models.ProbabilisticModel, dataset: trieste.data.Dataset | None = None, pending_points: trieste.types.TensorType | None = None) → trieste.acquisition.interface.AcquisitionFunction[source]#

Parameters:

model – The model.
dataset – The data from the observer. Must be populated.
pending_points – The points we penalize with respect to.

Returns:

The (log) expected improvement penalized with respect to the pending points.

Raises:

tf.errors.InvalidArgumentError – If the dataset is empty.

update_acquisition_function(function: trieste.acquisition.interface.AcquisitionFunction, model: trieste.models.ProbabilisticModel, dataset: trieste.data.Dataset | None = None, pending_points: trieste.types.TensorType | None = None, new_optimization_step: bool = True) → trieste.acquisition.interface.AcquisitionFunction[source]#

Parameters:

function – The acquisition function to update.
model – The model.
dataset – The data from the observer. Must be populated.
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.
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.

Returns:

The updated acquisition function.

class PenalizedAcquisition(base_acquisition_function: trieste.acquisition.interface.AcquisitionFunction, penalization: trieste.acquisition.interface.PenalizationFunction)[source]#

Class representing a penalized acquisition function.

Parameters:

base_acquisition_function – Base (unpenalized) acquisition function.
penalization – Penalization function.

class local_penalizer(model: trieste.models.ProbabilisticModel, pending_points: trieste.types.TensorType, lipschitz_constant: trieste.types.TensorType, eta: trieste.types.TensorType)[source]#

Bases: trieste.acquisition.interface.UpdatablePenalizationFunction

An UpdatablePenalizationFunction builds and updates a penalization function. Defining a penalization function that can be updated avoids having to retrace on every call.

Initialize the local penalizer.

Parameters:

model – The model over the specified dataset.
pending_points – The points we penalize with respect to.
lipschitz_constant – The estimated Lipschitz constant of the objective function.
eta – The estimated global minima.

Returns:

The local penalization function. This function will raise ValueError or InvalidArgumentError if used with a batch size greater than one.

update(pending_points: trieste.types.TensorType, lipschitz_constant: trieste.types.TensorType, eta: trieste.types.TensorType) → None[source]#: Update the local penalizer with new variable values.

class soft_local_penalizer(model: trieste.models.ProbabilisticModel, pending_points: trieste.types.TensorType, lipschitz_constant: trieste.types.TensorType, eta: trieste.types.TensorType)[source]#

Bases: local_penalizer

Return the soft local penalization function used for single-objective greedy batch Bayesian optimization in [GonzalezDHL16].

Soft penalization returns the probability that a candidate point does not belong in the exclusion zones of the pending points. For model posterior mean \(\mu\), model posterior variance \(\sigma^2\), current “best” function value \(\eta\), and an estimated Lipschitz constant \(L\),the penalization from a set of pending point \(x'\) on a candidate point \(x\) is given by .. math:: phi(x, x’) = frac{1}{2}textrm{erfc}(-z) where \(z = \frac{1}{\sqrt{2\sigma^2(x')}}(L||x'-x|| + \eta - \mu(x'))\).

The penalization from a set of pending points is just product of the individual penalizations. See [GonzalezDHL16] for a full derivation.

Parameters:

model – The model over the specified dataset.
pending_points – The points we penalize with respect to.
lipschitz_constant – The estimated Lipschitz constant of the objective function.
eta – The estimated global minima.

Returns:

The local penalization function. This function will raise ValueError or InvalidArgumentError if used with a batch size greater than one.

Initialize the local penalizer.

Parameters:

model – The model over the specified dataset.
pending_points – The points we penalize with respect to.
lipschitz_constant – The estimated Lipschitz constant of the objective function.
eta – The estimated global minima.

Returns:

The local penalization function. This function will raise ValueError or InvalidArgumentError if used with a batch size greater than one.

__call__(x: trieste.types.TensorType) → trieste.types.TensorType[source]#: Call penalization function..

class hard_local_penalizer(model: trieste.models.ProbabilisticModel, pending_points: trieste.types.TensorType, lipschitz_constant: trieste.types.TensorType, eta: trieste.types.TensorType)[source]#

Bases: local_penalizer

Return the hard local penalization function used for single-objective greedy batch Bayesian optimization in [ARC+19].

Hard penalization is a stronger penalizer than soft penalization and is sometimes more effective See [ARC+19] for details. Our implementation follows theirs, with the penalization from a set of pending points being the product of the individual penalizations.

Parameters:

model – The model over the specified dataset.
pending_points – The points we penalize with respect to.
lipschitz_constant – The estimated Lipschitz constant of the objective function.
eta – The estimated global minima.

Returns:

The local penalization function. This function will raise ValueError or InvalidArgumentError if used with a batch size greater than one.

Initialize the local penalizer.

Parameters:

model – The model over the specified dataset.
pending_points – The points we penalize with respect to.
lipschitz_constant – The estimated Lipschitz constant of the objective function.
eta – The estimated global minima.

Returns:

The local penalization function. This function will raise ValueError or InvalidArgumentError if used with a batch size greater than one.

__call__(x: trieste.types.TensorType) → trieste.types.TensorType[source]#: Call penalization function..

class FantasizerModelType[source]#

Bases: trieste.models.FastUpdateModel, trieste.models.interfaces.SupportsPredictJoint, trieste.models.interfaces.SupportsPredictY, trieste.models.interfaces.SupportsGetKernel, trieste.models.interfaces.SupportsGetObservationNoise, typing_extensions.Protocol

The model requirements for the Fantasizer acquisition function.

class FantasizerModelStack(model_with_event_size: tuple[ProbabilisticModelType, int], *models_with_event_sizes: tuple[ProbabilisticModelType, int])[source]#

Bases: trieste.models.interfaces.PredictJointModelStack, trieste.models.interfaces.PredictYModelStack, trieste.models.ModelStack[FantasizerModelType]

A stack of models FantasizerModelType models. Note that this delegates predict_joint and predict_y but none of the other methods.

The order of individual models specified at __init__() determines the order of the ModelStack output dimensions.

Parameters:

model_with_event_size – The first model, and the size of its output events. Note: This is a separate parameter to models_with_event_sizes simply so that the method signature requires at least one model. It is not treated specially.
*models_with_event_sizes – The other models, and sizes of their output events.

class Fantasizer(base_acquisition_function_builder: trieste.acquisition.interface.AcquisitionFunctionBuilder[trieste.models.interfaces.SupportsPredictJoint] | trieste.acquisition.interface.SingleModelAcquisitionBuilder[trieste.models.interfaces.SupportsPredictJoint] | None = None, fantasize_method: str = 'KB')[source]#

Bases: trieste.acquisition.interface.GreedyAcquisitionFunctionBuilder[FantasizerModelOrStack]

Builder of the acquisition function maker for greedily collecting batches. Fantasizer allows us to perform batch Bayesian optimization with any standard (non-batch) acquisition function.

Here, every time a query point is chosen by maximising an acquisition function, its corresponding observation is “fantasized”, and the models are conditioned further on this new artificial data.

This implies that the models need to predict what their updated predictions would be given new data, see FastUpdateModel. These equations are for instance in closed form for the GPR model, see [CGE14] (eqs. 8-10) for details.

There are several ways to “fantasize” data: the “kriging believer” heuristic (KB, see [GLRC10b]) uses the mean of the model as observations. “sample” uses samples from the model.

Parameters:

base_acquisition_function_builder – The acquisition function builder to use. Defaults to ExpectedImprovement.
fantasize_method – The following options are available: “KB” and “sample”. See class docs for more details.

Raises:

tf.errors.InvalidArgumentError – If fantasize_method is not “KB” or “sample”.

prepare_acquisition_function(models: Mapping[trieste.types.Tag, FantasizerModelOrStack], datasets: Mapping[trieste.types.Tag, trieste.data.Dataset] | None = None, pending_points: trieste.types.TensorType | None = None) → trieste.acquisition.interface.AcquisitionFunction[source]#

Parameters:

models – The models over each tag.
datasets – The data from the observer (optional).
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.

Returns:

An acquisition function.

update_acquisition_function(function: trieste.acquisition.interface.AcquisitionFunction, models: Mapping[trieste.types.Tag, FantasizerModelOrStack], datasets: Mapping[trieste.types.Tag, trieste.data.Dataset] | None = None, pending_points: trieste.types.TensorType | None = None, new_optimization_step: bool = True) → trieste.acquisition.interface.AcquisitionFunction[source]#

Parameters:

function – The acquisition function to update.
models – The models over each tag.
datasets – The data from the observer (optional).
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.
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.

Returns:

The updated acquisition function.

_generate_fantasized_data(fantasize_method: str, model: FantasizerModelOrStack, pending_points: trieste.types.TensorType) → trieste.data.Dataset[source]#

Generates “fantasized” data at pending_points depending on the chosen heuristic: - KB (kriging believer) uses the mean prediction of the models - sample uses samples from the GP posterior.

Parameters:

fantasize_method – the following options are available: “KB” and “sample”.
model – a model with predict method
dataset – past data
pending_points – points at which to fantasize data

Returns:

a fantasized dataset

class _fantasized_model(model: FantasizerModelType, fantasized_data: trieste.data.Dataset)[source]#

Bases: trieste.models.interfaces.SupportsPredictJoint, trieste.models.interfaces.SupportsGetKernel, trieste.models.interfaces.SupportsGetObservationNoise, trieste.models.interfaces.SupportsPredictY

Creates a new model from an existing one and additional data. This new model posterior is conditioned on both current model data and the additional one.

Parameters:

model – a model, must be of class FastUpdateModel
fantasized_data – additional dataset to condition on

Raises:

NotImplementedError – If model is not of class FastUpdateModel.

update_fantasized_data(fantasized_data: trieste.data.Dataset) → None[source]#

Parameters:: fantasized_data – new additional dataset to condition on

predict(query_points: trieste.types.TensorType) → tuple[trieste.types.TensorType, trieste.types.TensorType][source]#

This function wraps conditional_predict_f. It cannot directly call conditional_predict_f, since it does not accept query_points with rank > 2. We use map_fn to allow leading dimensions for query_points.

Parameters:: query_points – shape […*, N, d]
Returns:: mean, shape […*, …, N, L] and cov, shape […*, …, N, L], where … are the leading dimensions of fantasized_data

predict_joint(query_points: trieste.types.TensorType) → tuple[trieste.types.TensorType, trieste.types.TensorType][source]#

This function wraps conditional_predict_joint. It cannot directly call conditional_predict_joint, since it does not accept query_points with rank > 2. We use map_fn to allow leading dimensions for query_points.

Parameters:: query_points – shape […*, N, D]
Returns:: mean, shape […*, …, N, L] and cov, shape […*, …, L, N, N], where … are the leading dimensions of fantasized_data

sample(query_points: trieste.types.TensorType, num_samples: int) → trieste.types.TensorType[source]#

This function wraps conditional_predict_f_sample. It cannot directly call conditional_predict_joint, since it does not accept query_points with rank > 2. We use map_fn to allow leading dimensions for query_points.

Parameters:

query_points – shape […*, N, D]
num_samples – number of samples.

Returns:

samples of shape […*, …, S, N, L], where … are the leading dimensions of fantasized_data

predict_y(query_points: trieste.types.TensorType) → tuple[trieste.types.TensorType, trieste.types.TensorType][source]#

This function wraps conditional_predict_y. It cannot directly call conditional_predict_joint, since it does not accept query_points with rank > 2. We use tf.map_fn to allow leading dimensions for query_points.

Parameters:: query_points – shape […*, N, D]
Returns:: mean, shape […*, …, N, L] and var, shape […*, …, N, L], where … are the leading dimensions of fantasized_data

get_observation_noise() → trieste.types.TensorType[source]#

Return the variance of observation noise.

Returns:: The observation noise.

get_kernel() → gpflow.kernels.Kernel[source]#: Return the kernel of the model. :return: The kernel.

log(dataset: trieste.data.Dataset | None = None) → None[source]#

Log model-specific information at a given optimization step.

Parameters:: dataset – Optional data that can be used to log additional data-based model summaries.

_broadcast_predict(query_points: trieste.types.TensorType, fun: Callable[[trieste.types.TensorType], tuple[trieste.types.TensorType, trieste.types.TensorType]]) → tuple[trieste.types.TensorType, trieste.types.TensorType][source]#

Utility function that allows leading dimensions for query_points when fun only accepts rank 2 tensors. It works by flattening query_points into a rank 3 tensor, evaluate fun(query_points) through tf.map_fn, then restoring the leading dimensions.

Parameters:

query_points – shape […*, N, D]
fun – callable that returns two tensors (e.g. a predict function)

Returns:

two tensors (e.g. mean and variance) with shape […*, …]

_get_leading_dim_and_flatten(query_points: trieste.types.TensorType) → tuple[trieste.types.TensorType, trieste.types.TensorType][source]#

Parameters:: query_points – shape […*, N, D]
Returns:: leading_dim = ….*, query_points_flatten, shape [B, N, D]

_restore_leading_dim(x: trieste.types.TensorType, leading_dim: trieste.types.TensorType) → trieste.types.TensorType[source]#

“Un-flatten” the first dimension of x to leading_dim

Parameters:

x – shape [B, …]
leading_dim – […*]

Returns:

shape […*, …]