Accelerating Python Array Computations with Numba: A Practical Guide

Python is inherently slower for computational tasks, so optimizing code or using scientific libraries like NumPy and SciPy is common. However, when a custom algorithm is needed without resorting to low‑level extensions, Numba offers a solution.

What the article covers

Why NumPy alone can be insufficient

Basic usage of Numba

How Numba impacts code performance

When NumPy "can't help"

Consider sorting a very large array in increasing order. A simple in‑place conversion function is shown:

<code>[1, 2, 1, 3, 3, 5, 4, 6] → [1, 2, 2, 3, 3, 5, 5, 6]</code>

Implementation:

<code>def monotonically_increasing(a):
    max_value = 0
    for i in range(len(a)):
        if a[i] > max_value:
            max_value = a[i]
        a[i] = max_value
</code>

While NumPy excels at vectorized operations, this loop‑based scenario loses its advantage, taking about 2.5 seconds for a ten‑million‑element array.

Speeding up with Numba

Numba is a JIT compiler for Python, especially effective for NumPy array loops. Adding two lines of code yields a dramatic speedup:

<code>from numba import njit

@njit
def monotonically_increasing(a):
    max_value = 0
    for i in range(len(a)):
        if a[i] > max_value:
            max_value = a[i]
        a[i] = max_value
</code>

Runtime drops from 2.5 s to 0.19 s without changing the algorithm.

NumPy also provides numpy.maximum.accumulate , which reduces the runtime further to about 0.03 s.

Runtime

Python

for

loop

2560 ms

Numba

for

loop

190 ms

np.maximum.accumulate

30 ms

Introduction to Numba

If a needed function is absent from NumPy or SciPy, developers often resort to low‑level languages, increasing complexity. Numba lets you stay in Python while achieving compiled‑speed performance, supports fast iteration, and can target GPUs.

Numba parses the code, infers input types, and compiles specialized machine code on the fly. Different input types (e.g., unsigned 64‑bit integers vs. floats) produce different compiled versions.

Some drawbacks of Numba

Compilation overhead

The first call to a Numba‑decorated function incurs compilation time. Using IPython’s %time command illustrates the cost:

<code>In [1]: from numba import njit

In [2]: @njit
   ...: def add(a, b):
   ...:     return a + b

In [3]: %time add(1, 2)
CPU times: user 320 ms, sys: 117 ms, total: 437 ms
Wall time: 207 ms

In [4]: %time add(1, 2)
CPU times: user 17 µs, sys: 0 ns, total: 17 µs
Wall time: 24.3 µs
</code>

Subsequent calls are much faster, but changing input types (e.g., to floats) triggers recompilation, adding latency again.

<code>In [8]: %time add(1.5, 2.5)
CPU times: user 40.3 ms, sys: 1.14 ms, total: 41.5 ms
Wall time: 41 ms

In [9]: %time add(1.5, 2.5)
CPU times: user 16 µs, sys: 3 µs, total: 19 µs
Wall time: 26 µs
</code>

Simple arithmetic does not require Numba; the example merely highlights compilation cost.

Differences from pure Python/NumPy implementations

Numba implements a subset of Python and NumPy APIs, which can lead to missing features, performance variations, or bugs. Error messages from failed compilations may also be hard to interpret.

Comparing Numba with other options

Using only NumPy/SciPy: fast for vectorized operations but ineffective for loop‑heavy code.

Writing extensions in low‑level languages: maximum performance but requires abandoning Python.

Using Numba: accelerates Python loops with minimal code changes, though some Python/NumPy features may be unsupported.

Conclusion

Numba is easy to try; for any slow Python for‑loop performing mathematical work, adding a few lines of Numba code can dramatically improve execution speed.