numba_vs_cython.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This simple notebook is an update of the notebook by Jake Vanderplas at https://jakevdp.github.io/blog/2013/06/15/numba-vs-cython-take-2/ to run on Python 3.5. It compares usage of cython and numba illustrating the ability to vastly increase speed of computation with some simple tricks. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A very pithy description of this is that `X` contains $1000$ points in 3D space. The Euclidian distance between any two points can be calculated (vector length). This can be represented by a 1000 $\\times$ 1000 array containing all possible pairs. Calculating this in Python is slow. This can be easily sped up by \"precompiling\" the code. \n",
    "\n",
    "For more details on how to migrate to compiled code see [Hans Lantangen's tutorial](http://hplgit.github.io/primer.html/doc/pub/cython/cython-solarized.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The Cython extension is already loaded. To reload it, use:\n",
      "  %reload_ext Cython\n"
     ]
    }
   ],
   "source": [
    "# We are going to load the Cython extension for Jupyter. This allows us to \n",
    "# put a few things into Python code. Cython then turns it into C code, \n",
    "# compiles it, and runs it in the background. \n",
    "\n",
    "%load_ext Cython\n",
    "\n",
    "import numpy as np\n",
    "X = np.random.random((1000, 3))\n",
    "D = np.empty((1000, 1000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pure python version\n",
    "\n",
    "def pairwise_python(X, D):\n",
    "    M = X.shape[0]\n",
    "    N = X.shape[1]\n",
    "    for i in range(M):\n",
    "        for j in range(M):\n",
    "            d = 0.0\n",
    "            for k in range(N):\n",
    "                tmp = X[i, k] - X[j, k]\n",
    "                d += tmp * tmp\n",
    "            D[i, j] = np.sqrt(d)\n",
    "    return D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# numba version\n",
    "\n",
    "import numpy as np\n",
    "from numba import double\n",
    "from numba.decorators import jit\n",
    "\n",
    "# The next line is the only \"change to the code\"\n",
    "@jit     # (arg_types=[double[:,:], double[:,:]])\n",
    "def pairwise_numba(X, D):\n",
    "    M = X.shape[0]\n",
    "    N = X.shape[1]\n",
    "    for i in range(M):\n",
    "        for j in range(M):\n",
    "            d = 0.0\n",
    "            for k in range(N):\n",
    "                tmp = X[i, k] - X[j, k]\n",
    "                d += tmp * tmp\n",
    "            D[i, j] = np.sqrt(d)\n",
    "    return D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "# The above declares that this cell is Cython code.\n",
    "\n",
    "# Load some cheats that help cython\n",
    "cimport cython\n",
    "\n",
    "# Load the sqrt function, but with cimport, also import knowledge of the variable \n",
    "# type used (it's a floating point)\n",
    "from libc.math cimport sqrt\n",
    "\n",
    "# Tell cython not to check if my indices called ae actually inside the array\n",
    "# This is a good move if you are ABSOLUTELY sure of it. It's a bad move otherwise. \n",
    "@cython.boundscheck(False)\n",
    "\n",
    "# This disallows use of negative indices. \n",
    "@cython.wraparound(False)\n",
    "def pairwise_cython(double[:, ::1] X, double[:, ::1] D):\n",
    "    cdef int M = X.shape[0]\n",
    "    cdef int N = X.shape[1]\n",
    "    cdef double tmp, d\n",
    "    for i in range(M):\n",
    "        for j in range(M):\n",
    "            d = 0.0\n",
    "            for k in range(N):\n",
    "                tmp = X[i, k] - X[j, k]\n",
    "                d += tmp * tmp\n",
    "            D[i, j] = sqrt(d)\n",
    "    return D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.83 s ± 139 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit pairwise_python(X, D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.72 ms ± 92.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit pairwise_numba(X, D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.82 ms ± 123 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit pairwise_cython(X, D)"
   ]
  }
 ],
 "metadata": {
  "author": "mes",
  "hide_input": false,
  "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.7"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 0,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}