examples/5-jit-compiler.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accelerating code with Numba"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy.typing import NDArray\n",
    "import numba"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba and JIT"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Numba](https://numba.pydata.org/numba-doc/latest/user/5minguide.html) is a just-in-time compiler for Python that works best on code that uses NumPy arrays, Numpy functions, and loops. When a call is made to a Numba decorated function (with `@numba.jit()`) it is compiled to machine code “just-in-time” for execution and all or part of your code can subsequently run at native machine code speed!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus, using Numba inside Python classes or with Python objects like DataFrames is not very useful... It is better to use it with standalone functions that do mainly numerical calculations."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: the determinant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def determinant(matrix: NDArray) -> float:\n",
    "    \"\"\"Calculate the determinant of a matrix.\n",
    "\n",
    "    Args:\n",
    "        matrix (np.ndarray): a square matrix.\n",
    "\n",
    "    Returns:\n",
    "        float: the determinant of matrix.\n",
    "    \"\"\"\n",
    "    # Check the ndarray is a square matrix\n",
    "    assert len(matrix.shape) == 2\n",
    "    assert matrix.shape[0] == matrix.shape[1]\n",
    "\n",
    "    dim = matrix.shape[0]\n",
    "\n",
    "    # Convert the matrix to upper triangular form\n",
    "    for col in range(0, dim - 1):\n",
    "        for row in range(col + 1, dim):\n",
    "            if matrix[row, col] != 0.0:\n",
    "                coef = matrix[row, col] / matrix[col, col]\n",
    "                matrix[row, col:dim] = matrix[row, col:dim] - coef * matrix[col, col:dim]\n",
    "\n",
    "    return np.prod(np.diag(matrix))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The behaviour of the `nopython=True` compilation mode is to essentially compile the decorated function so that it will run entirely without the involvement of the Python interpreter. This is the recommended and best-practice way to use the Numba jit decorator as it leads to the best performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True)\n",
    "def jit_determinant(matrix: NDArray) -> float:\n",
    "    \"\"\"Calculate the determinant of a matrix faster using a just in time compiler.\n",
    "    The behaviour of the `nopython=True` compilation mode is to essentially compile the decorated\n",
    "    function so that it will run entirely without the involvement of the Python interpreter.\n",
    "\n",
    "    Args:\n",
    "        matrix (np.ndarray): a square matrix.\n",
    "\n",
    "    Returns:\n",
    "        float: the determinant of matrix.\n",
    "    \"\"\"\n",
    "    # Check the ndarray is a square matrix\n",
    "    assert len(matrix.shape) == 2\n",
    "    assert matrix.shape[0] == matrix.shape[1]\n",
    "\n",
    "    dim = matrix.shape[0]\n",
    "\n",
    "    # Convert the matrix to upper triangular form\n",
    "    for col in range(0, dim - 1):\n",
    "        for row in range(col + 1, dim):\n",
    "            if matrix[row, col] != 0.0:\n",
    "                coef = matrix[row, col] / matrix[col, col]\n",
    "                matrix[row, col:dim] = matrix[row, col:dim] - coef * matrix[col, col:dim]\n",
    "\n",
    "    return np.prod(np.diag(matrix))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Timing"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the jit compiled determinant is much faster than the interpreted version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Timing the basic determinant:\n",
      "300 µs ± 2.06 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
      "\n",
      "Timing the JIT determinant:\n",
      "17 µs ± 298 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "matrix = np.random.randn(40, 40)\n",
    "\n",
    "print(\"Timing the basic determinant:\")\n",
    "%timeit determinant(matrix)\n",
    "\n",
    "print(\"\\nTiming the JIT determinant:\")\n",
    "%timeit jit_determinant(matrix)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "quantum",
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}