Developer API Example on Hartmann6¶

from ax import (
    ChoiceParameter,
    ComparisonOp,
    Experiment,
    FixedParameter,
    Metric,
    Objective,
    OptimizationConfig,
    OrderConstraint,
    OutcomeConstraint,
    ParameterType,
    RangeParameter,
    SearchSpace,
    SumConstraint,
)
from ax.modelbridge.registry import Models
from ax.utils.notebook.plotting import init_notebook_plotting, render

init_notebook_plotting()

from ax.metrics.l2norm import L2NormMetric
from ax.metrics.hartmann6 import Hartmann6Metric


hartmann_search_space = SearchSpace(
    parameters=[
        RangeParameter(
            name=f"x{i}", parameter_type=ParameterType.FLOAT, lower=0.0, upper=1.0
        )
        for i in range(6)
    ]
)

choice_param = ChoiceParameter(
    name="choice", values=["foo", "bar"], parameter_type=ParameterType.STRING
)
fixed_param = FixedParameter(
    name="fixed", value=[True], parameter_type=ParameterType.BOOL
)

sum_constraint = SumConstraint(
    parameters=[
        hartmann_search_space.parameters["x0"],
        hartmann_search_space.parameters["x1"],
    ],
    is_upper_bound=True,
    bound=5.0,
)

order_constraint = OrderConstraint(
    lower_parameter=hartmann_search_space.parameters["x0"],
    upper_parameter=hartmann_search_space.parameters["x1"],
)

param_names = [f"x{i}" for i in range(6)]
optimization_config = OptimizationConfig(
    objective=Objective(
        metric=Hartmann6Metric(name="hartmann6", param_names=param_names),
        minimize=True,
    ),
    outcome_constraints=[
        OutcomeConstraint(
            metric=L2NormMetric(name="l2norm", param_names=param_names, noise_sd=0.2),
            op=ComparisonOp.LEQ,
            bound=1.25,
            relative=False,
        )
    ],
)

from ax import Runner


class MyRunner(Runner):
    def run(self, trial):
        trial_metadata = {"name": str(trial.index)}
        return trial_metadata

exp = Experiment(
    name="test_hartmann",
    search_space=hartmann_search_space,
    optimization_config=optimization_config,
    runner=MyRunner(),
)

from ax.modelbridge.registry import Models

NUM_SOBOL_TRIALS = 5
NUM_BOTORCH_TRIALS = 15

print(f"Running Sobol initialization trials...")
sobol = Models.SOBOL(search_space=exp.search_space)

for i in range(NUM_SOBOL_TRIALS):
    # Produce a GeneratorRun from the model, which contains proposed arm(s) and other metadata
    generator_run = sobol.gen(n=1)
    # Add generator run to a trial to make it part of the experiment and evaluate arm(s) in it
    trial = exp.new_trial(generator_run=generator_run)
    # Start trial run to evaluate arm(s) in the trial
    trial.run()
    # Mark trial as completed to record when a trial run is completed
    # and enable fetching of data for metrics on the experiment
    # (by default, trials must be completed before metrics can fetch their data,
    # unless a metric is explicitly configured otherwise)
    trial.mark_completed()

for i in range(NUM_BOTORCH_TRIALS):
    print(
        f"Running BO trial {i + NUM_SOBOL_TRIALS + 1}/{NUM_SOBOL_TRIALS + NUM_BOTORCH_TRIALS}..."
    )
    # Reinitialize GP+EI model at each step with updated data.
    gpei = Models.BOTORCH_MODULAR(experiment=exp, data=exp.fetch_data())
    generator_run = gpei.gen(n=1)
    trial = exp.new_trial(generator_run=generator_run)
    trial.run()
    trial.mark_completed()

print("Done!")

trial_data = exp.fetch_trials_data([NUM_SOBOL_TRIALS + NUM_BOTORCH_TRIALS - 1])
trial_data.df

exp.fetch_data().df

import numpy as np
from ax.plot.trace import optimization_trace_single_method

# `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.
objective_means = np.array([[trial.objective_mean for trial in exp.trials.values()]])
best_objective_plot = optimization_trace_single_method(
    y=np.minimum.accumulate(objective_means, axis=1),
    optimum=-3.32237,  # Known minimum objective for Hartmann6 function.
)
render(best_objective_plot)

from ax import Data
import pandas as pd


class BoothMetric(Metric):
    def fetch_trial_data(self, trial):
        records = []
        for arm_name, arm in trial.arms_by_name.items():
            params = arm.parameters
            records.append(
                {
                    "arm_name": arm_name,
                    "metric_name": self.name,
                    "trial_index": trial.index,
                    # in practice, the mean and sem will be looked up based on trial metadata
                    # but for this tutorial we will calculate them
                    "mean": (params["x1"] + 2 * params["x2"] - 7) ** 2
                    + (2 * params["x1"] + params["x2"] - 5) ** 2,
                    "sem": 0.0,
                }
            )
        return Data(df=pd.DataFrame.from_records(records))

    def is_available_while_running(self) -> bool:
        return True

from ax.storage.registry_bundle import RegistryBundle

bundle = RegistryBundle(
    metric_clss={BoothMetric: None, L2NormMetric: None, Hartmann6Metric: None},
    runner_clss={MyRunner: None},
)

from ax.storage.json_store.load import load_experiment
from ax.storage.json_store.save import save_experiment

save_experiment(exp, "experiment.json", encoder_registry=bundle.encoder_registry)

loaded_experiment = load_experiment(
    "experiment.json", decoder_registry=bundle.decoder_registry
)

from ax.storage.sqa_store.db import (
    init_engine_and_session_factory,
    get_engine,
    create_all_tables,
)
from ax.storage.sqa_store.load import load_experiment
from ax.storage.sqa_store.save import save_experiment

init_engine_and_session_factory(url="sqlite:///foo3.db")

engine = get_engine()
create_all_tables(engine)

from ax.storage.sqa_store.sqa_config import SQAConfig

exp.name = "new"

sqa_config = SQAConfig(
    json_encoder_registry=bundle.encoder_registry,
    json_decoder_registry=bundle.decoder_registry,
    metric_registry=bundle.metric_registry,
    runner_registry=bundle.runner_registry,
)

save_experiment(exp, config=sqa_config)

load_experiment(exp.name, config=sqa_config)