{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Service API Example on Hartmann6\n",
"\n",
"The Ax Service API is designed to allow the user to control scheduling of trials and data computation while having an easy to use interface with Ax.\n",
"\n",
"The user iteratively:\n",
"- Queries Ax for candidates\n",
"- Schedules / deploys them however they choose\n",
"- Computes data and logs to Ax\n",
"- Repeat"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
""
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"[INFO 08-29 13:00:04] ipy_plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.\n"
]
}
],
"source": [
"import numpy as np\n",
"\n",
"from ax.plot.contour import plot_contour\n",
"from ax.plot.trace import optimization_trace_single_method\n",
"from ax.service.ax_client import AxClient\n",
"from ax.metrics.branin import branin\n",
"from ax.utils.measurement.synthetic_functions import hartmann6\n",
"from ax.utils.notebook.plotting import render, init_notebook_plotting\n",
"\n",
"init_notebook_plotting()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Initialize client\n",
"\n",
"Create a client object to interface with Ax APIs. By default this runs locally without storage."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"ax = AxClient()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. Set up experiment\n",
"An experiment consists of a **search space** (parameters and parameter constraints) and **optimization configuration** (objective name, minimization setting, and outcome constraints). Note that:\n",
"- Only `name`, `parameters`, and `objective_name` arguments are required.\n",
"- Dictionaries in `parameters` have the following required keys: \"name\" - parameter name, \"type\" - parameter type (\"range\", \"choice\" or \"fixed\"), \"bounds\" for range parameters, \"values\" for choice parameters, and \"value\" for fixed parameters.\n",
"- Dictionaries in `parameters` can optionally include \"value_type\" (\"int\", \"float\", \"bool\" or \"str\"), \"log_scale\" flag for range parameters, and \"is_ordered\" flag for choice parameters.\n",
"- `parameter_constraints` should be a list of strings of form \"p1 >= p2\" or \"p1 + p2 <= some_bound\".\n",
"- `outcome_constraints` should be a list of strings of form \"constrained_metric <= some_bound\"."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[INFO 08-29 13:00:04] ax.service.utils.dispatch: Using Bayesian Optimization generation strategy: GenerationStrategy(name='Sobol+GPEI', steps=[Sobol for 6 arms, GPEI for subsequent arms], generated 0 arm(s) so far). Iterations after 6 will take longer to generate due to model-fitting.\n"
]
}
],
"source": [
"ax.create_experiment(\n",
" name=\"hartmann_test_experiment\",\n",
" parameters=[\n",
" {\n",
" \"name\": \"x1\",\n",
" \"type\": \"range\",\n",
" \"bounds\": [0.0, 1.0],\n",
" \"value_type\": \"float\", # Optional, defaults to inference from type of \"bounds\".\n",
" \"log_scale\": False, # Optional, defaults to False.\n",
" },\n",
" {\n",
" \"name\": \"x2\",\n",
" \"type\": \"range\",\n",
" \"bounds\": [0.0, 1.0],\n",
" },\n",
" {\n",
" \"name\": \"x3\",\n",
" \"type\": \"range\",\n",
" \"bounds\": [0.0, 1.0],\n",
" },\n",
" {\n",
" \"name\": \"x4\",\n",
" \"type\": \"range\",\n",
" \"bounds\": [0.0, 1.0],\n",
" },\n",
" {\n",
" \"name\": \"x5\",\n",
" \"type\": \"range\",\n",
" \"bounds\": [0.0, 1.0],\n",
" },\n",
" {\n",
" \"name\": \"x6\",\n",
" \"type\": \"range\",\n",
" \"bounds\": [0.0, 1.0],\n",
" },\n",
" ],\n",
" objective_name=\"hartmann6\",\n",
" minimize=True, # Optional, defaults to False.\n",
" parameter_constraints=[\"x1 + x2 <= 2.0\"], # Optional.\n",
" outcome_constraints=[\"l2norm <= 1.25\"], # Optional.\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. Define how to evaluate trials\n",
"When using Ax a service, evaluation of parameterizations suggested by Ax is done either locally or, more commonly, using an external scheduler. Below is a dummy evaluation function that outputs data for two metrics \"hartmann6\" and \"l2norm\". Note that all returned metrics correspond to either the `objective_name` set on experiment creation or the metric names mentioned in `outcome_constraints`."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"def evaluate(parameters):\n",
" x = np.array([parameters.get(f\"x{i+1}\") for i in range(6)])\n",
" # In our case, standard error is 0, since we are computing a synthetic function.\n",
" return {\"hartmann6\": (hartmann6(x), 0.0), \"l2norm\": (np.sqrt((x ** 2).sum()), 0.0)}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Result of the evaluation should generally be a mapping of the format: `{metric_name -> (mean, SEM)}`. If there is only one metric in the experiment – the objective – then evaluation function can return a single tuple of mean and SEM, in which case Ax will assume that evaluation corresponds to the objective. It can also return only the mean as a float, in which case Ax will treat SEM as unknown and use a model that can infer it. For more details on evaluation function, refer to the \"Trial Evaluation\" section in the docs."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4. Run optimization loop\n",
"With the experiment set up, we can start the optimization loop.\n",
"\n",
"At each step, the user queries the client for a new trial then submits the evaluation of that trial back to the client.\n",
"\n",
"Note that Ax auto-selects an appropriate optimization algorithm based on the search space. For more advance use cases that require a specific optimization algorithm, pass a `generation_strategy` argument into the `AxClient` constructor. Note that when Bayesian Optimization is used, generating new trials may take a few minutes."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 1/30...\n",
"Running trial 2/30...\n",
"Running trial 3/30...\n",
"Running trial 4/30...\n",
"Running trial 5/30...\n",
"Running trial 6/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 7/30...\n",
"Running trial 8/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 9/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 10/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 11/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 12/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 13/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 14/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 15/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 16/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 17/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 18/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 19/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 20/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 21/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 22/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 23/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 24/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 25/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 26/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 27/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 28/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 29/30...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running trial 30/30...\n"
]
}
],
"source": [
"for i in range(30):\n",
" print(f\"Running trial {i+1}/30...\")\n",
" parameters, trial_index = ax.get_next_trial()\n",
" # Local evaluation here can be replaced with deployment to external system.\n",
" ax.complete_trial(trial_index=trial_index, raw_data=evaluate(parameters))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5. Retrieve best parameters\n",
"\n",
"Once it's complete, we can access the best parameters found, as well as the corresponding metric values."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'x1': 0.41169705380325305,\n",
" 'x2': 0.8953323870029478,\n",
" 'x3': 0.33053037379919475,\n",
" 'x4': 0.5667668159515467,\n",
" 'x5': 0.37275383692522,\n",
" 'x6': 1.5474149616864724e-13}"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"best_parameters, values = ax.get_best_parameters()\n",
"best_parameters"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'hartmann6': -3.065546295723029, 'l2norm': 1.2411824154369033}"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"means, covariances = values\n",
"means"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For comparison, Hartmann6 minimum:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"-3.32237"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"hartmann6.fmin"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 6. Plot the response surface and optimization trace\n",
"Here we arbitrarily select \"x1\" and \"x2\" as the two parameters to plot for both metrics, \"hartmann6\" and \"l2norm\"."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"render(\n",
" plot_contour(\n",
" model=ax.generation_strategy.model, param_x='x1', param_y='x2', metric_name='hartmann6'\n",
" )\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can also plot the optimization trace, showing the progression of finding the point with the optimal objective:"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"# `plot_single_method` expects a 2-d array of means, because it expects to average means from multiple \n",
"# optimization runs, so we wrap out best objectives array in another array.\n",
"best_objectives = np.array([[trial.objective_mean for trial in ax.experiment.trials.values()]])\n",
"best_objective_plot = optimization_trace_single_method(\n",
" y=np.minimum.accumulate(best_objectives, axis=1),\n",
" optimum=hartmann6.fmin,\n",
" title=\"Model performance vs. # of iterations\",\n",
" ylabel=\"Hartmann6\",\n",
")\n",
"render(best_objective_plot)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Special Cases"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Evaluation failure**: should any optimization iterations fail during evaluation, `log_trial_failure` will ensure that the same trial is not proposed again."
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
"_, trial_index = ax.get_next_trial()\n",
"ax.log_trial_failure(trial_index=trial_index)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Adding custom trials**: should there be need to evaluate a specific parameterization, `attach_trial` will add it to the experiment."
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"({'x1': 9.0, 'x2': 9.0, 'x3': 9.0, 'x4': 9.0, 'x5': 9.0, 'x6': 9.0}, 31)"
]
},
"execution_count": 12,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ax.attach_trial(parameters={\"x1\": 9.0, \"x2\": 9.0, \"x3\": 9.0, \"x4\": 9.0, \"x5\": 9.0, \"x6\": 9.0})"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Need to run many trials in parallel**: for optimal results and optimization efficiency, we strongly recommend sequential optimization (generating a few trials, then waiting for them to be completed with evaluation data). However, if your use case needs to dispatch many trials in parallel before they are updated with data and you are running into the *\"All trials for current model have been generated, but not enough data has been observed to fit next model\"* error, instantiate `AxClient` as `AxClient(enforce_sequential_optimization=False)`."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "python3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.3"
}
},
"nbformat": 4,
"nbformat_minor": 2
}