{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Gradient checks\n",
    "\n",
    "It is best practice to do gradient checks before and after gradient-based optimization.\n",
    "\n",
    "1. Find suitable tolerances to use during optimization. Importantly, test your gradients using the settings you will use later on.\n",
    "2. At the optimum the values should be close to 0, except for parameters with active bounds. \n",
    "3. Gradient checks can help you identify inconsistencies and errors, especially when using custom gradient calculation or objectives.\n",
    "\n",
    "Here we show, how to use the gradient check methods that are implemented in pyPESTO, using the finite differences (FD) method as a comparison. There is a trade-off between the quality of the approximation and numerical noise, so it is recommended to try different FD step sizes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import benchmark_models_petab as models\n",
    "import numpy as np\n",
    "\n",
    "import pypesto.optimize as optimize\n",
    "import pypesto.petab\n",
    "\n",
    "np.random.seed(2)\n",
    "\n",
    "import pandas as pd\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Set up an example problem\n",
    "\n",
    "Create the pypesto problem and a random vector of parameter values.  \n",
    "Here, we use the startpoint sampling method to generate random parameter vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "\n",
    "model_name = \"Boehm_JProteomeRes2014\"\n",
    "petab_problem = models.get_problem(model_name)\n",
    "\n",
    "importer = pypesto.petab.PetabImporter(petab_problem)\n",
    "pypesto_problem = importer.create_problem(verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "startpoints = pypesto_problem.get_startpoints(n_starts=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gradient check before optimization\n",
    "\n",
    "Perform a gradient check at the location of one of the random parameter vectors. `check_grad` compares the gradients obtained by the finite differences (FD) method and the objective gradient. You can modify the finite differences step size via the argument `eps`.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pypesto_problem.objective.check_grad(\n",
    "    x=startpoints[0],\n",
    "    eps=1e-5,  # default\n",
    "    verbosity=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Explanation of the gradient check result columns:\n",
    "\n",
    "- `grad`: Objective gradient\n",
    "- `fd_f`: FD forward difference\n",
    "- `fd_b`: FD backward difference\n",
    "- `fd_c`: Approximation of FD central difference (reusing the information from `fd_f` and `fd_b`)\n",
    "- `fd_err`: Deviation between forward and backward differences `fd_f`, `fd_b`\n",
    "- `abs_err`: Absolute error between `grad` and the central FD gradient `fd_c`\n",
    "- `rel_err` Relative error between `grad` and the central FD gradient `fd_c`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If there are fixed parameters in your vector you might invoke an error due to the dimension mismatch. Use the helper method `Problem.get_reduced_vector` to get the reduced vector with only free (estimated) parameters.  \n",
    "Here we set a smaller FD step size `eps = 1e-6` and observe that the errors change:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameter_vector = pypesto_problem.get_reduced_vector(startpoints[0])\n",
    "\n",
    "pypesto_problem.objective.check_grad(\n",
    "    x=parameter_vector,\n",
    "    eps=1e-6,\n",
    "    verbosity=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The method `check_grad_multi_eps` calls the `check_grad` method multiple times with different settings for the FD step size and reports the setting that results in the smallest error. \n",
    "You can supply a list of FD step sizes to be tested via the `multi_eps` argument (or use the default ones), and use the `label` argument to switch between the FD, or absolute or relative error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gc = pypesto_problem.objective.check_grad_multi_eps(\n",
    "    x=parameter_vector,\n",
    "    verbosity=0,\n",
    "    label=\"rel_err\",  # default\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the pandas style methods to visualise the results of the gradient check, e.g.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def highlight_value_above_threshold(x, threshold=1):\n",
    "    return [\"color: darkorange\" if xi > threshold else None for xi in x]\n",
    "\n",
    "\n",
    "def highlight_gradient_check(gc: pd.DataFrame):\n",
    "    return (\n",
    "        gc.style.apply(\n",
    "            highlight_value_above_threshold,\n",
    "            subset=[\"fd_err\"],\n",
    "        )\n",
    "        .background_gradient(\n",
    "            cmap=sns.light_palette(\"purple\", as_cmap=True),\n",
    "            subset=[\"abs_err\"],\n",
    "        )\n",
    "        .background_gradient(\n",
    "            cmap=sns.light_palette(\"red\", as_cmap=True),\n",
    "            subset=[\"rel_err\"],\n",
    "        )\n",
    "        .background_gradient(\n",
    "            cmap=sns.color_palette(\"viridis\", as_cmap=True),\n",
    "            subset=[\"eps\"],\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "highlight_gradient_check(gc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are consistently large discrepancies between forward and backward FD and a large relative error for the parameter `k_exp_hetero`.  \n",
    "\n",
    "Ideally, all gradients would agree, but especially at not-so-smooth points of the objective, like (local) optima, large FD errors can occur.\n",
    "It is recommended to check gradients over a lot of random points and check if there are consistently large errors for specific parameters.  \n",
    "\n",
    "Below we perform a gradient check for another random point and observe small errors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameter_vector = startpoints[1]\n",
    "\n",
    "gc = pypesto_problem.objective.check_grad_multi_eps(\n",
    "    x=parameter_vector,\n",
    "    verbosity=0,\n",
    "    label=\"rel_err\",  # default\n",
    ")\n",
    "highlight_gradient_check(gc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Gradient check after optimization\n",
    "\n",
    "Next, we do optimization and perform a gradient check at a local optimum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "\n",
    "result = optimize.minimize(\n",
    "    problem=pypesto_problem,\n",
    "    optimizer=optimize.ScipyOptimizer(),\n",
    "    n_starts=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(Local) optima can be points with weird gradients. At a steep optimum, the `fd_err` is expected to be high.  \n",
    "\n",
    "At the local optimum shown below, the `sd_pSTAT5B_rel` forward and backward FD have opposite signs and are quite large, resulting in a substantial `fd_err`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameter vector at the local optimum, obtained from optimization\n",
    "parameter_vector = pypesto_problem.get_reduced_vector(\n",
    "    result.optimize_result[0].x\n",
    ")\n",
    "\n",
    "highlight_gradient_check(\n",
    "    gc=pypesto_problem.objective.check_grad_multi_eps(\n",
    "        x=parameter_vector,\n",
    "        verbosity=0,\n",
    "        label=\"rel_err\",  # default\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### How to \"fix\" my gradients?\n",
    "\n",
    "- Find suitable simulation tolerances.\n",
    "\n",
    "Specific to the petab-amici-pipeline:\n",
    "\n",
    "- Check the simulation logs for Warnings and Errors.\n",
    "- Consider switching between forward and adjoint sensitivity algorithms.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "developer",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}