Skip to content

Commit

Permalink
Ask-tell tutorial
Browse files Browse the repository at this point in the history 
Summary: Tutorial demonstrating ask-tell usage for minimizing hartmann6. Advanced features (constraints, MOO, storage, partial data, etc) not covered.

Differential Revision: D69322316
  • Loading branch information
@mpolson64 @facebook-github-bot
mpolson64 authored and facebook-github-bot committed Feb 19, 2025
1 parent 420c090 commit b9d6678
Show file tree
Hide file tree
Showing 2 changed files with 329 additions and 1 deletion.
322 changes: 322 additions & 0 deletions tutorials/ask_tell/ask_tell.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,322 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Ask-tell Optimization with Ax\n",
"\n",
"We often encounter complex optimization problems in which we wish to tune the values of some parameters to improve metric performance, but have little to no knowledge of their effect on each other.\n",
"We call this class of problems “black-box optimization” and it appears across many disciplines including machine learning, robotics, materials science, and chemistry.\n",
"These problems present a unique challenge, and can be made even more challenging if evaluations are expensive to conduct, time-consuming, or noisy.\n",
"\n",
"We can use Ax to efficiently conduct an experiment in which we \"ask\" for candidate points to evaluate, \"tell\" Ax the results, and repeat.\n",
"We'll uses Ax's `Client`, a tool for managing the state of our experiment, and we'll learn how to define an optimization problem, configure an experiment, run trials, analyze results, and persist the experiment for later use using the `Client`.\n",
"\n",
"Because Ax is a black box optimizer, we can use it to optimize any arbitrary function. In this example we will minimize the [Hartmann6 function](https://www.sfu.ca/~ssurjano/hart6.html), a complicated 6-dimensional function with multiple local minima -- a challenging test case for optimization algorithms that is commonly used in the global optimization literature.\n",
"Looking at its analytic form we can see that it would be incredibly challenging to efficiently find the global minimum either by manual trial-and-error or traditional design of experiments like grid-search or random-search.\n",
"\n",
"$$f(\\mathbf{x})=-\\sum_{i=1}^4 \\alpha_i \\exp \\left(-\\sum_{j=1}^6 A_{i j}\\left(x_j-P_{i j}\\right)^2\\right)$$\n",
"\n",
"where\n",
"\n",
"$$\\alpha=(1.0,1.2,3.0,3.2)^T$$\n",
"\n",
"$$\n",
"\\mathbf{A}=\\left(\\begin{array}{cccccc}10 & 3 & 17 & 3.50 & 1.7 & 8 \\\\ 0.05 & 10 & 17 & 0.1 & 8 & 14 \\\\ 3 & 3.5 & 1.7 & 10 & 17 & 8 \\\\ 17 & 8 & 0.05 & 10 & 0.1 & 14\\end{array}\\right)\n",
"$$\n",
"\n",
"\n",
"$$\n",
"\\mathbf{P}=10^{-4}\\left(\\begin{array}{cccccc}1312 & 1696 & 5569 & 124 & 8283 & 5886 \\\\ 2329 & 4135 & 8307 & 3736 & 1004 & 9991 \\\\ 2348 & 1451 & 3522 & 2883 & 3047 & 6650 \\\\ 4047 & 8828 & 8732 & 5743 & 1091 & 381\\end{array}\\right)\n",
"$$\n",
"\n",
"\n",
"### Learning Objectives\n",
"- Understand the basic concepts of black box optimization\n",
"- Learn how to define an optimization problem using Ax\n",
"- Configure and run an experiment using Ax's `Client`\n",
"- Analyze the results of the optimization\n",
"\n",
"### Prerequisites\n",
"\n",
"* Familiarity with Python and basic programming concepts\n",
"* Understanding of adaptive experimentation and Bayesian optimization (see [Introduction to Adaptive Experimentation](#) and [Introduction to Bayesian Optimization](#))\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 1: Import Necessary Modules\n",
"\n",
"First, ensure you have all the necessary imports:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"from ax.preview.api.client import Client\n",
"from ax.preview.api.configs import (\n",
" ExperimentConfig,\n",
" RangeParameterConfig,\n",
" ParameterType,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 2: Initialize the Client\n",
"\n",
"Create an instance of the `Client` to manage the state of your experiment."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"client = Client()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 3: Configure the Experiment\n",
"\n",
"The `Client` expects a series of `Config`s which define how the experiment will be run.\n",
"\n",
"The Hartmann6 is usually evaluated on the hypercube $x_i \\in (0, 1)$, so we will define six identical `RangeParameterConfig`s with the appropriate bounds, and add these to an `ExperimentConfig` along with other metadata about the experiment.\n",
"\n",
"You can also specify additional features like parameter constraints and parameter scaling. For more on configuring experiments, see this [recipe](#)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Define six float parameters for the Hartmann6 function\n",
"parameters = [\n",
" RangeParameterConfig(\n",
" name=f\"x{i + 1}\", parameter_type=ParameterType.FLOAT, bounds=(0, 1)\n",
" )\n",
" for i in range(6)\n",
"]\n",
"\n",
"# Create an experiment configuration\n",
"experiment_config = ExperimentConfig(\n",
" name=\"hartmann6_experiment\",\n",
" parameters=parameters,\n",
" # The following arguments are optional\n",
" description=\"Optimization of the Hartmann6 function\",\n",
" owner=\"developer\",\n",
")\n",
"\n",
"# Apply the experiment configuration to the client\n",
"client.configure_experiment(experiment_config=experiment_config)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 4: Configure Optimization\n",
"Now, we must set up the optimization objective in `Client`, where `objective` is a string that specifies single, scalarized, and multi-objective goals via SymPy string parsing.\n",
"For example, `\"-loss\"` means minimize `loss` (single objective), `\"ne1 + ne2\"` means maximize the sum of `ne1` and `ne2` (scalarization), and `\"-ne, qps\"` means simultaneously minimize `ne` and maximize `qps` (multi-objective). For more information, see the [string parsing recipe](#)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"metric_name = \"hartmann6\" # this name is used during the optimization loop in Step 5\n",
"objective = f\"-{metric_name}\" # minimization is specified by the negative sign\n",
"\n",
"client.configure_optimization(objective=objective)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 5: Run Trials\n",
"Here, we will configure the ask-tell loop.\n",
"\n",
"We begin by defining our Hartmann6 function as written above.\n",
"Remember, this is just an example problem and any Python function can be substituted here.\n",
"\n",
"Then we will iteratively call `client.get_next_trials` to \"ask\" Ax for a parameterization to evaluate, call `hartmann6` using those parameters, and \"tell\" Ax the result using `client.complete_trial`.\n",
"\n",
"This loop will run multiple trials to optimize the function.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Hartmann6 function\n",
"def hartmann6(x1, x2, x3, x4, x5, x6):\n",
" alpha = np.array([1.0, 1.2, 3.0, 3.2])\n",
" A = np.array([\n",
" [10, 3, 17, 3.5, 1.7, 8],\n",
" [0.05, 10, 17, 0.1, 8, 14],\n",
" [3, 3.5, 1.7, 10, 17, 8],\n",
" [17, 8, 0.05, 10, 0.1, 14]\n",
" ])\n",
" P = 10**-4 * np.array([\n",
" [1312, 1696, 5569, 124, 8283, 5886],\n",
" [2329, 4135, 8307, 3736, 1004, 9991],\n",
" [2348, 1451, 3522, 2883, 3047, 6650],\n",
" [4047, 8828, 8732, 5743, 1091, 381]\n",
" ])\n",
"\n",
" outer = 0.0\n",
" for i in range(4):\n",
" inner = 0.0\n",
" for j, x in enumerate([x1, x2, x3, x4, x5, x6]):\n",
" inner += A[i, j] * (x - P[i, j])**2\n",
" outer += alpha[i] * np.exp(-inner)\n",
" return -outer\n",
"\n",
"hartmann6(0.1, 0.45, 0.8, 0.25, 0.552, 1.0)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Optimization Loop\n",
"\n",
"Having configured the experiment and optimization settings and defined the Python function, proceed with the optimization loop, where several trials are run sequentially."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Number of trials to run\n",
"\n",
"num_trials = 30\n",
"\n",
"# Run trials\n",
"\n",
"for _ in range(num_trials):\n",
" trials = client.get_next_trials(\n",
" maximum_trials=1\n",
" ) # We will request just one trial at a time in this example\n",
" for trial_index, parameters in trials.items():\n",
" x1 = parameters[\"x1\"]\n",
" x2 = parameters[\"x2\"]\n",
" x3 = parameters[\"x3\"]\n",
" x4 = parameters[\"x4\"]\n",
" x5 = parameters[\"x5\"]\n",
" x6 = parameters[\"x6\"]\n",
"\n",
" result = hartmann6(x1, x2, x3, x4, x5, x6)\n",
"\n",
" # Set raw_data as a dictionary with metric names as keys and results as values\n",
"\n",
" raw_data = {metric_name: result}\n",
"\n",
" # Complete the trial with the result\n",
"\n",
" client.complete_trial(trial_index=trial_index, raw_data=raw_data)\n",
" print(f\"Completed trial {trial_index} with {raw_data=}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 6: Analyze Results\n",
"\n",
"After running trials, you can analyze the results.\n",
"Most commonly this means extracting the parameterization from the best performing trial you conducted."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"best_parameters, prediction, index, name = client.get_best_parameterization()\n",
"print(\"Best Parameters:\", best_parameters)\n",
"print(\"Prediction (mean, variance):\", prediction)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 7: Compute Analyses\n",
"\n",
"Ax can also produce a number of analyses to help interpret the results of the experiment via `client.compute_analyses`.\n",
"Users can manually select which analyses to run, or can allow Ax to select which would be most relevant.\n",
"In this case Ax selects the following:\n",
"* **Parrellel Coordinates Plot** shows which parameterizations were evaluated and what metric values were observed -- this is useful for getting a high level overview of how thoroughly the search space was explored and which regions tend to produce which outcomes\n",
"* **Interaction Analysis Plot** shows which parameters have the largest affect on the function and plots the most important parameters as 1 or 2 dimensional surfaces\n",
"* **Summary** lists all trials generated along with their parameterizations, observations, and miscellaneous metadata"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"client.compute_analyses(display=True) # By default Ax will display the AnalysisCards produced by compute_analyses"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Conclusion\n",
"\n",
"This tutorial demonstrates how to use Ax's `Client` for ask-tell optimization of Python functions using the Hartmann6 function as an example.\n",
"You can adjust the function and parameters to suit your specific optimization problem."
]
}
],
"metadata": {
"fileHeader": "",
"fileUid": "9dfaed34-de4d-42ed-8755-25343d677ef0",
"isAdHoc": false,
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"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.12.4"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
8 changes: 7 additions & 1 deletion website/tutorials.json
View file
Original file line number Diff line number Diff line change
@@ -1,2 +1,8 @@
{
}
"Basic usage": [
{
"id": "ask_tell",
"title": "Ask-tell Optimization with Ax"
}
]
}

0 comments on commit b9d6678

Please sign in to comment.