Decision Maker

decijax.decision_maker

Decision makers driving the ask-tell optimisation loop.

AbstractDecisionMaker dataclass

AbstractDecisionMaker(
    search_space: AbstractSearchSpace,
    model_builders: dict[str, AbstractModelBuilder],
    datasets: dict[str, Dataset],
    key: KeyArray,
    batch_size: int,
    post_ask: list[Callable],
    post_tell: list[Callable],
)

Bases: ABC

Abstract base class to handle the core decision making loop.

The decision making loop is split into two key steps, ask and tell. The ask step is typically used to decide which point to query next. The tell step is typically used to update models and datasets with newly queried points. These steps can be combined in a 'run' loop which alternates between asking which point to query next and telling the decision maker about the newly queried point having evaluated the black-box function of interest at this point.

Attributes:

• search_space (AbstractSearchSpace) –

  Search space over which we can evaluate the function(s) of interest.

• model_builders (dict[str, AbstractModelBuilder]) –

  dictionary of model builders, which are used to (re)fit models throughout the decision making loop. These objects are typically referred to as models in the model-based decision making literature. Tags are used to distinguish between models. In a typical Bayesian optimisation setup one of the tags will be OBJECTIVE, defined in decision_making.utils.

• datasets (dict[str, Dataset]) –

  dictionary of datasets, which are augmented with observations throughout the decision making loop. These are the canonical record of observations, in their original (untransformed) space, and are used to refit the models via the model_builders. Tags are used to distinguish datasets, and correspond to tags in model_builders.

• key (KeyArray) –

  JAX PRNG key owned by the decision maker, acting as the source of randomness for the initial model fit and for each step of run.

• batch_size (int) –

  Number of points to query at each step of the decision making loop. Note that SinglePointAcquisitionFunctions are only capable of generating one point to be queried at each iteration of the decision making loop.

• post_ask (list[Callable]) –

  List of functions to be executed after each ask step.

• post_tell (list[Callable]) –

  List of functions to be executed after each tell step.

ask abstractmethod

ask(key: KeyArray) -> Float[Array, 'B D']

Get the point(s) to be queried next.

Parameters:

• key (KeyArray) –

  JAX PRNG key for controlling random state.

Returns:

• Float[Array, 'B D'] –

  Point(s) to be queried next.

tell

tell(
    observation_datasets: Mapping[str, Dataset],
    key: KeyArray,
)

Add newly observed data to datasets and refit the corresponding models.

Parameters:

• observation_datasets (Mapping[str, Dataset]) –

  dictionary of datasets containing new observations. Tags are used to distinguish datasets, and correspond to tags in model_builders and self.datasets.

• key (KeyArray) –

  JAX PRNG key for controlling random state.

run

run(
    n_steps: int,
    black_box_function_evaluator: FunctionEvaluator,
) -> Mapping[str, Dataset]

Run the decision making loop continuously for n_steps.

This is broken down into three main steps: 1. Call the ask method to get the point to be queried next. 2. Call the black_box_function_evaluator to evaluate the black box functions of interest at the point chosen to be queried. 3. Call the tell method to update the datasets and posteriors with the newly observed data.

In addition to this, after the ask step, the functions in the post_ask list are executed, taking as arguments the decision maker and the point chosen to be queried next. Similarly, after the tell step, the functions in the post_tell list are executed, taking the decision maker as the sole argument.

Parameters:

• n_steps (int) –

  Number of steps to run the decision making loop for.

• black_box_function_evaluator (FunctionEvaluator) –

  Function evaluator which evaluates the black box functions of interest at supplied points.

Returns:

• Mapping[str, Dataset] –

  Dictionary of datasets containing the observations made throughout the

• Mapping[str, Dataset] –

  decision making loop, as well as the initial data supplied when

• Mapping[str, Dataset] –

  initialising the DecisionMaker.

AcquisitionDrivenDecisionMaker dataclass

AcquisitionDrivenDecisionMaker(
    search_space: AbstractSearchSpace,
    model_builders: dict[str, AbstractModelBuilder],
    datasets: dict[str, Dataset],
    key: KeyArray,
    batch_size: int,
    post_ask: list[Callable],
    post_tell: list[Callable],
    acquisition_function_builder: AbstractAcquisitionFunctionBuilder,
    acquisition_maximizer: AbstractAcquisitionMaximizer,
)

Bases: AbstractDecisionMaker

Class which handles the core decision making loop in a model-based setup.

In this setup we use surrogate model(s) for the function(s) of interest, and define an acquisition function which characterises how useful it would be to query a given point within the search space given the data we have observed so far. This can then be used to decide which point(s) to query next.

The decision making loop is split into two key steps, ask and tell. The ask step forms a AcquisitionFunction from the current posteriors and datasets and returns the point which maximises it. It also stores the formed acquisition function under the attribute self.current_acquisition_function so that it can be called, for instance for plotting, after the ask function has been called. The tell step adds a newly queried point to the datasets and updates the posteriors.

This can be run as a typical ask-tell loop, or the run method can be used to run the decision making loop for a fixed number of steps. Moreover, the run method executes the functions in post_ask and post_tell after each ask and tell step respectively. This enables the user to add custom functionality, such as the ability to plot values of interest during the optimization process.

Attributes:

• acquisition_function_builder (AbstractAcquisitionFunctionBuilder) –

  Object which builds acquisition functions from posteriors and datasets, to decide where to query next. In a typical Bayesian optimisation setup the point chosen to be queried next is the point which maximizes the acquisition function.

• acquisition_maximizer (AbstractAcquisitionMaximizer) –

  Object which maximizes acquisition functions over the search space.

ask

ask(key: KeyArray) -> Float[Array, 'B D']

Form acquisition function(s) and return the point(s) which maximise them.

This method also stores the acquisition function(s) in self.current_acquisition_functions so that they can be accessed after the ask function has been called. This is useful for non-deterministic acquisition functions, which may differ between calls to ask if the key argument gets changed.

Note that in general SinglePointAcquisitionFunctions are only capable of generating one point to be queried at each iteration of the decision making loop (i.e. self.batch_size must be 1). However, Thompson sampling can be used in a batched setting by drawing a batch of different samples from the GP posterior. This is done by calling build_acquisition_function with different keys sequentially, and optimising each of these individual samples in sequence in order to obtain self.batch_size points to query next.

Parameters:

• key (KeyArray) –

  JAX PRNG key for controlling random state.

Returns:

• Float[Array, 'B D'] –

  Point(s) to be queried next.