{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# JIT compilation and parallelization with Numba\n",
    "\n",
    "This is a topical tutorial. If you are new to iminuit, you should go through the basic tutorial first. \n",
    "\n",
    "The time that iminuit requires to return a result, if that time is perceptible, is usually dominated by the execution time of the cost function. To get good performance, it recommended to use array arthimetic and scipy and numpy functions in the body of the cost function. Python loops should be avoided, but if they are unavoidable, [numba](https://numba.pydata.org/) can help. Numba can also parallelize numerical calculations to make full use of multi-core CPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install matplotlib numpy numba scipy iminuit\n",
    "from iminuit import Minuit\n",
    "import numpy as np\n",
    "import numba as nb\n",
    "import math\n",
    "from scipy.stats import expon, norm\n",
    "from matplotlib import pyplot as plt\n",
    "from argparse import Namespace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate some toy data\n",
    "\n",
    "The standard fit in particle physics is the fit of a peak over some smooth background. We generate a Gaussian peak over exponential background, using scipy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 8512 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "np.random.seed(1)  # fix seed\n",
    "\n",
    "# true parameters for signal and background\n",
    "truth = Namespace(n_sig=1000, f_bkg=10, sig=(5.0, 0.5), bkg=(0.0, 4.0))\n",
    "n_bkg = truth.n_sig * truth.f_bkg\n",
    "\n",
    "# make a data set\n",
    "x = np.empty(truth.n_sig + n_bkg)\n",
    "\n",
    "# fill m variables\n",
    "x[:truth.n_sig] = norm(*truth.sig).rvs(truth.n_sig)\n",
    "x[truth.n_sig:] = expon(*truth.bkg).rvs(n_bkg)\n",
    "\n",
    "# cut a range in x\n",
    "xrange = np.array((1.0, 9.0))\n",
    "ma = (xrange[0] < x) & (x < xrange[1])\n",
    "x = x[ma]\n",
    "\n",
    "plt.hist((x[truth.n_sig:], x[:truth.n_sig]), bins=50, stacked=True, label=(\"background\", \"signal\"))\n",
    "plt.xlabel(\"x\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ideal starting values for iminuit\n",
    "start = np.array((truth.n_sig, n_bkg, truth.sig[0], truth.sig[1], truth.bkg[1]))\n",
    "\n",
    "\n",
    "# iminuit instance factory, will be called a lot in the benchmarks blow\n",
    "def m_init(fcn):\n",
    "    return Minuit.from_array_func(fcn,\n",
    "                                  start,\n",
    "                                  limit=((0, None), (0, None), None, (0, None),\n",
    "                                         (0, None)),\n",
    "                                  name=(\"ns\", \"nb\", \"mu\", \"sigma\", \"lambd\"),\n",
    "                                  errordef=Minuit.LIKELIHOOD,\n",
    "                                  pedantic=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-46424.886640130324"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# extended likelihood (https://doi.org/10.1016/0168-9002(90)91334-8)\n",
    "# this version uses numpy and scipy and array arithmetic\n",
    "def nll(par):\n",
    "    n_sig, n_bkg, mu, sigma, lambd = par\n",
    "    s = norm(mu, sigma)\n",
    "    b = expon(0, lambd)\n",
    "    # normalisation factors are needed for pdfs, since x range is restricted\n",
    "    sn = s.cdf(xrange)\n",
    "    bn = b.cdf(xrange)\n",
    "    sn = sn[1] - sn[0]\n",
    "    bn = bn[1] - bn[0]\n",
    "    return (n_sig + n_bkg) - np.sum(np.log(s.pdf(x) / sn * n_sig + b.pdf(x) / bn * n_bkg))\n",
    "\n",
    "\n",
    "nll(start)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "334 ms ± 5.66 ms per loop (mean ± std. dev. of 3 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 3 -n 1\n",
    "m = m_init(nll) # setup time is negligible\n",
    "m.migrad();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "About 2 seconds on this computer. Let's see whether we can beat that. The code above is already pretty fast, because numpy and scipy routines are fast, and we spend most of the time in those. But these implementations do not parallelize the execution and are not optimised for this particular CPU, unlike numba-jitted functions.\n",
    "\n",
    "To use numba, in theory we just need to put the `njit` decorator on top of the function, but often that doesn't work out of the box. numba understands many numpy functions, but no scipy. We must evaluate the code that uses scipy in 'object mode', which is numba-speak for calling into the Python interpreter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-46424.88664013033"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# first attempt to use numba\n",
    "@nb.njit(parallel=True)\n",
    "def jitted_nll_simple(par):\n",
    "    n_sig, n_bkg, mu, sigma, lambd = par\n",
    "    with nb.objmode(spdf=\"float64[:]\", bpdf=\"float64[:]\", sn=\"float64\", bn=\"float64\"):\n",
    "        s = norm(mu, sigma)\n",
    "        b = expon(0, lambd)\n",
    "        # normalisation factors are needed for pdfs, since x range is restricted\n",
    "        sn = np.diff(s.cdf(xrange))[0]\n",
    "        bn = np.diff(b.cdf(xrange))[0]\n",
    "        spdf = s.pdf(x)\n",
    "        bpdf = b.pdf(x)\n",
    "    no = n_sig + n_bkg\n",
    "    return no - np.sum(np.log(spdf / sn * n_sig + bpdf / bn * n_bkg))\n",
    "\n",
    "jitted_nll_simple(start) # test and warm-up JIT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "368 ms ± 9.69 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 3 -n 10\n",
    "m = m_init(jitted_nll_simple)\n",
    "m.migrad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is only marginally faster. Let's break the original function down by parts to see why."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.22 ms ± 89.7 µs per loop (mean ± std. dev. of 5 runs, 50 loops each)\n",
      "1.11 ms ± 43.8 µs per loop (mean ± std. dev. of 5 runs, 50 loops each)\n",
      "246 µs ± 3.17 µs per loop (mean ± std. dev. of 5 runs, 50 loops each)\n"
     ]
    }
   ],
   "source": [
    "# let's time the body of the function\n",
    "n_sig, n_bkg, mu, sigma, lambd = start\n",
    "s = norm(mu, sigma)\n",
    "b = expon(0, lambd)\n",
    "# normalisation factors are needed for pdfs, since x range is restricted\n",
    "sn = np.diff(s.cdf(xrange))[0]\n",
    "bn = np.diff(b.cdf(xrange))[0]\n",
    "spdf = s.pdf(x)\n",
    "bpdf = b.pdf(x)\n",
    "no = n_sig + n_bkg\n",
    "# no - np.sum(np.log(spdf / sn * n_sig + bpdf / bn * n_bkg))\n",
    "\n",
    "%timeit -r 5 -n 50 norm(*start[2:4]).pdf(x)\n",
    "%timeit -r 5 -n 50 expon(0, start[4]).pdf(x)\n",
    "%timeit -r 5 -n 50 np.sum(np.log(spdf / sn * n_sig + bpdf / bn * n_bkg))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most of the time is spend in those parts that numba could not accelerate and the total time is dominated by the slowest part.\n",
    "\n",
    "This, unfortunately, means we have to do much more manual work to make the function faster, since we have to replace the scipy routines with Python code that numba can accelerate and run in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-46424.886640130324"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kwd = {\"parallel\": True, \"fastmath\": True}\n",
    "\n",
    "@nb.njit(**kwd)\n",
    "def sum_log(fs, spdf, fb, bpdf):\n",
    "    return np.sum(np.log(fs * spdf + fb * bpdf))\n",
    "\n",
    "\n",
    "@nb.njit(**kwd)\n",
    "def norm_pdf(x, mu, sigma):\n",
    "    invs = 1.0 / sigma\n",
    "    z = (x - mu) * invs\n",
    "    invnorm = 1 / np.sqrt(2 * np.pi) * invs\n",
    "    return np.exp(-0.5 * z**2) * invnorm\n",
    "\n",
    "\n",
    "@nb.njit(**kwd)\n",
    "def nb_erf(x):\n",
    "    y = np.empty_like(x)\n",
    "    for i in nb.prange(len(x)):\n",
    "        y[i] = math.erf(x[i])\n",
    "    return y\n",
    "\n",
    "\n",
    "@nb.njit(**kwd)\n",
    "def norm_cdf(x, mu, sigma):\n",
    "    invs = 1.0 / (sigma * np.sqrt(2))\n",
    "    z = (x - mu) * invs\n",
    "    return 0.5 * (1 + nb_erf(z))\n",
    "\n",
    "\n",
    "@nb.njit(**kwd)\n",
    "def expon_pdf(x, lambd):\n",
    "    inv_lambd = 1.0 / lambd\n",
    "    return inv_lambd * np.exp(-inv_lambd * x)\n",
    "\n",
    "\n",
    "@nb.njit(**kwd)\n",
    "def expon_cdf(x, lambd):\n",
    "    inv_lambd = 1.0 / lambd\n",
    "    return 1.0 - np.exp(-inv_lambd * x)\n",
    "\n",
    "\n",
    "def jitted_nll(par):\n",
    "    n_sig, n_bkg, mu, sigma, lambd = par\n",
    "    # normalisation factors are needed for pdfs, since x range is restricted\n",
    "    sn = norm_cdf(xrange, mu, sigma)\n",
    "    bn = expon_cdf(xrange, lambd)\n",
    "    sn = sn[1] - sn[0]\n",
    "    bn = bn[1] - bn[0]\n",
    "    spdf = norm_pdf(x, mu, sigma)\n",
    "    bpdf = expon_pdf(x, lambd)\n",
    "    no = n_sig + n_bkg\n",
    "    return no - sum_log(n_sig / sn, spdf, n_bkg / bn, bpdf)\n",
    "\n",
    "\n",
    "jitted_nll(start)  # test and warm-up JIT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "93.9 ms ± 5.73 ms per loop (mean ± std. dev. of 5 runs, 50 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 5 -n 50\n",
    "m = m_init(jitted_nll)\n",
    "m.migrad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We managed to get a **4x** speed improvement over the initial code, which is impressive, but it cost us a lot of developer time. By putting these functions into a library, however, we only have to pay the developer cost once. We close with a breakdown of how numba accelerated the parts of `jitted_nll`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "300 µs ± 121 µs per loop (mean ± std. dev. of 5 runs, 50 loops each)\n",
      "156 µs ± 6.58 µs per loop (mean ± std. dev. of 5 runs, 50 loops each)\n",
      "257 µs ± 136 µs per loop (mean ± std. dev. of 5 runs, 50 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit -r 5 -n 50 norm_pdf(x, *start[2:4])\n",
    "%timeit -r 5 -n 50 expon_pdf(x, start[4])\n",
    "%timeit -r 5 -n 50 sum_log(n_sig / sn, spdf, n_bkg / bn, bpdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the total speed improvement mostly came from the acceleration of the normal and exponential pdfs, which got more than 4x faster, while the final sum was only accelerated by a factor of about two."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
