The Numerical Siblings: NumPy and Numba — Different Lives

If you work with Python for data science, statistics, machine learning, or scientific computing, you will likely encounter NumPy. But you may not have heard of Numba.

I once developed an algorithm for a descriptor and expected it to run faster than some existing descriptors. Surprisingly, it was actually slower, even though the implementation already made extensive use of NumPy's vectorized operations. I decided to investigate why, and the problem soon became clear: the algorithm performed a large number of repeated numerical operations within loops. Although the individual NumPy operations were optimized, repeatedly calling them from Python introduced considerable overhead, especially as the number of iterations increased. Then I discovered Numba. By using Numba to compile the computationally intensive portions of the algorithm, I was able to reduce much of this overhead, and the same algorithm subsequently became considerably faster. That experience made me curious about the relationship between NumPy and Numba and why they can produce such different performance outcomes.

The names sound remarkably similar, and both are used in numerical computing. So, are NumPy and Numba competing tools? Not at all. Like siblings, they have different roles and live different lives, but their differences are precisely what makes them complementary. NumPy excels at numerical and array operations, while Numba focuses on accelerating computationally intensive Python code.

NumPy

NumPy is a fundamental package for scientific computing in Python. It provides a powerful multidimensional array object, together with a broad collection of routines for efficient numerical operations, including mathematical and logical computations, array reshaping and manipulation, sorting and selection, input and output, discrete Fourier transforms, linear algebra, statistical operations, random simulation, and much more (NumPy Developers, n.d.). Consider a simple use case of NumPy involving a function that accepts a NumPy array and calculates the product of all its elements:

import numpy as np

array = np.array([[2, 3, 4],
                  [5, 6, 7]])

def product_2d(array):
    return np.prod(array)

Numba

Numba, on the other hand, is a Python compiler designed for numerical and array-based functions, allowing computationally intensive applications to be accelerated while keeping the code in Python. It uses the Low Level Virtual Machine (LLVM) compiler infrastructure to generate optimized machine code from Python code. With a few simple annotations, numerical and array-oriented programs can be Just-In-Time (JIT) compiled to achieve performance comparable to code written in C, C++, and Fortran, without requiring a switch to another programming language or Python interpreter (Numba Developers, n.d.).

Numba's key features include on-the-fly generation of machine code, support for both CPU and GPU hardware, and integration with the Python scientific-computing ecosystem through NumPy (Numba Developers, n.d.). A simple example of a Numba-optimized version of the code written earlier is shown below:

import numba

array = np.array([[2, 3, 4],
                  [5, 6, 7]])

@numba.jit
def product_2d(array):
    a, b = array.shape
    product = 1.0

    for i in range(a):
        for j in range(b):
            product *= array[i, j]

    return product

Here, the function retains the familiar Python loop structure, while Numba can compile the numerical operations into optimized machine code. Although no immediate advantage is evident here, using Numba can provide substantial performance improvements when the function involves computationally intensive loops.

Integrating NumPy and Numba

Numba comes in handy when an algorithm requires several custom calculations for each element, involves applying different conditions, updating multiple quantities, or interacting with neighbouring elements. This becomes particularly useful when such operations are repeated over large arrays. For example, consider an operation that repeatedly applies a mathematical transformation to every element of an array.

import numpy as np

array = np.random.rand(1_000_000)

def numpy_version(array):
    result = array.copy()

    for _ in range(10):
        result = np.sin(result) + np.sqrt(result)

    return result

The NumPy implementation is vectorized, but each iteration creates intermediate arrays for the different operations. When the array is very large and the operation is repeated many times, this can increase memory usage and memory traffic. However, Numba could be used to solve this problem.

import numpy as np
import numba

array = np.random.rand(1_000_000)

@numba.njit
def numba_version(array):
    result = array.copy()

    for _ in range(10):
        for i in range(len(result)):
            result[i] = np.sin(result[i]) + np.sqrt(result[i])

    return result

Here, Numba compiles the loops into machine code. Instead of repeatedly constructing intermediate arrays for the vectorized operations, the computation can be performed element by element within the compiled loop.

References

Numba Developers. (n.d.). Overview — Numba 0.64.0dev0 documentation. https://numba.pydata.org/numba-doc/dev/user/overview.html

NumPy Developers. (n.d.). NumPy documentation. https://numpy.org/doc/stable/