Python 各种numpy花式索引方法的性能，也包括numba

因为对于我的程序来说，快速索引

Numpy

数组是非常必要的，而且考虑到性能，花式索引没有很好的声誉，所以我决定做一些测试。特别是由于

Numba

发展非常快，我尝试了哪些方法可以很好地与Numba配合使用

作为输入，我一直在使用以下阵列进行小型阵列测试：

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

的重复数）：

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

-模式下使用jitting，

cach

ing和

nogil

I修饰了

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，很明显，小数组使用布尔掩码进行索引的效果最好（与

ndarray.take（idx）

相比，大约是2倍），而对于较大的数组

ndarray.take（idx）

的性能最好，在这种情况下，大约是布尔索引的6倍。盈亏平衡点的数组大小约为

1000

单元格，索引数组大小约为

20

单元格。
对于包含

1e5

元素和

5e3

索引数组大小的数组，

ndarray.take（idx）

将比布尔掩码索引快约10倍。所以，布尔索引似乎随着数组大小的增加而显著减慢，但在达到某个数组大小阈值后，它会稍微加快

对于numba jitted函数，除布尔掩码索引外，所有索引函数都有一个小的加速。简单的花式索引在这里效果最好，但仍然比没有jitting的布尔屏蔽慢。
对于较大的数组，布尔掩码索引比其他方法慢得多，甚至比非jitted版本还要慢。其他三种方法的性能都相当好，比非JIT版本快15%左右

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

编辑：
很抱歉，我忘了问我的问题，事实上我已经问过了。我只是在工作日快结束的时候快速输入这个，然后完全忘记了。。。 你知道比我测试的更好更快的方法吗？使用Cython，我的计时介于Numba和Python之间。
由于索引数组只需预定义一次，并且在长时间的迭代中不需要修改就可以使用，因此任何预先定义索引过程的方法都是非常好的。为此，我考虑使用跨步。但我无法预先定义一组自定义的步幅。是否可以使用步幅将预定义视图放入内存

编辑2:
我想我将把关于预定义常量索引数组的问题转移到一个新的更具体的问题上，该数组将在同一个值数组（其中只有值改变，而不是形状）上使用数百万次。这个问题太笼统了，也许我对这个问题的表述有点误导。我一打开新问题，就会在这里发布链接

---

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

我将其限制为纯索引，并省略了

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

s越多，结果数组也会越大，这也会占用更多的时间

作为此分支预测的示例，请使用此示例（可能因不同的系统/编译器而异）：

因此

True

和

False

的分布将严重影响使用布尔掩码的运行时，即使它们包含相同数量的

True

s！对于

压缩

-功能，同样的效果也是可见的

对于整数数组索引（同样也是

np.take

），另一个效果将是可见的：缓存位置。在您的情况下，索引是随机分布的，因此您的计算机必须对“处理器缓存”进行大量的“RAM”

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

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)

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)