{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3.1 Ax Platform Basic Usage\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/AccelerationConsortium/ac-microcourses/blob/main/docs/courses/hello-world/1.3.2-ax-service-api-basic.ipynb)\n",
    "\n",
    "This notebook uses the Ax Platform's Service API to perform Bayesian optimization. Here, we use an analytic function called [the Branin function](https://chat.openai.com/share/63e1255e-6d7d-4dab-81a3-7913509f2425), 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# only install if we are running in colab\n",
    "import sys\n",
    "IN_COLAB = 'google.colab' in sys.modules\n",
    "if IN_COLAB:\n",
    "    %pip install ax-platform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO 01-05 19:01:31] 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.\n",
      "[INFO 01-05 19:01:31] 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.\n",
      "[INFO 01-05 19:01:31] 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.\n",
      "[INFO 01-05 19:01:31] 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=[]).\n",
      "[INFO 01-05 19:01:31] ax.modelbridge.dispatch_utils: Using Models.BOTORCH_MODULAR since there are more ordered parameters than there are categories for the unordered categorical parameters.\n",
      "[INFO 01-05 19:01:31] 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\n",
      "[INFO 01-05 19:01:31] ax.modelbridge.dispatch_utils: calculated num_initialization_trials=5\n",
      "[INFO 01-05 19:01:31] ax.modelbridge.dispatch_utils: num_completed_initialization_trials=0 num_remaining_initialization_trials=5\n",
      "[INFO 01-05 19:01:31] ax.modelbridge.dispatch_utils: `verbose`, `disable_progbar`, and `jit_compile` are not yet supported when using `choose_generation_strategy` with ModularBoTorchModel, dropping these arguments.\n",
      "[INFO 01-05 19:01:31] 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.\n",
      "[INFO 01-05 19:01:31] ax.service.ax_client: Generated new trial 0 with parameters {'x1': 8.767505, 'x2': 9.342964}.\n",
      "[INFO 01-05 19:01:31] ax.service.ax_client: Completed trial 0 with data: {'branin': (56.664856, None)}.\n",
      "[INFO 01-05 19:01:31] ax.service.ax_client: Generated new trial 1 with parameters {'x1': 5.274806, 'x2': 9.295323}.\n",
      "[INFO 01-05 19:01:31] ax.service.ax_client: Completed trial 1 with data: {'branin': (80.666377, None)}.\n",
      "[INFO 01-05 19:01:31] ax.service.ax_client: Generated new trial 2 with parameters {'x1': 7.285222, 'x2': 14.839929}.\n",
      "[INFO 01-05 19:01:31] ax.service.ax_client: Completed trial 2 with data: {'branin': (199.542398, None)}.\n",
      "[INFO 01-05 19:01:31] ax.service.ax_client: Generated new trial 3 with parameters {'x1': -2.007367, 'x2': 1.319069}.\n",
      "[INFO 01-05 19:01:31] ax.service.ax_client: Completed trial 3 with data: {'branin': (76.437848, None)}.\n",
      "[INFO 01-05 19:01:31] ax.service.ax_client: Generated new trial 4 with parameters {'x1': 7.15399, 'x2': 12.869959}.\n",
      "[INFO 01-05 19:01:31] ax.service.ax_client: Completed trial 4 with data: {'branin': (151.775043, None)}.\n",
      "[INFO 01-05 19:01:32] ax.service.ax_client: Generated new trial 5 with parameters {'x1': 9.809306, 'x2': 5.469131}.\n",
      "[INFO 01-05 19:01:32] ax.service.ax_client: Completed trial 5 with data: {'branin': (8.125139, None)}.\n",
      "[INFO 01-05 19:01:34] ax.service.ax_client: Generated new trial 6 with parameters {'x1': 9.689023, 'x2': 0.1538}.\n",
      "[INFO 01-05 19:01:34] ax.service.ax_client: Completed trial 6 with data: {'branin': (7.249581, None)}.\n",
      "[INFO 01-05 19:01:35] ax.service.ax_client: Generated new trial 7 with parameters {'x1': 6.949056, 'x2': 2.269606}.\n",
      "[INFO 01-05 19:01:35] ax.service.ax_client: Completed trial 7 with data: {'branin': (18.741474, None)}.\n",
      "[INFO 01-05 19:01:36] ax.service.ax_client: Generated new trial 8 with parameters {'x1': 10.0, 'x2': 2.707994}.\n",
      "[INFO 01-05 19:01:36] ax.service.ax_client: Completed trial 8 with data: {'branin': (2.030143, None)}.\n",
      "[INFO 01-05 19:01:39] ax.service.ax_client: Generated new trial 9 with parameters {'x1': 10.0, 'x2': 3.702972}.\n",
      "[INFO 01-05 19:01:39] ax.service.ax_client: Completed trial 9 with data: {'branin': (2.433162, None)}.\n",
      "[INFO 01-05 19:01:40] ax.service.ax_client: Generated new trial 10 with parameters {'x1': -5.0, 'x2': 9.827796}.\n",
      "[INFO 01-05 19:01:40] ax.service.ax_client: Completed trial 10 with data: {'branin': (66.886933, None)}.\n",
      "[INFO 01-05 19:01:42] ax.service.ax_client: Generated new trial 11 with parameters {'x1': 9.243517, 'x2': 3.333914}.\n",
      "[INFO 01-05 19:01:42] ax.service.ax_client: Completed trial 11 with data: {'branin': (1.570388, None)}.\n",
      "[INFO 01-05 19:01:43] ax.service.ax_client: Generated new trial 12 with parameters {'x1': 9.42046, 'x2': 2.522998}.\n",
      "[INFO 01-05 19:01:43] ax.service.ax_client: Completed trial 12 with data: {'branin': (0.400643, None)}.\n",
      "[INFO 01-05 19:01:45] ax.service.ax_client: Generated new trial 13 with parameters {'x1': -5.0, 'x2': 15.0}.\n",
      "[INFO 01-05 19:01:45] ax.service.ax_client: Completed trial 13 with data: {'branin': (17.5083, None)}.\n",
      "[INFO 01-05 19:01:47] ax.service.ax_client: Generated new trial 14 with parameters {'x1': -2.237367, 'x2': 15.0}.\n",
      "[INFO 01-05 19:01:48] ax.service.ax_client: Completed trial 14 with data: {'branin': (27.030623, None)}.\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "from ax.service.ax_client import AxClient, ObjectiveProperties\n",
    "\n",
    "obj1_name = \"branin\"\n",
    "\n",
    "def branin(x1, x2):\n",
    "    y = float(\n",
    "        (x2 - 5.1 / (4 * math.pi**2) * x1**2 + 5.0 / math.pi * x1 - 6.0) ** 2\n",
    "        + 10 * (1 - 1.0 / (8 * math.pi)) * math.cos(x1)\n",
    "        + 10\n",
    "    )\n",
    "\n",
    "    return y\n",
    "\n",
    "ax_client = AxClient()\n",
    "ax_client.create_experiment(\n",
    "    parameters=[\n",
    "        {\"name\": \"x1\", \"type\": \"range\", \"bounds\": [-5.0, 10.0]},\n",
    "        {\"name\": \"x2\", \"type\": \"range\", \"bounds\": [0.0, 15.0]},\n",
    "    ],\n",
    "    objectives={\n",
    "        obj1_name: ObjectiveProperties(minimize=True),\n",
    "    },\n",
    ")\n",
    "\n",
    "for _ in range(15):\n",
    "    parameters, trial_index = ax_client.get_next_trial()\n",
    "    results = branin(\n",
    "        parameters[\"x1\"],\n",
    "        parameters[\"x2\"],\n",
    "    )\n",
    "    ax_client.complete_trial(trial_index=trial_index, raw_data=results)\n",
    "\n",
    "best_parameters, metrics = ax_client.get_best_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'x1': 9.42045994053506, 'x2': 2.522998247072662} ({'branin': 0.6237735107536508}, {'branin': {'branin': 0.27541560401879245}})\n"
     ]
    }
   ],
   "source": [
    "print(best_parameters, metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After running this, return to the [Bayesian optimization tutorial notebook](./1.3-bayesian-optimization.ipynb)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ac-microcourses",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}