numba.md

Numba and Jit

Joseph P. Vantassel, jpvantassel.com

One-Minute Summary

Numba is a Python package for improving the performance of your Python code.

Numba increases the performance of your Python code by transforming it into machine code (not C, C++, or Fortran but directly into machine code) using what is termed just-in-time compilation (jit).

You will see the best performance improvement with Numba and jit when your function performs mathematical type operations and little-to-no improvement otherwise.

Basic Syntax

import numba as nb
import numpy as np

@nb.jit(nopython=True)
def smooth_jit(array):
  sarray = np.empty_like(array)
  for index in range(1, len(array)-1):
    sarray[index] = (0.5*array[index-1] + array[index] + 0.5*array[index+1])/2
  return sarray

What is Numba

Numba is an open-source just-in-time (jit) compilation tool for Python code. Functions improved using Numba and jit can see increases in their performance to make them comparable to C and Fortran. You can use Numba and jit with only minimal changes to your existing Python code and with no need to write (or compile ahead of time) anything but pure Python.

What is Jit Compilation

Jit compilation stands for just-in-time compilation and refers to the process of compiling parts (e.g., functions) of an interpreted language (e.g., Python) to machine code when the function is called (i.e., at runtime). The parts of your Python code which will benefit the most from jit compilation are those which are called repeatedly (since Numba is smart enough to cache compiled functions) and those which involve mathematical operations (as these are the easiest to compile and optimize).

Basic Syntax

Numba and jit includes various functionality including compiling code for CPU, GPU, and multi-threading. For now we will consider only the basic syntax for compiling for CPU as this is the most common use of Numba and jit.

Jit decorator

Below is a simple example to demonstrate how to use jit for optimizing a Python function using the @nb.jit decorator.

import numba as nb                   # Import numba to get jit decorator
import numpy as np                   # Import numpy to get ndarray

@nb.jit(nopython=True)               # Decorator invokes jit, with kwarg such
                                     # that if types can't be optimized, error
def smooth_jit(array):               # Simple triangular smoothing function
  sarray = np.empty_like(array)
  for index in range(1, len(array)-1):
    sarray[index] = (0.5*array[index-1] + array[index] + 0.5*array[index+1])/2
  return sarray

The use of Numba and jit for this simple function improves performance on my machine by approximately a factor a 4. This is an impressive speed improvement considering it only involved two additional lines of pure Python code.

Sources

Numba