trieste.bayesian_optimizer#

This module contains the BayesianOptimizer class, used to perform Bayesian optimization.

Module Contents#

StateType[source]#

Unbound type variable.

SearchSpaceType[source]#

Type variable bound to SearchSpace.

ProbabilisticModelType[source]#

Covariant type variable bound to ProbabilisticModel.

TrainableProbabilisticModelType[source]#

Contravariant type variable bound to TrainableProbabilisticModel.

EarlyStopCallback[source]#

Early stop callback type, generic in the model and state types.

class Record[source]#

Bases: Generic[StateType, ProbabilisticModelType]

Container to record the state of each step of the optimization process.

datasets: Mapping[trieste.types.Tag, trieste.data.Dataset][source]#

The known data from the observer.

models: Mapping[trieste.types.Tag, ProbabilisticModelType][source]#

The models over the datasets.

acquisition_state: StateType | None[source]#

The acquisition state.

property dataset: trieste.data.Dataset[source]#

The dataset when there is just one dataset.

property model: ProbabilisticModelType[source]#

The model when there is just one dataset.

save(path: pathlib.Path | str) FrozenRecord[StateType, ProbabilisticModelType][source]#

Save the record to disk. Will overwrite any existing file at the same path.

class FrozenRecord[source]#

Bases: Generic[StateType, ProbabilisticModelType]

A Record container saved on disk.

Note that records are saved via pickling and are therefore neither portable nor secure. Only open frozen records generated on the same system.

path: pathlib.Path[source]#

The path to the pickled Record.

load() Record[StateType, ProbabilisticModelType][source]#

Load the record into memory.

property datasets: Mapping[trieste.types.Tag, trieste.data.Dataset][source]#

The known data from the observer.

property models: Mapping[trieste.types.Tag, ProbabilisticModelType][source]#

The models over the datasets.

property acquisition_state: StateType | None[source]#

The acquisition state.

property dataset: trieste.data.Dataset[source]#

The dataset when there is just one dataset.

property model: ProbabilisticModelType[source]#

The model when there is just one dataset.

class OptimizationResult[source]#

Bases: Generic[StateType, ProbabilisticModelType]

The final result, and the historical data of the optimization process.

final_result: trieste.utils.Result[Record[StateType, ProbabilisticModelType]][source]#

The final result of the optimization process. This contains either a Record or an exception.

history: list[Record[StateType, ProbabilisticModelType] | FrozenRecord[StateType, ProbabilisticModelType]][source]#

The history of the Records from each step of the optimization process. These Records are created at the start of each loop, and as such will never include the final_result. The records may be either in memory or on disk.

static step_filename(step: int, num_steps: int) str[source]#

Default filename for saved optimization steps.

astuple() tuple[trieste.utils.Result[Record[StateType, ProbabilisticModelType]], list[Record[StateType, ProbabilisticModelType] | FrozenRecord[StateType, ProbabilisticModelType]]][source]#

Note: In contrast to the standard library function dataclasses.astuple(), this method does not deepcopy instance attributes.

Returns:

The final_result and history as a 2-tuple.

property is_ok: bool[source]#

True if the final result contains a Record.

property is_err: bool[source]#

True if the final result contains an exception.

try_get_final_datasets() Mapping[trieste.types.Tag, trieste.data.Dataset][source]#

Convenience method to attempt to get the final data.

Returns:

The final data, if the optimization completed successfully.

Raises:

Exception – If an exception occurred during optimization.

try_get_final_dataset() trieste.data.Dataset[source]#

Convenience method to attempt to get the final data for a single dataset run.

Returns:

The final data, if the optimization completed successfully.

Raises:
  • Exception – If an exception occurred during optimization.

  • ValueError – If the optimization was not a single dataset run.

try_get_optimal_point() tuple[trieste.types.TensorType, trieste.types.TensorType, trieste.types.TensorType][source]#

Convenience method to attempt to get the optimal point for a single dataset, single objective run.

Returns:

Tuple of the optimal query point, observation and its index.

try_get_final_models() Mapping[trieste.types.Tag, ProbabilisticModelType][source]#

Convenience method to attempt to get the final models.

Returns:

The final models, if the optimization completed successfully.

Raises:

Exception – If an exception occurred during optimization.

try_get_final_model() ProbabilisticModelType[source]#

Convenience method to attempt to get the final model for a single model run.

Returns:

The final model, if the optimization completed successfully.

Raises:
  • Exception – If an exception occurred during optimization.

  • ValueError – If the optimization was not a single model run.

property loaded_history: list[Record[StateType, ProbabilisticModelType]][source]#

The history of the optimization process loaded into memory.

save_result(path: pathlib.Path | str) None[source]#

Save the final result to disk. Will overwrite any existing file at the same path.

save(base_path: pathlib.Path | str) None[source]#

Save the optimization result to disk. Will overwrite existing files at the same path.

classmethod from_path(base_path: pathlib.Path | str) OptimizationResult[StateType, ProbabilisticModelType][source]#

Load a previously saved OptimizationResult.

class BayesianOptimizer(observer: trieste.observer.Observer, search_space: SearchSpaceType)[source]#

Bases: Generic[SearchSpaceType]

This class performs Bayesian optimization, the data-efficient optimization of an expensive black-box objective function over some search space. Since we may not have access to the objective function itself, we speak instead of an observer that observes it.

Parameters:
  • observer – The observer of the objective function.

  • search_space – The space over which to search. Must be a SearchSpace.

__repr__() str[source]#

Return repr(self).

optimize(num_steps: int, datasets: Mapping[trieste.types.Tag, trieste.data.Dataset], models: Mapping[trieste.types.Tag, trieste.models.TrainableProbabilisticModel], *, track_state: bool = True, track_path: pathlib.Path | str | None = None, fit_model: bool = True, fit_initial_model: bool = True, early_stop_callback: EarlyStopCallback[trieste.models.TrainableProbabilisticModel, object] | None = None, start_step: int = 0) OptimizationResult[None, trieste.models.TrainableProbabilisticModel][source]#
optimize(num_steps: int, datasets: Mapping[trieste.types.Tag, trieste.data.Dataset], models: Mapping[trieste.types.Tag, TrainableProbabilisticModelType], acquisition_rule: trieste.acquisition.rule.AcquisitionRule[trieste.types.TensorType, SearchSpaceType, TrainableProbabilisticModelType], *, track_state: bool = True, track_path: pathlib.Path | str | None = None, fit_model: bool = True, fit_initial_model: bool = True, early_stop_callback: EarlyStopCallback[TrainableProbabilisticModelType, object] | None = None, start_step: int = 0) OptimizationResult[object, TrainableProbabilisticModelType]
optimize(num_steps: int, datasets: Mapping[trieste.types.Tag, trieste.data.Dataset], models: Mapping[trieste.types.Tag, TrainableProbabilisticModelType], acquisition_rule: trieste.acquisition.rule.AcquisitionRule[trieste.types.TensorType, SearchSpaceType, TrainableProbabilisticModelType], *, track_state: bool = True, track_path: pathlib.Path | str | None = None, fit_model: bool = True, fit_initial_model: bool = True, early_stop_callback: EarlyStopCallback[TrainableProbabilisticModelType, object] | None = None, start_step: int = 0) OptimizationResult[object, TrainableProbabilisticModelType]
optimize(num_steps: int, datasets: Mapping[trieste.types.Tag, trieste.data.Dataset], models: Mapping[trieste.types.Tag, TrainableProbabilisticModelType], acquisition_rule: trieste.acquisition.rule.AcquisitionRule[trieste.types.State[StateType | None, trieste.types.TensorType], SearchSpaceType, TrainableProbabilisticModelType], acquisition_state: StateType | None = None, *, track_state: bool = True, track_path: pathlib.Path | str | None = None, fit_model: bool = True, fit_initial_model: bool = True, early_stop_callback: EarlyStopCallback[TrainableProbabilisticModelType, StateType] | None = None, start_step: int = 0) OptimizationResult[StateType, TrainableProbabilisticModelType]
optimize(num_steps: int, datasets: Mapping[trieste.types.Tag, trieste.data.Dataset], models: Mapping[trieste.types.Tag, TrainableProbabilisticModelType], acquisition_rule: trieste.acquisition.rule.AcquisitionRule[trieste.types.State[StateType | None, trieste.types.TensorType], SearchSpaceType, TrainableProbabilisticModelType], acquisition_state: StateType | None = None, *, track_state: bool = True, track_path: pathlib.Path | str | None = None, fit_model: bool = True, fit_initial_model: bool = True, early_stop_callback: EarlyStopCallback[TrainableProbabilisticModelType, StateType] | None = None, start_step: int = 0) OptimizationResult[StateType, TrainableProbabilisticModelType]
optimize(num_steps: int, datasets: trieste.data.Dataset, models: trieste.models.TrainableProbabilisticModel, *, track_state: bool = True, track_path: pathlib.Path | str | None = None, fit_model: bool = True, fit_initial_model: bool = True, early_stop_callback: EarlyStopCallback[trieste.models.TrainableProbabilisticModel, object] | None = None, start_step: int = 0) OptimizationResult[None, trieste.models.TrainableProbabilisticModel]
optimize(num_steps: int, datasets: trieste.data.Dataset, models: TrainableProbabilisticModelType, acquisition_rule: trieste.acquisition.rule.AcquisitionRule[trieste.types.TensorType, SearchSpaceType, TrainableProbabilisticModelType], *, track_state: bool = True, track_path: pathlib.Path | str | None = None, fit_model: bool = True, fit_initial_model: bool = True, early_stop_callback: EarlyStopCallback[TrainableProbabilisticModelType, object] | None = None, start_step: int = 0) OptimizationResult[object, TrainableProbabilisticModelType]
optimize(num_steps: int, datasets: trieste.data.Dataset, models: TrainableProbabilisticModelType, acquisition_rule: trieste.acquisition.rule.AcquisitionRule[trieste.types.TensorType, SearchSpaceType, TrainableProbabilisticModelType], *, track_state: bool = True, track_path: pathlib.Path | str | None = None, fit_model: bool = True, fit_initial_model: bool = True, early_stop_callback: EarlyStopCallback[TrainableProbabilisticModelType, object] | None = None, start_step: int = 0) OptimizationResult[object, TrainableProbabilisticModelType]
optimize(num_steps: int, datasets: trieste.data.Dataset, models: TrainableProbabilisticModelType, acquisition_rule: trieste.acquisition.rule.AcquisitionRule[trieste.types.State[StateType | None, trieste.types.TensorType], SearchSpaceType, TrainableProbabilisticModelType], acquisition_state: StateType | None = None, *, track_state: bool = True, track_path: pathlib.Path | str | None = None, fit_model: bool = True, fit_initial_model: bool = True, early_stop_callback: EarlyStopCallback[TrainableProbabilisticModelType, StateType] | None = None, start_step: int = 0) OptimizationResult[StateType, TrainableProbabilisticModelType]
optimize(num_steps: int, datasets: trieste.data.Dataset, models: TrainableProbabilisticModelType, acquisition_rule: trieste.acquisition.rule.AcquisitionRule[trieste.types.State[StateType | None, trieste.types.TensorType], SearchSpaceType, TrainableProbabilisticModelType], acquisition_state: StateType | None = None, *, track_state: bool = True, track_path: pathlib.Path | str | None = None, fit_model: bool = True, fit_initial_model: bool = True, early_stop_callback: EarlyStopCallback[TrainableProbabilisticModelType, StateType] | None = None, start_step: int = 0) OptimizationResult[StateType, TrainableProbabilisticModelType]

Attempt to find the minimizer of the observer in the search_space (both specified at __init__()). This is the central implementation of the Bayesian optimization loop.

For each step in num_steps, this method:
  • Finds the next points with which to query the observer using the acquisition_rule’s acquire() method, passing it the search_space, datasets, models, and current acquisition state.

  • Queries the observer once at those points.

  • Updates the datasets and models with the data from the observer.

If any errors are raised during the optimization loop, this method will catch and return them instead and print a message (using absl at level absl.logging.ERROR). If track_state is enabled, then in addition to the final result, the history of the optimization process will also be returned. If track_path is also set, then the history and final result will be saved to disk rather than all being kept in memory.

Type hints:
  • The acquisition_rule must use the same type of SearchSpace as specified in __init__().

  • The acquisition_state must be of the type expected by the acquisition_rule. Any acquisition state in the optimization result will also be of this type.

Parameters:
  • num_steps – The number of optimization steps to run.

  • datasets – The known observer query points and observations for each tag.

  • models – The model to use for each Dataset in datasets.

  • acquisition_rule – The acquisition rule, which defines how to search for a new point on each optimization step. Defaults to EfficientGlobalOptimization with default arguments. Note that if the default is used, this implies the tags must be OBJECTIVE, the search space can be any SearchSpace, and the acquisition state returned in the OptimizationResult will be None.

  • acquisition_state – The acquisition state to use on the first optimization step. This argument allows the caller to restore the optimization process from an existing Record.

  • track_state – If True, this method saves the optimization state at the start of each step. Models and acquisition state are copied using copy.deepcopy.

  • track_path – If set, the optimization state is saved to disk at this path, rather than being copied in memory.

  • fit_model – If False then we never fit the model during BO (e.g. if we are using a rule that doesn’t rely on the models and don’t want to waste computation).

  • fit_initial_model – If False then we assume that the initial models have already been optimized on the datasets and so do not require optimization before the first optimization step.

  • early_stop_callback – An optional callback that is evaluated with the current datasets, models and optimization state before every optimization step. If this returns True then the optimization loop is terminated early.

  • start_step – The step number to start with. This number is removed from num_steps and is useful for restarting previous computations.

Returns:

An OptimizationResult. The final_result element contains either the final optimization data, models and acquisition state, or, if an exception was raised while executing the optimization loop, it contains the exception raised. In either case, the history element is the history of the data, models and acquisition state at the start of each optimization step (up to and including any step that fails to complete). The history will never include the final optimization result.

Raises:

ValueError

If any of the following are true:

  • num_steps is negative.

  • the keys in datasets and models do not match

  • datasets or models are empty

  • the default acquisition_rule is used and the tags are not OBJECTIVE.

continue_optimization(num_steps: int, optimization_result: OptimizationResult[StateType, TrainableProbabilisticModelType], *args: Any, **kwargs: Any) OptimizationResult[StateType, TrainableProbabilisticModelType][source]#

Continue a previous optimization that either failed, was terminated early, or which you simply wish to run for more steps.

Parameters:
  • num_steps – The total number of optimization steps, including any that have already been run.

  • optimization_result – The optimization result from which to extract the datasets, models and acquisition state. If the result was successful then the final result is used; otherwise the last record in the history is used. The size of the history is used to determine how many more steps are required.

  • args – Any more positional arguments to pass on to optimize.

  • kwargs – Any more keyword arguments to pass on to optimize.

Returns:

An OptimizationResult. The history will contain both the history from optimization_result (including the final_result if that was successful) and any new records.

write_summary_init(observer: trieste.observer.Observer, search_space: trieste.space.SearchSpace, acquisition_rule: trieste.acquisition.rule.AcquisitionRule[trieste.types.TensorType | trieste.types.State[StateType | None, trieste.types.TensorType], SearchSpaceType, TrainableProbabilisticModelType], datasets: Mapping[trieste.types.Tag, trieste.data.Dataset], models: Mapping[trieste.types.Tag, trieste.models.TrainableProbabilisticModel], num_steps: int) None[source]#

Write initial BO loop TensorBoard summary.

write_summary_initial_model_fit(datasets: Mapping[trieste.types.Tag, trieste.data.Dataset], models: Mapping[trieste.types.Tag, trieste.models.ProbabilisticModel], model_fitting_timer: trieste.utils.Timer) None[source]#

Write TensorBoard summary for the model fitting to the initial data.

observation_plot_init(datasets: Mapping[trieste.types.Tag, trieste.data.Dataset]) dict[trieste.types.Tag, pandas.DataFrame][source]#

Initialise query point pairplot dataframes with initial observations. Also logs warnings if pairplot dependencies are not installed.

write_summary_observations(datasets: Mapping[trieste.types.Tag, trieste.data.Dataset], models: Mapping[trieste.types.Tag, trieste.models.ProbabilisticModel], tagged_output: Mapping[trieste.types.Tag, trieste.types.TensorType], model_fitting_timer: trieste.utils.Timer, observation_plot_dfs: MutableMapping[trieste.types.Tag, pandas.DataFrame]) None[source]#

Write TensorBoard summary for the current step observations.

write_summary_query_points(datasets: Mapping[trieste.types.Tag, trieste.data.Dataset], models: Mapping[trieste.types.Tag, trieste.models.ProbabilisticModel], search_space: trieste.space.SearchSpace, query_points: trieste.types.TensorType, query_point_generation_timer: trieste.utils.Timer, query_plot_dfs: MutableMapping[int, pandas.DataFrame]) None[source]#

Write TensorBoard summary for the current step query points.

stop_at_minimum(minimum: tensorflow.Tensor | None = None, minimizers: tensorflow.Tensor | None = None, minimum_atol: float = 0, minimum_rtol: float = 0.05, minimizers_atol: float = 0, minimizers_rtol: float = 0.05, objective_tag: trieste.types.Tag = OBJECTIVE, minimum_step_number: int | None = None) EarlyStopCallback[trieste.models.TrainableProbabilisticModel, object][source]#

Generate an early stop function that terminates a BO loop when it gets close enough to the given objective minimum and/or minimizer points.

Parameters:
  • minimum – Optional minimum to stop at, with shape [1].

  • minimizers – Optional minimizer points to stop at, with shape [N, D].

  • minimum_atol – Absolute tolerance for minimum.

  • minimum_rtol – Relative tolerance for minimum.

  • minimizers_atol – Absolute tolerance for minimizer point.

  • minimizers_rtol – Relative tolerance for minimizer point.

  • objective_tag – The tag for the objective data.

  • minimum_step_number – Minimum step number to stop at.

Returns:

An early stop function that terminates if we get close enough to both the minimum and any of the minimizer points.