Using a custom BoTorch model

Open in GitHubRun in Google Colab

Using a custom BoTorch model with Ax

In this tutorial, we illustrate how to use a custom BoTorch model within Ax's botorch_modular API. This allows us to harness the convenience of Ax for running Bayesian Optimization loops while maintaining full flexibility in modeling.

Acquisition functions and their optimizers can be swapped out in much the same fashion. See for example the tutorial for Implementing a custom acquisition function.

If you want to do something non-standard, or would like to have full insight into every aspect of the implementation, please see this tutorial for how to write your own full optimization loop in BoTorch.

# Install dependencies if we are running in colab
import sys
import plotly.io as pio
if 'google.colab' in sys.modules:
    pio.renderers.default = "colab"
    %pip install botorch ax
else:
    # Ax uses Plotly to produce interactive plots. These are great for viewing and analysis,
    # though they also lead to large file sizes, which is not ideal for files living in GH.
    # Changing the default to `png` strips the interactive components to get around this.
    pio.renderers.default = "png"

import os
from contextlib import contextmanager, nullcontext

from ax.utils.testing.mock import mock_botorch_optimize_context_manager

SMOKE_TEST = os.environ.get("SMOKE_TEST")
NUM_EVALS = 10 if SMOKE_TEST else 25

Implementing the custom model

For this tutorial, we implement a very simple GPyTorch ExactGP model that uses an RBF kernel (with ARD) and infers a homoskedastic noise level.

Model definition is straightforward. Here we implement a GPyTorch ExactGP that inherits from GPyTorchModel; together these two superclasses add all the API calls that BoTorch expects in its various modules.

Note: BoTorch allows implementing any custom model that follows the Model API. For more information, please see the Model Documentation.

from typing import Optional

from botorch.models.gpytorch import GPyTorchModel
from gpytorch.distributions import MultivariateNormal
from gpytorch.kernels import RBFKernel, ScaleKernel
from gpytorch.likelihoods import GaussianLikelihood
from gpytorch.means import ConstantMean
from gpytorch.models import ExactGP
from torch import Tensor


class SimpleCustomGP(ExactGP, GPyTorchModel):

    _num_outputs = 1  # to inform GPyTorchModel API

    def __init__(self, train_X, train_Y, train_Yvar: Optional[Tensor] = None):
        # NOTE: This ignores train_Yvar and uses inferred noise instead.
        # squeeze output dim before passing train_Y to ExactGP
        super().__init__(train_X, train_Y.squeeze(-1), GaussianLikelihood())
        self.mean_module = ConstantMean()
        self.covar_module = ScaleKernel(
            base_kernel=RBFKernel(ard_num_dims=train_X.shape[-1]),
        )
        self.to(train_X)  # make sure we're on the right device/dtype

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return MultivariateNormal(mean_x, covar_x)

Instantiate a BoTorchGenerator in Ax

A BoTorchGenerator in Ax encapsulates both the surrogate -- which Ax calls a Surrogate and BoTorch calls a Model -- and an acquisition function. Here, we will only specify the custom surrogate and let Ax choose the default acquisition function.

Note that the Model (e.g., the SimpleCustomGP) must implement construct_inputs, as this is used to construct the inputs required for instantiating a Model instance from the experiment data.

from ax.models.torch.botorch_modular.model import BoTorchGenerator
from ax.models.torch.botorch_modular.surrogate import Surrogate, SurrogateSpec
from ax.models.torch.botorch_modular.utils import ModelConfig

ax_model = BoTorchGenerator(
    surrogate=Surrogate(
        surrogate_spec=SurrogateSpec(
            model_configs=[
                ModelConfig(
                    # The model class to use
                    botorch_model_class=SimpleCustomGP,
                    # Optional, MLL class with which to optimize model parameters
                    # mll_class=ExactMarginalLogLikelihood,
                    # Optional, dictionary of keyword arguments to model constructor
                    # model_options={}
                    # Passing in `None` to disable the default set of input transforms
                    # constructed in Ax, since the model doesn't support transforms.
                    input_transform_classes=None,
                )
            ]
        )
    ),
    # Optional, acquisition function class to use - see custom acquisition tutorial
    # botorch_acqf_class=qLogExpectedImprovement,
)

Combine with a ModelBridge

Models in Ax require a ModelBridge to interface with Experiments. A ModelBridge takes the inputs supplied by the Experiment and converts them to the inputs expected by the Model. For a BoTorchGenerator, we use TorchModelBridge. The Modular BoTorch interface creates the BoTorchGenerator and the TorchModelBridge in a single step, as follows:

from ax.modelbridge.registry import Generators
model_bridge = Generators.BOTORCH_MODULAR(
    experiment=experiment,
    data=data,
    surrogate=Surrogate(SimpleCustomGP),
    # Optional, will use default if unspecified
    # botorch_acqf_class=qLogNoisyExpectedImprovement,  
)
# To generate a trial
trial = model_bridge.gen(1)

Using the custom model in Ax to optimize the Branin function

We will demonstrate this with both the Service API (simpler, easier to use) and the Developer API (advanced, more customizable).

Optimization with Ax's Service API

A detailed tutorial on the Service API can be found here.

In order to customize the way the candidates are created in the Service API, we need to construct a new GenerationStrategy and pass it into AxClient.

from ax.generation_strategy.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Generators


gs = GenerationStrategy(
    steps=[
        # Quasi-random initialization step
        GenerationStep(
            model=Generators.SOBOL,
            num_trials=5,  # How many trials should be produced from this generation step
        ),
        # Bayesian optimization step using the custom acquisition function
        GenerationStep(
            model=Generators.BOTORCH_MODULAR,
            num_trials=-1,  # No limitation on how many trials should be produced from this step
            # For `BOTORCH_MODULAR`, we pass in kwargs to specify what surrogate or acquisition function to use.
            model_kwargs={
                "surrogate_spec": SurrogateSpec(
                    model_configs=[ModelConfig(botorch_model_class=SimpleCustomGP, input_transform_classes=None)]
                ),
            },
        ),
    ]
)

Setting up the experiment

In order to use the GenerationStrategy we just created, we will pass it into the AxClient.

import torch
from ax.service.ax_client import AxClient
from ax.service.utils.instantiation import ObjectiveProperties
from botorch.test_functions import Branin


# Initialize the client - AxClient offers a convenient API to control the experiment
ax_client = AxClient(generation_strategy=gs)
# Setup the experiment
ax_client.create_experiment(
    name="branin_test_experiment",
    parameters=[
        {
            "name": "x1",
            "type": "range",
            # It is crucial to use floats for the bounds, i.e., 0.0 rather than 0.
            # Otherwise, the parameter would be inferred as an integer range.
            "bounds": [-5.0, 10.0],
        },
        {
            "name": "x2",
            "type": "range",
            "bounds": [0.0, 15.0],
        },
    ],
    objectives={
        "branin": ObjectiveProperties(minimize=True),
    },
)
# Setup a function to evaluate the trials
branin = Branin()


def evaluate(parameters):
    x = torch.tensor([[parameters.get(f"x{i+1}") for i in range(2)]])
    # The GaussianLikelihood used by our model infers an observation noise level,
    # so we pass an sem value of NaN to indicate that observation noise is unknown
    return {"branin": (branin(x).item(), float("nan"))}

Output:

[INFO 03-31 08:33:40] ax.service.ax_client: Starting optimization with verbose logging. To disable logging, set the verbose_logging argument to False. Note that float values in the logs are rounded to 6 decimal points.
[INFO 03-31 08:33:40] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x1. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 03-31 08:33:40] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x2. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 03-31 08:33:40] ax.service.utils.instantiation: Created search space: SearchSpace(parameters=[RangeParameter(name='x1', parameter_type=FLOAT, range=[-5.0, 10.0]), RangeParameter(name='x2', parameter_type=FLOAT, range=[0.0, 15.0])], parameter_constraints=[]).

Running the BO loop

The next cell sets up a decorator solely to speed up the testing of the notebook in SMOKE_TEST mode. You can safely ignore this cell and the use of the decorator throughout the tutorial.

if SMOKE_TEST:
    fast_smoke_test = mock_botorch_optimize_context_manager
else:
    fast_smoke_test = nullcontext

# Set a seed for reproducible tutorial output
torch.manual_seed(0);

with fast_smoke_test():
    for i in range(NUM_EVALS):
        parameters, trial_index = ax_client.get_next_trial()
        # Local evaluation here can be replaced with deployment to external system.
        ax_client.complete_trial(trial_index=trial_index, raw_data=evaluate(parameters))

Output:

[INFO 03-31 08:33:40] ax.service.ax_client: Generated new trial 0 with parameters {'x1': 0.62583, 'x2': 14.359564} using model Sobol.
[INFO 03-31 08:33:40] ax.service.ax_client: Completed trial 0 with data: {'branin': (104.365417, nan)}.
[INFO 03-31 08:33:40] ax.service.ax_client: Generated new trial 1 with parameters {'x1': 3.166217, 'x2': 3.867106} using model Sobol.
[INFO 03-31 08:33:40] ax.service.ax_client: Completed trial 1 with data: {'branin': (2.996862, nan)}.
[INFO 03-31 08:33:40] ax.service.ax_client: Generated new trial 2 with parameters {'x1': 9.560105, 'x2': 10.718323} using model Sobol.
[INFO 03-31 08:33:40] ax.service.ax_client: Completed trial 2 with data: {'branin': (66.530624, nan)}.
[INFO 03-31 08:33:40] ax.service.ax_client: Generated new trial 3 with parameters {'x1': -3.878664, 'x2': 0.117947} using model Sobol.
[INFO 03-31 08:33:40] ax.service.ax_client: Completed trial 3 with data: {'branin': (198.850861, nan)}.
[INFO 03-31 08:33:40] ax.service.ax_client: Generated new trial 4 with parameters {'x1': -2.362858, 'x2': 8.855021} using model Sobol.
[INFO 03-31 08:33:40] ax.service.ax_client: Completed trial 4 with data: {'branin': (5.811776, nan)}.
[INFO 03-31 08:33:41] ax.service.ax_client: Generated new trial 5 with parameters {'x1': 3.810427, 'x2': 3.331499} using model BoTorch.
[INFO 03-31 08:33:41] ax.service.ax_client: Completed trial 5 with data: {'branin': (4.778038, nan)}.
[INFO 03-31 08:33:41] ax.service.ax_client: Generated new trial 6 with parameters {'x1': 4.467948, 'x2': 6.061272} using model BoTorch.
[INFO 03-31 08:33:41] ax.service.ax_client: Completed trial 6 with data: {'branin': (28.775272, nan)}.
[INFO 03-31 08:33:41] ax.service.ax_client: Generated new trial 7 with parameters {'x1': -0.217518, 'x2': 7.060189} using model BoTorch.
[INFO 03-31 08:33:41] ax.service.ax_client: Completed trial 7 with data: {'branin': (19.876953, nan)}.
[INFO 03-31 08:33:42] ax.service.ax_client: Generated new trial 8 with parameters {'x1': -5.0, 'x2': 9.812188} using model BoTorch.
[INFO 03-31 08:33:42] ax.service.ax_client: Completed trial 8 with data: {'branin': (67.116913, nan)}.
[INFO 03-31 08:33:42] ax.service.ax_client: Generated new trial 9 with parameters {'x1': -1.258478, 'x2': 9.744528} using model BoTorch.
[INFO 03-31 08:33:42] ax.service.ax_client: Completed trial 9 with data: {'branin': (15.312768, nan)}.
[INFO 03-31 08:33:42] ax.service.ax_client: Generated new trial 10 with parameters {'x1': 2.754369, 'x2': 0.669268} using model BoTorch.
[INFO 03-31 08:33:42] ax.service.ax_client: Completed trial 10 with data: {'branin': (4.82246, nan)}.
[INFO 03-31 08:33:43] ax.service.ax_client: Generated new trial 11 with parameters {'x1': 2.957732, 'x2': 2.447017} using model BoTorch.
[INFO 03-31 08:33:43] ax.service.ax_client: Completed trial 11 with data: {'branin': (0.560318, nan)}.
[INFO 03-31 08:33:43] ax.service.ax_client: Generated new trial 12 with parameters {'x1': 10.0, 'x2': 0.0} using model BoTorch.
[INFO 03-31 08:33:43] ax.service.ax_client: Completed trial 12 with data: {'branin': (10.960894, nan)}.
[INFO 03-31 08:33:43] ax.service.ax_client: Generated new trial 13 with parameters {'x1': 7.712431, 'x2': 0.0} using model BoTorch.
[INFO 03-31 08:33:43] ax.service.ax_client: Completed trial 13 with data: {'branin': (13.341015, nan)}.
[INFO 03-31 08:33:44] ax.service.ax_client: Generated new trial 14 with parameters {'x1': 10.0, 'x2': 3.669306} using model BoTorch.
[INFO 03-31 08:33:44] ax.service.ax_client: Completed trial 14 with data: {'branin': (2.387161, nan)}.
[INFO 03-31 08:33:44] ax.service.ax_client: Generated new trial 15 with parameters {'x1': 9.019555, 'x2': 2.729907} using model BoTorch.
[INFO 03-31 08:33:44] ax.service.ax_client: Completed trial 15 with data: {'branin': (1.506727, nan)}.
[INFO 03-31 08:33:45] ax.service.ax_client: Generated new trial 16 with parameters {'x1': -3.480847, 'x2': 15.0} using model BoTorch.
[INFO 03-31 08:33:45] ax.service.ax_client: Completed trial 16 with data: {'branin': (4.535527, nan)}.
[INFO 03-31 08:33:45] ax.service.ax_client: Generated new trial 17 with parameters {'x1': -3.018028, 'x2': 12.036105} using model BoTorch.
[INFO 03-31 08:33:45] ax.service.ax_client: Completed trial 17 with data: {'branin': (0.474244, nan)}.
[INFO 03-31 08:33:45] ax.service.ax_client: Generated new trial 18 with parameters {'x1': 9.492114, 'x2': 2.37747} using model BoTorch.
[INFO 03-31 08:33:45] ax.service.ax_client: Completed trial 18 with data: {'branin': (0.443647, nan)}.
[INFO 03-31 08:33:46] ax.service.ax_client: Generated new trial 19 with parameters {'x1': 3.892624, 'x2': 0.0} using model BoTorch.
[INFO 03-31 08:33:46] ax.service.ax_client: Completed trial 19 with data: {'branin': (6.08623, nan)}.
[INFO 03-31 08:33:47] ax.service.ax_client: Generated new trial 20 with parameters {'x1': -4.329194, 'x2': 15.0} using model BoTorch.
[INFO 03-31 08:33:47] ax.service.ax_client: Completed trial 20 with data: {'branin': (6.506817, nan)}.
[INFO 03-31 08:33:47] ax.service.ax_client: Generated new trial 21 with parameters {'x1': 3.146139, 'x2': 2.24539} using model BoTorch.
[INFO 03-31 08:33:47] ax.service.ax_client: Completed trial 21 with data: {'branin': (0.398666, nan)}.
[INFO 03-31 08:33:48] ax.service.ax_client: Generated new trial 22 with parameters {'x1': -3.113413, 'x2': 11.940744} using model BoTorch.
[INFO 03-31 08:33:48] ax.service.ax_client: Completed trial 22 with data: {'branin': (0.472794, nan)}.
[INFO 03-31 08:33:49] ax.service.ax_client: Generated new trial 23 with parameters {'x1': 9.46678, 'x2': 2.82058} using model BoTorch.
[INFO 03-31 08:33:49] ax.service.ax_client: Completed trial 23 with data: {'branin': (0.502407, nan)}.
[INFO 03-31 08:33:49] ax.service.ax_client: Generated new trial 24 with parameters {'x1': -3.161055, 'x2': 12.338703} using model BoTorch.
[INFO 03-31 08:33:49] ax.service.ax_client: Completed trial 24 with data: {'branin': (0.399991, nan)}.
[INFO 11-07 08:26:05] ax.service.ax_client: Generated new trial 3 with parameters {'x1': -3.878664, 'x2': 0.117947} using model Sobol.
[INFO 11-07 08:26:05] ax.service.ax_client: Completed trial 3 with data: {'branin': (198.850861, nan)}.
[INFO 11-07 08:26:05] ax.service.ax_client: Generated new trial 4 with parameters {'x1': -2.362858, 'x2': 8.855021} using model Sobol.
[INFO 11-07 08:26:05] ax.service.ax_client: Completed trial 4 with data: {'branin': (5.811776, nan)}.
[INFO 11-07 08:26:07] ax.service.ax_client: Generated new trial 5 with parameters {'x1': 2.562432, 'x2': 4.925782} using model BoTorch.
[INFO 11-07 08:26:07] ax.service.ax_client: Completed trial 5 with data: {'branin': (6.611189, nan)}.
[INFO 11-07 08:26:07] ax.service.ax_client: Generated new trial 6 with parameters {'x1': 5.50005, 'x2': 4.949873} using model BoTorch.
[INFO 11-07 08:26:07] ax.service.ax_client: Completed trial 6 with data: {'branin': (31.211433, nan)}.
[INFO 11-07 08:26:08] ax.service.ax_client: Generated new trial 7 with parameters {'x1': -2.300231, 'x2': 4.436402} using model BoTorch.
[INFO 11-07 08:26:08] ax.service.ax_client: Completed trial 7 with data: {'branin': (38.505764, nan)}.
[INFO 11-07 08:26:08] ax.service.ax_client: Generated new trial 8 with parameters {'x1': -1.583362, 'x2': 7.318469} using model BoTorch.
[INFO 11-07 08:26:08] ax.service.ax_client: Completed trial 8 with data: {'branin': (12.206194, nan)}.
[INFO 11-07 08:26:09] ax.service.ax_client: Generated new trial 9 with parameters {'x1': -5.0, 'x2': 9.066302} using model BoTorch.
[INFO 11-07 08:26:09] ax.service.ax_client: Completed trial 9 with data: {'branin': (78.675331, nan)}.
[INFO 11-07 08:26:09] ax.service.ax_client: Generated new trial 10 with parameters {'x1': 0.787884, 'x2': 6.879815} using model BoTorch.
[INFO 11-07 08:26:09] ax.service.ax_client: Completed trial 10 with data: {'branin': (20.990005, nan)}.
[INFO 11-07 08:26:10] ax.service.ax_client: Generated new trial 11 with parameters {'x1': 1.60023, 'x2': 0.584966} using model BoTorch.
[INFO 11-07 08:26:10] ax.service.ax_client: Completed trial 11 with data: {'branin': (19.951, nan)}.
[INFO 11-07 08:26:10] ax.service.ax_client: Generated new trial 12 with parameters {'x1': 10.0, 'x2': 0.0} using model BoTorch.
[INFO 11-07 08:26:10] ax.service.ax_client: Completed trial 12 with data: {'branin': (10.960894, nan)}.
[INFO 11-07 08:26:11] ax.service.ax_client: Generated new trial 13 with parameters {'x1': 7.38266, 'x2': 0.0} using model BoTorch.
[INFO 11-07 08:26:11] ax.service.ax_client: Completed trial 13 with data: {'branin': (16.027073, nan)}.
[INFO 11-07 08:26:11] ax.service.ax_client: Generated new trial 14 with parameters {'x1': 4.173322, 'x2': 0.0} using model BoTorch.
[INFO 11-07 08:26:11] ax.service.ax_client: Completed trial 14 with data: {'branin': (7.656268, nan)}.
[INFO 11-07 08:26:12] ax.service.ax_client: Generated new trial 15 with parameters {'x1': -3.935855, 'x2': 15.0} using model BoTorch.
[INFO 11-07 08:26:12] ax.service.ax_client: Completed trial 15 with data: {'branin': (3.810518, nan)}.
[INFO 11-07 08:26:12] ax.service.ax_client: Generated new trial 16 with parameters {'x1': -3.321259, 'x2': 12.38287} using model BoTorch.
[INFO 11-07 08:26:12] ax.service.ax_client: Completed trial 16 with data: {'branin': (0.660087, nan)}.
[INFO 11-07 08:26:13] ax.service.ax_client: Generated new trial 17 with parameters {'x1': 10.0, 'x2': 3.666754} using model BoTorch.
[INFO 11-07 08:26:13] ax.service.ax_client: Completed trial 17 with data: {'branin': (2.383767, nan)}.
[INFO 11-07 08:26:14] ax.service.ax_client: Generated new trial 18 with parameters {'x1': 9.34166, 'x2': 2.5446} using model BoTorch.
[INFO 11-07 08:26:14] ax.service.ax_client: Completed trial 18 with data: {'branin': (0.450308, nan)}.
[INFO 11-07 08:26:14] ax.service.ax_client: Generated new trial 19 with parameters {'x1': 3.076019, 'x2': 2.418569} using model BoTorch.
[INFO 11-07 08:26:14] ax.service.ax_client: Completed trial 19 with data: {'branin': (0.426966, nan)}.
[INFO 11-07 08:26:15] ax.service.ax_client: Generated new trial 20 with parameters {'x1': 9.537424, 'x2': 2.493842} using model BoTorch.
[INFO 11-07 08:26:15] ax.service.ax_client: Completed trial 20 with data: {'branin': (0.4648, nan)}.
[INFO 11-07 08:26:16] ax.service.ax_client: Generated new trial 21 with parameters {'x1': -3.360749, 'x2': 15.0} using model BoTorch.
[INFO 11-07 08:26:16] ax.service.ax_client: Completed trial 21 with data: {'branin': (5.432912, nan)}.
[INFO 11-07 08:26:17] ax.service.ax_client: Generated new trial 22 with parameters {'x1': 9.516079, 'x2': 2.791557} using model BoTorch.
[INFO 11-07 08:26:17] ax.service.ax_client: Completed trial 22 with data: {'branin': (0.494746, nan)}.
[INFO 11-07 08:26:19] ax.service.ax_client: Generated new trial 23 with parameters {'x1': 3.202976, 'x2': 2.439512} using model BoTorch.
[INFO 11-07 08:26:19] ax.service.ax_client: Completed trial 23 with data: {'branin': (0.460872, nan)}.
[INFO 11-07 08:26:20] ax.service.ax_client: Generated new trial 24 with parameters {'x1': 9.625609, 'x2': 2.470825} using model BoTorch.
[INFO 11-07 08:26:20] ax.service.ax_client: Completed trial 24 with data: {'branin': (0.622846, nan)}.
[INFO 11-07 08:26:21] ax.service.ax_client: Generated new trial 25 with parameters {'x1': -3.235781, 'x2': 12.32664} using model BoTorch.
[INFO 11-07 08:26:21] ax.service.ax_client: Completed trial 25 with data: {'branin': (0.471375, nan)}.
[INFO 11-07 08:26:22] ax.service.ax_client: Generated new trial 26 with parameters {'x1': 9.466124, 'x2': 2.301119} using model BoTorch.
[INFO 11-07 08:26:22] ax.service.ax_client: Completed trial 26 with data: {'branin': (0.449765, nan)}.
[W 241107 08:26:24 optimize:576] Optimization failed in gen_candidates_scipy with the following warning(s):
  [OptimizationWarning('Optimization failed within scipy.optimize.minimize with status 2 and message ABNORMAL_TERMINATION_IN_LNSRCH.'), OptimizationWarning('Optimization failed within scipy.optimize.minimize with status 2 and message ABNORMAL_TERMINATION_IN_LNSRCH.')]
  Trying again with a new set of initial conditions.
[INFO 11-07 08:26:25] ax.service.ax_client: Generated new trial 27 with parameters {'x1': 2.97826, 'x2': 2.43746} using model BoTorch.
[INFO 11-07 08:26:25] ax.service.ax_client: Completed trial 27 with data: {'branin': (0.526684, nan)}.
[INFO 11-07 08:26:27] ax.service.ax_client: Generated new trial 28 with parameters {'x1': -3.286554, 'x2': 12.040548} using model BoTorch.
[INFO 11-07 08:26:27] ax.service.ax_client: Completed trial 28 with data: {'branin': (0.84146, nan)}.
[W 241107 08:26:28 optimize:576] Optimization failed in gen_candidates_scipy with the following warning(s):
  [OptimizationWarning('Optimization failed within scipy.optimize.minimize with status 2 and message ABNORMAL_TERMINATION_IN_LNSRCH.')]
  Trying again with a new set of initial conditions.
[INFO 11-07 08:26:29] ax.service.ax_client: Generated new trial 29 with parameters {'x1': 9.459437, 'x2': 2.554713} using model BoTorch.
[INFO 11-07 08:26:29] ax.service.ax_client: Completed trial 29 with data: {'branin': (0.406186, nan)}.

Viewing the evaluated trials

ax_client.get_trials_data_frame()

	trial_index	arm_name	trial_status	generation_node	branin	x1	x2
0	0	0_0	COMPLETED	GenerationStep_0	104.365	0.62583	14.3596
1	1	1_0	COMPLETED	GenerationStep_0	2.99686	3.16622	3.86711
2	2	2_0	COMPLETED	GenerationStep_0	66.5306	9.56011	10.7183
3	3	3_0	COMPLETED	GenerationStep_0	198.851	-3.87866	0.117947
4	4	4_0	COMPLETED	GenerationStep_0	5.81178	-2.36286	8.85502
5	5	5_0	COMPLETED	GenerationStep_1	4.77804	3.81043	3.3315
6	6	6_0	COMPLETED	GenerationStep_1	28.7753	4.46795	6.06127
7	7	7_0	COMPLETED	GenerationStep_1	19.877	-0.217518	7.06019
8	8	8_0	COMPLETED	GenerationStep_1	67.1169	-5	9.81219
9	9	9_0	COMPLETED	GenerationStep_1	15.3128	-1.25848	9.74453
10	10	10_0	COMPLETED	GenerationStep_1	4.82246	2.75437	0.669268
11	11	11_0	COMPLETED	GenerationStep_1	0.560318	2.95773	2.44702
12	12	12_0	COMPLETED	GenerationStep_1	10.9609	10	0
13	13	13_0	COMPLETED	GenerationStep_1	13.341	7.71243	0
14	14	14_0	COMPLETED	GenerationStep_1	2.38716	10	3.66931
15	15	15_0	COMPLETED	GenerationStep_1	1.50673	9.01956	2.72991
16	16	16_0	COMPLETED	GenerationStep_1	4.53553	-3.48085	15
17	17	17_0	COMPLETED	GenerationStep_1	0.474244	-3.01803	12.0361
18	18	18_0	COMPLETED	GenerationStep_1	0.443647	9.49211	2.37747
19	19	19_0	COMPLETED	GenerationStep_1	6.08623	3.89262	0
20	20	20_0	COMPLETED	GenerationStep_1	6.50682	-4.32919	15
21	21	21_0	COMPLETED	GenerationStep_1	0.398666	3.14614	2.24539
22	22	22_0	COMPLETED	GenerationStep_1	0.472794	-3.11341	11.9407
23	23	23_0	COMPLETED	GenerationStep_1	0.502407	9.46678	2.82058
24	24	24_0	COMPLETED	GenerationStep_1	0.399991	-3.16106	12.3387

parameters, values = ax_client.get_best_parameters()
print(f"Best parameters: {parameters}")
print(f"Corresponding mean: {values[0]}, covariance: {values[1]}")

Output:

Best parameters: {'x1': 3.14613865412671, 'x2': 2.2453901311712436}
Corresponding mean: {'branin': 0.36734785691998795}, covariance: {'branin': {'branin': 0.08062404785076419}}

Plotting the response surface and optimization progress

from ax.utils.notebook.plotting import render

render(ax_client.get_contour_plot())

Output:

[INFO 03-31 08:33:49] ax.service.ax_client: Retrieving contour plot with parameter 'x1' on X-axis and 'x2' on Y-axis, for metric 'branin'. Remaining parameters are affixed to the middle of their range.

best_parameters, values = ax_client.get_best_parameters()
best_parameters, values[0]

Output:

({'x1': 3.14613865412671, 'x2': 2.2453901311712436},
{'branin': 0.36734785691998795})

render(ax_client.get_optimization_trace(objective_optimum=0.397887))

Optimization with the Developer API

A detailed tutorial on the Service API can be found here.

Set up the Experiment in Ax

We need 3 inputs for an Ax Experiment:

• A search space to optimize over;
• An optimization config specifiying the objective / metrics to optimize, and optional outcome constraints;
• A runner that handles the deployment of trials. For a synthetic optimization problem, such as here, this only returns simple metadata about the trial.

import pandas as pd
import torch
from ax.core import (
    Data,
    Experiment,
    Metric,
    Objective,
    OptimizationConfig,
    ParameterType,
    RangeParameter,
    Runner,
    SearchSpace,
)
from ax.utils.common.result import Ok
from botorch.test_functions import Branin


branin_func = Branin()

# For our purposes, the metric is a wrapper that structures the function output.
class BraninMetric(Metric):
    def fetch_trial_data(self, trial):
        records = []
        for arm_name, arm in trial.arms_by_name.items():
            params = arm.parameters
            tensor_params = torch.tensor([params["x1"], params["x2"]])
            records.append(
                {
                    "arm_name": arm_name,
                    "metric_name": self.name,
                    "trial_index": trial.index,
                    "mean": branin_func(tensor_params),
                    "sem": float(
                        "nan"
                    ),  # SEM (observation noise) - NaN indicates unknown
                }
            )
        return Ok(value=Data(df=pd.DataFrame.from_records(records)))


# Search space defines the parameters, their types, and acceptable values.
search_space = SearchSpace(
    parameters=[
        RangeParameter(
            name="x1", parameter_type=ParameterType.FLOAT, lower=-5, upper=10
        ),
        RangeParameter(
            name="x2", parameter_type=ParameterType.FLOAT, lower=0, upper=15
        ),
    ]
)

optimization_config = OptimizationConfig(
    objective=Objective(
        metric=BraninMetric(name="branin_metric", lower_is_better=True),
        minimize=True,  # This is optional since we specified `lower_is_better=True`
    )
)


class MyRunner(Runner):
    def run(self, trial):
        trial_metadata = {"name": str(trial.index)}
        return trial_metadata


exp = Experiment(
    name="branin_experiment",
    search_space=search_space,
    optimization_config=optimization_config,
    runner=MyRunner(),
)

Run the BO loop

First, we use the Sobol generator to create 5 (quasi-) random initial point in the search space. Ax controls objective evaluations via Trials.

• We generate a Trial using a generator run, e.g., Sobol below. A Trial specifies relevant metadata as well as the parameters to be evaluated. At this point, the Trial is at the CANDIDATE stage.
• We run the Trial using Trial.run(). In our example, this serves to mark the Trial as RUNNING. In an advanced application, this can be used to dispatch the Trial for evaluation on a remote server.
• Once the Trial is done running, we mark it as COMPLETED. This tells the Experiment that it can fetch the Trial data.

A Trial supports evaluation of a single parameterization. For parallel evaluations, see BatchTrial.

from ax.modelbridge.registry import Generators


sobol = Generators.SOBOL(experiment=exp)

for i in range(5):
    trial = exp.new_trial(generator_run=sobol.gen(1))
    trial.run()
    trial.mark_completed()

Once the initial (quasi-) random stage is completed, we can use our SimpleCustomGP with the default acquisition function chosen by Ax to run the BO loop.

with fast_smoke_test():
    for i in range(NUM_EVALS - 5):
        model_bridge = Generators.BOTORCH_MODULAR(
            experiment=exp,
            data=exp.fetch_data(),
            surrogate_spec=SurrogateSpec(
                model_configs=[ModelConfig(SimpleCustomGP, input_transform_classes=None)]
            ),
        )
        trial = exp.new_trial(generator_run=model_bridge.gen(1))
        trial.run()
        trial.mark_completed()

View the trials attached to the Experiment.

exp.trials

Output:

{0: Trial(experiment_name='branin_experiment', index=0, status=TrialStatus.COMPLETED, arm=Arm(name='0_0', parameters={'x1': 4.109151065349579, 'x2': 5.9993501007556915})),
1: Trial(experiment_name='branin_experiment', index=1, status=TrialStatus.COMPLETED, arm=Arm(name='1_0', parameters={'x1': -4.912690562196076, 'x2': 14.837858285754919})),
2: Trial(experiment_name='branin_experiment', index=2, status=TrialStatus.COMPLETED, arm=Arm(name='2_0', parameters={'x1': -0.015767342410981655, 'x2': 0.5538329249247909})),
3: Trial(experiment_name='branin_experiment', index=3, status=TrialStatus.COMPLETED, arm=Arm(name='3_0', parameters={'x1': 7.1351942559704185, 'x2': 8.564675766974688})),
4: Trial(experiment_name='branin_experiment', index=4, status=TrialStatus.COMPLETED, arm=Arm(name='4_0', parameters={'x1': 9.74812142085284, 'x2': 1.8847176153212786})),
5: Trial(experiment_name='branin_experiment', index=5, status=TrialStatus.COMPLETED, arm=Arm(name='5_0', parameters={'x1': 9.00982331764432, 'x2': 1.974644600140861})),
6: Trial(experiment_name='branin_experiment', index=6, status=TrialStatus.COMPLETED, arm=Arm(name='6_0', parameters={'x1': 9.416891796432306, 'x2': 13.736227637346953})),
7: Trial(experiment_name='branin_experiment', index=7, status=TrialStatus.COMPLETED, arm=Arm(name='7_0', parameters={'x1': -5.0, 'x2': 9.617154053196709})),
8: Trial(experiment_name='branin_experiment', index=8, status=TrialStatus.COMPLETED, arm=Arm(name='8_0', parameters={'x1': 9.963216595175087, 'x2': 3.7275431893048223})),
9: Trial(experiment_name='branin_experiment', index=9, status=TrialStatus.COMPLETED, arm=Arm(name='9_0', parameters={'x1': 8.36482615740155, 'x2': 3.0255373408915975})),
10: Trial(experiment_name='branin_experiment', index=10, status=TrialStatus.COMPLETED, arm=Arm(name='10_0', parameters={'x1': 10.0, 'x2': 0.0})),
11: Trial(experiment_name='branin_experiment', index=11, status=TrialStatus.COMPLETED, arm=Arm(name='11_0', parameters={'x1': -1.180060444974955, 'x2': 15.0})),
12: Trial(experiment_name='branin_experiment', index=12, status=TrialStatus.COMPLETED, arm=Arm(name='12_0', parameters={'x1': 10.0, 'x2': 2.474806316978115})),
13: Trial(experiment_name='branin_experiment', index=13, status=TrialStatus.COMPLETED, arm=Arm(name='13_0', parameters={'x1': 4.0400993337053395, 'x2': 0.0})),
14: Trial(experiment_name='branin_experiment', index=14, status=TrialStatus.COMPLETED, arm=Arm(name='14_0', parameters={'x1': 2.8293397538768144, 'x2': 2.2709475656684064})),
15: Trial(experiment_name='branin_experiment', index=15, status=TrialStatus.COMPLETED, arm=Arm(name='15_0', parameters={'x1': 2.740836662866605, 'x2': 0.0})),
16: Trial(experiment_name='branin_experiment', index=16, status=TrialStatus.COMPLETED, arm=Arm(name='16_0', parameters={'x1': 1.98680512174904, 'x2': 5.161171856075616})),
17: Trial(experiment_name='branin_experiment', index=17, status=TrialStatus.COMPLETED, arm=Arm(name='17_0', parameters={'x1': 9.46424657610246, 'x2': 2.539772241686965})),
18: Trial(experiment_name='branin_experiment', index=18, status=TrialStatus.COMPLETED, arm=Arm(name='18_0', parameters={'x1': 3.0615875339559713, 'x2': 2.524267679900479})),
19: Trial(experiment_name='branin_experiment', index=19, status=TrialStatus.COMPLETED, arm=Arm(name='19_0', parameters={'x1': 2.893659054922397, 'x2': 2.7205056210251})),
20: Trial(experiment_name='branin_experiment', index=20, status=TrialStatus.COMPLETED, arm=Arm(name='20_0', parameters={'x1': 3.193841898810341, 'x2': 2.2501131193481934})),
21: Trial(experiment_name='branin_experiment', index=21, status=TrialStatus.COMPLETED, arm=Arm(name='21_0', parameters={'x1': 9.472925630193352, 'x2': 2.7321785500790843})),
22: Trial(experiment_name='branin_experiment', index=22, status=TrialStatus.COMPLETED, arm=Arm(name='22_0', parameters={'x1': 2.56991388890152, 'x2': 15.0})),
23: Trial(experiment_name='branin_experiment', index=23, status=TrialStatus.COMPLETED, arm=Arm(name='23_0', parameters={'x1': -3.7100004536976137, 'x2': 15.0})),
24: Trial(experiment_name='branin_experiment', index=24, status=TrialStatus.COMPLETED, arm=Arm(name='24_0', parameters={'x1': -2.8067243700185616, 'x2': 11.572259851903533}))}

View the evaluation data about these trials.

exp.fetch_data().df

	arm_name	metric_name	mean	sem	trial_index
0	0_0	branin_metric	23.5445	nan	0
1	1_0	branin_metric	16.3152	nan	1
2	2_0	branin_metric	49.536	nan	2
3	3_0	branin_metric	70.2538	nan	3
4	4_0	branin_metric	1.6638	nan	4
5	5_0	branin_metric	1.24256	nan	5
6	6_0	branin_metric	127.363	nan	6
7	7_0	branin_metric	70.0318	nan	7
8	8_0	branin_metric	2.33546	nan	8
9	9_0	branin_metric	6.99407	nan	9
10	10_0	branin_metric	10.9609	nan	10
11	11_0	branin_metric	61.8482	nan	11
12	12_0	branin_metric	2.22208	nan	12
13	13_0	branin_metric	6.83765	nan	13
14	14_0	branin_metric	0.929893	nan	14
15	15_0	branin_metric	7.96183	nan	15
16	16_0	branin_metric	9.40782	nan	16
17	17_0	branin_metric	0.406343	nan	17
18	18_0	branin_metric	0.463215	nan	18
19	19_0	branin_metric	0.751143	nan	19
20	20_0	branin_metric	0.411232	nan	20
21	21_0	branin_metric	0.455786	nan	21
22	22_0	branin_metric	151.668	nan	22
23	23_0	branin_metric	3.64286	nan	23
24	24_0	branin_metric	0.938913	nan	24

Plot results

We can use convenient Ax utilities for plotting the results.

import numpy as np
from ax.plot.trace import optimization_trace_single_method


# `plot_single_method` expects a 2-d array of means, because it expects to average means from multiple
# optimization runs, so we wrap out best objectives array in another array.
objective_means = np.array([[trial.objective_mean for trial in exp.trials.values()]])
best_objective_plot = optimization_trace_single_method(
    y=np.minimum.accumulate(objective_means, axis=1),
    optimum=0.397887,  # Known minimum objective for Branin function.
)
render(best_objective_plot)