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.
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
init_notebook_plotting()
[INFO 09-23 20:29:20] ax.utils.notebook.plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.
[INFO 09-23 20:29:20] ax.utils.notebook.plotting: Please see (https://ax.dev/tutorials/visualizations.html#Fix-for-plots-that-are-not-rendering) if visualizations are not rendering.
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
)
/tmp/ipykernel_2335/1805422181.py:1: AxParameterWarning: `is_ordered` is not specified for `ChoiceParameter` "choice". Defaulting to `True` since there are exactly two choices.. To override this behavior (or avoid this warning), specify `is_ordered` during `ChoiceParameter` construction. Note that choice parameters with exactly 2 choices are always considered ordered and that the user-supplied `is_ordered` has no effect in this particular case. /tmp/ipykernel_2335/1805422181.py:1: AxParameterWarning: `sort_values` is not specified for `ChoiceParameter` "choice". Defaulting to `False` for parameters of `ParameterType` STRING. To override this behavior (or avoid this warning), specify `sort_values` during `ChoiceParameter` construction.
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"],
)
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,
)
],
)
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
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(),
)
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!")
Running Sobol initialization trials... Running BO trial 6/20...
Running BO trial 7/20...
Running BO trial 8/20...
Running BO trial 9/20...
Running BO trial 10/20...
Running BO trial 11/20...
Running BO trial 12/20...
Running BO trial 13/20...
Running BO trial 14/20...
Running BO trial 15/20...
Running BO trial 16/20...
Running BO trial 17/20...
Running BO trial 18/20...
Running BO trial 19/20...
Running BO trial 20/20...
Done!
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 | -2.553538 | 0.0 | 19 | 10000 | -2.553538 |
1 | 19_0 | l2norm | 1.041480 | 0.2 | 19 | 10000 | 1.041480 |
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.085655 | 0.0 | 0 | 10000 | -0.085655 |
1 | 0_0 | l2norm | 1.588772 | 0.2 | 0 | 10000 | 1.588772 |
2 | 1_0 | hartmann6 | -0.035950 | 0.0 | 1 | 10000 | -0.035950 |
3 | 1_0 | l2norm | 1.486780 | 0.2 | 1 | 10000 | 1.486780 |
4 | 2_0 | hartmann6 | -0.448604 | 0.0 | 2 | 10000 | -0.448604 |
5 | 2_0 | l2norm | 1.177263 | 0.2 | 2 | 10000 | 1.177263 |
6 | 3_0 | hartmann6 | -0.145877 | 0.0 | 3 | 10000 | -0.145877 |
7 | 3_0 | l2norm | 1.011089 | 0.2 | 3 | 10000 | 1.011089 |
8 | 4_0 | hartmann6 | -0.416899 | 0.0 | 4 | 10000 | -0.416899 |
9 | 4_0 | l2norm | 0.868664 | 0.2 | 4 | 10000 | 0.868664 |
10 | 5_0 | hartmann6 | -0.967230 | 0.0 | 5 | 10000 | -0.967230 |
11 | 5_0 | l2norm | 0.726146 | 0.2 | 5 | 10000 | 0.726146 |
12 | 6_0 | hartmann6 | -0.419545 | 0.0 | 6 | 10000 | -0.419545 |
13 | 6_0 | l2norm | 0.486646 | 0.2 | 6 | 10000 | 0.486646 |
14 | 7_0 | hartmann6 | -0.737130 | 0.0 | 7 | 10000 | -0.737130 |
15 | 7_0 | l2norm | 1.025983 | 0.2 | 7 | 10000 | 1.025983 |
16 | 8_0 | hartmann6 | -0.603645 | 0.0 | 8 | 10000 | -0.603645 |
17 | 8_0 | l2norm | 0.709729 | 0.2 | 8 | 10000 | 0.709729 |
18 | 9_0 | hartmann6 | -1.098661 | 0.0 | 9 | 10000 | -1.098661 |
19 | 9_0 | l2norm | 0.533604 | 0.2 | 9 | 10000 | 0.533604 |
20 | 10_0 | hartmann6 | -1.154395 | 0.0 | 10 | 10000 | -1.154395 |
21 | 10_0 | l2norm | 1.153987 | 0.2 | 10 | 10000 | 1.153987 |
22 | 11_0 | hartmann6 | -0.215039 | 0.0 | 11 | 10000 | -0.215039 |
23 | 11_0 | l2norm | 1.134288 | 0.2 | 11 | 10000 | 1.134288 |
24 | 12_0 | hartmann6 | -0.996909 | 0.0 | 12 | 10000 | -0.996909 |
25 | 12_0 | l2norm | 0.770433 | 0.2 | 12 | 10000 | 0.770433 |
26 | 13_0 | hartmann6 | -1.032950 | 0.0 | 13 | 10000 | -1.032950 |
27 | 13_0 | l2norm | 0.881829 | 0.2 | 13 | 10000 | 0.881829 |
28 | 14_0 | hartmann6 | -1.092444 | 0.0 | 14 | 10000 | -1.092444 |
29 | 14_0 | l2norm | 0.997956 | 0.2 | 14 | 10000 | 0.997956 |
30 | 15_0 | hartmann6 | -0.999576 | 0.0 | 15 | 10000 | -0.999576 |
31 | 15_0 | l2norm | 1.015494 | 0.2 | 15 | 10000 | 1.015494 |
32 | 16_0 | hartmann6 | -1.635377 | 0.0 | 16 | 10000 | -1.635377 |
33 | 16_0 | l2norm | 1.123209 | 0.2 | 16 | 10000 | 1.123209 |
34 | 17_0 | hartmann6 | -1.859101 | 0.0 | 17 | 10000 | -1.859101 |
35 | 17_0 | l2norm | 0.910136 | 0.2 | 17 | 10000 | 0.910136 |
36 | 18_0 | hartmann6 | -2.155254 | 0.0 | 18 | 10000 | -2.155254 |
37 | 18_0 | l2norm | 0.677451 | 0.2 | 18 | 10000 | 0.677451 |
38 | 19_0 | hartmann6 | -2.553538 | 0.0 | 19 | 10000 | -2.553538 |
39 | 19_0 | l2norm | 1.347050 | 0.2 | 19 | 10000 | 1.347050 |
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)
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
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)
Experiment(new)
Total runtime of script: 1 minutes, 30.17 seconds.