3.1 Ax Platform Basic Usage

This notebook uses the Ax Platform’s Service API to perform Bayesian optimization. Here, we use an analytic function called the Branin function, which is often used for optimization benchmarking. It takes two parameters as inputs, and returns one value as the output. The task here is to optimize the parameter inputs to minimize the function value.

# only install if we are running in colab
import sys
IN_COLAB = 'google.colab' in sys.modules
if IN_COLAB:
    %pip install ax-platform

import math
from ax.service.ax_client import AxClient, ObjectiveProperties

obj1_name = "branin"

def branin(x1, x2):
    y = float(
        (x2 - 5.1 / (4 * math.pi**2) * x1**2 + 5.0 / math.pi * x1 - 6.0) ** 2
        + 10 * (1 - 1.0 / (8 * math.pi)) * math.cos(x1)
        + 10
    )

    return y

ax_client = AxClient()
ax_client.create_experiment(
    parameters=[
        {"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
        {"name": "x2", "type": "range", "bounds": [0.0, 15.0]},
    ],
    objectives={
        obj1_name: ObjectiveProperties(minimize=True),
    },
)

for _ in range(15):
    parameters, trial_index = ax_client.get_next_trial()
    results = branin(
        parameters["x1"],
        parameters["x2"],
    )
    ax_client.complete_trial(trial_index=trial_index, raw_data=results)

best_parameters, metrics = ax_client.get_best_parameters()

[WARNING 07-19 14:05:57] ax.service.utils.with_db_settings_base: Ax currently requires a sqlalchemy version below 2.0. This will be addressed in a future release. Disabling SQL storage in Ax for now, if you would like to use SQL storage please install Ax with mysql extras via `pip install ax-platform[mysql]`.

[INFO 07-19 14:05:57] 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 07-19 14:05:57] 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 07-19 14:05:57] 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 07-19 14:05:57] ax.service.utils.instantiation: Created search space: SearchSpace(parameters=[RangeParameter(name='x1', parameter_type=FLOAT, range=[-5.0, 10.0]), RangeParameter(name='x2', parameter_type=FLOAT, range=[0.0, 15.0])], parameter_constraints=[]).

[INFO 07-19 14:05:57] ax.modelbridge.dispatch_utils: Using Models.BOTORCH_MODULAR since there is at least one ordered parameter and there are no unordered categorical parameters.

[INFO 07-19 14:05:57] ax.modelbridge.dispatch_utils: Calculating the number of remaining initialization trials based on num_initialization_trials=None max_initialization_trials=None num_tunable_parameters=2 num_trials=None use_batch_trials=False

[INFO 07-19 14:05:57] ax.modelbridge.dispatch_utils: calculated num_initialization_trials=5

[INFO 07-19 14:05:57] ax.modelbridge.dispatch_utils: num_completed_initialization_trials=0 num_remaining_initialization_trials=5

[INFO 07-19 14:05:57] 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 07-19 14:05:57] ax.modelbridge.dispatch_utils: Using Bayesian Optimization generation strategy: GenerationStrategy(name='Sobol+BoTorch', steps=[Sobol for 5 trials, BoTorch for subsequent trials]). Iterations after 5 will take longer to generate due to model-fitting.

[INFO 07-19 14:05:57] ax.service.ax_client: Generated new trial 0 with parameters {'x1': 7.46005, 'x2': 11.050261} using model Sobol.

[INFO 07-19 14:05:57] ax.service.ax_client: Completed trial 0 with data: {'branin': (108.433893, None)}.

[INFO 07-19 14:05:57] ax.service.ax_client: Generated new trial 1 with parameters {'x1': 9.75531, 'x2': 0.75612} using model Sobol.

[INFO 07-19 14:05:57] ax.service.ax_client: Completed trial 1 with data: {'branin': (4.965009, None)}.

[INFO 07-19 14:05:57] ax.service.ax_client: Generated new trial 2 with parameters {'x1': 6.443358, 'x2': 2.343874} using model Sobol.

[INFO 07-19 14:05:57] ax.service.ax_client: Completed trial 2 with data: {'branin': (21.005569, None)}.

[INFO 07-19 14:05:57] ax.service.ax_client: Generated new trial 3 with parameters {'x1': 5.215967, 'x2': 8.146018} using model Sobol.

[INFO 07-19 14:05:57] ax.service.ax_client: Completed trial 3 with data: {'branin': (62.69807, None)}.

[INFO 07-19 14:05:57] ax.service.ax_client: Generated new trial 4 with parameters {'x1': 8.841575, 'x2': 8.930721} using model Sobol.

[INFO 07-19 14:05:57] ax.service.ax_client: Completed trial 4 with data: {'branin': (49.646523, None)}.

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[2], line 27
ax_client.create_experiment(
   parameters=[
       {"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
   (...)
   },
)
for _ in range(15):
---> 27     parameters, trial_index = ax_client.get_next_trial()
   results = branin(
       parameters["x1"],
       parameters["x2"],
   )
   ax_client.complete_trial(trial_index=trial_index, raw_data=results)

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/utils/common/executils.py:163, in retry_on_exception.<locals>.func_wrapper.<locals>.actual_wrapper(*args, **kwargs)
           wait_interval = min(
               MAX_WAIT_SECONDS, initial_wait_seconds * 2 ** (i - 1)
           )
           time.sleep(wait_interval)
--> 163         return func(*args, **kwargs)
# If we are here, it means the retries were finished but
# The error was suppressed. Hence return the default value provided.
return default_return_on_suppression

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/service/ax_client.py:539, in AxClient.get_next_trial(self, ttl_seconds, force, fixed_features)
       raise OptimizationShouldStop(message=global_stopping_message)
try:
   trial = self.experiment.new_trial(
--> 539         generator_run=self._gen_new_generator_run(
           fixed_features=fixed_features
       ),
       ttl_seconds=ttl_seconds,
   )
except MaxParallelismReachedException as e:
   if self._early_stopping_strategy is not None:

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/service/ax_client.py:1790, in AxClient._gen_new_generator_run(self, n, fixed_features)
fixed_feats = (
   InstantiationBase.make_fixed_observation_features(
       fixed_features=fixed_features
   (...)
   else None
)
with with_rng_seed(seed=self._random_seed):
-> 1790     return not_none(self.generation_strategy).gen(
       experiment=self.experiment,
       n=n,
       pending_observations=self._get_pending_observation_features(
           experiment=self.experiment
       ),
       fixed_features=fixed_feats,
   )

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/modelbridge/generation_strategy.py:370, in GenerationStrategy.gen(self, experiment, data, n, pending_observations, **kwargs)
def gen(
   self,
   experiment: Experiment,
   (...)
   **kwargs: Any,
) -> GeneratorRun:
   """Produce the next points in the experiment. Additional kwargs passed to
   this method are propagated directly to the underlying model's `gen`, along
   with the `model_gen_kwargs` set on the current generation node.
   (...)
           resuggesting points that are currently being evaluated.
   """
--> 370     return self._gen_multiple(
       experiment=experiment,
       num_generator_runs=1,
       data=data,
       n=n,
       pending_observations=pending_observations,
       **kwargs,
   )[0]

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/modelbridge/generation_strategy.py:683, in GenerationStrategy._gen_multiple(self, experiment, num_generator_runs, data, n, pending_observations, **model_gen_kwargs)
for _ in range(num_generator_runs):
   try:
--> 683         generator_run = self._curr.gen(
           n=n,
           pending_observations=pending_observations,
           arms_by_signature_for_deduplication=experiment.arms_by_signature,
           **model_gen_kwargs,
       )
   except DataRequiredError as err:
       # Model needs more data, so we log the error and return
       # as many generator runs as we were able to produce, unless
       # no trials were produced at all (in which case its safe to raise).
       if len(generator_runs) == 0:

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/modelbridge/generation_node.py:712, in GenerationStep.gen(self, n, pending_observations, max_gen_draws_for_deduplication, arms_by_signature_for_deduplication, **model_gen_kwargs)
def gen(
   self,
   n: Optional[int] = None,
   (...)
   **model_gen_kwargs: Any,
) -> GeneratorRun:
--> 712     gr = super().gen(
       n=n,
       pending_observations=pending_observations,
       max_gen_draws_for_deduplication=max_gen_draws_for_deduplication,
       arms_by_signature_for_deduplication=arms_by_signature_for_deduplication,
       **model_gen_kwargs,
   )
   gr._generation_step_index = self.index
   return gr

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/modelbridge/generation_node.py:272, in GenerationNode.gen(self, n, pending_observations, max_gen_draws_for_deduplication, arms_by_signature_for_deduplication, **model_gen_kwargs)
# Keep generating until each of `generator_run.arms` is not a duplicate
# of a previous arm, if `should_deduplicate is True`
while should_generate_run:
--> 272     generator_run = self._gen(
       n=n,
       pending_observations=pending_observations,
       **model_gen_kwargs,
   )
   should_generate_run = (
       self.should_deduplicate
       and arms_by_signature_for_deduplication
   (...)
       )
   )
   n_gen_draws += 1

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/modelbridge/generation_node.py:334, in GenerationNode._gen(self, n, pending_observations, **model_gen_kwargs)
if n is None and model_spec.model_gen_kwargs:
   # If `n` is not specified, ensure that the `None` value does not
   # override the one set in `model_spec.model_gen_kwargs`.
   n = model_spec.model_gen_kwargs.get("n", None)
--> 334 return model_spec.gen(
   n=n,
   # For `pending_observations`, prefer the input to this function, as
   # `pending_observations` are dynamic throughout the experiment and thus
   # unlikely to be specified in `model_spec.model_gen_kwargs`.
   pending_observations=pending_observations,
   **model_gen_kwargs,
)

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/modelbridge/model_spec.py:221, in ModelSpec.gen(self, **model_gen_kwargs)
fitted_model = self.fitted_model
model_gen_kwargs = consolidate_kwargs(
   kwargs_iterable=[
       self.model_gen_kwargs,
   (...)
   keywords=get_function_argument_names(fitted_model.gen),
)
--> 221 return fitted_model.gen(**model_gen_kwargs)

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/modelbridge/base.py:786, in ModelBridge.gen(self, n, search_space, optimization_config, pending_observations, fixed_features, model_gen_options)
base_gen_args = self._get_transformed_gen_args(
   search_space=search_space,
   optimization_config=optimization_config,
   pending_observations=pending_observations,
   fixed_features=fixed_features,
)
# Apply terminal transform and gen
--> 786 gen_results = self._gen(
   n=n,
   search_space=base_gen_args.search_space,
   optimization_config=base_gen_args.optimization_config,
   pending_observations=base_gen_args.pending_observations,
   fixed_features=base_gen_args.fixed_features,
   model_gen_options=model_gen_options,
)
observation_features = gen_results.observation_features
best_obsf = gen_results.best_observation_features

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/modelbridge/torch.py:721, in TorchModelBridge._gen(self, n, search_space, pending_observations, fixed_features, model_gen_options, optimization_config)
observation_features = self._array_to_observation_features(
   X=gen_results.points.detach().cpu().clone().numpy(),
   candidate_metadata=gen_results.candidate_metadata,
)
try:
--> 721     xbest = not_none(self.model).best_point(
       search_space_digest=search_space_digest,
       torch_opt_config=torch_opt_config,
   )
except NotImplementedError:
   xbest = None

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/models/torch/botorch_modular/model.py:79, in single_surrogate_only.<locals>.impl(self, *args, **kwargs)
if len(self._surrogates) != 1:
   raise NotImplementedError(
       f"{f.__name__} not implemented for multi-surrogate case. Found "
       f"{self.surrogates=}."
   )
---> 79 return f(self, *args, **kwargs)

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/models/torch/botorch_modular/model.py:479, in BoTorchModel.best_point(self, search_space_digest, torch_opt_config)
@copy_doc(TorchModel.best_point)
@single_surrogate_only
def best_point(
   (...)
   torch_opt_config: TorchOptConfig,
) -> Optional[Tensor]:
   try:
--> 479         return self.surrogate.best_in_sample_point(
           search_space_digest=search_space_digest,
           torch_opt_config=torch_opt_config,
       )[0]
   except ValueError:
       return None

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/models/torch/botorch_modular/surrogate.py:622, in Surrogate.best_in_sample_point(self, search_space_digest, torch_opt_config, options)
if torch_opt_config.is_moo:
   raise NotImplementedError(
       "Best observed point is incompatible with MOO problems."
   )
--> 622 best_point_and_observed_value = best_in_sample_point(
   Xs=self.Xs,
   model=self,
   bounds=search_space_digest.bounds,
   objective_weights=torch_opt_config.objective_weights,
   outcome_constraints=torch_opt_config.outcome_constraints,
   linear_constraints=torch_opt_config.linear_constraints,
   fixed_features=torch_opt_config.fixed_features,
   risk_measure=torch_opt_config.risk_measure,
   options=options,
)
if best_point_and_observed_value is None:
   raise ValueError("Could not obtain best in-sample point.")

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/ax/models/model_utils.py:464, in best_in_sample_point(Xs, model, bounds, objective_weights, outcome_constraints, linear_constraints, fixed_features, risk_measure, options)
   raise UnsupportedError(f"Unknown best point method {method}.")
i = np.argmax(utility)
--> 464 if utility[i] == -np.Inf:
   return None
else:

File ~/checkouts/readthedocs.org/user_builds/ac-microcourses/envs/stable/lib/python3.11/site-packages/numpy/__init__.py:397, in __getattr__(attr)
   raise AttributeError(__former_attrs__[attr])
if attr in __expired_attributes__:
--> 397     raise AttributeError(
       f"`np.{attr}` was removed in the NumPy 2.0 release. "
       f"{__expired_attributes__[attr]}"
   )
if attr == "chararray":
   warnings.warn(
       "`np.chararray` is deprecated and will be removed from "
       "the main namespace in the future. Use an array with a string "
       "or bytes dtype instead.", DeprecationWarning, stacklevel=2)

AttributeError: `np.Inf` was removed in the NumPy 2.0 release. Use `np.inf` instead.

print(best_parameters, metrics)

{'x1': 9.42045994053506, 'x2': 2.522998247072662} ({'branin': 0.6237735107536508}, {'branin': {'branin': 0.27541560401879245}})

After running this, return to the Bayesian optimization tutorial notebook.