This tutorial illustrates the core visualization utilities available in Ax.
import numpy as np
from ax.modelbridge.cross_validation import cross_validate
from ax.plot.contour import interact_contour
from ax.plot.diagnostic import interact_cross_validation
from ax.plot.scatter import interact_fitted, plot_objective_vs_constraints, tile_fitted
from ax.plot.slice import plot_slice
from ax.service.ax_client import AxClient, ObjectiveProperties
from ax.utils.measurement.synthetic_functions import hartmann6
from ax.utils.notebook.plotting import init_notebook_plotting, render
init_notebook_plotting()
[INFO 09-23 20:30:52] ax.utils.notebook.plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.
[INFO 09-23 20:30:52] ax.utils.notebook.plotting: Please see (https://ax.dev/tutorials/visualizations.html#Fix-for-plots-that-are-not-rendering) if visualizations are not rendering.
The vizualizations require an experiment object and a model fit on the evaluated data. The routine below is a copy of the Service API tutorial, so the explanation here is omitted. Retrieving the experiment and model objects for each API paradigm is shown in the respective tutorials
noise_sd = 0.1
param_names = [f"x{i+1}" for i in range(6)] # x1, x2, ..., x6
def noisy_hartmann_evaluation_function(parameterization):
x = np.array([parameterization.get(p_name) for p_name in param_names])
noise1, noise2 = np.random.normal(0, noise_sd, 2)
return {
"hartmann6": (hartmann6(x) + noise1, noise_sd),
"l2norm": (np.sqrt((x**2).sum()) + noise2, noise_sd),
}
ax_client = AxClient()
ax_client.create_experiment(
name="test_visualizations",
parameters=[
{
"name": p_name,
"type": "range",
"bounds": [0.0, 1.0],
}
for p_name in param_names
],
objectives={"hartmann6": ObjectiveProperties(minimize=True)},
outcome_constraints=["l2norm <= 1.25"],
)
[INFO 09-23 20:30:52] 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 09-23 20:30:52] 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 09-23 20:30:52] 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 09-23 20:30:52] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x3. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 09-23 20:30:52] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x4. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 09-23 20:30:52] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x5. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 09-23 20:30:52] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x6. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 09-23 20:30:52] ax.service.utils.instantiation: Created search space: SearchSpace(parameters=[RangeParameter(name='x1', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x2', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x3', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x4', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x5', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x6', parameter_type=FLOAT, range=[0.0, 1.0])], parameter_constraints=[]).
[INFO 09-23 20:30:52] ax.modelbridge.dispatch_utils: Using Models.BOTORCH_MODULAR since there is at least one ordered parameter and there are no unordered categorical parameters.
[INFO 09-23 20:30:52] ax.modelbridge.dispatch_utils: Calculating the number of remaining initialization trials based on num_initialization_trials=None max_initialization_trials=None num_tunable_parameters=6 num_trials=None use_batch_trials=False
[INFO 09-23 20:30:52] ax.modelbridge.dispatch_utils: calculated num_initialization_trials=12
[INFO 09-23 20:30:52] ax.modelbridge.dispatch_utils: num_completed_initialization_trials=0 num_remaining_initialization_trials=12
[INFO 09-23 20:30:52] ax.modelbridge.dispatch_utils: `verbose`, `disable_progbar`, and `jit_compile` are not yet supported when using `choose_generation_strategy` with ModularBoTorchModel, dropping these arguments.
[INFO 09-23 20:30:52] ax.modelbridge.dispatch_utils: Using Bayesian Optimization generation strategy: GenerationStrategy(name='Sobol+BoTorch', steps=[Sobol for 12 trials, BoTorch for subsequent trials]). Iterations after 12 will take longer to generate due to model-fitting.
for i in range(20):
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=noisy_hartmann_evaluation_function(parameters)
)
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 0 with parameters {'x1': 0.301966, 'x2': 0.115237, 'x3': 0.662966, 'x4': 0.787345, 'x5': 0.514551, 'x6': 0.35244} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 0 with data: {'hartmann6': (0.007817, 0.1), 'l2norm': (1.20666, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 1 with parameters {'x1': 0.690867, 'x2': 0.848207, 'x3': 0.003923, 'x4': 0.416611, 'x5': 0.060394, 'x6': 0.503174} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 1 with data: {'hartmann6': (-0.086721, 0.1), 'l2norm': (1.156036, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 2 with parameters {'x1': 0.802228, 'x2': 0.481028, 'x3': 0.78, 'x4': 0.746598, 'x5': 0.473398, 'x6': 0.981709} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 2 with data: {'hartmann6': (-0.262696, 0.1), 'l2norm': (1.868653, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 3 with parameters {'x1': 0.190671, 'x2': 0.748055, 'x3': 0.433947, 'x4': 0.109519, 'x5': 0.960888, 'x6': 0.131726} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 3 with data: {'hartmann6': (-0.331837, 0.1), 'l2norm': (1.280234, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 4 with parameters {'x1': 0.084531, 'x2': 0.327964, 'x3': 0.187647, 'x4': 0.239736, 'x5': 0.808278, 'x6': 0.40264} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 4 with data: {'hartmann6': (-0.024598, 0.1), 'l2norm': (0.981802, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 5 with parameters {'x1': 0.907026, 'x2': 0.591085, 'x3': 0.596445, 'x4': 0.619006, 'x5': 0.258409, 'x6': 0.74211} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 5 with data: {'hartmann6': (-0.239878, 0.1), 'l2norm': (1.579705, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 6 with parameters {'x1': 0.584728, 'x2': 0.208281, 'x3': 0.36767, 'x4': 0.29665, 'x5': 0.204627, 'x6': 0.775324} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 6 with data: {'hartmann6': (-1.572804, 0.1), 'l2norm': (1.037927, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 7 with parameters {'x1': 0.406765, 'x2': 0.945157, 'x3': 0.963481, 'x4': 0.909565, 'x5': 0.721408, 'x6': 0.111605} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 7 with data: {'hartmann6': (-1.027357, 0.1), 'l2norm': (1.942456, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 8 with parameters {'x1': 0.47679, 'x2': 0.379172, 'x3': 0.25488, 'x4': 0.508789, 'x5': 0.11077, 'x6': 0.247573} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 8 with data: {'hartmann6': (-0.460301, 0.1), 'l2norm': (0.912411, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 9 with parameters {'x1': 0.529784, 'x2': 0.647298, 'x3': 0.912214, 'x4': 0.1305, 'x5': 0.564808, 'x6': 0.896629} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 9 with data: {'hartmann6': (-0.060766, 0.1), 'l2norm': (1.820705, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 10 with parameters {'x1': 0.977014, 'x2': 0.03012, 'x3': 0.184233, 'x4': 0.956763, 'x5': 0.901647, 'x6': 0.617313} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 10 with data: {'hartmann6': (0.029059, 0.1), 'l2norm': (1.650103, 0.1)}.
/tmp/tmp.QqcA7fo0ui/Ax-main/ax/modelbridge/cross_validation.py:463: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction. [INFO 09-23 20:30:53] ax.service.ax_client: Generated new trial 11 with parameters {'x1': 0.029488, 'x2': 0.761991, 'x3': 0.529554, 'x4': 0.342866, 'x5': 0.41452, 'x6': 0.269069} using model Sobol.
[INFO 09-23 20:30:53] ax.service.ax_client: Completed trial 11 with data: {'hartmann6': (-0.221124, 0.1), 'l2norm': (1.341961, 0.1)}.
[INFO 09-23 20:30:57] ax.service.ax_client: Generated new trial 12 with parameters {'x1': 0.515304, 'x2': 0.147016, 'x3': 0.270404, 'x4': 0.252918, 'x5': 0.125455, 'x6': 0.866776} using model BoTorch.
[INFO 09-23 20:30:57] ax.service.ax_client: Completed trial 12 with data: {'hartmann6': (-1.096015, 0.1), 'l2norm': (1.124111, 0.1)}.
[INFO 09-23 20:31:03] ax.service.ax_client: Generated new trial 13 with parameters {'x1': 0.616218, 'x2': 0.22177, 'x3': 0.398197, 'x4': 0.193171, 'x5': 0.17117, 'x6': 0.686749} using model BoTorch.
[INFO 09-23 20:31:03] ax.service.ax_client: Completed trial 13 with data: {'hartmann6': (-1.235589, 0.1), 'l2norm': (1.208472, 0.1)}.
[INFO 09-23 20:31:12] ax.service.ax_client: Generated new trial 14 with parameters {'x1': 0.517555, 'x2': 0.184403, 'x3': 0.334269, 'x4': 0.30403, 'x5': 0.282924, 'x6': 0.776514} using model BoTorch.
[INFO 09-23 20:31:12] ax.service.ax_client: Completed trial 14 with data: {'hartmann6': (-1.975216, 0.1), 'l2norm': (1.019098, 0.1)}.
[INFO 09-23 20:31:17] ax.service.ax_client: Generated new trial 15 with parameters {'x1': 0.466466, 'x2': 0.105748, 'x3': 0.196728, 'x4': 0.196277, 'x5': 0.289378, 'x6': 1.0} using model BoTorch.
[INFO 09-23 20:31:17] ax.service.ax_client: Completed trial 15 with data: {'hartmann6': (-0.711773, 0.1), 'l2norm': (1.124486, 0.1)}.
[INFO 09-23 20:31:26] ax.service.ax_client: Generated new trial 16 with parameters {'x1': 0.432326, 'x2': 0.144975, 'x3': 0.298308, 'x4': 0.391652, 'x5': 0.344003, 'x6': 0.753454} using model BoTorch.
[INFO 09-23 20:31:26] ax.service.ax_client: Completed trial 16 with data: {'hartmann6': (-2.17812, 0.1), 'l2norm': (1.165593, 0.1)}.
[INFO 09-23 20:31:31] ax.service.ax_client: Generated new trial 17 with parameters {'x1': 0.325259, 'x2': 0.10506, 'x3': 0.309941, 'x4': 0.051872, 'x5': 0.386782, 'x6': 0.757427} using model BoTorch.
[INFO 09-23 20:31:31] ax.service.ax_client: Completed trial 17 with data: {'hartmann6': (-1.606619, 0.1), 'l2norm': (0.995786, 0.1)}.
[INFO 09-23 20:31:35] ax.service.ax_client: Generated new trial 18 with parameters {'x1': 0.479433, 'x2': 0.190328, 'x3': 0.13175, 'x4': 0.471527, 'x5': 0.368483, 'x6': 0.745157} using model BoTorch.
[INFO 09-23 20:31:35] ax.service.ax_client: Completed trial 18 with data: {'hartmann6': (-1.508482, 0.1), 'l2norm': (0.977963, 0.1)}.
[INFO 09-23 20:31:37] ax.service.ax_client: Generated new trial 19 with parameters {'x1': 0.425636, 'x2': 0.021245, 'x3': 0.328259, 'x4': 0.532462, 'x5': 0.35233, 'x6': 0.736052} using model BoTorch.
[INFO 09-23 20:31:37] ax.service.ax_client: Completed trial 19 with data: {'hartmann6': (-1.303175, 0.1), 'l2norm': (1.042761, 0.1)}.
The plot below shows the response surface for hartmann6
metric as a function of the x1
, x2
parameters.
The other parameters are fixed in the middle of their respective ranges, which in this example is 0.5 for all of them.
# this could alternately be done with `ax.plot.contour.plot_contour`
render(ax_client.get_contour_plot(param_x="x1", param_y="x2", metric_name="hartmann6"))
[INFO 09-23 20:31:38] ax.service.ax_client: Retrieving contour plot with parameter 'x1' on X-axis and 'x2' on Y-axis, for metric 'hartmann6'. Remaining parameters are affixed to the middle of their range.
The plot below allows toggling between different pairs of parameters to view the contours.
model = ax_client.generation_strategy.model
render(interact_contour(model=model, metric_name="hartmann6"))