trieste.models#

This package contains the primary interfaces for probabilistic models, ProbabilisticModel and its trainable subclass TrainableProbabilisticModel. It also contains tooling for creating TrainableProbabilisticModels from config.

Package Contents#

class ModelConfig[source]#

This class is a specification for building a TrainableProbabilisticModel. It is not meant to be used by itself, it implements methods that facilitate building a Trieste model as a dictionary of model and optimizer arguments with create_model().

model :Any#

The low-level model to pass to the TrainableProbabilisticModel wrapper registered with the model via ModelRegistry. The model has to be one of the supported models, that is, registered via ModelRegistry. We use type Any here as this can be either a model that is supported by default (for example, GPflow- or GPflux-based models) or a user-defined model that has been registered.

model_args :dict[str, Any]#

The keyword arguments to pass to the model wrapper TrainableProbabilisticModel registered with the model via ModelRegistry.

optimizer :Any#

The low-level optimizer to pass to the Optimizer wrapper registered with the model via ModelRegistry, with which to train the model (by minimizing its loss function). The model has to be one of the supported models, that is, registered via ModelRegistry. We use type Any here as this can be either an optimizer that is supported by default (for example, GPflow or TensorFlow) or a user-defined optimizer that has been registered.

optimizer_args :dict[str, Any]#

The keyword arguments to pass to the optimizer wrapper Optimizer registered with the model via ModelRegistry.

build_model(self) → trieste.models.interfaces.TrainableProbabilisticModel#

Builds a Trieste model from the model and optimizer configuration.

class ModelRegistry[source]#

This is a registry of all supported models with their corresponding wrappers, and model optimizers with their corresponding wrapppers.

A single entry per model and optimizer is maintained, if same model is registered again it will overwrite the previous entry. Registry is primarily used by ModelConfig and create_model() to facilitate building models by specifying a simple dictionary of model and optimizer arguments.

Note that you do not need to register your custom model if you will provide an instance of TrainableProbabilisticModel directly to the BayesianOptimizer. Registering is required if you intend to build your custom model as a dictionary of arguments for the wrapper and the optimizer, or as a ModelConfig.

classmethod get_model_wrapper(cls, model_type: Type[Any]) → Type[trieste.models.interfaces.TrainableProbabilisticModel]#

Get a Trieste model wrapper for a given model type.

Parameters

model_type – The model type.

Returns

The wrapper which builds a model.

classmethod get_optimizer_wrapper(cls, optimizer_type: Type[Any]) → Type[trieste.models.optimizer.Optimizer]#

Get a Trieste model optimizer wrapper for a given optimizer type.

Parameters

optimizer_type – The optimizer type.

Returns

The optimizer wrapper to be used with the optimizer type.

classmethod register_model(cls, model_type: Type[Any], wrapper_type: Type[trieste.models.interfaces.TrainableProbabilisticModel]) → None#

Register a new model type. Note that this will overwrite a registry entry if the model has already been registered.

Parameters

  • model_type – The model type.

  • wrapper_type – The model wrapper to be used with the model type.

classmethod register_optimizer(cls, optimizer_type: Type[Any], wrapper_type: Type[trieste.models.optimizer.Optimizer]) → None#

Register a new optimizer type. Note that this will overwrite a registry entry if the optimizer has already been registered.

Parameters

  • optimizer_type – The optimizer type.

  • wrapper_type – The optimier wrapper to be used with the optimizer type.

classmethod get_registered_models(cls) → Iterable[Any]#

Provides a generator with all supported model types.

classmethod get_registered_optimizers(cls) → Iterable[Any]#

Provides a generator with all supported optimizer types.

ModelSpec[source]#

Type alias for any type that can be used to fully specify a model.

create_model(config: ModelSpec) → trieste.models.interfaces.TrainableProbabilisticModel[source]#

Build a model in a flexible way by providing a dictionary of model and optimizer arguments, a ModelConfig, or a TrainableProbabilisticModel. This function is primarily used by BayesianOptimizer to build a model.

Parameters

config – A configuration for building a Trieste model.

Returns

A Trieste model built according to config.

class FastUpdateModel[source]#

Bases: ProbabilisticModel, typing_extensions.Protocol

A model with the ability to predict based on (possibly fantasized) supplementary data.

abstract conditional_predict_f(self, query_points: trieste.types.TensorType, additional_data: trieste.data.Dataset) → tuple[trieste.types.TensorType, trieste.types.TensorType]#

Return the mean and variance of the independent marginal distributions at each point in query_points, given an additional batch of (possibly fantasized) data.

Parameters

  • query_points – The points at which to make predictions, of shape [M, D].

  • additional_data – Dataset with query_points with shape […, N, D] and observations with shape […, N, L]

Returns

The mean and variance of the independent marginal distributions at each point in query_points, with shape […, L, M, M].

abstract conditional_predict_joint(self, query_points: trieste.types.TensorType, additional_data: trieste.data.Dataset) → tuple[trieste.types.TensorType, trieste.types.TensorType]#
Parameters

  • query_points – The points at which to make predictions, of shape [M, D].

  • additional_data – Dataset with query_points with shape […, N, D] and observations with shape […, N, L]

Returns

The mean and covariance of the joint marginal distribution at each batch of points in query_points, with shape […, L, M, M].

abstract conditional_predict_f_sample(self, query_points: trieste.types.TensorType, additional_data: trieste.data.Dataset, num_samples: int) → trieste.types.TensorType#

Return num_samples samples from the independent marginal distributions at query_points, given an additional batch of (possibly fantasized) data.

Parameters

  • query_points – The points at which to sample, with shape […, N, D].

  • additional_data – Dataset with query_points with shape […, N, D] and observations with shape […, N, L]

  • num_samples – The number of samples at each point.

Returns

The samples. For a predictive distribution with event shape E, this has shape […, S, N] + E, where S is the number of samples.

abstract conditional_predict_y(self, query_points: trieste.types.TensorType, additional_data: trieste.data.Dataset) → tuple[trieste.types.TensorType, trieste.types.TensorType]#

Return the mean and variance of the independent marginal distributions at each point in query_points for the observations, including noise contributions, given an additional batch of (possibly fantasized) data.

Note that this is not supported by all models.

Parameters

  • query_points – The points at which to make predictions, of shape [M, D].

  • additional_data – Dataset with query_points with shape […, N, D] and observations with shape […, N, L]

Returns

The mean and variance of the independent marginal distributions at each point in query_points.

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

Bases: ProbabilisticModel, Generic[ProbabilisticModelType]

A ModelStack is a wrapper around a number of ProbabilisticModels of type ProbabilisticModelType. It combines the outputs of each model for predictions and sampling.

Note: Only supports vector outputs (i.e. with event shape [E]). Outputs for any two models are assumed independent. Each model may itself be single- or multi-output, and any one multi-output model may have dependence between its outputs. When we speak of event size in this class, we mean the output dimension for a given ProbabilisticModel, whether that is the ModelStack itself, or one of the subsidiary ProbabilisticModels within the ModelStack. Of course, the event size for a ModelStack will be the sum of the event sizes of each subsidiary model.

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.

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

query_points – The points at which to make predictions, of shape […, D].

Returns

The predictions from all the wrapped models, concatenated along the event axis in the same order as they appear in __init__(). If the wrapped models have predictive distributions with event shapes [\(E_i\)], the mean and variance will both have shape […, \(\sum_i E_i\)].

sample(self, query_points: trieste.types.TensorType, num_samples: int) → trieste.types.TensorType#
Parameters

  • query_points – The points at which to sample, with shape […, N, D].

  • num_samples – The number of samples at each point.

Returns

The samples from all the wrapped models, concatenated along the event axis. For wrapped models with predictive distributions with event shapes [\(E_i\)], this has shape […, S, N, \(\sum_i E_i\)], where S is the number of samples.

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

query_points – The points at which to make predictions, of shape […, D].

Returns

The predictions from all the wrapped models, concatenated along the event axis in the same order as they appear in __init__(). If the wrapped models have predictive distributions with event shapes [\(E_i\)], the mean and variance will both have shape […, \(\sum_i E_i\)].

Raises

NotImplementedError – If any of the models don’t implement predict_y.

log(self) → None#

Log model-specific information at a given optimization step.

class ProbabilisticModel[source]#

Bases: typing_extensions.Protocol

A probabilistic model.

NOTE: This and its subclasses are defined as Protocols rather than ABCs in order to allow acquisition functions to depend on the intersection of different model types. As a result, it is also possible to pass models to acquisition functions that don’t explicitly inherit from this class, as long as they implement all the necessary methods. This may change in future if https://github.com/python/typing/issues/213 is implemented.

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

Return the mean and variance of the independent marginal distributions at each point in query_points.

This is essentially a convenience method for predict_joint(), where non-event dimensions of query_points are all interpreted as broadcasting dimensions instead of batch dimensions, and the covariance is squeezed to remove redundant nesting.

Parameters

query_points – The points at which to make predictions, of shape […, D].

Returns

The mean and variance of the independent marginal distributions at each point in query_points. For a predictive distribution with event shape E, the mean and variance will both have shape […] + E.

abstract sample(self, query_points: trieste.types.TensorType, num_samples: int) → trieste.types.TensorType#

Return num_samples samples from the independent marginal distributions at query_points.

Parameters

  • query_points – The points at which to sample, with shape […, N, D].

  • num_samples – The number of samples at each point.

Returns

The samples. For a predictive distribution with event shape E, this has shape […, S, N] + E, where S is the number of samples.

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

Return the mean and variance of the independent marginal distributions at each point in query_points for the observations, including noise contributions.

Note that this is not supported by all models.

Parameters

query_points – The points at which to make predictions, of shape […, D].

Returns

The mean and variance of the independent marginal distributions at each point in query_points. For a predictive distribution with event shape E, the mean and variance will both have shape […] + E.

log(self) → None#

Log model-specific information at a given optimization step.

ProbabilisticModelType[source]#

Contravariant type variable bound to ProbabilisticModel. This is used to specify classes such as samplers and acquisition function builders that take models as input parameters and might ony support models with certain features.

class ReparametrizationSampler(sample_size: int, model: ProbabilisticModelType)[source]#

Bases: abc.ABC, Generic[ProbabilisticModelType]

These samplers employ the reparameterization trick to draw samples from a ProbabilisticModel‘s predictive distribution across a discrete set of points. See [WHD18] for details.

Note that our TrainableModelStack currently assumes that all ReparametrizationSampler constructors have only these inputs and so will not work with more complicated constructors.

Parameters

  • sample_size – The desired number of samples.

  • model – The model to sample from.

Raises

ValueError (or InvalidArgumentError) – If sample_size is not positive.

abstract sample(self, at: trieste.types.TensorType, *, jitter: float = DEFAULTS.JITTER) → trieste.types.TensorType#
Parameters

  • at – Input points that define the sampler of shape [N, D].

  • jitter – The size of the jitter to use when stabilizing the Cholesky decomposition of the covariance matrix.

Returns

Samples of shape [sample_size, D].

reset_sampler(self) → None#

Reset the sampler so that new samples are drawn at the next sample() call.

class TrainableModelStack[source]#

Bases: ModelStack[TrainableProbabilisticModel], TrainableProbabilisticModel

A TrainableModelStack is a wrapper around a number of TrainableProbabilisticModels. It delegates training data to each model for updates and optimization.

TrainableProbabilisticModels within the TrainableModelStack. Of course, the event size for a TrainableModelStack will be the sum of the event sizes of each subsidiary model.

Initialize self. See help(type(self)) for accurate signature.

update(self, dataset: trieste.data.Dataset) → None#

Update all the wrapped models on their corresponding data. The data for each model is extracted by splitting the observations in dataset along the event axis according to the event sizes specified at __init__().

Parameters

dataset – The query points and observations for all the wrapped models.

optimize(self, dataset: trieste.data.Dataset) → None#

Optimize all the wrapped models on their corresponding data. The data for each model is extracted by splitting the observations in dataset along the event axis according to the event sizes specified at __init__().

Parameters

dataset – The query points and observations for all the wrapped models.

class TrainableProbabilisticModel[source]#

Bases: ProbabilisticModel, typing_extensions.Protocol

A trainable probabilistic model.

abstract update(self, dataset: trieste.data.Dataset) → None#

Update the model given the specified dataset. Does not train the model.

Parameters

dataset – The data with which to update the model.

abstract optimize(self, dataset: trieste.data.Dataset) → None#

Optimize the model objective with respect to (hyper)parameters given the specified dataset.

Parameters

dataset – The data with which to train the model.

TrajectoryFunction[source]#

Type alias for trajectory functions. These have essentially the same behavior as an AcquisitionFunction but have additional sampling properties.

An TrajectoryFunction evaluates a batch of B samples, each across different sets of N query points (of dimension D) i.e. takes input of shape [N, B, D] and returns shape [N, B].

A key property of these trajectory functions is that the same sample draw is evaluated for all queries. This property is known as consistency.

class TrajectoryFunctionClass[source]#

Bases: abc.ABC

An TrajectoryFunctionClass is a trajectory function represented using a class rather than as a standalone function. Using a class to represent a trajectory function makes it easier to update and resample without having to retrace the function.

abstract __call__(self, x: trieste.types.TensorType) → trieste.types.TensorType#

Call trajectory function.

class TrajectorySampler(model: ProbabilisticModelType)[source]#

Bases: abc.ABC, Generic[ProbabilisticModelType]

This class builds functions that approximate a trajectory sampled from an underlying ProbabilisticModel.

Unlike the ReparametrizationSampler, a TrajectorySampler provides consistent samples (i.e ensuring that the same sample draw is used for all evaluations of a particular trajectory function).

Parameters

model – The model to sample from.

abstract get_trajectory(self) → TrajectoryFunction#

Sample a batch of B trajectories. Note that the batch size B is determined by the first call of the TrajectoryFunction. To change the batch size of a TrajectoryFunction after initialization, you must recall get_trajecotry().

Returns

A trajectory function representing an approximate trajectory from the model, taking an input of shape [N, B, D] and returning shape [N, B].

resample_trajectory(self, trajectory: TrajectoryFunction) → TrajectoryFunction#

A TrajectoryFunction can often be efficiently updated in-place to provide a new sample without retracing. Note that if the underlying ProbabilisticModel has been updated, then we must call update_trajectory() to get a new sample from the new model.

Efficient implementations of a TrajectorySampler will have a custom method here to allow in-place resampling. However, the default behavior is just to make a new trajectory from scratch.

Parameters

trajectory – The trajectory function to be resampled.

Returns

The new resampled trajectory function.

update_trajectory(self, trajectory: TrajectoryFunction) → TrajectoryFunction#

Update a TrajectoryFunction to reflect an update in its underlying ProbabilisticModel and resample accordingly.

Efficient implementations will have a custom method here to allow in-place resampling and updating. However, the default behavior is just to make a new trajectory from scratch.

Parameters

trajectory – The trajectory function to be resampled.

Returns

The new trajectory function updated for a new model