README.md

numba vs numpy performance benchmark

A small benchmark that compares the wall-clock performance of equivalent numeric operations written with plain numpy and with numba @njit-compiled code.

What it measures

On a large 1-D float array (10M elements) it compares five operations:

Operation	Description
numpy sum	np.ndarray.sum() (native/vectorized)
numba sum	plain compiled loop summing the array
numba sum par	compiled parallel prange loop
numpy loop	x*2 + 1 via a pure Python loop (slow!)
numba loop	x*2 + 1 via a compiled loop

The two key points it demonstrates:

1. Numba speeds up even native numpy ops (e.g. sum) by JIT-compiling a tight loop, giving ~13x here.
2. Numba crushes Python-style loops numpy cannot vectorize (the numpy loop case), giving ~80x here.

How to run

Using uv:

uv venv --python 3.12
uv pip install numpy numba
uv run python numba_vs_numpy_bench.py

Or with plain pip:

python -m venv .venv
source .venv/bin/activate
pip install numpy numba
python numba_vs_numpy_bench.py

Requires Python 3.9+ and a compatible numpy/numba pair.

Example output

Array size: 10000000
----------------------------------------
numpy sum               156.977 ms
numba sum                12.225 ms
numba sum par             4.598 ms
numpy loop             4116.068 ms
numba loop               48.758 ms
----------------------------------------
numba sum speedup vs numpy:   12.84x
numba par sum speedup:        34.14x
numba loop vs numpy loop:     84.42x

Notes

• Results vary by hardware, Python version, and numpy/numba versions.
• A warm-up call is made before timing so numba's one-time JIT compilation is excluded from the measurements.
• The numpy loop case is intentionally un-vectorized to show where numba's real strength lies. In practice you would vectorize it, but that is not always possible in real code.

numba_vs_numpy_bench.py

"""Benchmark comparing the performance of numpy vs numba (just-in-time) execution.

This script measures the wall-clock time of equivalent numeric operations
implemented with plain numpy and with @njit-compiled numba functions, on a
large 1-D float array.

The goal is to illustrate where numba provides a speedup over numpy:
  * Operations numpy can vectorize natively (e.g. ``sum``) - moderate speedup.
  * Element-by-element Python-style loops that numpy cannot vectorize - large
    speedup, because the loop body is compiled to native code.
"""

import time

import numpy as np
from numba import njit, prange


def numpy_sum(n):
    """Sum a random array of ``n`` values using numpy's native ``sum``.

    Args:
        n (int): Number of elements in the array.

    Returns:
        float: Sum of the elements.
    """
    a = np.random.rand(n)
    return a.sum()


@njit
def numba_sum(a):
    """Sum an array with a compiled numba loop.

    Args:
        a (np.ndarray): Input 1-D float array.

    Returns:
        float: Sum of the elements.
    """
    s = 0.0
    for i in range(a.shape[0]):
        s += a[i]
    return s


@njit(parallel=True)
def numba_sum_parallel(a):
    """Sum an array using numba's parallel ``prange`` loop.

    Args:
        a (np.ndarray): Input 1-D float array.

    Returns:
        float: Sum of the elements.
    """
    s = 0.0
    for i in prange(a.shape[0]):
        s += a[i]
    return s


def numpy_loop(a):
    """Apply ``x*2 + 1`` element-wise using a pure Python loop over numpy.

    This is deliberately NOT vectorized, so it is very slow. It represents
    code numpy cannot optimize and is the case where numba shines.

    Args:
        a (np.ndarray): Input 1-D float array.

    Returns:
        np.ndarray: Result of ``a[i] * 2 + 1`` for each element.
    """
    out = np.empty_like(a)
    for i in range(a.shape[0]):
        out[i] = a[i] * 2.0 + 1.0
    return out


@njit
def numba_loop(a):
    """Apply ``x*2 + 1`` element-wise with a compiled numba loop.

    Benchmark comparison to :func:`numpy_loop`.

    Args:
        a (np.ndarray): Input 1-D float array.

    Returns:
        np.ndarray: Result of ``a[i] * 2 + 1`` for each element.
    """
    out = np.empty_like(a)
    for i in range(a.shape[0]):
        out[i] = a[i] * 2.0 + 1.0
    return out


def bench(name, fn, *args, reps=5):
    """Benchmark ``fn`` and print its best wall-clock time over ``reps`` runs.

    A single warm-up call runs first to skip lazy initialization (e.g. numba
    compilation or JIT for the first invocation). The best (minimum) time is
    reported to reduce noise from background system activity.

    Args:
        name (str): Label printed in the results table.
        fn (callable): Function to benchmark.
        *args: Positional arguments passed to ``fn``.
        reps (int, optional): Number of timed repetitions. Defaults to 5.

    Returns:
        float: Best measured wall-clock time in seconds.
    """
    fn(*args)
    times = []
    for _ in range(reps):
        t0 = time.perf_counter()
        fn(*args)
        times.append(time.perf_counter() - t0)
    best = min(times)
    print(f"{name:<20} {best * 1e3:>10.3f} ms")
    return best


def main():
    """Run the full benchmark and print a speedup summary."""
    n = 10_000_000
    rng = np.random.default_rng(42)
    a = rng.random(n)

    print(f"Array size: {n}")
    print("-" * 40)

    # Warm up numba compilation so timing excludes JIT overhead.
    numba_sum(a)
    numba_sum_parallel(a)
    numba_loop(a)

    t_np = bench("numpy sum", numpy_sum, n)
    t_nb = bench("numba sum", numba_sum, a)
    t_nbp = bench("numba sum par", numba_sum_parallel, a)
    t_np_loop = bench("numpy loop", numpy_loop, a)
    t_nb_loop = bench("numba loop", numba_loop, a)

    print("-" * 40)
    print(f"numba sum speedup vs numpy:   {t_np / t_nb:.2f}x")
    print(f"numba par sum speedup:        {t_np / t_nbp:.2f}x")
    print(f"numba loop vs numpy loop:     {t_np_loop / t_nb_loop:.2f}x")


if __name__ == "__main__":
    main()