numba.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Numba - Just In Time Compilation for Python\n",
    "\n",
    "Some technologies are so easy to use, and the effects are so powerful, that they appear to be magic.  Numba is one such technology.  \n",
    "\n",
    "Let's start with Paul's prime sieve code from earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import  math\n",
    "def sieve_primes(n):\n",
    "    a = [True for x in range(n + 1)]\n",
    "    i = 2\n",
    "    while i <= math.sqrt(n):\n",
    "        if a[i]:\n",
    "            for j in range(i*i, n + 1, i):\n",
    "                a[j] = False\n",
    "        i += 1\n",
    "    return [i for i in range(2, len(a)) if a[i]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check it's working OK\n",
    "sieve_primes(30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Time it for all primes less than 5 million\n",
    "N = 5000000\n",
    "original_speed = %timeit -o sieve_primes(N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now with Numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import jit\n",
    "\n",
    "@jit\n",
    "def numba_sieve_primes(n):\n",
    "    a = [True for x in range(n + 1)]\n",
    "    i = 2\n",
    "    while i <= math.sqrt(n):\n",
    "        if a[i]:\n",
    "            for j in range(i*i, n + 1, i):\n",
    "                a[j] = False\n",
    "        i += 1\n",
    "    return [i for i in range(2, len(a)) if a[i]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check it's working OK\n",
    "numba_sieve_primes(30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Time it \n",
    "numba_speed = %timeit -o numba_sieve_primes(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Numba is {0:.1f} times faster than pure Python for the primes function'.format(original_speed.best/numba_speed.best))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**A significant speed up and all we did was type @jit**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What is Numba doing?\n",
    "#### Just in time compilation and compilation overhead\n",
    "\n",
    "By adding the @jit decorator to our Python function, we are asking Numba to try to compile the function to machine code. The compilation happens the first time the function is called, hence the name **jit - just in time** compilation.  The compiler in question is the powerful open source compiler framework, LLVM.\n",
    "\n",
    "This means that the first time a numba function is called there will be a delay due to the compilation but all subsequent calls will be fast.  Here's an explicit demonstration using `monte_carlo_pi`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import jit\n",
    "import random\n",
    "import time\n",
    "\n",
    "@jit\n",
    "def monte_carlo_pi(nsamples):\n",
    "    acc = 0\n",
    "    for i in range(nsamples):\n",
    "        x = random.random()\n",
    "        y = random.random()\n",
    "        if (x ** 2 + y ** 2) < 1.0:\n",
    "            acc += 1\n",
    "    return 4.0 * acc / nsamples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "pi = monte_carlo_pi(N)\n",
    "end = time.time()\n",
    "print(\"Elapsed (with compilation) = %s\" % (end - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "pi = monte_carlo_pi(N)\n",
    "end = time.time()\n",
    "print(\"Elapsed (no compilation) = %s\" % (end - start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have compilation out of the way, let's time it more robustly using timeit and compare with a version that doesn't make use of numba."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numba_time = %timeit -o monte_carlo_pi(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Plain Python version that doesn't use numba\n",
    "def python_monte_carlo_pi(nsamples):\n",
    "    acc = 0\n",
    "    for i in range(nsamples):\n",
    "        x = random.random()\n",
    "        y = random.random()\n",
    "        if (x ** 2 + y ** 2) < 1.0:\n",
    "            acc += 1\n",
    "    return 4.0 * acc / nsamples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "python_time = %timeit -o python_monte_carlo_pi(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Numba version of monte_carlo_pi is {0:.1f} times faster than plain Python'.format(python_time.best/numba_time.best))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba doesn't work on everything\n",
    "\n",
    "Numba has a lot of functionality but there is still a lot of Python that it does not know about and hence cannot compile to machine code.  The following function is given as an example in the numba documentation as one that Numba can't do anything with"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import jit\n",
    "import pandas as pd\n",
    "\n",
    "x = {'a': [1, 2, 3], 'b': [20, 30, 40]}\n",
    "\n",
    "@jit\n",
    "def use_pandas(a): \n",
    "    df = pd.DataFrame.from_dict(a) \n",
    "    df += 1                        \n",
    "    return df.cov()                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_pandas(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might expect an error message when numba is given a function it can't compile but instead everything seemed to have just worked.\n",
    "\n",
    "To dig deeper, let's create a version of this function without any numba decorations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def use_pandas_nonumba(a): \n",
    "    df = pd.DataFrame.from_dict(a) \n",
    "    df += 1                        \n",
    "    return df.cov()   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit use_pandas_nonumba(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit use_pandas(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The numba version is actually slightly slower!  Why?\n",
    "\n",
    "What is happening behind the scenes is that Numba tries to compile the function, realises it doesn't know anything about pandas and so gives up and just runs the function as a normal Python routine.  All of the checking adds a little overhead and we get the opposite of what we were trying to achieve -- slightly slower code instead of much faster code.\n",
    "\n",
    "You might want to know about Numba's failure. If so, you can use the `nopython` flag which tells Numba that you only want it to produce pure compiled functions -- no interaction with the python interpreter at all."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def use_pandas(a): \n",
    "    df = pd.DataFrame.from_dict(a) \n",
    "    df += 1                        \n",
    "    return df.cov()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This will no result in an error rather than a slightly slower version of the Python funcion\n",
    "use_pandas(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Going Parallel\n",
    "\n",
    "It is also possible to create parallel code using Numba.  The easiest way to proceed is to add `parallel=True` to the Numba decorator for a function that's already working in `nopython` mode and change `range` to its parallel equivalent `prange`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import prange\n",
    "\n",
    "@jit(parallel=True)\n",
    "def parallel_monte_carlo_pi(nsamples):\n",
    "    acc = 0\n",
    "    for i in prange(nsamples):\n",
    "        x = random.random()\n",
    "        y = random.random()\n",
    "        if (x ** 2 + y ** 2) < 1.0:\n",
    "            acc += 1\n",
    "    return 4.0 * acc / nsamples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numba_parallel_time = %timeit -o parallel_monte_carlo_pi(N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that Numba is doing certain things for us automatically to make this work. For example Numba ensures that in the loop `x`, `y` and `acc` are unique to each thread, and then performs a reduction operation at the end of the loop to combine the individual `acc` values into a single value.\n",
    "\n",
    "I recommend reading the [Numba documentation](https://numba.pydata.org/numba-doc/latest/user/parallel.html) to learn more about the automatic parallelization. In particular:\n",
    "\n",
    "- Numba attempts to automatically identify operations that can run in parallel, this requires no modifications to the user program.\n",
    "- When using `prange` the user is required to make sure the loop interations don't have dependencies, except for supported reductions. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall how long this took in serial mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numba_serial_time = %timeit -o monte_carlo_pi(N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and the original, plain Python version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "python_time = %timeit -o python_monte_carlo_pi(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('The parallel numba version was {0:.1f} times faster than the serial numba version and'.\n",
    "      format(numba_serial_time.best/numba_parallel_time.best))\n",
    "print('it was {0:.1f} times faster than the plain python version'.\n",
    "     format(python_time.best/numba_parallel_time.best))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numba also automatically ensures that each thread will get a different random number seed but, at the time of writing, I can't figure out how to determine what these seeds are.  \n",
    "\n",
    "It should also be noted that the underlying Random Number Generator used by **random.random** is [Mersenne Twister](https://en.wikipedia.org/wiki/Mersenne_Twister) which has a very large period.  Although it is likely that multiple random streams created with different seeds won't overlap, it isn't guaranteed \n",
    "\n",
    "A discussion on this topic can be found at https://github.com/numba/numba/issues/2486 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "1. Produce a parallel Numba version of the `sieve_primes` function.\n",
    "2. Try to write a parallel Numba version of the Asian option pricing routine below\n",
    "3. If you have any python code that does something of interest to you, try to write an accelerated version in Numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math # This is the standard Pyton math module. Not the numpy one\n",
    "import random # This is the standard Python random module.  Not the numpy one\n",
    "\n",
    "def Asian(so,k,r,v,t,m,n):\n",
    "    \"\"\"\n",
    "    I have not identified what the arguments mean since the original MATLAB code didn't either. \n",
    "    This doc-string will be updated once I learn what they are!\n",
    "    \"\"\"\n",
    "    dt = t/m\n",
    "    AsianPayoffSum = 0\n",
    "    for i in range(1,n+1):\n",
    "        s = so\n",
    "        stSum = so\n",
    "        at = so\n",
    "        for j in range(1,m+1):\n",
    "            st = s * math.exp(((r-v**2/2)*dt) + (v*random.normalvariate(0,1)*math.sqrt(dt)))\n",
    "            stSum = stSum + st \n",
    "            at = stSum/(j+1)\n",
    "            s = st\n",
    "        AsianPayoff = max(at-k,0);\n",
    "        AsianPayoffSum = AsianPayoffSum + AsianPayoff;\n",
    "    AsianCall = math.exp(-r*t)*(AsianPayoffSum/n)\n",
    "    return(AsianCall)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba - Next steps\n",
    "\n",
    "This session has provided just an outline of some of the functionality provided by Numba.  Some of the topics that haven't been covered include\n",
    "\n",
    "* GPU Support - Numba has strong support for [NVIDIA GPUs](http://numba.pydata.org/numba-doc/latest/cuda/index.html) and some [support for AMD](http://numba.pydata.org/numba-doc/latest/roc/index.html).\n",
    "* Numba has support for different [threading layers](http://numba.pydata.org/numba-doc/latest/user/threading-layer.html) including OpenMP and Intel Threading Building Blocks (TBB)\n",
    "* Numba can make use of the [Intel Short Vector Math Library (SVML)](http://numba.pydata.org/numba-doc/latest/user/performance-tips.html#intel-svml) for fast evaluation of trigonometric functions, square roots, exponentials and many more.\n",
    "* The code compiled by Numba makes use of many of your CPU's features including SIMD instructions.  Detailed diagnostics concerning the compilation can be obtained."
   ]
  }
 ],
 "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.3"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}