01 - Go fast with Numba.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Numba\n",
    "Numba can be installed both with `conda` or `pip`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install numba numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that Numba is designed to be using in conjunction with NumPy, so for other libraries like Pandas, it may not work as well."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How compliing makes it faster\n",
    "In this example below, we will see how we can make a function go faster by compiling it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import jit\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "x = np.arange(100).reshape(10, 10)\n",
    "\n",
    "def normal(a):\n",
    "    trace = 0.0\n",
    "    for i in range(a.shape[0]):\n",
    "        trace += np.tanh(a[i, i])\n",
    "    return a + trace\n",
    "\n",
    "@jit(nopython=True)\n",
    "def go_fast(a): # Function is compiled and runs in machine code\n",
    "    trace = 0.0\n",
    "    for i in range(a.shape[0]):\n",
    "        trace += np.tanh(a[i, i])\n",
    "    return a + trace\n",
    "\n",
    "# TIME WITHOUT COMPILING\n",
    "start = time.perf_counter()\n",
    "normal(x)\n",
    "end = time.perf_counter()\n",
    "print(\"Elapsed (without compilation) = {}s\".format((end - start)))\n",
    "\n",
    "# DO NOT REPORT THIS... COMPILATION TIME IS INCLUDED IN THE EXECUTION TIME!\n",
    "start = time.perf_counter()\n",
    "go_fast(x)\n",
    "end = time.perf_counter()\n",
    "print(\"Elapsed (with compilation) = {}s\".format((end - start)))\n",
    "\n",
    "# NOW THE FUNCTION IS COMPILED, RE-TIME IT EXECUTING FROM CACHE\n",
    "start = time.perf_counter()\n",
    "go_fast(x)\n",
    "end = time.perf_counter()\n",
    "print(\"Elapsed (after compilation) = {}s\".format((end - start)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the time that run the function with compliation for the 1st time is much slower than runing it without the complilation. However, after compliation, it is way faster.\n",
    "\n",
    "**Question: Why it is slower the 1st time?**\n",
    "\n",
    "So we always have a trade off. If the function is only used a couple of times, it may as will be better without the compliation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vector opteraions VS loop operations\n",
    "Before we moved on and talk more about how we can make use of Numba to make things faster, we should also knows that Numpy loves vectors. Let's compare the same trigonometric identity opeataion `cos(x)^2 + sin(x)^2` with vector operation and looping through it's elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.arange(1.e7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ident_np(x):\n",
    "    return np.cos(x) ** 2 + np.sin(x) ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ident_np(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ident_loops(x):\n",
    "    r = np.empty_like(x)\n",
    "    n = len(x)\n",
    "    for i in range(n):\n",
    "        r[i] = np.cos(x[i]) ** 2 + np.sin(x[i]) ** 2\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ident_loops(arr)  # warning: really slow, can take a few munites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the same operation is almost 10 times slower if it is elementwise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loop is where Numba really shine\n",
    "So it seems Numba loves repatitive operations, what happens if we speed up the above operations with `@njit`? (`@njit` is the short hand for `@jit(nopython=True)`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import njit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def ident_np_njit(x):\n",
    "    return np.cos(x) ** 2 + np.sin(x) ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compliting it\n",
    "ident_np_njit(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ident_np_njit(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def ident_loops_njit(x):\n",
    "    r = np.empty_like(x)\n",
    "    n = len(x)\n",
    "    for i in range(n):\n",
    "        r[i] = np.cos(x[i]) ** 2 + np.sin(x[i]) ** 2\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compiling it\n",
    "ident_loops_njit(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ident_loops_njit(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you see now, the different between the vetor operation and the elementwise looping operation is almost none. In other words, the `@njit` speed up the operation almost 10 folds.\n",
    "\n",
    "There are still improvement of the vector operations but it is not as much as the elementwise looping.\n",
    "\n",
    "**Exercise: Can you change the following into elementwise operation?**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def julia_fast(mesh, c=-1, num_iter=10, radius=2):\n",
    "\n",
    "    z = mesh.copy()\n",
    "    diverge_len = np.zeros(z.shape)\n",
    "\n",
    "    for i in range(num_iter):\n",
    "        conv_mask = np.abs(z) < radius\n",
    "        z[conv_mask] = np.square(z[conv_mask]) + c\n",
    "        diverge_len[conv_mask] += 1\n",
    "\n",
    "    return diverge_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Un-commend the line below to see how we do it.\n",
    "#%load element_op.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fast math\n",
    "In calculations, [IEEE 754 compliance](https://en.wikipedia.org/wiki/IEEE_754) is followed to make sure the calcuation is accurate to the precision specified. This will lock the order of the operation and limit the speed of some processed. When the precision is not a huge concern. We can relex the compliance and perform fastmath to gain some speed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.arange(1.e7)\n",
    "\n",
    "@njit(fastmath=False)\n",
    "def do_sum(A):\n",
    "    acc = 0.\n",
    "    # without fastmath, this loop must accumulate in strict order\n",
    "    for x in A:\n",
    "        acc += np.sqrt(x)\n",
    "    return acc\n",
    "\n",
    "@njit(fastmath=True)\n",
    "def do_sum_fast(A):\n",
    "    acc = 0.\n",
    "    # with fastmath, the reduction can be vectorized as floating point\n",
    "    # reassociation is permitted.\n",
    "    for x in A:\n",
    "        acc += np.sqrt(x)\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compile the functions\n",
    "do_sum(arr)\n",
    "do_sum_fast(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "do_sum(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "do_sum_fast(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you see, the fast math version is takes almost only half the time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Make use of multiple cores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit(parallel=True)\n",
    "def ident_parallel(x):\n",
    "    return np.cos(x) ** 2 + np.sin(x) ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compiling it\n",
    "ident_parallel(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ident_parallel(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Depending on how many cores your have on your computer, you will see carious degrees of improvement. On my 3.1 GHz Dual-Core Intel Core i7 early 2015 MacBook Pro, it is slightly faster than the one using njit and almost 7 times faster than without compilation. On my new 8 Core M2 2022 Macbook air, it is 10 times faster. \n",
    "\n",
    "**Question: What about yours?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we get a feeling of how Numba makes things faster and when Numba can makes things really fast, lets go to the second notebook and see what can go wrong and how to fix it."
   ]
  }
 ],
 "metadata": {
  "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}