{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# iminuit and automatic differentiation with JAX\n",
    "\n",
    "This is a topical tutorial. If you are new to iminuit, you should go through the basic tutorial first. \n",
    "\n",
    "iminuit's minimization algorithm MIGRAD uses a mix of gradient descent and Newton's method to find the minimum. Both require a first derivative, which MIGRAD usually computes numerically from finite differences. This requires many function evaluations and the gradient may not be accurate. As an alternative, iminuit also allows the user to compute the gradient and pass it to MIGRAD.\n",
    "\n",
    "Although computing derivatives is often straight-forward, it is usually too much hassle to do manually. Automatic differentiation (AD) is an interesting alternative, it allows one to compute exact derivatives efficiently for pure Python/numpy functions. We demonstrate automatic differentiation with the JAX module, which can not only compute derivatives, but also accelerates the computation of Python code (including the gradient code) with a just-in-time compiler.\n",
    "\n",
    "[Recommended read: Gentle introduction to AD](https://www.kaggle.com/borisettinger/gentle-introduction-to-automatic-differentiation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example A: Fit of a gaussian model to a histogram\n",
    "\n",
    "We fit a gaussian to a histogram using a maximum-likelihood approach based on Poisson statistics. This example is used to investigate how automatic differentiation can accelerate a typical fit in a counting experiment.\n",
    "\n",
    "To compare fits with and without passing an analytic gradient fairly, we use `Minuit.strategy = 0`, which prevents Minuit from automatically computing the Hesse matrix after the fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:26:37.436843Z",
     "start_time": "2020-02-21T10:26:37.432080Z"
    }
   },
   "outputs": [],
   "source": [
    "# !pip install jax jaxlib matplotlib\n",
    "\n",
    "from jax.config import config\n",
    "config.update(\"jax_enable_x64\", True) # enable float64 precision\n",
    "from jax import numpy as np # replacement for normal numpy\n",
    "from jax.scipy.special import erf # replacement for scipy.special.erf\n",
    "from jax import jit, grad, jacfwd, jacrev\n",
    "from iminuit import Minuit\n",
    "\n",
    "import numpy as onp # original numpy still needed, since jax does not cover full API\n",
    "from numpy.random import default_rng"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We generate some toy data and write the negative log-likelihood (nll) for a fit to binned data, assuming Poisson-distributed counts.\n",
    "\n",
    "**Note:** We write all statistical functions in pure Python code, to demonstrate Jax's ability to automatically differentiate and JIT compile this code. In practice, one should import JIT-able statistical distributions from jax.scipy.stats. The library versions can be expected to have fewer bugs and to be faster and more accurate than hand-written code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:26:37.594856Z",
     "start_time": "2020-02-21T10:26:37.585943Z"
    }
   },
   "outputs": [],
   "source": [
    "# generate some toy data\n",
    "rng = default_rng(seed=1)\n",
    "w, xe = onp.histogram(rng.normal(0, 1, 10000), bins=1000)\n",
    "\n",
    "# cdf of a normal distribution, needed to compute the expected counts per bin\n",
    "def cdf(x, mu, sigma):\n",
    "    z = (x - mu) / sigma\n",
    "    return 0.5 * (1 + erf(z / np.sqrt(2)))\n",
    "\n",
    "## better alternative for real code:\n",
    "# from jax.scipy.stats.norm import cdf\n",
    "\n",
    "\n",
    "def nll(par): # negative log-likelihood with constants stripped\n",
    "    amp = par[0]\n",
    "    rest = par[1:]\n",
    "    mu = amp * (cdf(xe[1:], *rest) - cdf(xe[:-1], *rest))\n",
    "    result = np.sum(mu - w * np.log(mu + 1e-100))\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check results from all combinations of using JIT and gradient and then compare the execution times.\n",
    "\n",
    "|      |     |     |\n",
    "|:----:|:---:|:---:|\n",
    "|      |~JIT~| JIT |\n",
    "|~grad~| m1  | m3  |\n",
    "| grad | m2  | m4  |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:26:37.890967Z",
     "start_time": "2020-02-21T10:26:37.886224Z"
    }
   },
   "outputs": [],
   "source": [
    "start_values = (1.5 * onp.sum(w), 1.0, 2.0)\n",
    "limits = ((0, None), None, (0, None))\n",
    "\n",
    "def make_and_run_minuit(fcn, grad=None):\n",
    "    m = Minuit.from_array_func(fcn, start_values, grad=grad, name=(\"amp\", \"mu\", \"sigma\"),\n",
    "                               limit=limits, pedantic=False, errordef=Minuit.LIKELIHOOD)\n",
    "    m.strategy = 0\n",
    "    fmin = m.migrad()\n",
    "    return m, fmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:26:38.532308Z",
     "start_time": "2020-02-21T10:26:38.368563Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hdembinski/.local/lib/python3.6/site-packages/jax/lib/xla_bridge.py:122: UserWarning: No GPU/TPU found, falling back to CPU.\n",
      "  warnings.warn('No GPU/TPU found, falling back to CPU.')\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td colspan=\"2\" title=\"Minimum value of function\">\n",
       "FCN = -1.936e+04\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"3\" title=\"No. of calls in last algorithm and total number of calls\">\n",
       "Ncalls = 66 (66 total)\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td colspan=\"2\" title=\"Estimated distance to minimum and target threshold\">\n",
       "EDM = 1.76e-08 (Goal: 0.0001)\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"3\" title=\"Increase in FCN which corresponds to 1 standard deviation\">\n",
       "up = 0.5\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" title=\"Validity of the migrad call\">\n",
       "Valid Min.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Validity of parameters\">\n",
       "Valid Param.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is EDM above goal EDM?\">\n",
       "Above EDM\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"2\" title=\"Did last migrad call reach max call limit?\">\n",
       "Reached call limit\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"2\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" title=\"Did Hesse fail?\">\n",
       "Hesse failed\n",
       "</td>\n",
       "<td align=\"center\" title=\"Has covariance matrix\">\n",
       "Has cov.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is covariance matrix accurate?\">\n",
       "Accurate\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is covariance matrix positive definite?\">\n",
       "Pos. def.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Was positive definiteness enforced by Minuit?\">\n",
       "Forced\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#FF7878;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "</tr>\n",
       "</table>\n",
       "<table>\n",
       "<tr style=\"background-color:#F4F4F4;\">\n",
       "<td/>\n",
       "<th title=\"Variable name\">\n",
       "Name\n",
       "</th>\n",
       "<th title=\"Value of parameter\">\n",
       "Value\n",
       "</th>\n",
       "<th title=\"Hesse error\">\n",
       "Hesse Error\n",
       "</th>\n",
       "<th title=\"Minos lower error\">\n",
       "Minos Error-\n",
       "</th>\n",
       "<th title=\"Minos upper error\">\n",
       "Minos Error+\n",
       "</th>\n",
       "<th title=\"Lower limit of the parameter\">\n",
       "Limit-\n",
       "</th>\n",
       "<th title=\"Upper limit of the parameter\">\n",
       "Limit+\n",
       "</th>\n",
       "<th title=\"Is the parameter fixed in the fit\">\n",
       "Fixed\n",
       "</th>\n",
       "</tr>\n",
       "<tr style=\"background-color:#FFFFFF;\">\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "amp\n",
       "</td>\n",
       "<td>\n",
       "10.0k\n",
       "</td>\n",
       "<td>\n",
       "0.1k\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "<tr style=\"background-color:#F4F4F4;\">\n",
       "<td>\n",
       "1\n",
       "</td>\n",
       "<td>\n",
       "mu\n",
       "</td>\n",
       "<td>\n",
       "-0.011\n",
       "</td>\n",
       "<td>\n",
       "0.010\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "<tr style=\"background-color:#FFFFFF;\">\n",
       "<td>\n",
       "2\n",
       "</td>\n",
       "<td>\n",
       "sigma\n",
       "</td>\n",
       "<td>\n",
       "999m\n",
       "</td>\n",
       "<td>\n",
       "7m\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "</table>\n"
      ],
      "text/plain": [
       "------------------------------------------------------------------\n",
       "| FCN = -1.936e+04              |      Ncalls=66 (66 total)      |\n",
       "| EDM = 1.76e-08 (Goal: 0.0001) |            up = 0.5            |\n",
       "------------------------------------------------------------------\n",
       "|  Valid Min.   | Valid Param.  | Above EDM | Reached call limit |\n",
       "------------------------------------------------------------------\n",
       "|     True      |     True      |   False   |       False        |\n",
       "------------------------------------------------------------------\n",
       "| Hesse failed  |   Has cov.    | Accurate  | Pos. def. | Forced |\n",
       "------------------------------------------------------------------\n",
       "|     False     |     True      |   False   |   True    | False  |\n",
       "------------------------------------------------------------------\n",
       "-------------------------------------------------------------------------------------------\n",
       "|   | Name  |   Value   | Hesse Err | Minos Err- | Minos Err+ | Limit-  | Limit+  | Fixed |\n",
       "-------------------------------------------------------------------------------------------\n",
       "| 0 | amp   |   10.0k   |   0.1k    |            |            |    0    |         |       |\n",
       "| 1 | mu    |  -0.011   |   0.010   |            |            |         |         |       |\n",
       "| 2 | sigma |   999m    |    7m     |            |            |    0    |         |       |\n",
       "-------------------------------------------------------------------------------------------"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1, fmin = make_and_run_minuit(nll)\n",
    "fmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:26:39.371830Z",
     "start_time": "2020-02-21T10:26:38.797460Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td colspan=\"2\" title=\"Minimum value of function\">\n",
       "FCN = -1.936e+04\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"3\" title=\"No. of calls in last algorithm and total number of calls\">\n",
       "Ncalls = 25 (25 total)\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td colspan=\"2\" title=\"Estimated distance to minimum and target threshold\">\n",
       "EDM = 4.95e-05 (Goal: 0.0001)\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"3\" title=\"Increase in FCN which corresponds to 1 standard deviation\">\n",
       "up = 0.5\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" title=\"Validity of the migrad call\">\n",
       "Valid Min.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Validity of parameters\">\n",
       "Valid Param.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is EDM above goal EDM?\">\n",
       "Above EDM\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"2\" title=\"Did last migrad call reach max call limit?\">\n",
       "Reached call limit\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"2\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" title=\"Did Hesse fail?\">\n",
       "Hesse failed\n",
       "</td>\n",
       "<td align=\"center\" title=\"Has covariance matrix\">\n",
       "Has cov.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is covariance matrix accurate?\">\n",
       "Accurate\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is covariance matrix positive definite?\">\n",
       "Pos. def.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Was positive definiteness enforced by Minuit?\">\n",
       "Forced\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#FF7878;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "</tr>\n",
       "</table>\n",
       "<table>\n",
       "<tr style=\"background-color:#F4F4F4;\">\n",
       "<td/>\n",
       "<th title=\"Variable name\">\n",
       "Name\n",
       "</th>\n",
       "<th title=\"Value of parameter\">\n",
       "Value\n",
       "</th>\n",
       "<th title=\"Hesse error\">\n",
       "Hesse Error\n",
       "</th>\n",
       "<th title=\"Minos lower error\">\n",
       "Minos Error-\n",
       "</th>\n",
       "<th title=\"Minos upper error\">\n",
       "Minos Error+\n",
       "</th>\n",
       "<th title=\"Lower limit of the parameter\">\n",
       "Limit-\n",
       "</th>\n",
       "<th title=\"Upper limit of the parameter\">\n",
       "Limit+\n",
       "</th>\n",
       "<th title=\"Is the parameter fixed in the fit\">\n",
       "Fixed\n",
       "</th>\n",
       "</tr>\n",
       "<tr style=\"background-color:#FFFFFF;\">\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "amp\n",
       "</td>\n",
       "<td>\n",
       "10.0k\n",
       "</td>\n",
       "<td>\n",
       "0.1k\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "<tr style=\"background-color:#F4F4F4;\">\n",
       "<td>\n",
       "1\n",
       "</td>\n",
       "<td>\n",
       "mu\n",
       "</td>\n",
       "<td>\n",
       "-0.011\n",
       "</td>\n",
       "<td>\n",
       "0.010\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "<tr style=\"background-color:#FFFFFF;\">\n",
       "<td>\n",
       "2\n",
       "</td>\n",
       "<td>\n",
       "sigma\n",
       "</td>\n",
       "<td>\n",
       "999m\n",
       "</td>\n",
       "<td>\n",
       "7m\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "</table>\n"
      ],
      "text/plain": [
       "------------------------------------------------------------------\n",
       "| FCN = -1.936e+04              |      Ncalls=25 (25 total)      |\n",
       "| EDM = 4.95e-05 (Goal: 0.0001) |            up = 0.5            |\n",
       "------------------------------------------------------------------\n",
       "|  Valid Min.   | Valid Param.  | Above EDM | Reached call limit |\n",
       "------------------------------------------------------------------\n",
       "|     True      |     True      |   False   |       False        |\n",
       "------------------------------------------------------------------\n",
       "| Hesse failed  |   Has cov.    | Accurate  | Pos. def. | Forced |\n",
       "------------------------------------------------------------------\n",
       "|     False     |     True      |   False   |   True    | False  |\n",
       "------------------------------------------------------------------\n",
       "-------------------------------------------------------------------------------------------\n",
       "|   | Name  |   Value   | Hesse Err | Minos Err- | Minos Err+ | Limit-  | Limit+  | Fixed |\n",
       "-------------------------------------------------------------------------------------------\n",
       "| 0 | amp   |   10.0k   |   0.1k    |            |            |    0    |         |       |\n",
       "| 1 | mu    |  -0.011   |   0.010   |            |            |         |         |       |\n",
       "| 2 | sigma |   999m    |    7m     |            |            |    0    |         |       |\n",
       "-------------------------------------------------------------------------------------------"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m2, fmin = make_and_run_minuit(nll, grad=grad(nll))\n",
    "fmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:26:39.510553Z",
     "start_time": "2020-02-21T10:26:39.373728Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td colspan=\"2\" title=\"Minimum value of function\">\n",
       "FCN = -1.936e+04\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"3\" title=\"No. of calls in last algorithm and total number of calls\">\n",
       "Ncalls = 25 (25 total)\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td colspan=\"2\" title=\"Estimated distance to minimum and target threshold\">\n",
       "EDM = 4.96e-05 (Goal: 0.0001)\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"3\" title=\"Increase in FCN which corresponds to 1 standard deviation\">\n",
       "up = 0.5\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" title=\"Validity of the migrad call\">\n",
       "Valid Min.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Validity of parameters\">\n",
       "Valid Param.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is EDM above goal EDM?\">\n",
       "Above EDM\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"2\" title=\"Did last migrad call reach max call limit?\">\n",
       "Reached call limit\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"2\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" title=\"Did Hesse fail?\">\n",
       "Hesse failed\n",
       "</td>\n",
       "<td align=\"center\" title=\"Has covariance matrix\">\n",
       "Has cov.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is covariance matrix accurate?\">\n",
       "Accurate\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is covariance matrix positive definite?\">\n",
       "Pos. def.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Was positive definiteness enforced by Minuit?\">\n",
       "Forced\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#FF7878;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "</tr>\n",
       "</table>\n",
       "<table>\n",
       "<tr style=\"background-color:#F4F4F4;\">\n",
       "<td/>\n",
       "<th title=\"Variable name\">\n",
       "Name\n",
       "</th>\n",
       "<th title=\"Value of parameter\">\n",
       "Value\n",
       "</th>\n",
       "<th title=\"Hesse error\">\n",
       "Hesse Error\n",
       "</th>\n",
       "<th title=\"Minos lower error\">\n",
       "Minos Error-\n",
       "</th>\n",
       "<th title=\"Minos upper error\">\n",
       "Minos Error+\n",
       "</th>\n",
       "<th title=\"Lower limit of the parameter\">\n",
       "Limit-\n",
       "</th>\n",
       "<th title=\"Upper limit of the parameter\">\n",
       "Limit+\n",
       "</th>\n",
       "<th title=\"Is the parameter fixed in the fit\">\n",
       "Fixed\n",
       "</th>\n",
       "</tr>\n",
       "<tr style=\"background-color:#FFFFFF;\">\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "amp\n",
       "</td>\n",
       "<td>\n",
       "10.0k\n",
       "</td>\n",
       "<td>\n",
       "0.1k\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "<tr style=\"background-color:#F4F4F4;\">\n",
       "<td>\n",
       "1\n",
       "</td>\n",
       "<td>\n",
       "mu\n",
       "</td>\n",
       "<td>\n",
       "-0.011\n",
       "</td>\n",
       "<td>\n",
       "0.010\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "<tr style=\"background-color:#FFFFFF;\">\n",
       "<td>\n",
       "2\n",
       "</td>\n",
       "<td>\n",
       "sigma\n",
       "</td>\n",
       "<td>\n",
       "999m\n",
       "</td>\n",
       "<td>\n",
       "7m\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "</table>\n"
      ],
      "text/plain": [
       "------------------------------------------------------------------\n",
       "| FCN = -1.936e+04              |      Ncalls=25 (25 total)      |\n",
       "| EDM = 4.96e-05 (Goal: 0.0001) |            up = 0.5            |\n",
       "------------------------------------------------------------------\n",
       "|  Valid Min.   | Valid Param.  | Above EDM | Reached call limit |\n",
       "------------------------------------------------------------------\n",
       "|     True      |     True      |   False   |       False        |\n",
       "------------------------------------------------------------------\n",
       "| Hesse failed  |   Has cov.    | Accurate  | Pos. def. | Forced |\n",
       "------------------------------------------------------------------\n",
       "|     False     |     True      |   False   |   True    | False  |\n",
       "------------------------------------------------------------------\n",
       "-------------------------------------------------------------------------------------------\n",
       "|   | Name  |   Value   | Hesse Err | Minos Err- | Minos Err+ | Limit-  | Limit+  | Fixed |\n",
       "-------------------------------------------------------------------------------------------\n",
       "| 0 | amp   |   10.0k   |   0.1k    |            |            |    0    |         |       |\n",
       "| 1 | mu    |  -0.011   |   0.010   |            |            |         |         |       |\n",
       "| 2 | sigma |   999m    |    7m     |            |            |    0    |         |       |\n",
       "-------------------------------------------------------------------------------------------"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m3, fmin = make_and_run_minuit(jit(nll), grad=grad(nll))\n",
    "fmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:26:40.573574Z",
     "start_time": "2020-02-21T10:26:40.229476Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td colspan=\"2\" title=\"Minimum value of function\">\n",
       "FCN = -1.936e+04\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"3\" title=\"No. of calls in last algorithm and total number of calls\">\n",
       "Ncalls = 25 (25 total)\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td colspan=\"2\" title=\"Estimated distance to minimum and target threshold\">\n",
       "EDM = 4.96e-05 (Goal: 0.0001)\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"3\" title=\"Increase in FCN which corresponds to 1 standard deviation\">\n",
       "up = 0.5\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" title=\"Validity of the migrad call\">\n",
       "Valid Min.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Validity of parameters\">\n",
       "Valid Param.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is EDM above goal EDM?\">\n",
       "Above EDM\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"2\" title=\"Did last migrad call reach max call limit?\">\n",
       "Reached call limit\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"2\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" title=\"Did Hesse fail?\">\n",
       "Hesse failed\n",
       "</td>\n",
       "<td align=\"center\" title=\"Has covariance matrix\">\n",
       "Has cov.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is covariance matrix accurate?\">\n",
       "Accurate\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is covariance matrix positive definite?\">\n",
       "Pos. def.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Was positive definiteness enforced by Minuit?\">\n",
       "Forced\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#FF7878;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "</tr>\n",
       "</table>\n",
       "<table>\n",
       "<tr style=\"background-color:#F4F4F4;\">\n",
       "<td/>\n",
       "<th title=\"Variable name\">\n",
       "Name\n",
       "</th>\n",
       "<th title=\"Value of parameter\">\n",
       "Value\n",
       "</th>\n",
       "<th title=\"Hesse error\">\n",
       "Hesse Error\n",
       "</th>\n",
       "<th title=\"Minos lower error\">\n",
       "Minos Error-\n",
       "</th>\n",
       "<th title=\"Minos upper error\">\n",
       "Minos Error+\n",
       "</th>\n",
       "<th title=\"Lower limit of the parameter\">\n",
       "Limit-\n",
       "</th>\n",
       "<th title=\"Upper limit of the parameter\">\n",
       "Limit+\n",
       "</th>\n",
       "<th title=\"Is the parameter fixed in the fit\">\n",
       "Fixed\n",
       "</th>\n",
       "</tr>\n",
       "<tr style=\"background-color:#FFFFFF;\">\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "amp\n",
       "</td>\n",
       "<td>\n",
       "10.0k\n",
       "</td>\n",
       "<td>\n",
       "0.1k\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "<tr style=\"background-color:#F4F4F4;\">\n",
       "<td>\n",
       "1\n",
       "</td>\n",
       "<td>\n",
       "mu\n",
       "</td>\n",
       "<td>\n",
       "-0.011\n",
       "</td>\n",
       "<td>\n",
       "0.010\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "<tr style=\"background-color:#FFFFFF;\">\n",
       "<td>\n",
       "2\n",
       "</td>\n",
       "<td>\n",
       "sigma\n",
       "</td>\n",
       "<td>\n",
       "999m\n",
       "</td>\n",
       "<td>\n",
       "7m\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "</table>\n"
      ],
      "text/plain": [
       "------------------------------------------------------------------\n",
       "| FCN = -1.936e+04              |      Ncalls=25 (25 total)      |\n",
       "| EDM = 4.96e-05 (Goal: 0.0001) |            up = 0.5            |\n",
       "------------------------------------------------------------------\n",
       "|  Valid Min.   | Valid Param.  | Above EDM | Reached call limit |\n",
       "------------------------------------------------------------------\n",
       "|     True      |     True      |   False   |       False        |\n",
       "------------------------------------------------------------------\n",
       "| Hesse failed  |   Has cov.    | Accurate  | Pos. def. | Forced |\n",
       "------------------------------------------------------------------\n",
       "|     False     |     True      |   False   |   True    | False  |\n",
       "------------------------------------------------------------------\n",
       "-------------------------------------------------------------------------------------------\n",
       "|   | Name  |   Value   | Hesse Err | Minos Err- | Minos Err+ | Limit-  | Limit+  | Fixed |\n",
       "-------------------------------------------------------------------------------------------\n",
       "| 0 | amp   |   10.0k   |   0.1k    |            |            |    0    |         |       |\n",
       "| 1 | mu    |  -0.011   |   0.010   |            |            |         |         |       |\n",
       "| 2 | sigma |   999m    |    7m     |            |            |    0    |         |       |\n",
       "-------------------------------------------------------------------------------------------"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m4, fmin = make_and_run_minuit(jit(nll), grad=jit(grad(nll)))\n",
    "fmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:26:45.031931Z",
     "start_time": "2020-02-21T10:26:40.674388Z"
    }
   },
   "outputs": [],
   "source": [
    "from timeit import timeit\n",
    "times = { \"no JIT, no grad\": \"m1\",\n",
    "          \"no JIT, grad\": \"m2\",\n",
    "          \"JIT, no grad\": \"m3\",\n",
    "          \"JIT, grad\": \"m4\" }\n",
    "for k, v in times.items():\n",
    "    t = timeit(f\"{v}.args[:] = start_values; {v}.migrad()\",\n",
    "               f\"from __main__ import {v}, start_values\", number=1)\n",
    "    times[k] = t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:26:45.142272Z",
     "start_time": "2020-02-21T10:26:45.033451Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Just using the JIT compiler is a 1.2-fold increase.\n",
      "Using JIT and gradient is a 28-fold increase.\n"
     ]
    }
   ],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "x = onp.fromiter(times.values(), dtype=float)\n",
    "xmin = np.min(x)\n",
    "\n",
    "y = -np.arange(4)\n",
    "plt.barh(y, x)\n",
    "for yi, k, v in zip(y, times, x):\n",
    "    plt.text(v, yi, f\"{v/xmin:.1f}x\")\n",
    "plt.yticks(y, times.keys())\n",
    "for loc in (\"top\", \"right\"):\n",
    "    plt.gca().spines[loc].set_visible(False)\n",
    "plt.xlabel(\"execution time / s\")\n",
    "print(f\"\"\"Just using the JIT compiler is a {x[0]/x[2]:.1f}-fold increase.\n",
    "Using JIT and gradient is a {x[0]/x[3]:.0f}-fold increase.\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Conclusions:\n",
    "\n",
    "1. As expected, the best results are obtained by JIT compiling the function and the gradient.\n",
    "\n",
    "2. JIT compiling the cost function with Jax is only a minor performance improvement (numba is able to do much better). \n",
    "\n",
    "3. JIT compiling the gradient is important. Using a pure-Python gradient decreases performance in this example.\n",
    "\n",
    "\n",
    "The gain from using a gradient is larger for functions with hundreds of parameters, as is common in machine learning. Human-made models often have less than 10 parameters, and then the gain is not so dramatic. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing covariance matrices with JAX\n",
    "\n",
    "Automatic differentiation gives us another way to compute uncertainties of fitted parameters. MINUIT compute the uncertainties with the HESSE algorithm by default, which computes the matrix of second derivates approximately using finite differences and inverts this.\n",
    "\n",
    "Let's compare the output of HESSE with the exact (within floating point precision) computation using automatic differentiation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:27:38.715871Z",
     "start_time": "2020-02-21T10:27:37.907690Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sigma[amp]  : HESSE =  100.0, JAX =  100.0\n",
      "sigma[mu]   : HESSE = 0.0100, JAX = 0.0100\n",
      "sigma[sigma]: HESSE = 0.0071, JAX = 0.0071\n"
     ]
    }
   ],
   "source": [
    "m4.hesse()\n",
    "cov_hesse = m4.np_covariance()\n",
    "\n",
    "def jax_covariance(par):\n",
    "    hessian = jacfwd(jacrev(nll))\n",
    "    return np.linalg.inv(hessian(par))\n",
    "\n",
    "par = m4.np_values()\n",
    "cov_jax = jax_covariance(par)\n",
    "\n",
    "print(f\"sigma[amp]  : HESSE = {cov_hesse[0, 0] ** 0.5:6.1f}, JAX = {cov_jax[0, 0] ** 0.5:6.1f}\")\n",
    "print(f\"sigma[mu]   : HESSE = {cov_hesse[1, 1] ** 0.5:6.4f}, JAX = {cov_jax[1, 1] ** 0.5:6.4f}\")\n",
    "print(f\"sigma[sigma]: HESSE = {cov_hesse[2, 2] ** 0.5:6.4f}, JAX = {cov_jax[2, 2] ** 0.5:6.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Success, HESSE and JAX give the same answer within the  relevant precision.\n",
    "\n",
    "**Note:** If you compute the covariance matrix in this way from a least-squares cost function, you must multiply it by 2.\n",
    "\n",
    "Let us compare the performance of HESSE with Jax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "73.1 ms ± 9.23 ms per loop (mean ± std. dev. of 3 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 1 -r 3\n",
    "m = Minuit.from_array_func(nll, par, pedantic=False)\n",
    "m.hesse()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "67.3 ms ± 7.39 ms per loop (mean ± std. dev. of 3 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 1 -r 3\n",
    "jax_covariance(par)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The computation with Jax is faster, but not by much. It is also more accurate (although the added precision is not relevant).\n",
    "\n",
    "Altogether, Minuit's HESSE algorithm still makes sense today. It has the advantage that it can process any function, while Jax cannot. Jax cannot differentiate a function that calls into C/C++ code or Cython code, for example.\n",
    "\n",
    "Final note: If we JIT compile `jax_covariance`, it greatly outperforms Minuit's HESSE algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "268 µs ± 77.4 µs per loop (mean ± std. dev. of 3 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 1 -r 3 jit_jax_covariance = jit(jax_covariance); jit_jax_covariance(par)\n",
    "jit_jax_covariance(par)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is much faster... but only because the compilation cost is excluded here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.5 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 1 -r 1\n",
    "# if we include the JIT compilation cost, the performance drops dramatically\n",
    "@jit\n",
    "def jax_covariance(par):\n",
    "    hessian = jacfwd(jacrev(nll))\n",
    "    return np.linalg.inv(hessian(par))\n",
    "\n",
    "jax_covariance(par)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With compilation cost included, it is much slower.\n",
    "\n",
    "Conclusion: Using the JIT compiler makes a lot of sense if the covariance matrix has to be computed repeatedly for the same cost function but different parameters, but this is not the common case in practice."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example B: Fit data points with uncertainties in x and y\n",
    "\n",
    "Let's say we have some data points $(x_i \\pm \\sigma_{x,i}, y_i \\pm \\sigma_{y,i})$ and we have a model $y=f(x)$ that we want to adapt to this data. If $\\sigma_{x,i}$ was zero, we could use the usual least-squares method, minimizing the sum of squared residuals $r^2_i = (y_i - f(x_i))^2 / \\sigma^2_{y,i}$. Here, we don't know where to evaluate $f(x)$, since the exact $x$-location is only known up to $\\sigma_{x,i}$.\n",
    "\n",
    "We can approximately extend the standard least-squares method to handle this case. We use that the uncertainty along the $x$-axis can be converted into an additional uncertainty along the $y$-axis with error propagation,\n",
    "$$\n",
    "f(x_i \\pm \\sigma_{x,i}) \\simeq f(x_i) \\pm f'(x_i)\\,\\sigma_{x,i}.\n",
    "$$\n",
    "Using this, we obtain modified squared residuals\n",
    "$$\n",
    "r^2_i = \\frac{(y_i - f(x_i))^2}{\\sigma^2_{y,i} + (f'(x_i) \\,\\sigma_{x,i})^2}.\n",
    "$$\n",
    "\n",
    "We demonstrate this with a fit of a polynomial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:25:43.510168Z",
     "start_time": "2020-02-21T10:25:43.371319Z"
    }
   },
   "outputs": [],
   "source": [
    "# polynomial model\n",
    "def f(x, par):\n",
    "    return np.polyval(par, x)\n",
    "\n",
    "\n",
    "# true polynomial f(x) = x^2 + 2 x + 3\n",
    "par_true = (1, 2, 3)\n",
    "\n",
    "\n",
    "# grad computes derivative with respect to the first argument\n",
    "f_prime = jit(grad(f))\n",
    "\n",
    "\n",
    "# checking first derivative f'(x) = 2 x + 2\n",
    "assert f_prime(0.0, par_true) == 2\n",
    "assert f_prime(1.0, par_true) == 4\n",
    "# ok!\n",
    "\n",
    "# generate toy data\n",
    "n = 30 \n",
    "data_x = np.linspace(-4, 7, n)\n",
    "data_y = f(data_x, par_true)\n",
    "\n",
    "rng = default_rng(seed=1)\n",
    "sigma_x = 0.5\n",
    "sigma_y = 5\n",
    "data_x += rng.normal(0, sigma_x, n)\n",
    "data_y += rng.normal(0, sigma_y, n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:25:43.646212Z",
     "start_time": "2020-02-21T10:25:43.512384Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10696 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.errorbar(data_x, data_y, sigma_y, sigma_x, fmt=\"o\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:25:44.032210Z",
     "start_time": "2020-02-21T10:25:43.648365Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(876.49545695, dtype=float64)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# define the cost function\n",
    "@jit\n",
    "def cost(par):\n",
    "    result = 0.0\n",
    "    for xi, yi in zip(data_x, data_y):\n",
    "        y_var = sigma_y ** 2 + (f_prime(xi, par) * sigma_x) ** 2\n",
    "        result += (yi - f(xi, par)) ** 2 / y_var\n",
    "    return result\n",
    "\n",
    "# test the jit-ed function\n",
    "cost(np.zeros(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:25:44.059729Z",
     "start_time": "2020-02-21T10:25:44.034029Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td colspan=\"2\" title=\"Minimum value of function\">\n",
       "FCN = 23.14\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"3\" title=\"No. of calls in last algorithm and total number of calls\">\n",
       "Ncalls = 91 (91 total)\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td colspan=\"2\" title=\"Estimated distance to minimum and target threshold\">\n",
       "EDM = 3.12e-05 (Goal: 0.0002)\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"3\" title=\"Increase in FCN which corresponds to 1 standard deviation\">\n",
       "up = 1.0\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" title=\"Validity of the migrad call\">\n",
       "Valid Min.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Validity of parameters\">\n",
       "Valid Param.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is EDM above goal EDM?\">\n",
       "Above EDM\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"2\" title=\"Did last migrad call reach max call limit?\">\n",
       "Reached call limit\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" colspan=\"2\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" title=\"Did Hesse fail?\">\n",
       "Hesse failed\n",
       "</td>\n",
       "<td align=\"center\" title=\"Has covariance matrix\">\n",
       "Has cov.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is covariance matrix accurate?\">\n",
       "Accurate\n",
       "</td>\n",
       "<td align=\"center\" title=\"Is covariance matrix positive definite?\">\n",
       "Pos. def.\n",
       "</td>\n",
       "<td align=\"center\" title=\"Was positive definiteness enforced by Minuit?\">\n",
       "Forced\n",
       "</td>\n",
       "</tr>\n",
       "<tr>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "True\n",
       "</td>\n",
       "<td align=\"center\" style=\"background-color:#92CCA6;\">\n",
       "False\n",
       "</td>\n",
       "</tr>\n",
       "</table>\n",
       "<table>\n",
       "<tr style=\"background-color:#F4F4F4;\">\n",
       "<td/>\n",
       "<th title=\"Variable name\">\n",
       "Name\n",
       "</th>\n",
       "<th title=\"Value of parameter\">\n",
       "Value\n",
       "</th>\n",
       "<th title=\"Hesse error\">\n",
       "Hesse Error\n",
       "</th>\n",
       "<th title=\"Minos lower error\">\n",
       "Minos Error-\n",
       "</th>\n",
       "<th title=\"Minos upper error\">\n",
       "Minos Error+\n",
       "</th>\n",
       "<th title=\"Lower limit of the parameter\">\n",
       "Limit-\n",
       "</th>\n",
       "<th title=\"Upper limit of the parameter\">\n",
       "Limit+\n",
       "</th>\n",
       "<th title=\"Is the parameter fixed in the fit\">\n",
       "Fixed\n",
       "</th>\n",
       "</tr>\n",
       "<tr style=\"background-color:#FFFFFF;\">\n",
       "<td>\n",
       "0\n",
       "</td>\n",
       "<td>\n",
       "x0\n",
       "</td>\n",
       "<td>\n",
       "1.25\n",
       "</td>\n",
       "<td>\n",
       "0.15\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "<tr style=\"background-color:#F4F4F4;\">\n",
       "<td>\n",
       "1\n",
       "</td>\n",
       "<td>\n",
       "x1\n",
       "</td>\n",
       "<td>\n",
       "1.5\n",
       "</td>\n",
       "<td>\n",
       "0.5\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "<tr style=\"background-color:#FFFFFF;\">\n",
       "<td>\n",
       "2\n",
       "</td>\n",
       "<td>\n",
       "x2\n",
       "</td>\n",
       "<td>\n",
       "1.6\n",
       "</td>\n",
       "<td>\n",
       "1.5\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "<td>\n",
       "\n",
       "</td>\n",
       "</tr>\n",
       "</table>\n"
      ],
      "text/plain": [
       "------------------------------------------------------------------\n",
       "| FCN = 23.14                   |      Ncalls=91 (91 total)      |\n",
       "| EDM = 3.12e-05 (Goal: 0.0002) |            up = 1.0            |\n",
       "------------------------------------------------------------------\n",
       "|  Valid Min.   | Valid Param.  | Above EDM | Reached call limit |\n",
       "------------------------------------------------------------------\n",
       "|     True      |     True      |   False   |       False        |\n",
       "------------------------------------------------------------------\n",
       "| Hesse failed  |   Has cov.    | Accurate  | Pos. def. | Forced |\n",
       "------------------------------------------------------------------\n",
       "|     False     |     True      |   True    |   True    | False  |\n",
       "------------------------------------------------------------------\n",
       "------------------------------------------------------------------------------------------\n",
       "|   | Name |   Value   | Hesse Err | Minos Err- | Minos Err+ | Limit-  | Limit+  | Fixed |\n",
       "------------------------------------------------------------------------------------------\n",
       "| 0 | x0   |   1.25    |   0.15    |            |            |         |         |       |\n",
       "| 1 | x1   |    1.5    |    0.5    |            |            |         |         |       |\n",
       "| 2 | x2   |    1.6    |    1.5    |            |            |         |         |       |\n",
       "------------------------------------------------------------------------------------------"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = Minuit.from_array_func(cost, np.zeros(3), pedantic=False)\n",
    "m.migrad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:25:44.063672Z",
     "start_time": "2020-02-21T10:25:44.061147Z"
    }
   },
   "outputs": [],
   "source": [
    "par = m.np_values()\n",
    "cov = m.np_covariance()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-21T10:25:44.566228Z",
     "start_time": "2020-02-21T10:25:44.065443Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 16572 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.errorbar(data_x, data_y, sigma_y, sigma_x, fmt=\"o\", label=\"data\")\n",
    "x = np.linspace(data_x[0], data_x[-1], 200)\n",
    "plt.plot(x, f(x, par), label=\"fit\")\n",
    "plt.legend()\n",
    "\n",
    "# check fit quality\n",
    "chi2 = m.fval\n",
    "ndof = len(data_y) - 3\n",
    "plt.title(f\"$\\\\chi^2 / n_\\\\mathrm{{dof}} = {chi2:.2f} / {ndof} = {chi2/ndof:.2f}$\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We obtained a good fit."
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
