更快的成员资格测试（numpy.isin 太慢）

Question

我正在做一个实时图像处理项目，需要更改一些参数。确切的细节并不重要，所以我只关注相关部分：

我有两个 Numpy 数组，其中一个是索引值列表（它有重复的索引），另一个是有效索引列表。

我尝试了np.intersect1D，但是它删除了重复的索引值，所以我尝试了np.isin 和np.in1d，这两个都需要大约 900 毫秒才能完成。

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

我也尝试使用集合和字典作为成员资格测试的查找表，但它们也相当慢（大约 1.2 秒）。

作为最后的努力，我在 Cython 中实现了一个二进制搜索算法，该算法在 validRange 的排序版本上运行，但由于某种原因它会随机崩溃内核（使用完全相同的输入，它会在不同的点失败）。诚然，我正在使用的数组相当大（索引数组有 4,706,878 个元素，有效范围有 527,076 个元素），但即便如此，intersect1D 还是设法完成（接近 400 毫秒）。

尽管这不是常见的操作，但在更新设置时将实时系统停止大约 900 毫秒是不可接受的。

有没有什么方法可以加快这个过程，使其接近 intersect1d 的 300 到 400 毫秒的时间，而不需要删除重复的元素？我环顾了谷歌并尝试了一些不同的建议，但它们并没有真正改善。

附：我知道有些人想对使用np.where(x>y) 发表评论，但validRange 不是连续的，为方便起见被称为。例如，实际值可能是：[1,2,5,8,15,20,21]，因此简单的条件无法解决此问题。

Answer 1

您可以使用布尔查找表快速计算np.isin 的结果，因为validRange 的值是相对较小的正有界整数。使用 np.bincount 可以非常高效地构建表格。仅当validRange 的大小明显小于index_array 时，此方法才非常快，这里显然就是这种情况。

这是一个例子：

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]

在我的机器上，对于 index_array 和 validRange 中的值，此实现比原始实现快 7 到 32 倍。

如果您想要更快的代码，您可以使用 Numba 来避免创建相当大的临时数组，并使用它轻松并行化操作。有了所有这些，您应该能够享受由此产生的真正实时计算。

---

请注意，如果需要，您可以在 Nvidia 卡上使用 CuPy 或在其他兼容 OpenCL 的设备上使用 ClPy，使用 GPU 来进一步加快计算速度。您可以使用简单的import cupy as np 来完成此操作。但是请注意，这些包在某些计算之间引入了相当大的延迟，这对于实时计算可能不是很好。尽管如此，与将计算置于迭代循环中时的初始实现相比，使用带有 GTX-1660S GPU 的 CuPy，我可以达到 150 倍的速度提升！