Faster membership test (numpy.isin too slow)

Question

I'm working on a real time image processing project and there is a requirement for changing some parameters. The exact details are not important, so I'll just stick to the relevant part:

I have two Numpy arrays, one of which is a list of index values (it has repeated indices) and another one is a list of valid indices.

I tried np.intersect1D, however that removes the duplicate index values, so then I tried np.isin and np.in1d, both of which take around 900ms to finish.

mask = np.isin(index_array,validRange)
new_index_array = index_array[mask]
new_coeff_array = coeff_array[mask]

I also tried using a set and a dictionary as look up tables for membership test, but they too were rather slow (around 1.2seconds).

As a last effort, I implemented a binary search algorithm in Cython, which operated over the sorted version of validRange but for some reason it would crash the kernel randomly (with the exact same inputs, it would fail at different points). Admittedly the arrays I'm working with are rather large (index array has 4,706,878 elements and valid range has 527,076 elements), but even so, intersect1D manages to finish (closer to 400ms).

Even though this is not a common operation, stopping a real time system for around 900ms while the settings are being updated is not really acceptable.

Is there any way to speed up this process to be closer to the 300 to 400ms time of the intersect1d, without removing the duplicate elements? I looked around google and tried a few different suggestions, but they didn't really improve things.

P.S. I know some want to comment about using np.where(x>y) but the validRange is not continuous and is called that for convenience. Actual values could for example be: [1,2,5,8,15,20,21], so a simple condition can't solve this.

user_12939557 · Answer

You can use a boolean lookup table to compute the result of np.isin very quickly since values of validRange are relatively-small positive bounded integers. The table can be built very efficiently using np.bincount. This method is very fast only when the size of validRange is significantly smaller than index_array which is clearly the case here.

Here is an example:

assert len(validRange) > 0

# Build the lookup table
lut = np.bincount(validRange) > 0
bound = len(lut)

# Pre-filter the invalid ranges and locate the value to check further
# If index_array contains only positive values, 
# then you can just use:  mask = index_array < bound
mask = np.logical_and(index_array >= 0, index_array < bound)
validRangeIdx = np.where(mask)

# Correct the mask by using the very fast LUT
mask[validRangeIdx] = lut[index_array[validRangeIdx]]

new_index_array = index_array[mask]
new_coeff_array = coeff_array[mask]

On my machine, this implementation is between 7 to 32 times faster than the original implementation regarding the values in both index_array and validRange.

If you want an even faster code, you can use Numba to avoid the creation of the quite-big temporary arrays and parallelize the operations easily with it. With all of this, you should be able to enjoy the resulting truly real-time computation.

Note that you can use your GPU to speed up the computation even more if needed using CuPy on Nvidia cards or ClPy on other OpenCL-compatible devices. You can do this with a simple import cupy as np. Note however that those packages introduces a quite big latency between some computations that may not be great for a real-time computation. Nevertheless, I can reach a speed up of 150 with CuPy with a GTX-1660S GPU compared to the initial implementation when the computation is put in an iterative loop!

Faster membership test (numpy.isin too slow)

1 Answers