为什么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

我们仍然比不连续的情况快一个数量级。 我已经准备好接受Victor的解释，即实际的罪魁祸首是正在制作副本（使用.copy()强制复制的时间是暗示性的）。 之后，短路是否发生并不重要。

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

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

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

并且，更重要的是：是否有一种解决方法，即某种方法可以强制numpy短路UPDATE 而不在非连续数组上制作副本 ？

Answer 1

该问题与使用步幅时阵列的内存对齐有关。 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)你计时数组a 5000个副本

执行此操作并比较这两个定时调用的结果：

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] ”问题的解决np.argmax ：

我创建实现的功能的小库argmax这是一个包装np.argmax ，但是当输入参数是与设置为-1跨距值一维数组布尔它增加了性能：

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

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

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

---

编辑：当处理不同于-1的步幅值（具有1D布尔数组）并且具有良好结果时，我扩展了库以提高性能argmax ： a[::2] ， a[::-3]等

试试看。