tutorials/information_theoretic_acquisition_functions.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "738920c6",
      "metadata": {},
      "source": [
        "# Information-theoretic acquisition functions"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c6b5d9a9",
      "metadata": {},
      "source": [
        "This notebook illustrates the use of some information-theoretic acquisition functions in BoTorch for single and multi-objective optimization. We present a single-objective example in section 1 and a multi-objective example in section 2. Before introducing these examples, we present an overview on the different approaches and how they are estimated."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "876c4359",
      "metadata": {},
      "source": [
        "## Notation\n",
        "\n",
        "We consider the problem of maximizing a function $f: \\mathbb{X} \\rightarrow \\mathbb{R}^M$. In the single-objective setting ($M=1$), the maximum is defined as usual with respect to the total ordering over the real numbers. In the multi-objective setting ($M>1$), the maximum is defined with respect to the Pareto partial ordering over vectors. By an abuse in notation, we denote the optimal set of inputs and outputs by\n",
        "\n",
        "$$\\mathbb{X}^* = \\text{arg}\\max_{\\mathbf{x} \\in \\mathbb{X}} f(\\mathbf{x}) \\subseteq \\mathbb{X} \\quad \\text{and} \\quad \\mathbb{Y}^* = f(\\mathbb{X}^*) = \\max_{\\mathbf{x} \\in \\mathbb{X}} f(\\mathbf{x}) \\subset \\mathbb{R}^M,$$\n",
        "\n",
        "respectively for both the single and multi-objective setting. We denote the collection of optimal input-output pairs by $(\\mathbb{X}^*, \\mathbb{Y}^*)$."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1499a1cd",
      "metadata": {},
      "source": [
        "## Information-theoretic acquisition functions\n",
        "\n",
        "Information-theoretic (IT) acquisition functions work by quantifying the utility of an input $\\mathbf{x} \\in \\mathbb{X}$ based on how \"informative\" the corresponding observation $\\mathbf{y} \\in \\mathbb{R}^M$ will be in learning more about the distribution of some statistic of the function $S(f)$. Here, we define the notion of information via the mutual information ($\\text{MI}$):\n",
        "\n",
        "\\begin{equation}\n",
        "    \\alpha^{\\text{IT}}(\\mathbf{x}|D_n) \n",
        "    = \\text{MI}(\\mathbf{y}; S(f)| \\mathbf{x}, D_n) \n",
        "    = H[p(\\mathbf{y}|D_n)] - \\mathbb{E}_{p(S(f)|D_n)}[H[p(\\mathbf{y}| \\mathbf{x}, D_n, S(f)]],\n",
        "\\end{equation}\n",
        "\n",
        "where $D_n = \\{(\\mathbf{x}_t, \\mathbf{y}_t)\\}_{t=1,\\dots,n}$ denotes the data set of sampled inputs and observations and the function $H$ denotes the differential entropy $H[p(\\mathbf{x})] = - \\int p(\\mathbf{x}) \\log(p(\\mathbf{x})) d\\mathbf{x}$. The main difference between existing information-theoretic acquisition functions in the literature is the choice of statistic $S$ and the modelling assumptions that are made in order to estimate the resulting acquisition function. In this notebook, we focus on three particular cases of information-theoretic acquisition functions:"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7218d85d",
      "metadata": {},
      "source": [
        "### Predictive Entropy Search (PES)\n",
        "\n",
        "The PES acquisition function [1] considers the problem of learning more about the distribution of the optimal inputs: $S(f) = \\mathbb{X}^*$.\n",
        "\n",
        "\\begin{equation}\n",
        "\\alpha^{\\text{PES}}(\\mathbf{x}|D_n) \n",
        "= \\text{MI}(\\mathbf{y}; \\mathbb{X}^*| \\mathbf{x}, D_n) \n",
        "= H[p(\\mathbf{y}|D_n)] - \\mathbb{E}_{p(\\mathbb{X}^*|D_n)}[H[p(\\mathbf{y}| \\mathbf{x}, D_n, \\mathbb{X}^*)]].\n",
        "\\end{equation}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c7b1d071",
      "metadata": {},
      "source": [
        "### Max-value Entropy Search (MES)\n",
        "\n",
        "The MES acquisition function [2] considers the problem of learning more about the distribution of the optimal outputs: $S(f) = \\mathbb{Y}^*$.\n",
        "\n",
        "\\begin{equation}\n",
        "\\alpha^{\\text{MES}}(\\mathbf{x}|D_n) \n",
        "= \\text{MI}(\\mathbf{y}; \\mathbb{Y}^*| \\mathbf{x}, D_n) \n",
        "= H[p(\\mathbf{y}|D_n)] - \\mathbb{E}_{p(\\mathbb{Y}^*|D_n)}[H[p(\\mathbf{y}| \\mathbf{x}, D_n, \\mathbb{Y}^*)]].\n",
        "\\end{equation}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3c8e68ed",
      "metadata": {},
      "source": [
        "### Joint Entropy Search (JES)\n",
        "\n",
        "The JES acquisition function [3] considers the problem of learning more about the distribution of the optimal inputs and outputs: $S(f) = (\\mathbb{X}^*, \\mathbb{Y}^*)$.\n",
        "\n",
        "\\begin{equation}\n",
        "\\alpha^{\\text{JES}}(\\mathbf{x}|D_n) \n",
        "= \\text{MI}(\\mathbf{y}; (\\mathbb{X}^*, \\mathbb{Y}^*)| \\mathbf{x}, D_n) \n",
        "= H[p(\\mathbf{y}|D_n)] - \\mathbb{E}_{p((\\mathbb{X}^*, \\mathbb{Y}^*)|D_n)}[H[p(\\mathbf{y}| \\mathbf{x}, D_n, (\\mathbb{X}^*, \\mathbb{Y}^*))]].\n",
        "\\end{equation}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f14bbab1",
      "metadata": {},
      "source": [
        "## Estimation\n",
        "\n",
        "In order to estimate the three acquistion functions listed above, we make two simplfying assumptions:\n",
        "\n",
        "**[Assumption 1]** We assume an independent Gaussian process prior on each objective function.\n",
        "\n",
        "**[Assumption 2]** We assume a Gaussian observation likelihood."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c69dfe94",
      "metadata": {},
      "source": [
        "### First term\n",
        "\n",
        "Under the modelling assumptions, the first term in each of the acquisition functions is an entropy of a Gaussian random variable, which is analytically tractable.\n",
        "\n",
        "### Second term\n",
        "\n",
        "The second term in each of the acquisition functions is an expectation of an entropy over an intractable distribution. The expectation can be estimated using Monte Carlo, whilst the entropy has to be approximated using different strategies such as moment-matching.\n",
        "\n",
        "**Monte Carlo.** To sample from the distribution over the optimal points, we can first (approximately) sample a collection of posterior paths $f_j \\sim p(f|D_n)$ and then optimize them to obtain the sample of optimal points $(\\mathbb{X}^*_j, \\mathbb{Y}^*_j)$ for $j=1,\\dots,J$. \n",
        "\n",
        "**PES entropy estimate.** In `qPredictiveEntropySearch` and `qMultiObjectivePredictiveEntropySearch`, we approximate the entropy term arising in PES using the expectation propagation strategy described in [4]. In particular, we first relax the global optimality condition:\n",
        "\n",
        "\\begin{align}\n",
        "    H[p(\\mathbf{y}| \\mathbf{x}, D_n, \\mathbb{X}^*)]\n",
        "    &\\overset{(1)}{=} H[p(\\mathbf{y}| \\mathbf{x}, D_n, f(\\mathbb{X}) \\preceq f(\\mathbb{X}^*))]\n",
        "    \\\\\\\\\n",
        "    &\\overset{(2)}{\\leq} H[p(\\mathbf{y}| \\mathbf{x}, D_n, f(X_n \\cup \\{\\mathbf{x}\\}) \\preceq f(\\mathbb{X}^*))].\n",
        "\\end{align}\n",
        "\n",
        "(1) This statement follows from the observation that conditioning on the optimal points $\\mathbb{X}^*$ is equivalent to knowing that all points lie below the objective values at the optimal inputs: $f(\\mathbb{X}) \\preceq f(\\mathbb{X}^*)$. \n",
        "\n",
        "(2) We replace the global optimality condition with the local optimality condition: $f(X_n \\cup \\{\\mathbf{x}\\}) \\preceq f(\\mathbb{X}^*)$, where $X_n = \\{\\mathbf{x}_t\\}_{t=1,\\dots,n}$. . The upper bound follows from the standard result that conditioning on more information only decreases the entropy: $H(A|B) \\leq H(A)$ for any random variables $A$ and $B$.\n",
        "\n",
        "We then estimate the resulting lower bound of the PES acquisition function by approximating the intractable distribution $p(\\mathbf{y}| \\mathbf{x}, D_n, f(X_n \\cup \\{\\mathbf{x}\\}) \\preceq f(\\mathbb{X}^*))$ with a product of Gaussian random variables, which is fitted via an iterative moment-matching procedure known as expectation propagation. The entropy of this resulting distribution can then be computed analytically.\n",
        "\n",
        "**MES and JES entropy estimate.** In `qLowerBoundMultiObjectiveMaxValueEntropySearch`, `qLowerBoundJointEntropySearch` and `qLowerBoundMultiObjectiveJointEntropySearch`, we approximate the entropy term arising in MES and JES using the strategies described in [3]. These estimates rely on different upper bounds of the entropy term, which results in different lower bounds for the mutual information. These estimates are motivated by the following chain inequalities for the entropy in the JES expression:\n",
        "\n",
        "\\begin{align}\n",
        "    H[p(\\mathbf{y}| \\mathbf{x}, D_n, (\\mathbb{X}^*, \\mathbb{Y}^*))]\n",
        "    &\\overset{(1)}{=} H[p(\\mathbf{y}| \\mathbf{x}, D_n \\cup (\\mathbb{X}^*, \\mathbb{Y}^*), f(\\mathbb{X}) \\preceq \\mathbb{Y}^*)]\n",
        "    \\\\\\\\\n",
        "    &\\overset{(2)}{\\leq} H[p(\\mathbf{y}| \\mathbf{x}, D_n \\cup (\\mathbb{X}^*, \\mathbb{Y}^*), f(\\mathbf{x}) \\preceq \\mathbb{Y}^*)]\n",
        "    \\\\\\\\\n",
        "    &\\overset{(3)}{\\leq} H[\\mathcal{N}(\\mathbf{y}| \\mathbf{m}_{(\\mathbf{x}, (\\mathbb{X}^*, \\mathbb{Y}^*))}, \\mathbf{V}_{(\\mathbf{x}, (\\mathbb{X}^*, \\mathbb{Y}^*))})]\n",
        "    \\\\\\\\\n",
        "    &\\overset{(4)}{\\leq} H[\\mathcal{N}(\\mathbf{y}| \\mathbf{m}_{(\\mathbf{x}, (\\mathbb{X}^*, \\mathbb{Y}^*))}, \\text{diag}(\\mathbf{V}_{(\\mathbf{x}, (\\mathbb{X}^*, \\mathbb{Y}^*))}))],\n",
        "\\end{align}\n",
        "\n",
        "where \n",
        "\n",
        "\\begin{align}\n",
        "    \\mathbf{m}_{(\\mathbf{x}, (\\mathbb{X}^*, \\mathbb{Y}^*))} = \\mathbb{E}[p(\\mathbf{y}| \\mathbf{x}, D_n \\cup (\\mathbb{X}^*, \\mathbb{Y}^*), f(\\mathbf{x}) \\preceq \\mathbb{Y}^*)]\n",
        "\\end{align}\n",
        "\n",
        "\\begin{align}\n",
        "    \\mathbf{V}_{(\\mathbf{x}, (\\mathbb{X}^*, \\mathbb{Y}^*))} = \\mathbb{C}\\text{ov}[p(\\mathbf{y}| \\mathbf{x}, D_n \\cup (\\mathbb{X}^*, \\mathbb{Y}^*), f(\\mathbf{x}) \\preceq \\mathbb{Y}^*)].\n",
        "\\end{align}\n",
        "\n",
        "(1) This statement follows from the observation that conditioning on the optimal points $(\\mathbb{X}^*, \\mathbb{Y}^*)$ is equivalent to knowing that $\\mathbb{X}^*$ maps to $\\mathbb{Y}^*$ and that all points lie below the Pareto front, $f(\\mathbb{X}) \\preceq f(\\mathbb{X}^*) = \\mathbb{Y}^*$. \n",
        "\n",
        "(2) We replace the global optimality condition with the local optimality condition: $f(\\mathbf{x}) \\preceq \\mathbb{Y}^*$. The upper bound follows from the standard result that conditioning on more information only decreases the entropy: $H(A|B) \\leq H(A)$ for any random variables $A$ and $B$.\n",
        "\n",
        "(3) We upper bound the entropy using the standard result that the multivariate Gaussian distribution has the maximum entropy over all distributions supported on $\\mathbb{R}^M$ with the same first two moments.\n",
        "\n",
        "(4) We upper bound the entropy by again using the standard result that conditioning on more information only decreases the entropy.\n",
        "\n",
        "**(Conditioning)** A similar chain of inequalities can be obtained for the entropy in the MES term by replacing the augmented data set $D_n \\cup (\\mathbb{X}^*, \\mathbb{Y}^*)$ with the original data set $D_n$. The only real difference between the JES and MES estimate is whether we condition on the extra samples $(\\mathbb{X}^*_j, \\mathbb{Y}^*_j)$ or not for $j=1,\\dots,J$. As a result of this conditioning, the JES estimate can be more expensive than the MES estimate.\n",
        "\n",
        "**(Noiseless setting)** When the observations are exact, $\\mathbf{y} = f(\\mathbf{x})$, then the entropy term in (2) can be computed exactly. By setting `estimation_type=\"0\"`, we use this estimate. In the setting where there is observation noise, the estimate also includes an ad-hoc correction which can be useful (more details in the appendix of [3]).\n",
        "\n",
        "**(Monte Carlo)** The entropy term in (2) can be estimated using Monte Carlo because the distribution has a tractable density under the assumptions. By setting `estimation_type=\"MC\"`, we use this Monte Carlo estimate.\n",
        "\n",
        "**(Lower bound)** The entropy term in (3) and (4) can be computed exactly. By setting `estimation_type=\"LB\"`, we use this lower bound estimate in (3). By setting `estimation_type=\"LB2\"`, we use lower bound estimate in (4)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "658c9cc3",
      "metadata": {},
      "source": [
        "### Batch\n",
        "\n",
        "For the batch setting, the first term is again analytically tractable. The second term can be estimated using Monte Carlo, whilst the entropy term again has to be estimated.\n",
        "\n",
        "**PES entropy estimate.** In `qPredictiveEntropySearch` and `qMultiObjectivePredictiveEntropySearch`, the entropy term is again approximated using expectation propagation. In particular, we approximate $p(Y| X, D_n, f(X_n \\cup X) \\preceq f(\\mathbb{X}^*))$ with a product of Gaussian random variables. \n",
        "\n",
        "**MES and JES entropy estimate** In `qLowerBoundMultiObjectiveMaxValueEntropySearch`, `qLowerBoundJointEntropySearch` and `qLowerBoundMultiObjectiveJointEntropySearch`, we approximate a lower bound to the MES and JES acquisition function:\n",
        "\n",
        "\\begin{equation}\n",
        "\\alpha^{\\text{LB-MES}}(X|D_n) \n",
        "= \\text{MI}(Y; \\mathbb{Y}^*| X, D_n) \n",
        "= H[p(Y|D_n)] - \\sum_{\\mathbf{x} \\in X} \\mathbb{E}_{p(\\mathbb{Y}^*|D_n)}[H[p(\\mathbf{y}| \\mathbf{x}, D_n, \\mathbb{Y}^*)]],\n",
        "\\end{equation}\n",
        "\n",
        "\\begin{equation}\n",
        "\\alpha^{\\text{LB-JES}}(X|D_n) \n",
        "= \\text{MI}(Y; (\\mathbb{X}^*, \\mathbb{Y}^*)| X, D_n) \n",
        "= H[p(Y|D_n)] - \\sum_{\\mathbf{x} \\in X} \\mathbb{E}_{p((\\mathbb{X}^*, \\mathbb{Y}^*)|D_n)}[H[p(\\mathbf{y}| \\mathbf{x}, D_n, (\\mathbb{X}^*, \\mathbb{Y}^*))]].\n",
        "\\end{equation}\n",
        "\n",
        "The advantage of these expressions is that it allows us to take advantage of the existing entropy estimates for the sequential setting."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1ce75823",
      "metadata": {},
      "source": [
        "## References"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "262131d3",
      "metadata": {},
      "source": [
        "[1] J.M. Hernández-Lobato, M.W. Hoffman and Z. Ghahramani, [**Predictive Entropy Search for Efficient Global Optimization of Black-box Functions**](https://arxiv.org/abs/1406.2541), NeurIPS, 2014.\n",
        "\n",
        "[2] Z. Wang and S. Jegelka, [**Max-value Entropy Search for Efficient Bayesian Optimization**](https://arxiv.org/abs/1703.01968), ICML, 2017.\n",
        "\n",
        "[3] B. Tu, A. Gandy, N. Kantas and B. Shafei, [**Joint Entropy Search for Multi-Objective Bayesian Optimization**](https://arxiv.org/abs/2210.02905), NeurIPS, 2022.\n",
        "\n",
        "[4] C. Hvarfner, F. Hutter and N. Nardi, [**Joint Entropy Search for Maximally-Informed Bayesian Optimization**](https://arxiv.org/abs/2206.04771), NeurIPS, 2022.\n",
        "\n",
        "[5] E. Garrido-Merchán and D. Hernández-Lobato, [**Predictive Entropy Search for Multi-objective Bayesian Optimization with Constraints**](https://www.sciencedirect.com/science/article/abs/pii/S0925231219308525), Neurocomputing, 2019."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7490ac1c",
      "metadata": {},
      "source": [
        "# 1. Single-objective example "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5c3e4976",
      "metadata": {},
      "source": [
        "In this section, we present a simple example in one-dimension with one objective to illustrate the use of these acquisition functions. We first define the objective function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "908e289f",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import torch\n",
        "from botorch.fit import fit_gpytorch_mll\n",
        "from botorch.models.gp_regression import SingleTaskGP\n",
        "from botorch.models.transforms.outcome import Standardize\n",
        "from botorch.utils.sampling import draw_sobol_samples\n",
        "from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood\n",
        "\n",
        "SMOKE_TEST = os.environ.get(\"SMOKE_TEST\")\n",
        "tkwargs = {\"dtype\": torch.double, \"device\": \"cpu\"}\n",
        "\n",
        "\n",
        "def f(x):\n",
        "    p1 = torch.cos(torch.pi * x)\n",
        "    p2 = 10 * torch.sin(torch.pi * x)\n",
        "    p3 = 2 * torch.sin(2 * torch.pi * x)\n",
        "    p4 = 2 * torch.sin(6 * torch.pi * x)\n",
        "    return p1 + p2 + p3 + p4\n",
        "\n",
        "\n",
        "bounds = torch.tensor([[0.0], [1.0]], **tkwargs)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0df52007",
      "metadata": {},
      "source": [
        "We now generate some data and then fit the Gaussian process model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5770f703",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "torch.manual_seed(0)\n",
        "np.random.seed(0)\n",
        "n = 5\n",
        "train_X = draw_sobol_samples(bounds=bounds, n=n, q=1, seed=12345678).squeeze(-2)\n",
        "train_Y = f(train_X)\n",
        "\n",
        "\n",
        "def fit_model(train_X, train_Y, num_outputs):\n",
        "    model = SingleTaskGP(train_X, train_Y, outcome_transform=Standardize(m=num_outputs))\n",
        "    mll = ExactMarginalLogLikelihood(model.likelihood, model)\n",
        "    fit_gpytorch_mll(mll)\n",
        "    return model\n",
        "\n",
        "\n",
        "model = fit_model(train_X=train_X, train_Y=train_Y, num_outputs=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0b6b02f9",
      "metadata": {},
      "source": [
        "We now plot the objective function and the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "877a342b",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "X = torch.linspace(bounds[0, 0], bounds[1, 0], 1000, **tkwargs)\n",
        "mean_fX = model.posterior(X).mean.squeeze(-1).detach().numpy()\n",
        "std_fX = torch.sqrt(model.posterior(X).variance).squeeze(-1).detach().numpy()\n",
        "\n",
        "plt.scatter(train_X, train_Y, color=\"k\", label=\"Observations\")\n",
        "plt.plot(X, f(X), color=\"k\", linewidth=2, label=\"Objective function\")\n",
        "plt.plot(X, mean_fX, color=\"dodgerblue\", linewidth=3, label=\"Posterior model\")\n",
        "plt.fill_between(\n",
        "    X, (mean_fX + 3 * std_fX), (mean_fX - 3 * std_fX), alpha=0.2, color=\"dodgerblue\"\n",
        ")\n",
        "plt.xlabel(\"x\", fontsize=15)\n",
        "plt.ylabel(\"y\", fontsize=15)\n",
        "plt.legend(fontsize=15)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6ec0e247",
      "metadata": {},
      "source": [
        "To compute the information-theoretic acquisition functions, we first need to get some Monte Carlo samples of the optimal inputs and outputs. The method `sample_optimal_points` generates `num_samples` approximate samples of the Gaussian process model and optimizes them sequentially using an optimizer. In the single-objective setting, the number of optimal points (`num_points`) should be set to one. For simplicitly, we consider optimization via random search. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "79e93848",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "from botorch.acquisition.utils import get_optimal_samples\n",
        "\n",
        "num_samples = 12\n",
        "\n",
        "optimal_inputs, optimal_outputs = get_optimal_samples(\n",
        "    model, bounds=bounds, num_optima=num_samples\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "620d538a",
      "metadata": {},
      "source": [
        "We now initialize the information-theoretic acquisition functions. The PES can simply be initialized using just the optimal set of inputs. For the MES and JES acquisition function, we also have to specify the region of integration, which is $\\{\\mathbf{y}: \\mathbf{y} \\preceq \\mathbb{Y}^*\\}$ for a maximization problem. This is done by providing a Tensor of bounds, which is obtained via the method `compute_sample_box_decomposition`.\n",
        "\n",
        "Note that for the MES algorithm, we use the multi-objective implementation `qLowerBoundMultiObjectiveMaxValueEntropySearch`, which implements all the estimation types into one acquisition function. BoTorch alreadys supports many other strategies to estimate the single-objective MES algorithms in `botorch.acquisition.max_value_entropy`, which is described in the other complementary notebooks."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "320b07cc",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "from botorch.acquisition.joint_entropy_search import qJointEntropySearch\n",
        "from botorch.acquisition.max_value_entropy_search import qLowerBoundMaxValueEntropy\n",
        "from botorch.acquisition.predictive_entropy_search import qPredictiveEntropySearch\n",
        "\n",
        "pes = qPredictiveEntropySearch(model=model, optimal_inputs=optimal_inputs)\n",
        "\n",
        "# Here we use the lower bound estimates for the MES and JES\n",
        "# Note that the single-objective MES interface is slightly different,\n",
        "# as it utilizes the Gumbel max-value approximation internally and\n",
        "# therefore does not take the max values as input.\n",
        "mes_lb = qLowerBoundMaxValueEntropy(\n",
        "    model=model,\n",
        "    candidate_set=torch.rand(1000, 1),\n",
        ")\n",
        "jes_lb = qJointEntropySearch(\n",
        "    model=model,\n",
        "    optimal_inputs=optimal_inputs,\n",
        "    optimal_outputs=optimal_outputs,\n",
        "    estimation_type=\"LB\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ec4692e9",
      "metadata": {},
      "source": [
        "To illustrate the acquisition functions, we evaluate it over the whole input space and plot it. As described in [3], the JES should be an upper bound to both the PES and MES, although the estimates might not be."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "382e37f4",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "# the acquisition function call takes a three-dimensional tensor\n",
        "fwd_X = X.unsqueeze(-1).unsqueeze(-1)\n",
        "\n",
        "# make the acquisition functions live on the same scale\n",
        "scale_acqvals = True\n",
        "\n",
        "pes_X = pes(fwd_X).detach().numpy()\n",
        "mes_lb_X = mes_lb(fwd_X).detach().numpy()\n",
        "jes_lb_X = jes_lb(fwd_X).detach().numpy()\n",
        "\n",
        "if scale_acqvals:\n",
        "    pes_X = pes_X / pes_X.max()\n",
        "    mes_lb_X = mes_lb_X / mes_lb_X.max()\n",
        "    jes_lb_X = jes_lb_X / jes_lb_X.max()\n",
        "\n",
        "plt.plot(X, pes_X, color=\"mediumseagreen\", linewidth=3, label=\"PES\")\n",
        "plt.plot(X, mes_lb_X, color=\"crimson\", linewidth=3, label=\"MES-LB\")\n",
        "plt.plot(X, jes_lb_X, color=\"dodgerblue\", linewidth=3, label=\"JES-LB\")\n",
        "\n",
        "plt.vlines(\n",
        "    X[pes_X.argmax()], 0, 1, color=\"mediumseagreen\", linewidth=1.5, linestyle=\"--\"\n",
        ")\n",
        "plt.vlines(X[mes_lb_X.argmax()], 0, 1, color=\"crimson\", linewidth=1.5, linestyle=\":\")\n",
        "plt.vlines(\n",
        "    X[jes_lb_X.argmax()], 0, 1, color=\"dodgerblue\", linewidth=1.5, linestyle=\"--\"\n",
        ")\n",
        "plt.legend(fontsize=15)\n",
        "plt.xlabel(\"$x$\", fontsize=15)\n",
        "plt.ylabel(r\"$\\alpha(x)$\", fontsize=15)\n",
        "plt.title(\"Entropy-based acquisition functions\", fontsize=15)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3ce0f584",
      "metadata": {},
      "source": [
        "To maximize the acquisition function in a standard Bayesian optimization loop, we can use the standard optimization routines. Note that the PES acquisition function might not be differentiable since some operations that may arise during expectation propagation are not differentiable. Therefore, we use a finite difference approach to optimize this acquisition function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f7f639bb",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "from botorch.optim import optimize_acqf\n",
        "\n",
        "# Use finite difference for PES\n",
        "candidate, acq_value = optimize_acqf(\n",
        "    acq_function=pes,\n",
        "    bounds=bounds,\n",
        "    q=1,\n",
        "    num_restarts=4,\n",
        "    raw_samples=256,\n",
        "    options={\"with_grad\": False},\n",
        ")\n",
        "print(\"PES: candidate={}, acq_value={}\".format(candidate, acq_value))\n",
        "\n",
        "candidate, acq_value = optimize_acqf(\n",
        "    acq_function=mes_lb,\n",
        "    bounds=bounds,\n",
        "    q=1,\n",
        "    num_restarts=4,\n",
        "    raw_samples=256,\n",
        ")\n",
        "print(\"MES-LB: candidate={}, acq_value={}\".format(candidate, acq_value))\n",
        "\n",
        "candidate, acq_value = optimize_acqf(\n",
        "    acq_function=jes_lb,\n",
        "    bounds=bounds,\n",
        "    q=1,\n",
        "    num_restarts=4,\n",
        "    raw_samples=256,\n",
        ")\n",
        "print(\"JES-LB: candidate={}, acq_value={}\".format(candidate, acq_value))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e95f0846",
      "metadata": {},
      "source": [
        "# 2. Multi-objective batch example "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "57237806",
      "metadata": {},
      "source": [
        "In this section, we illustrate a simple multi-objective example. First we generate some data and fit the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fabc86e9",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "from botorch.acquisition.multi_objective.utils import (\n",
        "    compute_sample_box_decomposition,\n",
        "    random_search_optimizer,\n",
        "    sample_optimal_points,\n",
        ")\n",
        "from botorch.test_functions.multi_objective import ZDT1\n",
        "\n",
        "d = 4\n",
        "M = 2\n",
        "n = 8\n",
        "\n",
        "if SMOKE_TEST:\n",
        "    q = 2\n",
        "else:\n",
        "    q = 4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "34787908",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "problem = ZDT1(dim=d, num_objectives=M, noise_std=0, negate=True)\n",
        "bounds = problem.bounds.to(**tkwargs)\n",
        "\n",
        "train_X = draw_sobol_samples(bounds=bounds, n=n, q=1, seed=123).squeeze(-2)\n",
        "train_Y = problem(train_X)\n",
        "\n",
        "model = fit_model(train_X=train_X, train_Y=train_Y, num_outputs=M)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b7174710",
      "metadata": {},
      "source": [
        "We now obtain Monte Carlo samples of the optimal inputs and outputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "56bd5f5a",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "num_pareto_samples = 8\n",
        "num_pareto_points = 8\n",
        "\n",
        "# We set the parameters for the random search\n",
        "optimizer_kwargs = {\n",
        "    \"pop_size\": 500,\n",
        "    \"max_tries\": 10,\n",
        "}\n",
        "\n",
        "ps, pf = sample_optimal_points(\n",
        "    model=model,\n",
        "    bounds=bounds,\n",
        "    num_samples=num_pareto_samples,\n",
        "    num_points=num_pareto_points,\n",
        "    optimizer=random_search_optimizer,\n",
        "    optimizer_kwargs=optimizer_kwargs,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "76c35b23",
      "metadata": {},
      "source": [
        "We initialize the acquisition functions as before."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2c7dfaf0",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "from botorch.acquisition.multi_objective.joint_entropy_search import (\n",
        "    qLowerBoundMultiObjectiveJointEntropySearch,\n",
        ")\n",
        "from botorch.acquisition.multi_objective.max_value_entropy_search import (\n",
        "    qLowerBoundMultiObjectiveMaxValueEntropySearch,\n",
        ")\n",
        "from botorch.acquisition.multi_objective.predictive_entropy_search import (\n",
        "    qMultiObjectivePredictiveEntropySearch,\n",
        ")\n",
        "\n",
        "pes = qMultiObjectivePredictiveEntropySearch(model=model, pareto_sets=ps)\n",
        "\n",
        "# Compute the box-decomposition\n",
        "hypercell_bounds = compute_sample_box_decomposition(pf)\n",
        "\n",
        "# # Here we use the lower bound estimates for the MES and JES\n",
        "mes_lb = qLowerBoundMultiObjectiveMaxValueEntropySearch(\n",
        "    model=model,\n",
        "    hypercell_bounds=hypercell_bounds,\n",
        "    estimation_type=\"LB\",\n",
        ")\n",
        "\n",
        "jes_lb = qLowerBoundMultiObjectiveJointEntropySearch(\n",
        "    model=model,\n",
        "    pareto_sets=ps,\n",
        "    pareto_fronts=pf,\n",
        "    hypercell_bounds=hypercell_bounds,\n",
        "    estimation_type=\"LB\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6a6d071b",
      "metadata": {},
      "source": [
        "We now optimize the batch acquistion functions. For the batch PES, we optimize the batch acquisition function directly. Whereas for the MES and JES we use a sequential optimization strategy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ceac58f5",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Use finite difference for PES. This may take some time\n",
        "candidates, acq_values = optimize_acqf(\n",
        "    acq_function=pes,\n",
        "    bounds=bounds,\n",
        "    q=q,\n",
        "    num_restarts=4,\n",
        "    raw_samples=512,\n",
        "    options={\"with_grad\": False},\n",
        ")\n",
        "print(\"PES: \\ncandidates={}\".format(candidates))\n",
        "\n",
        "# Sequentially greedy optimization\n",
        "candidates, acq_values = optimize_acqf(\n",
        "    acq_function=mes_lb,\n",
        "    bounds=bounds,\n",
        "    q=q,\n",
        "    num_restarts=4,\n",
        "    raw_samples=512,\n",
        "    sequential=True,\n",
        ")\n",
        "print(\"MES-LB: \\ncandidates={}\".format(candidates))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d9281308",
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "# Sequentially greedy optimization\n",
        "candidates, acq_values = optimize_acqf(\n",
        "    acq_function=jes_lb,\n",
        "    bounds=bounds,\n",
        "    q=q,\n",
        "    num_restarts=4,\n",
        "    raw_samples=512,\n",
        "    sequential=True,\n",
        ")\n",
        "print(\"JES-LB: \\ncandidates={}\".format(candidates))"
      ]
    }
  ],
  "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.9.16"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}