1-Introduction-to-Numba.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "<center>\n",
    "    \n",
    "# 1 - Introduction to `Numba`\n",
    "    \n",
    "<img src=\"imgs/numba_logo.png\" alt=\"dask\" width=\"300\"/>\n",
    "\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Two approaches to write scientific or numeric software:\n",
    "\n",
    "- Traditionally, scripting everything on **C** or **C++**, then writting **Python** wrappers as user-friendly interface with code (*bottom-up*).\n",
    "- Nowadays, scripting everything on **Python** and, only when needed for performance, we speed up code with **Cython** or **Numba** (*top-down*). \n",
    "\n",
    "<center>\n",
    "<img src=\"imgs/two_approaches.png\" alt=\"numba\" width=\"400\"/>\n",
    "   \n",
    "##### \"For day-to-day scientific data exploration, speed-of-development is primary, and speed-of-execution is often secondary.\"\n",
    "    \n",
    "Jake Vanderplas.\n",
    "\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Just-in-time Compilation with Numba\n",
    "\n",
    "<center>\n",
    "<img src=\"imgs/compiled_vs_interpreted.png\" alt=\"numba\" width=\"1000\"/>\n",
    "</center>\n",
    "\n",
    "\n",
    "- Numba compiles Python functions *on the fly* to machine code using LLVM\n",
    "- Easy to use: just *decorate* your Python function with `@numba.jit`.  \n",
    "- Compatibility with Numpy arrays\n",
    "- Enables parallelization (use all the CPU cores in your machine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import urllib.request\n",
    "\n",
    "import numba\n",
    "import numpy as np\n",
    "from skimage.color import rgb2gray\n",
    "from skimage.io import imread\n",
    "from skimage.util import img_as_float\n",
    "\n",
    "from utils import show_images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 0) Let's load some data! \n",
    "\n",
    "For this introduction we are going to use some image that fits in memory. We are going to load an RGB [picture of Manhattan](https://unsplash.com/photos/5ULk8EgE8tg) taken by Miltiadis Fragkidis. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download image --  (approx. 2Mb)\n",
    "\n",
    "url = \"https://unsplash.com/photos/5ULk8EgE8tg/download?ixid=MnwxMjA3fDB8MXxhbGx8fHx8fHx8fHwxNjc2OTAxMjYx&force=true\"\n",
    "manhattan_image = os.path.join(\"data\", \"manhattan.jpg\")\n",
    "urllib.request.urlretrieve(url, manhattan_image);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load image and convert to float\n",
    "\n",
    "img = img_as_float(imread(manhattan_image))\n",
    "\n",
    "print(f\"The image has shape {img.shape}\")\n",
    "print(\n",
    "    f\"The full image has {np.prod(img.shape[:2]) / 1e6 : .0f} MPix, and occupies {img.nbytes / 1e6:.0f} Mb in RAM\"\n",
    ")\n",
    "\n",
    "# RGB -> Grayscale\n",
    "img_gray = rgb2gray(img)\n",
    "\n",
    "# Plot\n",
    "show_images(images=[img, img], zoom=[False, True], titles=\"Original \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Uniform filter \n",
    "\n",
    "As part of one pipeline, we are interested in performing smoothing with a 2-d uniform filter, of the form:\n",
    "\n",
    "$$\n",
    "\\frac{1}{9} \\cdot \\begin{pmatrix} \n",
    "1 & 1 & 1 \\\\\n",
    "1 & 1 & 1 \\\\\n",
    "1 & 1 & 1 \\\\\n",
    "\\end{pmatrix} \n",
    "$$\n",
    "\n",
    "This filter is applied at each pixel of the image, via a convolution.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) Let's first attempt this with pure Python code and time the performance. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def uniform_filter(image):\n",
    "    out = np.zeros_like(image)\n",
    "    for i in range(1, image.shape[0] - 1):\n",
    "        for j in range(1, image.shape[1] - 1):\n",
    "            out[i, j] = (\n",
    "                image[i + -1, j + -1]\n",
    "                + image[i + -1, j + 0]\n",
    "                + image[i + -1, j + 1]\n",
    "                + image[i + 0, j + -1]\n",
    "                + image[i + 0, j + 0]\n",
    "                + image[i + 0, j + 1]\n",
    "                + image[i + 1, j + -1]\n",
    "                + image[i + 1, j + 0]\n",
    "                + image[i + 1, j + 1]\n",
    "            ) / 9\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time smooth_img_gray = uniform_filter(img_gray)\n",
    "\n",
    "show_images(\n",
    "    images=[img_gray, smooth_img_gray, img_gray - smooth_img_gray],\n",
    "    titles=[\"Original\", \"Smoothed\", \"Difference\"],\n",
    "    zoom=True,\n",
    "    cmap=[None, None, \"Accent\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On my machine, his took around 26 seconds to perform. \n",
    "\n",
    "Python is very slow at for-loops because it uses [dynamic typing](https://stackoverflow.com/a/1517670), and at each iteration, the types must be checked. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "`Numba` can look the code ahead at run-time and optimize for repeated or unneeded operations, which can result in speed-ups. This is known as Just-in-time (JIT) compilation. `Numba` can compile Python functions with a the `@numba.jit` decorator. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba\n",
    "\n",
    "\n",
    "@numba.jit(nopython=True)\n",
    "def numba_uniform_filter(image):\n",
    "    out = np.zeros_like(image)\n",
    "    for i in range(1, image.shape[0] - 1):\n",
    "        for j in range(1, image.shape[1] - 1):\n",
    "            out[i, j] = (\n",
    "                image[i + -1, j + -1]\n",
    "                + image[i + -1, j + 0]\n",
    "                + image[i + -1, j + 1]\n",
    "                + image[i + 0, j + -1]\n",
    "                + image[i + 0, j + 0]\n",
    "                + image[i + 0, j + 1]\n",
    "                + image[i + 1, j + -1]\n",
    "                + image[i + 1, j + 0]\n",
    "                + image[i + 1, j + 1]\n",
    "            ) / 9\n",
    "    return out\n",
    "\n",
    "\n",
    "# We could also reuse the previous function and compile it.\n",
    "# Feel free to try if this would yield the same functionality:\n",
    "\n",
    "# @numba.jit\n",
    "# def numba_uniform_filter(image):\n",
    "#     return uniform_filter(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The first time that we run the decorated Python function, Numba will compile it\n",
    "numba_uniform_filter(img_gray)\n",
    "# Time without compilation\n",
    "%timeit smooth_img_gray = numba_uniform_filter(img_gray)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On my machine, his took around 82 miliseconds to perform. \n",
    "\n",
    "That means that only adding one line of code (`@numba.jit`) enables us to speed up our pipeline **317 times faster**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Providing arguments to @JIT\n",
    "Numba allows us to provide some arguments that can accellerate the performance even more. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba\n",
    "\n",
    "\n",
    "@numba.jit(nopython=True, parallel=True, nogil=True, fastmath=True)\n",
    "def numba_uniform_filter(x):\n",
    "    out = np.zeros_like(x)\n",
    "    # rule of thumb -- parallalize outermost loop\n",
    "    for i in numba.prange(1, x.shape[0] - 1):\n",
    "        for j in range(1, x.shape[1] - 1):\n",
    "            out[i, j] = (\n",
    "                x[i + -1, j + -1]\n",
    "                + x[i + -1, j + 0]\n",
    "                + x[i + -1, j + 1]\n",
    "                + x[i + 0, j + -1]\n",
    "                + x[i + 0, j + 0]\n",
    "                + x[i + 0, j + 1]\n",
    "                + x[i + 1, j + -1]\n",
    "                + x[i + 1, j + 0]\n",
    "                + x[i + 1, j + 1]\n",
    "            ) / 9\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Time execution\n",
    "\n",
    "# 1) The first time that we run the decorated Python function, Numba will compile it\n",
    "numba_uniform_filter(img_gray)\n",
    "\n",
    "# 2) Time without compilation\n",
    "%timeit smooth_img_gray = numba_uniform_filter(img_gray)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare that with a Numpy implementation and with the Scipy's optimized `scipy.ndimage.uniform_filter`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def numpy_uniform_filter(image):\n",
    "    out = np.zeros_like(image)\n",
    "\n",
    "    out[1:-1, 1:-1] += image[2:, 2:]\n",
    "    out[1:-1, 1:-1] += image[2:, 1:-1]\n",
    "    out[1:-1, 1:-1] += image[2:, :-2]\n",
    "    out[1:-1, 1:-1] += image[1:-1, 2:]\n",
    "    out[1:-1, 1:-1] += image[1:-1, 1:-1]\n",
    "    out[1:-1, 1:-1] += image[1:-1, :-2]\n",
    "    out[1:-1, 1:-1] += image[:-2, 2:]\n",
    "    out[1:-1, 1:-1] += image[:-2, 1:-1]\n",
    "    out[1:-1, 1:-1] += image[:-2, :-2]\n",
    "    out /= 9\n",
    "    return out\n",
    "\n",
    "\n",
    "%timeit smooth_img_gray = numpy_uniform_filter(img_gray)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.ndimage\n",
    "\n",
    "%timeit smooth_img_gray = scipy.ndimage.uniform_filter(img_gray, size=3, mode=\"constant\", cval=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Stencils\n",
    "\n",
    "Stencils are a class of linear operators where output $x[i]$ is given by a weighted linear combination of its neighborhod: \n",
    "\n",
    "$$y[i] = \\sum_{k\\in\\mathcal{N_{i}}}\\alpha_{i} x[i-k]$$\n",
    "\n",
    "Notable examples include multi-dimensional convolution and correlation.\n",
    "\n",
    "**Numba** provides the [`@stencil` decorator](https://numba.pydata.org/numba-doc/latest/user/stencil.html) so that users may easily specify a stencil kernel and Numba then generates the looping code necessary to apply that kernel to some input array. \n",
    "\n",
    "Thus, the stencil decorator allows clearer, more concise code and in conjunction with the parallel jit option enables higher performance through parallelization of the stencil execution.\n",
    "\n",
    "#### Let's create a Numba stencil, and JIT compile it.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.stencil\n",
    "def _smooth_stencil(x):\n",
    "    return (\n",
    "        x[-1, -1]\n",
    "        + x[-1, 0]\n",
    "        + x[-1, 1]\n",
    "        + x[0, -1]\n",
    "        + x[0, 0]\n",
    "        + x[0, 1]\n",
    "        + x[1, -1]\n",
    "        + x[1, 0]\n",
    "        + x[1, 1]\n",
    "    ) / 9\n",
    "\n",
    "\n",
    "@numba.jit(nopython=True, parallel=True)\n",
    "def smooth_stencil(x):\n",
    "    return _smooth_stencil(x)\n",
    "\n",
    "\n",
    "# Time execution\n",
    "\n",
    "# 1) The first time that we run the decorated Python function, Numba will compile it\n",
    "smooth_stencil(img_gray)\n",
    "\n",
    "# 2) Time without compilation\n",
    "%timeit smooth_img_gray = smooth_stencil(img_gray)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Numba benchmarking\n",
    "\n",
    "| Architecture                  | Time      |   Ratio   |\n",
    "|-------------------------------|-----------|-----------|\n",
    "| Single CPU Core (Python)      |  26 s     |  x1       |\n",
    "| Single CPU Core (Numpy)       | 260 ms    |  x100     |\n",
    "| Single CPU Core (Scipy)       | 178 ms    |  x150     |\n",
    "| Single CPU Core (Numba)       |  82 ms    |  x300     |\n",
    "| 16 CPU Cores (parallel=True)  |  50 ms    |  x500     |\n",
    "| 16 CPU Cores (stencil)        |  25 ms    |  x1000    |\n",
    "\n",
    "\n",
    "### Some notes:\n",
    "The argument `nopython=True` in Numba's `@jit` compiler defines a compilation mode that generates code that does not access the Python C API. This compilation mode produces the highest performance code, but requires that the native types of all values in the function can be inferred. Unless otherwise instructed, the `@jit` decorator will automatically fall back to object mode if `nopython` mode cannot be used.\n",
    "\n",
    "As a side note, if compilation time is an issue, Numba JIT supports on-disk caching of compiled functions and also has an [Ahead-Of-Time](https://numba.readthedocs.io/en/stable/user/pycc.html) compilation mode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bonus:\n",
    "Other JIT [options](https://numba.readthedocs.io/en/stable/user/jit.html#compilation-options) to play with: \n",
    "- fastmath\n",
    "- nopython\n",
    "- cache\n",
    "- nogil\n",
    "\n",
    "# Bonus:\n",
    "[Numba's `@guvectorize`](https://numba.pydata.org/numba-doc/dev/user/vectorize.html)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}