tutorials/compare_mc_analytic_acquisition.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analytic and MC-based Expected Improvement (EI) acquisition\n",
    "\n",
    "In this tutorial, we compare the analytic and MC-based EI acquisition functions and show both `scipy`- and `torch`-based optimizers for optimizing the acquisition. This tutorial highlights the modularity of botorch and the ability to easily try different acquisition functions and accompanying optimization algorithms on the same fitted model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison of analytic and MC-based EI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from botorch.fit import fit_gpytorch_model\n",
    "from botorch.models import SingleTaskGP\n",
    "from botorch.test_functions import Hartmann\n",
    "from gpytorch.mlls import ExactMarginalLogLikelihood\n",
    "\n",
    "neg_hartmann6 = Hartmann(dim=6, negate=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we generate some random data and fit a SingleTaskGP for a 6-dimensional synthetic test function 'Hartmann6'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_x = torch.rand(10, 6)\n",
    "train_obj = neg_hartmann6(train_x).unsqueeze(-1)\n",
    "model = SingleTaskGP(train_X=train_x, train_Y=train_obj)\n",
    "mll = ExactMarginalLogLikelihood(model.likelihood, model)\n",
    "fit_gpytorch_model(mll);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize an analytic EI acquisition function on the fitted model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from botorch.acquisition import ExpectedImprovement\n",
    "\n",
    "best_value = train_obj.max()\n",
    "EI = ExpectedImprovement(model=model, best_f=best_value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we optimize the analytic EI acquisition function using 50 random restarts chosen from 100 initial raw samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from botorch.optim import optimize_acqf\n",
    "\n",
    "new_point_analytic, _ = optimize_acqf(\n",
    "    acq_function=EI,\n",
    "    bounds=torch.tensor([[0.0] * 6, [1.0] * 6]),\n",
    "    q=1,\n",
    "    num_restarts=20,\n",
    "    raw_samples=100,\n",
    "    options={},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.7490, 0.0804, 0.5910, 0.5143, 0.1870, 0.8092]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_point_analytic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's swap out the analytic acquisition function and replace it with an MC version. Note that we are in the `q = 1` case; for `q > 1`, an analytic version does not exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from botorch.acquisition import qExpectedImprovement\n",
    "from botorch.sampling import SobolQMCNormalSampler\n",
    "\n",
    "\n",
    "sampler = SobolQMCNormalSampler(num_samples=500, seed=0, resample=False)        \n",
    "MC_EI = qExpectedImprovement(\n",
    "    model, best_f=best_value, sampler=sampler\n",
    ")\n",
    "torch.manual_seed(seed=0) # to keep the restart conditions the same\n",
    "new_point_mc, _ = optimize_acqf(\n",
    "    acq_function=MC_EI,\n",
    "    bounds=torch.tensor([[0.0] * 6, [1.0] * 6]),\n",
    "    q=1,\n",
    "    num_restarts=20,\n",
    "    raw_samples=100,\n",
    "    options={},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.7490, 0.0803, 0.5912, 0.5144, 0.1870, 0.8093]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_point_mc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that the two generated points are close."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0002)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.norm(new_point_mc - new_point_analytic)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using a torch optimizer on a stochastic acquisition function\n",
    "\n",
    "We could also optimize using a `torch` optimizer. This is particularly useful for the case of a stochastic acquisition function, which we can obtain by setting `resample=True`. First, we illustrate the usage of `torch.optim.Adam`. In the code snippet below, `gen_batch_initial_candidates` uses a heuristic to select a set of restart locations, `gen_candidates_torch` is a wrapper to the `torch` optimizer for maximizing the acquisition value, and `get_best_candidates` finds the best result amongst the random restarts.\n",
    "\n",
    "Under the hood, `gen_candidates_torch` uses a convergence criterion based on exponential moving averages of the loss. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from botorch.generation import get_best_candidates, gen_candidates_torch\n",
    "from botorch.optim import gen_batch_initial_conditions\n",
    "\n",
    "resampler = SobolQMCNormalSampler(num_samples=500, seed=0, resample=True)        \n",
    "MC_EI_resample = qExpectedImprovement(\n",
    "    model, best_f=best_value, sampler=resampler\n",
    ")\n",
    "bounds = torch.tensor([[0.0] * 6, [1.0] * 6])\n",
    "\n",
    "batch_initial_conditions = gen_batch_initial_conditions(\n",
    "    acq_function=MC_EI_resample,\n",
    "    bounds=bounds,\n",
    "    q=1,\n",
    "    num_restarts=20,\n",
    "    raw_samples=100,\n",
    ")\n",
    "batch_candidates, batch_acq_values = gen_candidates_torch(\n",
    "    initial_conditions=batch_initial_conditions,\n",
    "    acquisition_function=MC_EI_resample,\n",
    "    lower_bounds=bounds[0],\n",
    "    upper_bounds=bounds[1],\n",
    "    optimizer=torch.optim.Adam,\n",
    "    options={\"maxiter\": 500},\n",
    ")\n",
    "new_point_torch_Adam = get_best_candidates(\n",
    "    batch_candidates=batch_candidates, batch_values=batch_acq_values\n",
    ").detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.7493, 0.0783, 0.5926, 0.5091, 0.1921, 0.8094]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_point_torch_Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0078)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.norm(new_point_torch_Adam - new_point_analytic)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By changing the `optimizer` parameter to `gen_candidates_torch`, we can also try `torch.optim.SGD`. Note that without the adaptive step size selection of Adam, basic SGD does worse job at optimizing without further manual tuning of the optimization parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_candidates, batch_acq_values = gen_candidates_torch(\n",
    "    initial_conditions=batch_initial_conditions,\n",
    "    acquisition_function=MC_EI_resample,\n",
    "    lower_bounds=bounds[0],\n",
    "    upper_bounds=bounds[1],\n",
    "    optimizer=torch.optim.SGD,\n",
    "    options={\"maxiter\": 500},\n",
    ")\n",
    "new_point_torch_SGD = get_best_candidates(\n",
    "    batch_candidates=batch_candidates, batch_values=batch_acq_values\n",
    ").detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.7583, 0.1894, 0.5593, 0.8277, 0.1477, 0.8988]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_point_torch_SGD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.3475)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.norm(new_point_torch_SGD - new_point_analytic)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}