fastpercentile: As-fast-as-possible median & percentile calculation on integer and float arrays

There's no reason why median and percentile calculations should be any slower than a np.max() call, and yet, on a 1-billion-element numpy array:

np.max                    :  0.080 seconds
np.median                 :  5.529 seconds
np.percentile             :  8.878 seconds

This package provides optimized versions of median and percentile calculations on numpy arrays, giving ~100× faster speeds than np.percentile:

fastpercentile.median     :  0.083 seconds
fastpercentile.percentile :  0.084 seconds

import time
import numpy as np
import fastpercentile

# A 1-billion-element uint16 volume (takes up ~2 GB of RAM)
arr = np.random.randint(0, 65536, size=1_000_000_000, dtype=np.uint16)
qs = [1, 50, 99, 99.9]
fastpercentile.median(arr)  # Compile (happens once on first call) before measuring runtimes

commands_to_run = """
np.max(arr)
np.median(arr)
np.percentile(arr, qs)
fastpercentile.median(arr)
fastpercentile.percentile(arr, qs)
"""

for command in commands_to_run.strip().splitlines():
    start = time.perf_counter()
    result = eval(command)
    print(f'{command.split("(")[0]:26s}: {time.perf_counter() - start:6.3f} seconds, result {result}')

The speed of the optimized calculation is no longer limited by CPU processing speed, but instead by data transfer bandwidth between the RAM and the CPU. Therefore, this runs literally as fast as your hardware could possibly allow.

Algorithm

For small-integer dtypes (int8, uint8, int16, uint16) the data only takes one of at most 65536 distinct values, so a single parallel pass of counting how many times each value occurs (that is, building a histogram of value occurrences) gives everything needed to compute any percentile. After the histogram is built, walking the cumulative count to find the bin holding each requested rank costs essentially nothing. The whole thing is a few hundred lines of numba, and additional memory usage is only ~16 MB regardless of input size (32 threads × one 65536-bin local table each, plus a final reduced histogram), so it adds no measurable RAM pressure on top of the input.

There is one deliberate difference from numpy.percentile on non-finite data: when a requested rank is integral (no interpolation needed) but a neighboring value is infinite, numpy.percentile still evaluates its interpolation term and returns NaN (0 * inf), whereas fastpercentile returns the exact order statistic — which is what numpy.median returns in the same situation. This is the same spirit as the narrow-signed-integer case above, where we avoid a numpy overflow.

Usage

import numpy as np
import fastpercentile

arr = np.random.randint(0, 65536, size=(305, 96, 69, 846), dtype=np.uint16)

# A scalar percentile
p99 = fastpercentile.percentile(arr, 99)

# The median (50th percentile)
m = fastpercentile.median(arr)

# Multiple percentiles in a single pass over the data
p1, p50, p99, p99_9 = fastpercentile.percentile(arr, [1, 50, 99, 99.9])

# Reduce along one or more axes, just like numpy.  Pass a single int,
# a tuple of ints, or None (the default, which reduces everything).
per_frame = fastpercentile.percentile(arr, 99, axis=(1, 2, 3))  # shape (305,)
keep = fastpercentile.percentile(arr, 50, axis=0, keepdims=True)  # shape (1, 96, 69, 846)

# Floats work too (float32 and float64).  By default any NaN
# propagates, just like numpy; use the nan-aware variants to ignore
# NaNs instead.
floats = np.random.standard_normal((1000, 1000)).astype(np.float32)
floats[floats > 3] = np.nan
m = fastpercentile.median(floats)                 # NaN, because NaNs are present
m = fastpercentile.nanmedian(floats)              # ignores the NaNs
p = fastpercentile.nanpercentile(floats, [25, 75], axis=1)  # shape (2, 1000)

# Or just grab the histogram if you want to do something else with it.
# (Only works for 8-bit and 16-bit values because 32-bit and 64-bit
# histograms don't fit in memory.)
hist = fastpercentile.histogram(arr)  # length 65536 for uint16

Results match numpy.percentile(arr, q, axis=axis, keepdims=keepdims) with the default 'linear' interpolation method (typically exact for integer inputs). In fact, on narrow signed dtypes numpy.percentile can overflow during its interpolation and return a value outside the data range; fastpercentile interpolates in float64 and stays correct.

When you reduce along an axis, the work is parallelized across the resulting slices (or within each slice when there are only a few), so reductions like a per-frame percentile over an image stack stay fast. And for arrays small enough that a full histogram would be mostly empty, fastpercentile simply sorts instead, so it never pays the histogram's fixed cost on a tiny input.

Installation

Option 1: pip install from PyPI:

pip install fastpercentile

Option 2: pip install directly from GitHub:

pip install git+https://github.com/jasper-tms/fastpercentile.git

Option 3: First git clone this repo and then pip install it from your clone:

cd ~/repos
git clone https://github.com/jasper-tms/fastpercentile.git
cd fastpercentile
pip install .

Notes on threading

fastpercentile uses every logical core on the machine by default (via numba.get_num_threads()). To limit it for a particular call, pass n_threads=N; to set it globally, use numba.set_num_threads(N) or the NUMBA_NUM_THREADS environment variable. On most systems the workload saturates DRAM bandwidth around nproc / 2 threads, so reserving a few cores for the rest of the machine costs little throughput.