The loop API is the most lightweight way to do optimization in Ax. The user makes one call to optimize, which performs all of the optimization under the hood and returns the optimized parameters.
For more customizability of the optimization procedure, consider the Service or Developer API.
import numpy as np
from ax.plot.contour import plot_contour
from ax.plot.trace import optimization_trace_single_method
from ax.service.managed_loop import optimize
from ax.metrics.branin import branin
from ax.utils.measurement.synthetic_functions import hartmann6
from ax.utils.notebook.plotting import render, init_notebook_plotting
init_notebook_plotting()
[INFO 04-30 16:00:39] ipy_plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.
First, we set up an evaluation function that returns evaluations for two metrics: "hartmann6" and "l2norm".
def hartmann_evaluation_function(parameterization):
x = np.array([parameterization.get(f"x{i+1}") for i in range(6)])
# In our case, standard error is 0, since we are computing a synthetic function.
return {"hartmann6": (hartmann6(x), 0.0), "l2norm": (np.sqrt((x ** 2).sum()), 0.0)}
The setup for the loop is fully compatible with JSON. The optimization algorithm is selected based on the properties of the problem search space.
best_parameters, values, experiment, model = optimize(
parameters=[
{
"name": "x1",
"type": "range",
"bounds": [0.0, 1.0],
"value_type": "float", # Optional, defaults to inference from type of "bounds".
"log_scale": False, # Optional, defaults to False.
},
{
"name": "x2",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x3",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x4",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x5",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x6",
"type": "range",
"bounds": [0.0, 1.0],
},
],
experiment_name="test",
objective_name="hartmann6",
evaluation_function=hartmann_evaluation_function,
minimize=True, # Optional, defaults to False.
parameter_constraints=["x1 + x2 <= 20"], # Optional.
outcome_constraints=["l2norm <= 1.25"], # Optional.
total_trials=30, # Optional.
)
[INFO 04-30 16:00:39] ax.service.utils.dispatch: Using Bayesian Optimization generation strategy. Iterations after 6 will take longer to generate due to model-fitting. [INFO 04-30 16:00:39] ax.service.managed_loop: Started full optimization with 30 steps. [INFO 04-30 16:00:39] ax.service.managed_loop: Running optimization trial 1... [INFO 04-30 16:00:39] ax.service.managed_loop: Running optimization trial 2... [INFO 04-30 16:00:39] ax.service.managed_loop: Running optimization trial 3... [INFO 04-30 16:00:39] ax.service.managed_loop: Running optimization trial 4... [INFO 04-30 16:00:39] ax.service.managed_loop: Running optimization trial 5... [INFO 04-30 16:00:39] ax.service.managed_loop: Running optimization trial 6... [INFO 04-30 16:00:39] ax.service.managed_loop: Running optimization trial 7... [INFO 04-30 16:03:22] ax.service.managed_loop: Running optimization trial 8... [INFO 04-30 16:06:32] ax.service.managed_loop: Running optimization trial 9... [INFO 04-30 16:09:53] ax.service.managed_loop: Running optimization trial 10... [INFO 04-30 16:14:00] ax.service.managed_loop: Running optimization trial 11... [INFO 04-30 16:18:21] ax.service.managed_loop: Running optimization trial 12... [INFO 04-30 16:22:16] ax.service.managed_loop: Running optimization trial 13... [INFO 04-30 16:24:24] ax.service.managed_loop: Running optimization trial 14... [INFO 04-30 16:28:02] ax.service.managed_loop: Running optimization trial 15... [INFO 04-30 16:31:10] ax.service.managed_loop: Running optimization trial 16... [INFO 04-30 16:35:04] ax.service.managed_loop: Running optimization trial 17... [INFO 04-30 16:39:13] ax.service.managed_loop: Running optimization trial 18... [INFO 04-30 16:43:41] ax.service.managed_loop: Running optimization trial 19... [INFO 04-30 16:47:26] ax.service.managed_loop: Running optimization trial 20... [INFO 04-30 16:51:29] ax.service.managed_loop: Running optimization trial 21... [INFO 04-30 16:52:26] ax.service.managed_loop: Running optimization trial 22... [INFO 04-30 16:54:26] ax.service.managed_loop: Running optimization trial 23... [INFO 04-30 16:59:06] ax.service.managed_loop: Running optimization trial 24... [INFO 04-30 17:00:50] ax.service.managed_loop: Running optimization trial 25... [INFO 04-30 17:03:20] ax.service.managed_loop: Running optimization trial 26... [INFO 04-30 17:05:30] ax.service.managed_loop: Running optimization trial 27... [INFO 04-30 17:07:03] ax.service.managed_loop: Running optimization trial 28... [INFO 04-30 17:11:53] ax.service.managed_loop: Running optimization trial 29... [INFO 04-30 17:13:45] ax.service.managed_loop: Running optimization trial 30...
And we can introspect optimization results:
best_parameters
{'x1': 0.2518118433201088,
'x2': 0.12875223997444202,
'x3': 0.432035773040578,
'x4': 0.2853897302458177,
'x5': 0.31579857892882346,
'x6': 0.6537879408811222}
means, covariances = values
means
{'l2norm': 0.9355511573766859, 'hartmann6': -3.272066359418134}
For comparison, minimum of Hartmann6 is:
hartmann6.fmin
-3.32237
Here we arbitrarily select "x1" and "x2" as the two parameters to plot for both metrics, "hartmann6" and "l2norm".
render(plot_contour(model=model, param_x='x1', param_y='x2', metric_name='hartmann6'))
render(plot_contour(model=model, param_x='x1', param_y='x2', metric_name='l2norm'))
We also plot optimization trace, which shows best hartmann6 objective value seen by each iteration of the optimization:
# `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.
best_objectives = np.array([[trial.objective_mean for trial in experiment.trials.values()]])
best_objective_plot = optimization_trace_single_method(
y=np.minimum.accumulate(best_objectives, axis=1),
optimum=hartmann6.fmin,
title="Model performance vs. # of iterations",
ylabel="Hartmann6",
)
render(best_objective_plot)