numba_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Numba\n",
    "Numba is good for interfacing your user-defined Python functions into very fast machine code (for any machine!). Numba's speed is in repeated applications of the same functions, as it compiles the function into machine code in the first pass and is very fast when calling it afterward. This applies to functions that do mostly scalar and vector operations and do not rely on specialized Python libraries and objects. If Numba can recognize the types/shapes of the inputs and the mapping to outputs, it could be made much faster. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from timeit import default_timer as timer\n",
    "import numpy as np\n",
    "from numba import jit\n",
    "from sys import getsizeof"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define random 2-d array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of array x (MB):  7.63\n"
     ]
    }
   ],
   "source": [
    "n = 10**3\n",
    "x = np.random.random((n,n))\n",
    "print('Size of array x (MB): ',round(getsizeof(x)/1024/1024,2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define arbitrary function that does some computation and looping:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def function(x):\n",
    "    y = np.zeros(x.shape)\n",
    "    for i in range(0,x.shape[0]):\n",
    "        for j in range(0,x.shape[1]):\n",
    "            y[j,i] = x[i,j] + x[j,i] - x[i,j]*x[j,i] # arbitrary symmetrization w/ Hadamard product\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Time one evaluation of the function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time of 1 function evaluation without Numba (seconds):  0.8829193459999998\n"
     ]
    }
   ],
   "source": [
    "s = timer()\n",
    "y = function(x)\n",
    "e = timer()\n",
    "normal_tfe = e - s\n",
    "print('Time of 1 function evaluation without Numba (seconds): ',normal_tfe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, redefine same function, except decorate with Numba's just-in-time compiler:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def function(x):\n",
    "    y = np.zeros(x.shape)\n",
    "    for i in range(0,x.shape[0]):\n",
    "        for j in range(0,x.shape[1]):\n",
    "            y[j,i] = x[i,j] + x[j,i] - x[i,j]*x[j,i] # arbitrary symmetrization w/ Hadamard product\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time of 1st function evaluation with Numba (seconds):  0.24062856900000007\n"
     ]
    }
   ],
   "source": [
    "s = timer()\n",
    "y = function(x)\n",
    "e = timer()\n",
    "print('Time of 1st function evaluation with Numba (seconds): ',e - s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is a larger speedup after the first application because Numba compiled the function in machine code to run it once. Let's test how much faster it got with this function that loops over the compiled function until it hits the time it took to run once in Python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loop_numba(x,normal_tfe):\n",
    "    i = 0 # counter\n",
    "    s,e = timer(),0 # start timer\n",
    "    while e-s < normal_tfe:\n",
    "        y = function(x)\n",
    "        e = timer()\n",
    "        i = i + 1\n",
    "    return i\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare the function evaluation times:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Numba-compiled function ran 119 times in the amount of time it took to run once in Python.\n"
     ]
    }
   ],
   "source": [
    "count_nfes_numba = loop_numba(x,normal_tfe)\n",
    "print('Numba-compiled function ran {} times in the amount of time it took to run once in Python.'.format(count_nfes_numba))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But of course, for this operation, we could just use Numpy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loop_numpy(x,normal_tfe):\n",
    "    i = 0 # counter\n",
    "    s,e = timer(),0 # start timer\n",
    "    while e-s < normal_tfe:\n",
    "        y = x + np.transpose(x) - np.multiply(x,np.transpose(x)) # same as what we coded in for-loops above\n",
    "        e = timer()\n",
    "        i = i + 1\n",
    "    return i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Numpy-vectorized function ran 118 times in the amount of time it took to run once in Python.\n"
     ]
    }
   ],
   "source": [
    "count_nfes_numpy = loop_numpy(x,normal_tfe)\n",
    "print('Numpy-vectorized function ran {} times in the amount of time it took to run once in Python.'.format(count_nfes_numpy))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, this process extends to more complex user-defined functions that are written mostly in Numpy, like a watershed model. As long as the functions: \n",
    "1. Accept and work in mostly scalars and vectors (not DataFrames, or other objects/classes)\n",
    "2. Do mostly looping and vector operations\n",
    "3. Avoid using dictionaries (in general), tuples (in operations), and sets (idk how)\n",
    "4. Avoid using other Python libraries besides Numpy/Scipy\n",
    "5. Are called many times\n",
    "6. Return scalars or vectors, not objects\n",
    "\n",
    "Then, the application can be sped-up with Numba. This comes at the cost of losing the flexibility of working with various Python objects and syntactical tricks, but most of the common things like if-else, logical comparisons, Python math, and lists/list comprehensions are still useable. \n",
    "\n",
    "Numba also has some support for automatic parallelization and kernelization of functions and loops. This means you can automatically distribute code and automatically optimize how the operations are distributed. This is an advantage over other approaches, such as the mpi4py or multiprocessing libraries, where we explicitly define how to parallelize code. Numba can also compile your functions to run on GPUs! But, the basic functionality is in the @jit decorator.\n"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}