Accelerating code with Numba¶

Imports¶

In [1]:

import numpy as np
from numpy.typing import NDArray
import numba

import numpy as np from numpy.typing import NDArray import numba

Numba and JIT¶

Numba 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!

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.

Example: the determinant¶

In [2]:

def determinant(matrix: NDArray) -> float:
    """Calculate the determinant of a matrix.

    Args:
        matrix (np.ndarray): a square matrix.

    Returns:
        float: the determinant of matrix.
    """
    # Check the ndarray is a square matrix
    assert len(matrix.shape) == 2
    assert matrix.shape[0] == matrix.shape[1]

    dim = matrix.shape[0]

    # Convert the matrix to upper triangular form
    for col in range(0, dim - 1):
        for row in range(col + 1, dim):
            if matrix[row, col] != 0.0:
                coef = matrix[row, col] / matrix[col, col]
                matrix[row, col:dim] = matrix[row, col:dim] - coef * matrix[col, col:dim]

    return np.prod(np.diag(matrix))

def determinant(matrix: NDArray) -> float: """Calculate the determinant of a matrix. Args: matrix (np.ndarray): a square matrix. Returns: float: the determinant of matrix. """ # Check the ndarray is a square matrix assert len(matrix.shape) == 2 assert matrix.shape[0] == matrix.shape[1] dim = matrix.shape[0] # Convert the matrix to upper triangular form for col in range(0, dim - 1): for row in range(col + 1, dim): if matrix[row, col] != 0.0: coef = matrix[row, col] / matrix[col, col] matrix[row, col:dim] = matrix[row, col:dim] - coef * matrix[col, col:dim] return np.prod(np.diag(matrix))

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.

In [3]:

@numba.jit(nopython=True)
def jit_determinant(matrix: NDArray) -> float:
    """Calculate the determinant of a matrix faster using a just in time compiler.
    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.

    Args:
        matrix (np.ndarray): a square matrix.

    Returns:
        float: the determinant of matrix.
    """
    # Check the ndarray is a square matrix
    assert len(matrix.shape) == 2
    assert matrix.shape[0] == matrix.shape[1]

    dim = matrix.shape[0]

    # Convert the matrix to upper triangular form
    for col in range(0, dim - 1):
        for row in range(col + 1, dim):
            if matrix[row, col] != 0.0:
                coef = matrix[row, col] / matrix[col, col]
                matrix[row, col:dim] = matrix[row, col:dim] - coef * matrix[col, col:dim]

    return np.prod(np.diag(matrix))

@numba.jit(nopython=True) def jit_determinant(matrix: NDArray) -> float: """Calculate the determinant of a matrix faster using a just in time compiler. 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. Args: matrix (np.ndarray): a square matrix. Returns: float: the determinant of matrix. """ # Check the ndarray is a square matrix assert len(matrix.shape) == 2 assert matrix.shape[0] == matrix.shape[1] dim = matrix.shape[0] # Convert the matrix to upper triangular form for col in range(0, dim - 1): for row in range(col + 1, dim): if matrix[row, col] != 0.0: coef = matrix[row, col] / matrix[col, col] matrix[row, col:dim] = matrix[row, col:dim] - coef * matrix[col, col:dim] return np.prod(np.diag(matrix))

Timing¶

As we can see, the jit compiled determinant is much faster than the interpreted version:

In [5]:

matrix = np.random.randn(40, 40)

print("Timing the basic determinant:")
%timeit determinant(matrix)

print("\nTiming the JIT determinant:")
%timeit jit_determinant(matrix)

matrix = np.random.randn(40, 40) print("Timing the basic determinant:") %timeit determinant(matrix) print("\nTiming the JIT determinant:") %timeit jit_determinant(matrix)

Timing the basic determinant:
300 µs ± 2.06 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Timing the JIT determinant:
17 µs ± 298 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)