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.

Subpackages#

Submodules#

Package Contents#

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(query_points: trieste.types.TensorType, additional_data: trieste.data.Dataset) → tuple[trieste.types.TensorType, trieste.types.TensorType][source]#

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(query_points: trieste.types.TensorType, additional_data: trieste.data.Dataset) → tuple[trieste.types.TensorType, trieste.types.TensorType][source]#
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(query_points: trieste.types.TensorType, additional_data: trieste.data.Dataset, num_samples: int) → trieste.types.TensorType[source]#

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(query_points: trieste.types.TensorType, additional_data: trieste.data.Dataset) → tuple[trieste.types.TensorType, trieste.types.TensorType][source]#

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(query_points: trieste.types.TensorType) → tuple[trieste.types.TensorType, trieste.types.TensorType][source]#
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(query_points: trieste.types.TensorType, num_samples: int) → trieste.types.TensorType[source]#
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.

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.

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 python/typing#213 is implemented.

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

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(query_points: trieste.types.TensorType, num_samples: int) → trieste.types.TensorType[source]#

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 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.

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(at: trieste.types.TensorType, *, jitter: float = DEFAULTS.JITTER) → trieste.types.TensorType[source]#
Parameters:

  • at – Where to sample the predictive distribution, with shape […, 1, D], for points of dimension D.

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

Returns:

The samples, of shape […, S, B, L], where S is the sample_size, B is the number of points per batch, and L is the number of latent model dimensions.

reset_sampler() → None[source]#

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

class SupportsCovarianceWithTopFidelity[source]#

Bases: ProbabilisticModel, typing_extensions.Protocol

A probabilistic model is multifidelity and has access to a method to calculate the covariance between a point and the same point at the top fidelity

abstract property num_fidelities: int#

The number of fidelities

abstract covariance_with_top_fidelity(query_points: trieste.types.TensorType) → trieste.types.TensorType[source]#

Calculate the covariance of the output at query_point and a given fidelity with the highest fidelity output at the same query_point.

Parameters:

query_points – The query points to calculate the covariance for, of shape [N, D+1], where the final column of the final dimension contains the fidelity of the query point

Returns:

The covariance with the top fidelity for the query_points, of shape [N, P]

class TrainableModelStack(model_with_event_size: tuple[ProbabilisticModelType, int], *models_with_event_sizes: tuple[ProbabilisticModelType, int])[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.

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.

update(dataset: trieste.data.Dataset) → None[source]#

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(dataset: trieste.data.Dataset) → Sequence[Any][source]#

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(dataset: trieste.data.Dataset) → None[source]#

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

Parameters:

dataset – The data with which to update the model.

abstract optimize(dataset: trieste.data.Dataset) → Any[source]#

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

Parameters:

dataset – The data with which to train the model.

Returns:

Any (optimizer-specific) optimization result.

TrajectoryFunction[source]#

Type alias for trajectory functions. These have similar behaviour to an AcquisitionFunction but have additional sampling properties and support multiple model outputs.

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, L], where L is the number of outputs of the model. Note that we require the L dimension to be present, even if there is only one output.

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__(x: trieste.types.TensorType) → trieste.types.TensorType[source]#

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() → TrajectoryFunction[source]#

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_trajectory().

Returns:

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

resample_trajectory(trajectory: TrajectoryFunction) → TrajectoryFunction[source]#

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(trajectory: TrajectoryFunction) → TrajectoryFunction[source]#

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