{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "47de02e1",
   "metadata": {},
   "source": [
    "# Running Tune experiments with AxSearch\n",
    "In this tutorial we introduce Ax, while running a simple Ray Tune experiment. Tune’s Search Algorithms integrate with Ax and, as a result, allow you to seamlessly scale up a Ax optimization process - without sacrificing performance.\n",
    "\n",
    "Ax is a platform for optimizing any kind of experiment, including machine learning experiments, A/B tests, and simulations. Ax can optimize discrete configurations (e.g., variants of an A/B test) using multi-armed bandit optimization, and continuous/ordered configurations (e.g. float/int parameters) using Bayesian optimization. Results of A/B tests and simulations with reinforcement learning agents often exhibit high amounts of noise. Ax supports state-of-the-art algorithms which work better than traditional Bayesian optimization in high-noise settings. Ax also supports multi-objective and constrained optimization which are common to real-world problems (e.g. improving load time without increasing data use). Ax belongs to the domain of  \"derivative-free\" and \"black-box\" optimization.\n",
    "\n",
    "In this example we minimize a simple objective to briefly demonstrate the usage of AxSearch with Ray Tune via `AxSearch`. It's useful to keep in mind that despite the emphasis on machine learning experiments, Ray Tune optimizes any implicit or explicit objective. Here we assume `ax-platform==0.2.4` library is installed withe python version >= 3.7. To learn more, please refer to the [Ax website](https://ax.dev/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "297d8b18",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# !pip install ray[tune]\n",
    "!pip install ax-platform==0.2.4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59b1e0d1",
   "metadata": {},
   "source": [
    "Click below to see all the imports we need for this example.\n",
    "You can also launch directly into a Binder instance to run this notebook yourself.\n",
    "Just click on the rocket symbol at the top of the navigation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbae6dbe",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "\n",
    "import ray\n",
    "from ray import tune\n",
    "from ray.tune.suggest.ax import AxSearch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b2b6af7",
   "metadata": {},
   "source": [
    "Let's start by defining a classic benchmark for global optimization.\n",
    "The form here is explicit for demonstration, yet it is typically a black-box.\n",
    "We artificially sleep for a bit (`0.02` seconds) to simulate a long-running ML experiment.\n",
    "This setup assumes that we're running multiple `step`s of an experiment and try to tune 6-dimensions of the `x` hyperparameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f7fbe0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def landscape(x):\n",
    "    \"\"\"\n",
    "    Hartmann 6D function containing 6 local minima.\n",
    "    It is a classic benchmark for developing global optimization algorithms.\n",
    "    \"\"\"\n",
    "    alpha = np.array([1.0, 1.2, 3.0, 3.2])\n",
    "    A = np.array(\n",
    "        [\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",
    "    )\n",
    "    P = 10 ** (-4) * np.array(\n",
    "        [\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",
    "    y = 0.0\n",
    "    for j, alpha_j in enumerate(alpha):\n",
    "        t = 0\n",
    "        for k in range(6):\n",
    "            t += A[j, k] * ((x[k] - P[j, k]) ** 2)\n",
    "        y -= alpha_j * np.exp(-t)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b1ae9df",
   "metadata": {},
   "source": [
    "Next, our `objective` function takes a Tune `config`, evaluates the `landscape` of our experiment in a training loop,\n",
    "and uses `tune.report` to report the `landscape` back to Tune."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c3f252e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(config):\n",
    "    for i in range(config[\"iterations\"]):\n",
    "        x = np.array([config.get(\"x{}\".format(i + 1)) for i in range(6)])\n",
    "        tune.report(\n",
    "            timesteps_total=i, landscape=landscape(x), l2norm=np.sqrt((x ** 2).sum())\n",
    "        )\n",
    "        time.sleep(0.02)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9982d95",
   "metadata": {},
   "source": [
    "Next we define a search space. The critical assumption is that the optimal hyperparamters live within this space. Yet, if the space is very large, then those hyperparamters may be difficult to find in a short amount of time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30f75f5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_space = {\n",
    "    \"iterations\":100,\n",
    "    \"x1\": tune.uniform(0.0, 1.0),\n",
    "    \"x2\": tune.uniform(0.0, 1.0),\n",
    "    \"x3\": tune.uniform(0.0, 1.0),\n",
    "    \"x4\": tune.uniform(0.0, 1.0),\n",
    "    \"x5\": tune.uniform(0.0, 1.0),\n",
    "    \"x6\": tune.uniform(0.0, 1.0)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "106d8578",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "ray.init(configure_logging=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "932f74e6",
   "metadata": {},
   "source": [
    "Now we define the search algorithm from `AxSearch`. If you want to constrain your parameters or even the space of outcomes, that can be easily done by passing the argumentsas below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34dd5c95",
   "metadata": {},
   "outputs": [],
   "source": [
    "algo = AxSearch(\n",
    "    parameter_constraints=[\"x1 + x2 <= 2.0\"],\n",
    "    outcome_constraints=[\"l2norm <= 1.25\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6d18a99",
   "metadata": {},
   "source": [
    "We also use `ConcurrencyLimiter` to constrain to 4 concurrent trials. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcd905ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "algo = tune.suggest.ConcurrencyLimiter(algo, max_concurrent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10fd5427",
   "metadata": {},
   "source": [
    "The number of samples is the number of hyperparameter combinations that will be tried out. This Tune run is set to `1000` samples.\n",
    "You can decrease this if it takes too long on your machine, or you can set a time limit easily through `stop` argument in `tune.run()` as we will show here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c53349a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 100\n",
    "stop_timesteps = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c661045",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# Reducing samples for smoke tests\n",
    "num_samples = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91076c5a",
   "metadata": {},
   "source": [
    "Finally, we run the experiment to find the global minimum of the provided landscape (which contains 5 false minima). The argument to metric, `\"landscape\"`, is provided via the `objective` function's `tune.report`. The experiment `\"min\"`imizes the \"mean_loss\" of the `landscape` by searching within `search_space` via `algo`, `num_samples` times or when `\"timesteps_total\": stop_timesteps`. This previous sentence is fully characterizes the search problem we aim to solve. With this in mind, notice how efficient it is to execute `tune.run()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f519d63",
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis = tune.run(\n",
    "    objective,\n",
    "    name=\"ax\",\n",
    "    metric=\"landscape\",\n",
    "    mode=\"min\",\n",
    "    search_alg=algo,\n",
    "    num_samples=num_samples,\n",
    "    config=search_space,\n",
    "    stop={\"timesteps_total\": stop_timesteps}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "860b53b0",
   "metadata": {},
   "source": [
    "And now we have the hyperparameters found to minimize the mean loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12906421",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Best hyperparameters found were: \", analysis.best_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68872424",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "ray.shutdown()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}