get_countour_plot() not plotting all trials

[davifebba]

Hello, I'm learning Ax and playing with the tutorial code of the Service API, but I noticed that render(ax_client.get_contour_plot()) does not display all trials over the response surface, but displays a total of (number of trials -1) samples. How to display the full set of samples on the response surface?

For example, running an optimization campaign for 10 trials:

from ax.service.ax_client import AxClient, ObjectiveProperties from ax.utils.measurement.synthetic_functions import hartmann6, Branin from ax.utils.notebook.plotting import init_notebook_plotting, render from ax.modelbridge.registry import Models from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy from ax.models.torch.botorch_modular.surrogate import Surrogate from botorch.acquisition.monte_carlo import ( qNoisyExpectedImprovement, ) from botorch.models.gp_regression import FixedNoiseGP

import numpy as np init_notebook_plotting()

gs = GenerationStrategy( steps=[ GenerationStep( # Initialization step # Which model to use for this step model=Models.SOBOL, # How many generator runs (each of which is then made a trial) to produce with this step num_trials=5, # How many trials generated from this step must be COMPLETED # before the next one min_trials_observed=5, ), GenerationStep( # BayesOpt step model=Models.BOTORCH_MODULAR, # No limit on how many generator runs will be produced num_trials=-1, model_kwargs={ # Kwargs to pass to BoTorchModel.__init__ "surrogate": Surrogate(FixedNoiseGP), "botorch_acqf_class": qNoisyExpectedImprovement, }, ), ] )

ax_client = AxClient(generation_strategy = gs)

ax_client.create_experiment( name="hartmann_test_experiment", 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], }, ], objectives={"hartmann6": ObjectiveProperties(minimize=True)}, parameter_constraints=["x1 + x2 <= 2.0"], # Optional. #outcome_constraints=["l2norm <= 1.25"], # Optional. )

def evaluate(parameters): x = np.array([parameters.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)}

for i in range(10): 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=evaluate(parameters))

render(ax_client.get_contour_plot())

[mgarrard]

Hey @davifebba, thanks for the question - i can't see the full data frame that shows the parameterizations, but my hunch is that one of the arms is out of sample -- likely one of the Sobol arms. The contour plots the in-sample trials. Could you paste the full dataframe or check if all the arms are in sample?

[davifebba]

Thanks for replying! Could you please advise on how to check if all the arms are in-sample? From another run of the script above, I have this data frame below:

[mgarrard]

Hi @davifebba, I recreated the notebook in colab and was able to see that you are right the last trial that was ran isn't being populated in the contour plots correctly right now. Will flag this as a bug and look into a solution further - just wanted to update you. In the meantime the plots should still be pretty informative, and the dataframes are also accurate.

[ligerlac]

The problem appears to be that get_contour_plot() displays the model's training data, which does not include the very last trial. I created a PR here: https://github.com/facebook/Ax/pull/2305

[mgarrard]

Hey folks, so sorry for the delay here, I think the best solution moving forward is to refit the model before doing the plotting.

For now this can be done with

ax_client.generation_strategy._fit_current_model(data=ax_client.experiment.lookup_data())

and we'll update the plotting code internally to handle that as well. cc @ligerlac @davifebba

Closing for now, but feel free to re-open if necessary