python - Performance of various numpy fancy indexing methods, also with numba

Question

Since for my program fast indexing of Numpy arrays is quite necessary and fancy indexing doesn't have a good reputation considering performance, I decided to make a few tests. Especially since Numba is developing quite fast, I tried which methods work well with numba.

As inputs I've been using the following arrays for my small-arrays-test:

import numpy as np
import numba as nb

x = np.arange(0, 100, dtype=np.float64)  # array to be indexed
idx = np.array((0, 4, 55, -1), dtype=np.int32)  # fancy indexing array
bool_mask = np.zeros(x.shape, dtype=np.bool)  # boolean indexing mask
bool_mask[idx] = True  # set same elements as in idx True
y = np.zeros(idx.shape, dtype=np.float64)  # output array
y_bool = np.zeros(bool_mask[bool_mask == True].shape, dtype=np.float64)  #bool output array (only for convenience)

And the following arrays for my large-arrays-test (y_bool needed here to cope with dupe numbers from randint):

x = np.arange(0, 1000000, dtype=np.float64)
idx = np.random.randint(0, 1000000, size=int(1000000/50))
bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[idx] = True
y = np.zeros(idx.shape, dtype=np.float64)
y_bool = np.zeros(bool_mask[bool_mask == True].shape, dtype=np.float64)

This yields the following timings without using numba:

%timeit x[idx]
#1.08 μs ± 21 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
#large arrays: 129 μs ± 3.45 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit x[bool_mask]
#482 ns ± 18.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
#large arrays: 621 μs ± 15.9 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.take(x, idx)
#2.27 μs ± 104 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
# large arrays: 112 μs ± 5.76 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit np.take(x, idx, out=y)
#2.65 μs ± 134 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
# large arrays: 134 μs ± 4.47 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit x.take(idx)
#919 ns ± 21.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 108 μs ± 1.71 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit x.take(idx, out=y)
#1.79 μs ± 40.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# larg arrays: 131 μs ± 2.92 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit np.compress(bool_mask, x)
#1.93 μs ± 95.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 618 μs ± 15.8 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.compress(bool_mask, x, out=y_bool)
#2.58 μs ± 167 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
# large arrays: 637 μs ± 9.88 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit x.compress(bool_mask)
#900 ns ± 82.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 628 μs ± 17.8 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit x.compress(bool_mask, out=y_bool)
#1.78 μs ± 59.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 628 μs ± 13.8 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.extract(bool_mask, x)
#5.29 μs ± 194 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
# large arrays: 641 μs ± 13 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

And with numba, using jitting in nopython-mode, caching and nogil I decorated the ways of indexing, which are supported by numba:

@nb.jit(nopython=True, cache=True, nogil=True)
def fancy(x, idx):
    x[idx]

@nb.jit(nopython=True, cache=True, nogil=True)
def fancy_bool(x, bool_mask):
    x[bool_mask]

@nb.jit(nopython=True, cache=True, nogil=True)
def taker(x, idx):
    np.take(x, idx)

@nb.jit(nopython=True, cache=True, nogil=True)
def ndtaker(x, idx):
    x.take(idx)

This yields the following results for small and large arrays:

%timeit fancy(x, idx)
#686 ns ± 25.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 84.7 μs ± 1.82 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit fancy_bool(x, bool_mask)
#845 ns ± 31 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 843 μs ± 14.2 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit taker(x, idx)
#814 ns ± 21.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 87 μs ± 1.52 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit ndtaker(x, idx)
#831 ns ± 24.5 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 85.4 μs ± 2.69 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

---

Summary

While for numpy without numba it is clear that small arrays are by far best indexed with boolean masks (about a factor 2 compared to ndarray.take(idx)), for larger arrays ndarray.take(idx) will perform best, in this case around 6 times faster than boolean indexing. The breakeven-point is at an array-size of around 1000 cells with and index-array-size of around 20 cells.
For arrays with 1e5 elements and 5e3 index array size, ndarray.take(idx) will be around 10 times faster than boolean mask indexing. So it seems that boolean indexing seems to slow down considerably with array size, but catches up a little after some array-size-threshold is reached.

For the numba jitted functions there is a small speedup for all indexing functions except for boolean mask indexing. Simple fancy indexing works best here, but is still slower than boolean masking without jitting.
For larger arrays boolean mask indexing is a lot slower than the other methods, and even slower than the non-jitted version. The three other methods all perform quite good and around 15% faster than the non-jitted version.

For my case with many arrays of different sizes, fancy indexing with numba is the best way to go. Perhaps some other people can also find some useful information in this quite lengthy post.

Edit:
I'm sorry that I forgot to ask my question, which I in fact have. I was just rapidly typing this at the end of my workday and completely forgot it... Well, do you know any better and faster method than those that I tested? Using Cython my timings were between Numba and Python.
As the index array is predefined once and used without alteration in long iterations, any way of pre-defining the indexing process would be great. For this I thought about using strides. But I wasn't able to pre-define a custom set of strides. Is it possible to get a predefined view into the memory using strides?

Edit 2:
I guess I'll move my question about predefined constant index arrays which will be used on the same value array (where only the values change but not the shape) for a few million times in iterations to a new and more specific question. This question was too general and perhaps I also formulated the question a little bit misleading. I'll post the link here as soon as I opened the new question!
Here is the link to the followup question.

Answer 1

Your summary isn't completely correct, you already did tests with differently sized arrays but one thing that you didn't do was to change the number of elements indexed.

I restricted it to pure indexing and omitted take (which effectively is integer array indexing) and compress and extract (because these are effectively boolean array indexing). The only difference for these are the constant factors. The constant factor for the methods take and compress will be less than the overhead for the numpy functions np.take and np.compress but otherwise the effects will be negligible for reasonably sized arrays.

Just let me present it with different numbers:

# ~ every 500th element
x = np.arange(0, 1000000, dtype=np.float64)
idx = np.random.randint(0, 1000000, size=int(1000000/500))  # changed the ratio!
bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[idx] = True

%timeit x[idx]
# 51.6 μs ± 2.02 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit x[bool_mask]
# 1.03 ms ± 37.1 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)


# ~ every 50th element
idx = np.random.randint(0, 1000000, size=int(1000000/50))  # changed the ratio!
bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[idx] = True

%timeit x[idx]
# 1.46 ms ± 55.1 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit x[bool_mask]
# 2.69 ms ± 154 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)


# ~ every 5th element
idx = np.random.randint(0, 1000000, size=int(1000000/5))  # changed the ratio!
bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[idx] = True

%timeit x[idx]
# 14.9 ms ± 495 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit x[bool_mask]
# 8.31 ms ± 181 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

So what happened here? It's simple: Integer array indexing only needs to access as many elements as there are values in the index-array. That means if there are few matches it will be quite fast but slow if there are many indices. Boolean array indexing, however, always needs to walk through the whole boolean array and check for "true" values. That means it should be roughly "constant" for the array.

But, wait, it's not really constant for boolean arrays and why does integer array indexing take longer (last case) than boolean array indexing even if it has to process ~5 times less elements?

That's where it gets more complicated. In this case the boolean array had True at random places which means that it will be subject to branch prediction failures. These will be more likely if True and False will have equal occurrences but at random places. That's why the boolean array indexing got slower - because the ratio of True to False got more equal and thus more "random". Also the result array will be larger if there are more Trues which also consumes more time.

As an example for this branch prediction thing use this as example (could differ with different system/compilers):

bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[:1000000//2] = True   # first half True, second half False
%timeit x[bool_mask]
# 5.92 ms ± 118 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[::2] = True   # True and False alternating
%timeit x[bool_mask]
# 16.6 ms ± 361 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[::2] = True
np.random.shuffle(bool_mask)  # shuffled
%timeit x[bool_mask]
# 18.2 ms ± 325 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

So the distribution of True and False will critically affect the runtime with boolean masks even if they contain the same amount of Trues! The same effect will be visible for the compress-functions.

For integer array indexing (and likewise np.take) another effect will be visible: cache locality. The indices in your case are randomly distributed, so your computer has to do a lot of "RAM" to "processor cache" loads because it's very unlikely two indices will be near to each other.

Compare this:

idx = np.random.randint(0, 1000000, size=int(1000000/5))
%timeit x[idx]
# 15.6 ms ± 703 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

idx = np.random.randint(0, 1000000, size=int(1000000/5))
idx = np.sort(idx)  # sort them
%timeit x[idx]
# 4.33 ms ± 366 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

By sorting the indices the chances immensely increased that the next value will already be in the cache and this can lead to huge speedups. That's a very important factor if you know that the indices will be sorted (for example if they were created by np.where they are sorted, which makes the result of np.where especially efficient for indexing).

So, it's not like integer array indexing is slower for small arrays and faster for large arrays it depends on much more factors. Both do have their use-cases and depending on the circumstances one might be (considerably) faster than the other.

---

Let me also talk a bit about the numba functions. First some general statements:

• cache won't make a difference, it just avoids recompiling the function. In interactive environments this is essentially useless. It's faster if you would package the functions in a module though.
• nogil by itself won't provide any speed boost. It will be faster if it's called in different threads because each function execution can release the GIL and then multiple calls can run in parallel.

Otherwise I don't know how numba effectivly implements these functions, however when you use NumPy features in numba it could be slower or faster - but even if it's faster it won't be much faster (except maybe for small arrays). Because if it could be made faster the NumPy developers would also implement it. My rule of thumb is: If you can do it (vectorized) with NumPy don't bother with numba. Only if you can't do it with vectorized NumPy functions or NumPy would use too many temporary arrays then numba will shine!