{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyVBMC Example 3: Output diagnostics and saving results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we demonstrate extended usage of PyVBMC. We will take a brief look at PyVBMC's diagnostic output, and show you how to save the results of optimization to disk.\n",
    "\n",
    "This notebook is Part 3 of a series of notebooks in which we present various example usages for VBMC with the PyVBMC package. The code used in this example is available as a script [here](https://github.com/acerbilab/pyvbmc/blob/main/examples/scripts/pyvbmc_example_3_full_code.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.stats as scs\n",
    "from scipy.optimize import minimize\n",
    "from pyvbmc import VBMC\n",
    "from pyvbmc.formatting import format_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Model definition and setup\n",
    "\n",
    "For demonstration purposes, we will run PyVBMC with a restricted budget of function evaluations, insufficient to achieve convergence. Then we will inspect the output diagnostics, and resume optimization.\n",
    "\n",
    "We use a higher-dimensional analogue of the same toy target function in Example 1, a broad [Rosenbrock's banana function](https://en.wikipedia.org/wiki/Rosenbrock_function) in $D = 4$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "D = 4  # A four-dimensional problem\n",
    "prior_mu = np.zeros(D)\n",
    "prior_var = 3 * np.ones(D)\n",
    "\n",
    "\n",
    "def log_prior(theta):\n",
    "    \"\"\"Multivariate normal prior on theta.\"\"\"\n",
    "    cov = np.diag(prior_var)\n",
    "    return scs.multivariate_normal(prior_mu, cov).logpdf(theta)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The likelihood function of your model will in general depend on the observed data. This data can be fixed as a global variable, as we did directly above for `prior_mu` and `prior_var`. It can also be defined by a default second argument: to PyVBMC there is no difference so long as the function can be called with only a single argument (the parameters `theta`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_likelihood(theta, data=np.ones(D)):\n",
    "    \"\"\"D-dimensional Rosenbrock's banana function.\"\"\"\n",
    "    # In this simple demo the data just translates the parameters:\n",
    "    theta = np.atleast_2d(theta)\n",
    "    theta = theta + data\n",
    "\n",
    "    x, y = theta[:, :-1], theta[:, 1:]\n",
    "    return -np.sum((x**2 - y) ** 2 + (x - 1) ** 2 / 100, axis=1)\n",
    "\n",
    "\n",
    "def log_joint(theta, data=np.ones(D)):\n",
    "    \"\"\"log-density of the joint distribution.\"\"\"\n",
    "    return log_likelihood(theta, data) + log_prior(theta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "LB = np.full(D, -np.inf)  # Lower bounds\n",
    "UB = np.full(D, np.inf)  # Upper bounds\n",
    "PLB = np.full(D, prior_mu - np.sqrt(prior_var))  # Plausible lower bounds\n",
    "PUB = np.full(D, prior_mu + np.sqrt(prior_var))  # Plausible upper bounds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a typical inference scenario, we recommend starting from a \"good\" point (i.e. one near the mode). We can run a  quick preliminary optimization, though a more extensive optimization would not harm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(41)\n",
    "x0 = np.random.uniform(PLB, PUB)  # Random point inside plausible box\n",
    "x0 = minimize(\n",
    "    lambda t: -log_joint(t),\n",
    "    x0,\n",
    "    bounds=[\n",
    "        (-np.inf, np.inf),\n",
    "        (-np.inf, np.inf),\n",
    "        (-np.inf, np.inf),\n",
    "        (-np.inf, np.inf),\n",
    "    ],\n",
    ").x\n",
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reshaping x0 to row vector.\n",
      "Reshaping lower bounds to (1, 4).\n",
      "Reshaping upper bounds to (1, 4).\n",
      "Reshaping plausible lower bounds to (1, 4).\n",
      "Reshaping plausible upper bounds to (1, 4).\n"
     ]
    }
   ],
   "source": [
    "# Limit number of function evaluations\n",
    "options = {\n",
    "    \"max_fun_evals\": 10 * D,\n",
    "}\n",
    "# We can specify either the log-joint, or the log-likelihood and log-prior.\n",
    "# In other words, the following lines are equivalent:\n",
    "vbmc = VBMC(\n",
    "    log_likelihood,\n",
    "    x0,\n",
    "    LB,\n",
    "    UB,\n",
    "    PLB,\n",
    "    PUB,\n",
    "    options=options,\n",
    "    log_prior=log_prior,\n",
    ")\n",
    "# vbmc = VBMC(\n",
    "#     log_joint,\n",
    "#     x0, LB, UB, PLB, PUB, options=options,\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(PyVBMC expects the bounds to be `(1, D)` row vectors, and the initial point(s) to be of shape `(n, D)`, but it will accept and re-shape vectors of shape `(D,)` as well.)\n",
    "\n",
    "## 2. Running the model and checking convergence diagnostics\n",
    "\n",
    "Now we run PyVBMC with a very small budget of 40 function evaluations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Beginning variational optimization assuming EXACT observations of the log-joint.\n",
      " Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action\n",
      "     0         10          -3.64         1.20    230538.53        2        inf     start warm-up\n",
      "     1         15          -3.40         2.05        15.56        2        inf     \n",
      "     2         20          -3.42         2.08        14.12        2        242     \n",
      "     3         25          35.04        70.57      2305.69        2   3.88e+04     \n",
      "     4         30          39.27        71.14      4950.88        2   8.28e+04     trim data\n",
      "     5         35          -3.39         1.05       227.02        2   3.93e+03     \n",
      "     6         40           3.33        12.07        27.91        2        528     \n",
      "   inf         40          -3.13         1.00         0.20       50   3.93e+03     finalize\n",
      "Inference terminated: reached maximum number of function evaluations options.max_fun_evals.\n",
      "Estimated ELBO: -3.130 +/-0.995.\n",
      "Caution: Returned variational solution may have not converged.\n"
     ]
    }
   ],
   "source": [
    "vp, results = vbmc.optimize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyVBMC is warning us that convergence is doubtful. We can look at the output for more information and diagnostics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n"
     ]
    }
   ],
   "source": [
    "print(results[\"success_flag\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`False` means that PyVBMC has not converged to a stable solution within the given number of function evaluations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    'function': '<function VBMC._init_log_joint.<locals>.log_joint at 0x7f00ea0d79d0>',\n",
      "    'problem_type': 'unconstrained',\n",
      "    'iterations': 6,\n",
      "    'func_count': 40,\n",
      "    'best_iter': 5,\n",
      "    'train_set_size': 31,\n",
      "    'components': 50,\n",
      "    'r_index': 3929.341498249543,\n",
      "    'convergence_status': 'no',\n",
      "    'overhead': nan,\n",
      "    'rng_state': 'rng',\n",
      "    'algorithm': 'Variational Bayesian Monte Carlo',\n",
      "    'version': '0.1.0',\n",
      "    'message': 'Inference terminated: reached maximum number of function evaluations options.max_fun_evals.',\n",
      "    'elbo': -3.129758475814559,\n",
      "    'elbo_sd': 0.995262203269921,\n",
      "    'success_flag': False,\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(format_dict(results))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the `info` dictionary:\n",
    "- the `convergence_status` field says 'no' (probable lack of convergence);\n",
    "- the reliability index `r_index` is 3.68, (should be less than 1).\n",
    "Our diagnostics tell us that this run has not converged, suggesting to increase the budget.\n",
    "\n",
    "Note that convergence to a solution does not mean that it is a _good_ solution. You should always check the returned variational posteriors, and ideally should compare across multiple runs of PyVBMC."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Saving results\n",
    "\n",
    "We can also save the `VBMC` instance to disk and reload it later, in order to check the results and convergence diagnostics, sample from the posterior, or resume the optimization from checkpoint etc. If you are only interested in the final (best) variational solution, as opposed to the full iteration history of the optimization, then you may wish to save only the final `VariationalPosterior` instead.\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-info\"><b>Note:</b> Some complex attributes of the <code style=\"background:LightBlue\">VBMC</code> instance — such as the stored function(s) representing the log joint density — may not behave as expected when saved and loaded by different Python minor versions (e.g. <code style=\"background:LightBlue\">3.9</code> and <code style=\"background:LightBlue\">3.10</code>), due to differing dependencies. The instance should still load, and its static data will remain, but if you plan to resume optimization as shown here, then we suggest you use the same version of Python to save and load the <code style=\"background:LightBlue\">VBMC</code> instance.</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we specify `overwrite=True`, since we don't care about overwriting our\n",
    "# test file. By default, `overwrite=False` and PyVBMC will raise an error if\n",
    "# the file already exists.\n",
    "vbmc.save(\"vbmc_test_save.pkl\", overwrite=True)\n",
    "# We could also save just the final variational posterior:\n",
    "# vbmc.vp.save(\"vp_test_save.pkl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.  Loading results and resuming the optimization process"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `VBMC.load(file)` class method will load a previously-saved `VBMC` instance from the specified `file`. We can load the instance saved above and resume the optimization process, increasing maximum number of function evaluations. The default budget is $50(d+2)$ evaluations, where $d$ is the dimension of the parameter space (though this example will not require the full budget). We can change this or other options by passing the `new_options` keyword to `VBMC.load(...)`. Here we increase `max_fun_evals` and resume the optimization process:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Beginning variational optimization assuming EXACT observations of the log-joint.\n",
      "Continuing optimization from previous state.\n",
      " Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action\n",
      "     7         45          -4.84         1.80         4.46        2        108     \n",
      "     8         50          -4.87         0.33         7.65        2        129     end warm-up\n",
      "     9         55          73.00        57.41        89.44        2   1.75e+03     \n",
      "    10         60          -9.39         4.61        20.33        2        629     \n",
      "    11         65          -4.70         0.12         2.32        3       54.7     \n",
      "    12         70          -4.53         0.02         0.12        4       2.58     \n",
      "    13         75          -4.42         0.01         0.04        5       1.13     \n",
      "    14         75          -4.12         0.22         0.19        6       4.82     rotoscale\n",
      "    15         80          -4.51         0.07         0.05        6       2.45     \n",
      "    16         85          -4.42         0.02         0.02        6      0.667     \n",
      "    17         90          -4.39         0.01         0.01        6      0.382     \n",
      "    18         95          -4.38         0.01         0.01        9      0.142     \n",
      "    19        100          -4.30         0.01         0.01       12      0.451     \n",
      "    20        105          -4.25         0.01         0.01       15      0.297     rotoscale, undo rotoscale\n",
      "    21        110          -4.21         0.00         0.01       18      0.271     \n",
      "    22        115          -4.17         0.00         0.00       21      0.158     \n",
      "    23        120          -4.19         0.00         0.00       23     0.0931     \n",
      "    24        125          -4.18         0.00         0.00       23     0.0491     \n",
      "    25        130          -4.17         0.00         0.00       24     0.0325     stable\n",
      "   inf        130          -4.15         0.00         0.00       50     0.0325     finalize\n",
      "Inference terminated: variational solution stable for options.tol_stable_count fcn evaluations.\n",
      "Estimated ELBO: -4.149 +/-0.001.\n"
     ]
    }
   ],
   "source": [
    "new_options = {\n",
    "    \"max_fun_evals\": 50 * (D + 2),\n",
    "}\n",
    "vbmc = VBMC.load(\n",
    "    \"vbmc_test_save.pkl\",\n",
    "    new_options=new_options,\n",
    "    iteration=None,  # the default: start from the last stored iteration.\n",
    "    set_random_state=False,  # the default: don't modify the random state\n",
    "    # (can be set to True for reproducibility).\n",
    ")\n",
    "vp, results = vbmc.optimize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    'function': '<function VBMC._init_log_joint.<locals>.log_joint at 0x7f00ea1301f0>',\n",
      "    'problem_type': 'unconstrained',\n",
      "    'iterations': 25,\n",
      "    'func_count': 130,\n",
      "    'best_iter': 25,\n",
      "    'train_set_size': 121,\n",
      "    'components': 50,\n",
      "    'r_index': 0.03246322965421961,\n",
      "    'convergence_status': 'probable',\n",
      "    'overhead': nan,\n",
      "    'rng_state': 'rng',\n",
      "    'algorithm': 'Variational Bayesian Monte Carlo',\n",
      "    'version': '0.1.0',\n",
      "    'message': 'Inference terminated: variational solution stable for options.tol_stable_count fcn evaluations.',\n",
      "    'elbo': -4.149058408953476,\n",
      "    'elbo_sd': 0.001492050324205448,\n",
      "    'success_flag': True,\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(format_dict(results))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the default budget of function evaluations, we can see that the `convergence_status` is 'probable' and the `r_index` is much less than 1, suggesting convergence has been acheived. We can save the result to file and load it later, e.g. to perform futher validation or to sample from the variational posterior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-1.82750929 -0.25320367  0.51048786  1.41698176]\n",
      " [-2.38874097  0.03977422 -0.35191954  0.32260624]\n",
      " [-1.26096266 -0.58007695 -1.03945047 -0.0098995 ]\n",
      " [-1.88089158 -0.59198002 -0.72811124 -1.00054502]\n",
      " [-0.60285619 -0.07388834 -1.15702975 -0.25769202]]\n",
      "[ 3 14 20  4  8]\n"
     ]
    }
   ],
   "source": [
    "vbmc.save(\"vbmc_test_save.pkl\", overwrite=True)\n",
    "vbmc = VBMC.load(\"vbmc_test_save.pkl\")\n",
    "\n",
    "samples, components = vbmc.vp.sample(5)\n",
    "# `samples` are samples drawn from the variational posterior.\n",
    "# `components` are the index of the mixture components each\n",
    "#  sample was drawn from.\n",
    "print(samples)\n",
    "print(components)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Conclusions\n",
    "\n",
    "In this notebook, we have given a brief overview of PyVBMC's output diagnostics, and shown how to save and load results and resume optimization from a specific iteration.\n",
    "\n",
    "In the next notebook, we will illustrate running PyVBMC multiple times in order to validate the results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Acknowledgments\n",
    "\n",
    "Work on the PyVBMC package was funded by the [Finnish Center for Artificial Intelligence FCAI](https://fcai.fi/)."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "cf2c3e35bb9d622e963fe7adafe5d3d77a0ee2382f730f35475e9a620896d84b"
  },
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}