Kwker

NumPy drop-in

kwker.numpy_ops makes NumPy's sorting functions run on Kwker. Your program keeps calling np.sort, np.argsort and the others. After install(), those calls go to Kwker, and the results are NumPy's.

PythonRuns on your machine.
import numpy as np
import kwker.numpy_ops as knp

knp.install()                      # from here on, np.sort and the others run Kwker
x = np.random.default_rng(0).integers(0, 1000, 1_000_000)
print(np.sort(x)[:5], np.argsort(x, kind="stable")[:3].dtype)
knp.uninstall()                    # NumPy's own functions back
Output
[0 0 0 0 0] int64

To switch it on for one block only, use the context manager:

PythonRuns on your machine.
import numpy as np
import kwker.numpy_ops as knp

x = np.random.default_rng(1).standard_normal(500_000)
with knp.accelerated():            # Kwker inside the block only
    order = np.argsort(x, kind="stable")
print(np.array_equal(order, np.argsort(x, kind="stable")))   # NumPy's own result, outside the block
Output
True

Which functions

Function With Kwker
np.sort the same values, any kind and axis
np.argsort the stable permutation, any axis: exactly NumPy's for kind="stable" or stable=True
np.partition, np.argpartition one kth: the value a full sort would put at kth moves there, smaller or equal values before it, larger or equal after; neither side sorted
np.lexsort NumPy's permutation, last key most significant
np.searchsorted NumPy's positions, both sides, without sorter=
np.unique the same values, and return_index, return_inverse and return_counts
np.intersect1d, np.union1d, np.setdiff1d, np.setxor1d the same results, for two arrays of one dtype

Kwker takes a call when the input is a NumPy array (or a list) of booleans, integers, float16, float32, float64, datetime64 or timedelta64 with at least 1,024 elements (np.lexsort: at least 16,384 keys over all its columns). Smaller arrays stay on NumPy, which is quicker for them. Every other call goes to NumPy unchanged. That includes masked arrays and other subclasses, strings, objects, complex numbers, structured order=, sorter= and several kth values.

When results can differ

Where NumPy defines the result, Kwker returns the same result. Where NumPy leaves it open, Kwker can return another valid answer:

NaN sorts last, as in NumPy.

What is not covered

install() replaces the functions in the numpy namespace. Array methods such as a.sort() and a.argsort() keep NumPy's code, and so does a function imported by name before install() (from numpy import sort). To speed those up, call np.sort(a) or the Kwker functions directly.

Measure a program both ways

python -m kwker audit --mode numpy runs your program with and without the drop-in and compares the time and the output. See Evaluate on your machine.

ShellOn your machine.
python -m kwker audit --mode numpy -- analysis.py --input data.npy

KWKER_NUMPY_MIN sets the smallest array Kwker takes (default 1024), and knp.set_min_size(n) does the same from code.

Next steps