Developer API Example on Hartmann6
The Developer API is suitable when the user wants maximal customization of the
optimization loop. This tutorial demonstrates optimization of a Hartmann6 function using
the Experiment construct. In this example, trials will be evaluated synchronously.
import sys
in_colab = 'google.colab' in sys.modules
if in_colab:
%pip install ax-platform
from ax import (
ChoiceParameter,
ComparisonOp,
Experiment,
FixedParameter,
Metric,
Objective,
OptimizationConfig,
OrderConstraint,
OutcomeConstraint,
ParameterType,
RangeParameter,
SearchSpace,
SumConstraint,
)
from ax.modelbridge.registry import Models
from ax.utils.notebook.plotting import init_notebook_plotting, render
import plotly.io as pio
init_notebook_plotting()
if in_colab:
pio.renderers.default = "colab"
1. Create Search Space
First, we define a search space, which defines the type and allowed range for the parameters.
from ax.metrics.l2norm import L2NormMetric
from ax.metrics.hartmann6 import Hartmann6Metric
hartmann_search_space = SearchSpace(
parameters=[
RangeParameter(
name=f"x{i}", parameter_type=ParameterType.FLOAT, lower=0.0, upper=1.0
)
for i in range(6)
]
)
Note that there are two other parameter classes, FixedParameter and ChoiceParameter. Although we won't use these in this example, you can create them as follows.
choice_param = ChoiceParameter(
name="choice", values=["foo", "bar"], parameter_type=ParameterType.STRING
)
fixed_param = FixedParameter(
name="fixed", value=[True], parameter_type=ParameterType.BOOL
)
Sum constraints enforce that the sum of a set of parameters is greater or less than some bound, and order constraints enforce that one parameter is smaller than the other. We won't use these either, but see two examples below.
sum_constraint = SumConstraint(
parameters=[
hartmann_search_space.parameters["x0"],
hartmann_search_space.parameters["x1"],
],
is_upper_bound=True,
bound=5.0,
)
order_constraint = OrderConstraint(
lower_parameter=hartmann_search_space.parameters["x0"],
upper_parameter=hartmann_search_space.parameters["x1"],
)
2. Create Optimization Config
Second, we define the optimization_config with an objective and
outcome_constraints.
When doing the optimization, we will find points that minimize the objective while
obeying the constraints (which in this case means l2norm < 1.25).
Note: we are using Hartmann6Metric and L2NormMetric here, which have built in
evaluation functions for testing. For creating your own cutom metrics, see
8. Defining custom metrics.
param_names = [f"x{i}" for i in range(6)]
optimization_config = OptimizationConfig(
objective=Objective(
metric=Hartmann6Metric(name="hartmann6", param_names=param_names),
minimize=True,
),
outcome_constraints=[
OutcomeConstraint(
metric=L2NormMetric(name="l2norm", param_names=param_names, noise_sd=0.2),
op=ComparisonOp.LEQ,
bound=1.25,
relative=False,
)
],
)
3. Define a Runner
Before an experiment can collect data, it must have a Runner attached. A runner handles the deployment of trials. A trial must be "run" before it can be evaluated.
Here, we have a dummy runner that does nothing. In practice, a runner might be in charge of pushing an experiment to production.
The only method that needs to be defined for runner subclasses is run, which performs
any necessary deployment logic, and returns a dictionary of resulting metadata. This
metadata can later be accessed through the trial's run_metadata property.
from ax import Runner
class MyRunner(Runner):
def run(self, trial):
trial_metadata = {"name": str(trial.index)}
return trial_metadata
4. Create Experiment
Next, we make an Experiment with our search space, runner, and optimization config.
exp = Experiment(
name="test_hartmann",
search_space=hartmann_search_space,
optimization_config=optimization_config,
runner=MyRunner(),
)
5. Perform Optimization
Run the optimization using the settings defined on the experiment. We will create 5 random sobol points for exploration followed by 15 points generated using the GPEI optimizer.
Instead of a member of the Models enum to produce generator runs, users can leverage a
GenerationStrategy. See the
Generation Strategy Tutorial for
more info.
from ax.modelbridge.registry import Models
NUM_SOBOL_TRIALS = 5
NUM_BOTORCH_TRIALS = 15
print(f"Running Sobol initialization trials...")
sobol = Models.SOBOL(search_space=exp.search_space)
for i in range(NUM_SOBOL_TRIALS):
# Produce a GeneratorRun from the model, which contains proposed arm(s) and other metadata
generator_run = sobol.gen(n=1)
# Add generator run to a trial to make it part of the experiment and evaluate arm(s) in it
trial = exp.new_trial(generator_run=generator_run)
# Start trial run to evaluate arm(s) in the trial
trial.run()
# Mark trial as completed to record when a trial run is completed
# and enable fetching of data for metrics on the experiment
# (by default, trials must be completed before metrics can fetch their data,
# unless a metric is explicitly configured otherwise)
trial.mark_completed()
for i in range(NUM_BOTORCH_TRIALS):
print(
f"Running BO trial {i + NUM_SOBOL_TRIALS + 1}/{NUM_SOBOL_TRIALS + NUM_BOTORCH_TRIALS}..."
)
# Reinitialize GP+EI model at each step with updated data.
gpei = Models.BOTORCH_MODULAR(experiment=exp, data=exp.fetch_data())
generator_run = gpei.gen(n=1)
trial = exp.new_trial(generator_run=generator_run)
trial.run()
trial.mark_completed()
print("Done!")
6. Inspect trials' data
Now we can inspect the Experiment's data by calling fetch_data(), which retrieves
evaluation data for all trials of the experiment.
To fetch trial data, we need to run it and mark it completed. For most metrics in Ax,
data is only available once the status of the trial is COMPLETED, since in real-worlds
scenarios, metrics can typically only be fetched after the trial finished running.
NOTE: Metrics classes may implement the is_available_while_running method. When this
method returns True, data is available when trials are either RUNNING or
COMPLETED. This can be used to obtain intermediate results from A/B test trials and
other online experiments, or when metric values are available immediately, like in the
case of synthetic problem metrics.
We can also use the fetch_trials_data function to get evaluation data for a specific
trials in the experiment, like so:
trial_data = exp.fetch_trials_data([NUM_SOBOL_TRIALS + NUM_BOTORCH_TRIALS - 1])
trial_data.df
| arm_name | metric_name | mean | sem | trial_index | n | frac_nonnull | |
|---|---|---|---|---|---|---|---|
| 0 | 19_0 | hartmann6 | -1.19736 | 0 | 19 | 10000 | -1.19736 |
| 1 | 19_0 | l2norm | 1.24024 | 0.2 | 19 | 10000 | 1.24024 |
The below call to exp.fetch_data() also attaches data to the last trial, which because
of the way we looped through Botorch trials in
5. Perform Optimization,
would otherwise not have data attached. This is necessary to get objective_means in
7. Plot results.
exp.fetch_data().df
| arm_name | metric_name | mean | sem | trial_index | n | frac_nonnull | |
|---|---|---|---|---|---|---|---|
| 0 | 0_0 | hartmann6 | -0.108356 | 0 | 0 | 10000 | -0.108356 |
| 1 | 0_0 | l2norm | 1.61962 | 0.2 | 0 | 10000 | 1.61962 |
| 2 | 1_0 | hartmann6 | -0.191654 | 0 | 1 | 10000 | -0.191654 |
| 3 | 1_0 | l2norm | 1.03261 | 0.2 | 1 | 10000 | 1.03261 |
| 4 | 2_0 | hartmann6 | -0.412359 | 0 | 2 | 10000 | -0.412359 |
| 5 | 2_0 | l2norm | 1.30781 | 0.2 | 2 | 10000 | 1.30781 |
| 6 | 3_0 | hartmann6 | -0.01257 | 0 | 3 | 10000 | -0.01257 |
| 7 | 3_0 | l2norm | 1.69188 | 0.2 | 3 | 10000 | 1.69188 |
| 8 | 4_0 | hartmann6 | -0.000315 | 0 | 4 | 10000 | -0.000315 |
| 9 | 4_0 | l2norm | 1.31454 | 0.2 | 4 | 10000 | 1.31454 |
| 10 | 5_0 | hartmann6 | -0.338911 | 0 | 5 | 10000 | -0.338911 |
| 11 | 5_0 | l2norm | 1.31222 | 0.2 | 5 | 10000 | 1.31222 |
| 12 | 6_0 | hartmann6 | -0.118221 | 0 | 6 | 10000 | -0.118221 |
| 13 | 6_0 | l2norm | 1.21758 | 0.2 | 6 | 10000 | 1.21758 |
| 14 | 7_0 | hartmann6 | -0.083487 | 0 | 7 | 10000 | -0.083487 |
| 15 | 7_0 | l2norm | 1.43442 | 0.2 | 7 | 10000 | 1.43442 |
| 16 | 8_0 | hartmann6 | -0.071166 | 0 | 8 | 10000 | -0.071166 |
| 17 | 8_0 | l2norm | 0.777986 | 0.2 | 8 | 10000 | 0.777986 |
| 18 | 9_0 | hartmann6 | -0.022054 | 0 | 9 | 10000 | -0.022054 |
| 19 | 9_0 | l2norm | 1.52327 | 0.2 | 9 | 10000 | 1.52327 |
| 20 | 10_0 | hartmann6 | -0.575221 | 0 | 10 | 10000 | -0.575221 |
| 21 | 10_0 | l2norm | 1.41935 | 0.2 | 10 | 10000 | 1.41935 |
| 22 | 11_0 | hartmann6 | -0.074926 | 0 | 11 | 10000 | -0.074926 |
| 23 | 11_0 | l2norm | 1.17233 | 0.2 | 11 | 10000 | 1.17233 |
| 24 | 12_0 | hartmann6 | -0.846021 | 0 | 12 | 10000 | -0.846021 |
| 25 | 12_0 | l2norm | 1.38625 | 0.2 | 12 | 10000 | 1.38625 |
| 26 | 13_0 | hartmann6 | -0.528232 | 0 | 13 | 10000 | -0.528232 |
| 27 | 13_0 | l2norm | 1.22178 | 0.2 | 13 | 10000 | 1.22178 |
| 28 | 14_0 | hartmann6 | -0.075541 | 0 | 14 | 10000 | -0.075541 |
| 29 | 14_0 | l2norm | 1.13524 | 0.2 | 14 | 10000 | 1.13524 |
| 30 | 15_0 | hartmann6 | -0.145934 | 0 | 15 | 10000 | -0.145934 |
| 31 | 15_0 | l2norm | 1.07899 | 0.2 | 15 | 10000 | 1.07899 |
| 32 | 16_0 | hartmann6 | -0.100319 | 0 | 16 | 10000 | -0.100319 |
| 33 | 16_0 | l2norm | 1.43435 | 0.2 | 16 | 10000 | 1.43435 |
| 34 | 17_0 | hartmann6 | -1.31714 | 0 | 17 | 10000 | -1.31714 |
| 35 | 17_0 | l2norm | 1.1739 | 0.2 | 17 | 10000 | 1.1739 |
| 36 | 18_0 | hartmann6 | -0.607066 | 0 | 18 | 10000 | -0.607066 |
| 37 | 18_0 | l2norm | 0.736811 | 0.2 | 18 | 10000 | 0.736811 |
| 38 | 19_0 | hartmann6 | -1.19736 | 0 | 19 | 10000 | -1.19736 |
| 39 | 19_0 | l2norm | 1.33775 | 0.2 | 19 | 10000 | 1.33775 |
7. Plot results
Now we can plot the results of our optimization:
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=-3.32237, # Known minimum objective for Hartmann6 function.
)
render(best_objective_plot)
8. Defining custom metrics
In order to perform an optimization, we also need to define an optimization config for the experiment. An optimization config is composed of an objective metric to be minimized or maximized in the experiment, and optionally a set of outcome constraints that place restrictions on how other metrics can be moved by the experiment.
In order to define an objective or outcome constraint, we first need to subclass Metric. Metrics are used to evaluate trials, which are individual steps of the experiment sequence. Each trial contains one or more arms for which we will collect data at the same time.
Our custom metric(s) will determine how, given a trial, to compute the mean and SEM of each of the trial's arms.
The only method that needs to be defined for most metric subclasses is
fetch_trial_data, which defines how a single trial is evaluated, and returns a pandas
dataframe.
The is_available_while_running method is optional and returns a boolean, specifying
whether the trial data can be fetched before the trial is complete. See
6. Inspect trials' data
for more details.
from ax import Data
import pandas as pd
class BoothMetric(Metric):
def fetch_trial_data(self, trial):
records = []
for arm_name, arm in trial.arms_by_name.items():
params = arm.parameters
records.append(
{
"arm_name": arm_name,
"metric_name": self.name,
"trial_index": trial.index,
# in practice, the mean and sem will be looked up based on trial metadata
# but for this tutorial we will calculate them
"mean": (params["x1"] + 2 * params["x2"] - 7) ** 2
+ (2 * params["x1"] + params["x2"] - 5) ** 2,
"sem": 0.0,
}
)
return Data(df=pd.DataFrame.from_records(records))
def is_available_while_running(self) -> bool:
return True
9. Save to JSON or SQL
At any point, we can also save our experiment to a JSON file. To ensure that our custom metrics and runner are saved properly, we first need to register them.
from ax.storage.registry_bundle import RegistryBundle
bundle = RegistryBundle(
metric_clss={BoothMetric: None, L2NormMetric: None, Hartmann6Metric: None},
runner_clss={MyRunner: None},
)
from ax.storage.json_store.load import load_experiment
from ax.storage.json_store.save import save_experiment
save_experiment(exp, "experiment.json", encoder_registry=bundle.encoder_registry)
loaded_experiment = load_experiment(
"experiment.json", decoder_registry=bundle.decoder_registry
)
To save our experiment to SQL, we must first specify a connection to a database and create all necessary tables.
from ax.storage.sqa_store.db import (
init_engine_and_session_factory,
get_engine,
create_all_tables,
)
from ax.storage.sqa_store.load import load_experiment
from ax.storage.sqa_store.save import save_experiment
init_engine_and_session_factory(url="sqlite:///foo3.db")
engine = get_engine()
create_all_tables(engine)
from ax.storage.sqa_store.sqa_config import SQAConfig
exp.name = "new"
sqa_config = SQAConfig(
json_encoder_registry=bundle.encoder_registry,
json_decoder_registry=bundle.decoder_registry,
metric_registry=bundle.metric_registry,
runner_registry=bundle.runner_registry,
)
save_experiment(exp, config=sqa_config)
load_experiment(exp.name, config=sqa_config)