为什么numpy不会在非连续数组上短路？

Question

考虑以下简单测试：

import numpy as np
from timeit import timeit

a = np.random.randint(0,2,1000000,bool)

让我们找到第一个True的索引

timeit(lambda:a.argmax(), number=1000)
# 0.000451055821031332

这相当快，因为​​numpy 短路了。

它也适用于连续切片，

timeit(lambda:a[1:-1].argmax(), number=1000)
# 0.0006490410305559635

但似乎不是在非连续的。我主要感兴趣的是找到最后一个True：

timeit(lambda:a[::-1].argmax(), number=1000)
# 0.3737605109345168

更新：我认为观察到的减速不是由于短路引起的假设是不准确的（感谢@Victor Ruiz）。的确，在 所有False 数组的最坏情况

b=np.zeros_like(a)
timeit(lambda:b.argmax(), number=1000)
# 0.04321779008023441

我们仍然比非连续的快一个数量级 案子。我已经准备好接受维克多的解释，即真正的罪魁祸首 正在制作副本（使用.copy() 强制复制的时间是 暗示）。之后是否真的不再重要 短路发生与否。

但其他步长 != 1 会产生类似的行为。

timeit(lambda:a[::2].argmax(), number=1000)
# 0.19192566303536296

问题：为什么numpy 在最后两个示例中没有短路UPDATE 不复制？

而且，更重要的是：是否有解决方法，即强制numpy 短路更新 不复制也在非连续数组上?

Answer 1

问题与使用 strides 时数组的内存对齐有关。 a[1:-1]、a[::-1] 被认为在内存中对齐，但 a[::2] 不要：

a = np.random.randint(0,2,1000000,bool)

print(a[1:-1].flags.c_contiguous) # True
print(a[::-1].flags.c_contiguous) # False
print(a[::2].flags.c_contiguous) # False

这解释了为什么np.argmax 在a[::2] 上运行缓慢（来自ndarrays 的文档）：

NumPy 中的一些算法可以处理任意步长的数组。然而，一些算法需要单段数组。当一个不规则的跨步数组被传入此类算法时，会自动生成一个副本。

np.argmax(a[::2]) 正在制作数组的副本。因此，如果您执行timeit(lambda: np.argmax(a[::2]), number=5000)，您正在计时数组的 5000 个副本a

执行这个并比较这两个计时调用的结果：

print(timeit(lambda: np.argmax(a[::2]), number=5000))

b = a[::2].copy()
print(timeit(lambda: np.argmax(b), number=5000))

---

编辑： 深入研究 numpy 的 C 中的源代码，我发现 argmax 函数的下划线实现 PyArray_ArgMax 在某些时候调用 PyArray_ContiguousFromAny 以确保给定的输入数组在内存中对齐（C 风格）

然后，如果数组的 dtype 是 bool，它会委托给 BOOL_argmax 函数。 查看它的代码，似乎 总是 应用了短路。

总结

• 为了避免被np.argmax 复制，请确保输入数组在内存中是连续的
• 当数据类型为布尔值时始终应用短路。

Answer 2

我对解决这个问题很感兴趣。因此，我提出了下一个解决方案，该解决方案可以避免由于np.argmax 的内部 ndarray 副本而导致的“a[::-1]”问题案例：

我创建了一个实现函数argmax 的小型库，该函数是np.argmax 的包装器，但是当输入参数是步幅值设置为-1 的一维布尔数组时，它提高了性能：

https://github.com/Vykstorm/numpy-bool-argmax-ext

对于这些情况，它使用低级 C 例程来查找具有最大值 (True) 的项的索引 k，从数组 a 的末尾开始。
然后你可以用len(a)-k-1 计算argmax(a[::-1])

低级方法不执行任何内部 ndarray 副本，因为它使用数组 a 进行操作，该数组已经是 C 连续并在内存中对齐。它也适用于短路

---

编辑： 我扩展了库以提高 argmax 的性能，同时在处理不同于 -1 的步幅值（使用 1D 布尔数组）时也取得了良好的效果：a[::2]、a[::-3] 等

试一试。