各种numpy花式索引方法的性能，以及numba的性能

Question

由于对我的程序来说，Numpy数组的快速索引是非常必要的，而且考虑到性能，花式索引没有良好的声誉，因此我决定进行一些测试。尤其是由于Numba它发展很快，我尝试了哪种方法与numba一起工作。

作为输入，我一直在我的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)

以下是我的大数组测试的以下数组（y_bool此处用于处理来自的重复数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)

在不使用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)

并使用numba，在nopython-mode中使用jitting
，caching和nogil我装饰了索引的方式，这些方式受以下支持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)

对于大型阵列，这将产生以下结果：

%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)

概要

对于没有numba的numpy，很明显，到目前为止，最好使用布尔掩码对小数组进行索引（与相比，大约是2
ndarray.take(idx)），而对于较大的数组ndarray.take(idx)将是最佳的，在这种情况下，比布尔索引快6倍。收支平衡点的大小是大约1000单元格的大小，而索引数组大小是大约20单元格的大小。
对于具有1e5元素和5e3索引数组大小的数组，ndarray.take(idx)将比布尔掩码索引 快10倍
左右。因此，布尔索引似乎随着数组大小而显着放慢，但是在达到某些数组大小阈值后会有所追赶。

对于numba jitted函数，除了布尔掩码索引之外，所有索引函数的速度都有所提高。简单的花式索引在这里效果最好，但是比不带jitting的布尔掩码慢。
对于较大的数组，布尔值掩码索引比其他方法慢很多，甚至比非Jitted版本慢。其他三种方法的性能都非常好，并且比非固定版本快15％。

对于具有许多不同大小的数组的情况，使用numba进行花式索引是最好的方法。也许其他人也可以在这篇冗长的文章中找到一些有用的信息。

编辑：
对不起，我忘了问我的问题，实际上我有。我只是在工作结束时迅速输入了这个信息，却完全忘记了……好吧，您知道有什么比我测试过的方法更好，更快的方法吗？使用Cython的时间介于Numba和Python之间。
由于索引数组是一次预定义的，并且在长时间迭代中无需更改就可以使用，因此任何预定义索引过程的方式都将非常有用。为此，我想到了大步前进。但是我无法预定义一组自定义的步骤。是否可以使用跨步将预定义视图导入内存？

编辑2：
我想我将有关预定义的常量索引数组的问题移到一个新的且更具体的问题上，该常量索引数组将在同一值数组（其中值仅更改但形状不更改）上使用数百万次。这个问题太笼统了，也许我也提出了一些误导性的问题。打开新问题后，我将在此处发布链接！

Answer 1

您的总结并不完全正确，您已经使用不同大小的数组进行了测试，但是您没有做的一件事就是更改索引元素的数量。

我将其限制为纯索引，并省略了take（实际上是整数数组索引）compress和extract（因为它们实际上是布尔数组索引）。这些的唯一区别是恒定因素。对这些方法的常数因子take和compress将小于开销为numpy的功能np.take和np.compress否则的影响可以忽略不计的合理大小的数组。

让我用不同的数字表示它：

# ~ 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)

那么这里发生了什么？很简单：整数数组索引只需要访问与索引数组中的值一样多的元素。这意味着如果匹配很少，那么如果索引很多，将会非常快但是很慢。但是，布尔数组索引始终需要遍历整个布尔数组并检查“真”值。这意味着该数组应该大致“恒定”。

但是，等等，对于布尔数组来说，它并不是真正恒定的，为什么整数数组索引需要比布尔数组索引花费更长的时间（最后一种情况），即使它必须处理的元素减少约5倍？

那就是它变得更加复杂的地方。在这种情况下，布尔数组True位于随机位置，这意味着它将经受 分支预测失败
。如果True并且False发生次数相等，但发生在随机位置，则发生这些事件的可能性更大。这就是为什么布尔数组索引越来越慢-
因为比True来False得到更加平等，从而更加“随机”。如果有更多True的，结果数组也会更大，这也会消耗更多的时间。

作为该分支预测事物的示例，请以以下示例为例（可能因不同的系统/编译器而异）：

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)

因此，即使布尔掩码包含相同数量的s
，Trueand的分布False也会严重影响运行时间True。compress-functions可以看到相同的效果。

对于整数数组索引（同样np.take），将看到另一个效果： cache locality
。您的情况下的索引是随机分布的，因此您的计算机必须执行很多“ RAM”到“处理器缓存”的加载，因为两个索引彼此之间不太可能相互接近。

比较一下：

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)

通过对索引进行排序，下一个值将已经存在于高速缓存中的机会大大增加，这可能导致巨大的加速。如果您知道索引将被排序，那么这是一个非常重要的因素（例如，如果索引是由np.where它们创建的，则np.where对索引进行排序特别有效）。

因此，这并不取决于整数数组索引对于较小的数组而言较慢，而对于较大的数组而言则取决于更多因素。两者都有其用例，并且根据情况，它们可能（相当）快于另一种。

让我也谈一谈numba函数。首先是一些一般性的陈述：

• cache不会有所不同，只是避免重新编译函数。在交互式环境中，这实际上是没有用的。但是，如果将功能打包在模块中会更快。
• nogil本身不会提供任何速度提升。如果在不同的线程中调用它将更快，因为每个函数执行可以释放GIL，然后多个调用可以并行运行。

否则我不知道numba如何有效地实现这些功能，但是当您在numba中使用NumPy功能时，它可能会变慢或变快-
但即使速度更快，它也不会变快（也许对于小型阵列而言）。因为如果可以更快地进行开发，NumPy开发人员也会实施它。我的经验法则是：如果可以使用NumPy（向量化）进行操作，请不要理会numba。只有当您无法使用向量化NumPy函数或NumPy做到这一点时，它将使用过多的临时数组，然后numba才会发光！