Skip to main content
Version: Next

Getting Started with Ax

Complex optimization problems where we wish to tune multiple parameters to improve metric performance, but the inter-parameter interactions are not fully understood, are common across various fields including machine learning, robotics, materials science, and chemistry. This category of problem is known as "black-box" optimization. The complexity of black-box optimization problems further increases if evaluations are expensive to conduct, time-consuming, or noisy.

We can use Ax to efficiently conduct an experiment in which we "ask" for candidate points to evaluate, "tell" Ax the results, and repeat. We'll uses Ax's Client, a tool for managing the state of our experiment, and we'll learn how to define an optimization problem, configure an experiment, run trials, analyze results, and persist the experiment for later use using the Client.

Because Ax is a black box optimizer, we can use it to optimize any arbitrary function. In this example we will minimize the Hartmann6 function, a complicated 6-dimensional function with multiple local minima. Hartmann6 is a challenging benchmark for optimization algorithms commonly used in the global optimization literature -- it tests the algorithm's ability to identify the true global minimum, rather than mistakenly converging on a local minimum. Looking at its analytic form we can see that it would be incredibly challenging to efficiently find the global minimum either by manual trial-and-error or traditional design of experiments like grid-search or random-search.

f(x)=i=14αiexp(j=16Aij(xjPij)2) f(\mathbf{x})=-\sum_{i=1}^4 \alpha_i \exp \left(-\sum_{j=1}^6 A_{i j}\left(x_j-P_{i j}\right)^2\right)

Learning Objectives

  • Understand the basic concepts of black box optimization
  • Learn how to define an optimization problem using Ax
  • Configure and run an experiment using Ax's Client
  • Analyze the results of the optimization

Prerequisites

Step 1: Import Necessary Modules

First, ensure you have all the necessary imports:

import numpy as np
from ax.api.client import Client
from ax.api.configs import RangeParameterConfig

Step 2: Initialize the Client

Create an instance of the Client to manage the state of your experiment.

client = Client()

Step 3: Configure the Experiment

The Client instance can be configured with a series of Configs that define how the experiment will be run.

The Hartmann6 problem is usually evaluated on the hypercube xi(0,1)x_i \in (0, 1), so we will define six identical RangeParameterConfigs with these bounds.

You may specify additional features like parameter constraints to further refine the search space and parameter scaling to help navigate parameters with nonuniform effects.

# Define six float parameters x1, x2, x3, ... for the Hartmann6 function, which is typically evaluated on the unit hypercube
parameters = [
RangeParameterConfig(
name="x1", parameter_type="float", bounds=(0, 1)
),
RangeParameterConfig(
name="x2", parameter_type="float", bounds=(0, 1)
),
RangeParameterConfig(
name="x3", parameter_type="float", bounds=(0, 1)
),
RangeParameterConfig(
name="x4", parameter_type="float", bounds=(0, 1)
),
RangeParameterConfig(
name="x5", parameter_type="float", bounds=(0, 1)
),
RangeParameterConfig(
name="x6", parameter_type="float", bounds=(0, 1)
),
]

client.configure_experiment(parameters=parameters)

Step 4: Configure Optimization

Now, we must configure the objective for this optimization, which we do using Client.configure_optimization. This method expects a string objective, an expression containing either a single metric to maximize, a linear combination of metrics to maximize, or a tuple of multiple metrics to jointly maximize. These expressions are parsed using SymPy. For example:

  • "score" would direct Ax to maximize a metric named score
  • "-loss" would direct Ax to Ax to minimize a metric named loss
  • "task_0 + 0.5 * task_1" would direct Ax to maximize the sum of two task scores, downweighting task_1 by a factor of 0.5
  • "score, -flops" would direct Ax to simultaneously maximize score while minimizing flops

See these recipes for more information on configuring objectives and outcome constraints.

metric_name = "hartmann6" # this name is used during the optimization loop in Step 5
objective = f"-{metric_name}" # minimization is specified by the negative sign

client.configure_optimization(objective=objective)

Step 5: Run Trials

Here, we will configure the ask-tell loop.

We begin by defining the Hartmann6 function as written above. Remember, this is just an example problem and any Python function can be substituted here.

# Hartmann6 function
def hartmann6(x1, x2, x3, x4, x5, x6):
alpha = np.array([1.0, 1.2, 3.0, 3.2])
A = np.array([
[10, 3, 17, 3.5, 1.7, 8],
[0.05, 10, 17, 0.1, 8, 14],
[3, 3.5, 1.7, 10, 17, 8],
[17, 8, 0.05, 10, 0.1, 14]
])
P = 10**-4 * np.array([
[1312, 1696, 5569, 124, 8283, 5886],
[2329, 4135, 8307, 3736, 1004, 9991],
[2348, 1451, 3522, 2883, 3047, 6650],
[4047, 8828, 8732, 5743, 1091, 381]
])

outer = 0.0
for i in range(4):
inner = 0.0
for j, x in enumerate([x1, x2, x3, x4, x5, x6]):
inner += A[i, j] * (x - P[i, j])**2
outer += alpha[i] * np.exp(-inner)
return -outer

hartmann6(0.1, 0.45, 0.8, 0.25, 0.552, 1.0)
Output:
np.float64(-0.4878737485613134)

Optimization Loop

We will iteratively call client.get_next_trials to "ask" Ax for a parameterization to evaluate, then call hartmann6 using those parameters, and finally "tell" Ax the result using client.complete_trial.

This loop will run multiple trials to optimize the function.

for _ in range(10): # Run 10 rounds of trials
# We will request three trials at a time in this example
trials = client.get_next_trials(max_trials=3)

for trial_index, parameters in trials.items():
x1 = parameters["x1"]
x2 = parameters["x2"]
x3 = parameters["x3"]
x4 = parameters["x4"]
x5 = parameters["x5"]
x6 = parameters["x6"]

result = hartmann6(x1, x2, x3, x4, x5, x6)

# Set raw_data as a dictionary with metric names as keys and results as values
raw_data = {metric_name: result}

# Complete the trial with the result
client.complete_trial(trial_index=trial_index, raw_data=raw_data)

print(f"Completed trial {trial_index} with {raw_data=}")
Output:
[WARNING 05-31 05:03:57] ax.api.client: 3 trials requested but only 2 could be generated.
Completed trial 0 with raw_data={'hartmann6': np.float64(-0.5053149917022333)}
Completed trial 1 with raw_data={'hartmann6': np.float64(-0.05310880076340437)}
Completed trial 2 with raw_data={'hartmann6': np.float64(-0.0027575806405279457)}
Completed trial 3 with raw_data={'hartmann6': np.float64(-1.1181996618940482)}
Completed trial 4 with raw_data={'hartmann6': np.float64(-0.2994968277034288)}
[WARNING 05-31 05:03:58] ax.api.client: 3 trials requested but only 1 could be generated.
Completed trial 5 with raw_data={'hartmann6': np.float64(-0.6199607419741057)}
Completed trial 6 with raw_data={'hartmann6': np.float64(-1.4171429433355645)}
Completed trial 7 with raw_data={'hartmann6': np.float64(-1.1860464764090044)}
Completed trial 8 with raw_data={'hartmann6': np.float64(-0.2886977847465286)}
Completed trial 9 with raw_data={'hartmann6': np.float64(-0.6033495951413292)}
Completed trial 10 with raw_data={'hartmann6': np.float64(-0.5731342396172782)}
Completed trial 11 with raw_data={'hartmann6': np.float64(-0.2143303396514023)}
Completed trial 12 with raw_data={'hartmann6': np.float64(-1.9297193854919958)}
Completed trial 13 with raw_data={'hartmann6': np.float64(-0.7576422475972197)}
Completed trial 14 with raw_data={'hartmann6': np.float64(-1.6511268198737103)}
Completed trial 15 with raw_data={'hartmann6': np.float64(-1.8101063659626961)}
Completed trial 16 with raw_data={'hartmann6': np.float64(-1.1807823364671957)}
Completed trial 17 with raw_data={'hartmann6': np.float64(-1.2341467968286983)}
Completed trial 18 with raw_data={'hartmann6': np.float64(-1.3994378547667)}
Completed trial 19 with raw_data={'hartmann6': np.float64(-2.516460519985058)}
Completed trial 20 with raw_data={'hartmann6': np.float64(-1.0792181277935151)}
Completed trial 21 with raw_data={'hartmann6': np.float64(-2.938633144766989)}
Completed trial 22 with raw_data={'hartmann6': np.float64(-3.0138871209561264)}
Completed trial 23 with raw_data={'hartmann6': np.float64(-2.473956511483838)}
Completed trial 24 with raw_data={'hartmann6': np.float64(-3.223634648212362)}
Completed trial 25 with raw_data={'hartmann6': np.float64(-2.9959900546169878)}
Completed trial 26 with raw_data={'hartmann6': np.float64(-2.9569657233654043)}

Step 6: Analyze Results

After running trials, you can analyze the results. Most commonly this means extracting the parameterization from the best performing trial you conducted.

Hartmann6 has a known global minimum of f(x)=3.322f(x*) = -3.322 at x=(0.201,0.150,0.477,0.273,0.312,0.657)x* = (0.201, 0.150, 0.477, 0.273, 0.312, 0.657). Ax is able to identify a point very near to this true optimum using just 30 evaluations. This is possible due to the sample-efficiency of Bayesian optimization, the optimization method we use under the hood in Ax.

best_parameters, prediction, index, name = client.get_best_parameterization()
print("Best Parameters:", best_parameters)
print("Prediction (mean, variance):", prediction)
Output:
Best Parameters: {'x1': 0.15915548492999046, 'x2': 0.27191153291213743, 'x3': 0.5192361185007736, 'x4': 0.2761008279597605, 'x5': 0.3493177058737067, 'x6': 0.6094043810697247}
Prediction (mean, variance): {'hartmann6': (np.float64(-3.00590874991418), np.float64(0.002086710923302185))}

Step 7: Compute Analyses

Ax can also produce a number of analyses to help interpret the results of the experiment via client.compute_analyses. Users can manually select which analyses to run, or can allow Ax to select which would be most relevant. In this case Ax selects the following:

  • Parrellel Coordinates Plot shows which parameterizations were evaluated and what metric values were observed -- this is useful for getting a high level overview of how thoroughly the search space was explored and which regions tend to produce which outcomes
  • Sensitivity Analysis Plot shows which parameters have the largest affect on the objective using Sobol Indicies
  • Slice Plot shows how the model predicts a single parameter effects the objective along with a confidence interval
  • Contour Plot shows how the model predicts a pair of parameters effects the objective as a 2D surface
  • Summary lists all trials generated along with their parameterizations, observations, and miscellaneous metadata
  • Cross Validation helps to visualize how well the surrogate model is able to predict out of sample points
# display=True instructs Ax to sort then render the resulting analyses
cards = client.compute_analyses(display=True)

Parallel Coordinates for hartmann6

The parallel coordinates plot displays multi-dimensional data by representing each parameter as a parallel axis. This plot helps in assessing how thoroughly the search space has been explored and in identifying patterns or clusterings associated with high-performing (good) or low-performing (bad) arms. By tracing lines across the axes, one can observe correlations and interactions between parameters, gaining insights into the relationships that contribute to the success or failure of different configurations within the experiment.

loading...

Sensitivity Analysis for hartmann6

Understand how each parameter affects hartmann6 according to a second-order sensitivity analysis.

loading...

x1 vs. hartmann6

The slice plot provides a one-dimensional view of predicted outcomes for hartmann6 as a function of a single parameter, while keeping all other parameters fixed at their status_quo value (or mean value if status_quo is unavailable). This visualization helps in understanding the sensitivity and impact of changes in the selected parameter on the predicted metric outcomes.

loading...

x2 vs. hartmann6

The slice plot provides a one-dimensional view of predicted outcomes for hartmann6 as a function of a single parameter, while keeping all other parameters fixed at their status_quo value (or mean value if status_quo is unavailable). This visualization helps in understanding the sensitivity and impact of changes in the selected parameter on the predicted metric outcomes.

loading...

x6 vs. hartmann6

The slice plot provides a one-dimensional view of predicted outcomes for hartmann6 as a function of a single parameter, while keeping all other parameters fixed at their status_quo value (or mean value if status_quo is unavailable). This visualization helps in understanding the sensitivity and impact of changes in the selected parameter on the predicted metric outcomes.

loading...

x5 vs. hartmann6

The slice plot provides a one-dimensional view of predicted outcomes for hartmann6 as a function of a single parameter, while keeping all other parameters fixed at their status_quo value (or mean value if status_quo is unavailable). This visualization helps in understanding the sensitivity and impact of changes in the selected parameter on the predicted metric outcomes.

loading...

x2, x5 vs. hartmann6

The contour plot visualizes the predicted outcomes for hartmann6 across a two-dimensional parameter space, with other parameters held fixed at their status_quo value (or mean value if status_quo is unavailable). This plot helps in identifying regions of optimal performance and understanding how changes in the selected parameters influence the predicted outcomes. Contour lines represent levels of constant predicted values, providing insights into the gradient and potential optima within the parameter space.

loading...

Summary for Experiment

High-level summary of the Trial-s in this Experiment

trial_indexarm_nametrial_statusgeneration_nodehartmann6x1x2x3x4x5x6
000_0COMPLETEDCenterOfSearchSpace-0.5053150.50.50.50.50.50.5
111_0COMPLETEDSobol-0.0531090.8720330.9400290.8007740.418220.3012420.185559
222_0COMPLETEDSobol-0.0027580.1116230.0454960.1881950.8554450.8286760.979904
333_0COMPLETEDSobol-1.11820.2555680.5189040.6068420.1837740.50050.545
444_0COMPLETEDSobol-0.2994970.5107750.4993740.4044790.6206930.0980650.370593
555_0COMPLETEDMBM-0.6199610.2446240.4302630.626340.0299990.9317720.665512
666_0COMPLETEDMBM-1.417140.1424640.5673660.6302050.211870.2671860.537635
777_0COMPLETEDMBM-1.186050.323120.4357970.4794830.1349680.3157750.397272
888_0COMPLETEDMBM-0.2886980.2218150.6670290.8217190.2346580.5586120.417136
999_0COMPLETEDMBM-0.603350.0503170.5301150.567210.1220860.0631210.536614
101010_0COMPLETEDMBM-0.5731340.0326960.2581210.5874250.2066030.0950340.321497
111111_0COMPLETEDMBM-0.214330.0733040.8841490.5542870.1190080.1059770.539749
121212_0COMPLETEDMBM-1.929720.1463080.4934850.6049910.1866160.3314450.601086
131313_0COMPLETEDMBM-0.7576420.1190120.5322250.5984130.1976610.3334680.314529
141414_0COMPLETEDMBM-1.651130.3612720.5096430.6337570.16890.3082470.661154
151515_0COMPLETEDMBM-1.8101100.3596240.5334630.0844530.3525530.64819
161616_0COMPLETEDMBM-1.1807800.3339990.9134770.1458940.3537050.641758
171717_0COMPLETEDMBM-1.2341500.4456910.1866030.1015460.3563830.659469
181818_0COMPLETEDMBM-1.3994400.5475560.5505950.1550170.3671940.605632
191919_0COMPLETEDMBM-2.516460.0408760.3465250.5414310.2656070.3475850.603833
202020_0COMPLETEDMBM-1.079220.3343510.5270230.5705180.034070.3681260.584111
212121_0COMPLETEDMBM-2.938630.1287230.3164720.5166030.2961760.3214580.667262
222222_0COMPLETEDMBM-3.013890.1591550.2719120.5192360.2761010.3493180.609404
232323_0COMPLETEDMBM-2.473960.0373210.354720.535590.3295580.3144050.70817
242424_0COMPLETEDMBM-3.223630.1941640.1371640.4861150.3259640.3307250.654245
252525_0COMPLETEDMBM-2.995990.1841320.1602230.3900740.3624360.3378050.634113
262626_0COMPLETEDMBM-2.956970.2081910.1099240.5868270.3615540.3168180.652197

Cross Validation for hartmann6

The cross-validation plot displays the model fit for each metric in the experiment. It employs a leave-one-out approach, where the model is trained on all data except one sample, which is used for validation. The plot shows the predicted outcome for the validation set on the y-axis against its actual value on the x-axis. Points that align closely with the dotted diagonal line indicate a strong model fit, signifying accurate predictions. Additionally, the plot includes 95% confidence intervals that provide insight into the noise in observations and the uncertainty in model predictions. A horizontal, flat line of predictions indicates that the model has not picked up on sufficient signal in the data, and instead is just predicting the mean.

loading...

Conclusion

This tutorial demonstrates how to use Ax's Client for ask-tell optimization of Python functions using the Hartmann6 function as an example. You can adjust the function and parameters to suit your specific optimization problem.