python - Numpy Argwhere 效率低下

Question

我正在学习像 argwhere 这样的方法和 nonzero在 NumPy 中。看来 numpy.nonzero(x)函数返回一维元组 ndarray对象，以便此函数的输出可用于索引。
我还没有准备好C nonzero的源代码因为我不知道如何找到它。但是，我想 nonzero函数将构造一个 m通过 ndim ndarray对象(对于某些 m 取决于输入)将保存 a 的非零元素的索引.为了验证这个猜测是否正确，我尝试了:

import numpy as np
from numpy.random import Generator, PCG64
rg = Generator(PCG64())
x = rg.integers(0,2,(10000,10000))
y = np.nonzero(x)
print(y[0].base is y[1].base)
z = y[0].base
print(type(z),z.shape)
print(np.array_equal(z[:,0].reshape(-1),y[0]))
print(np.array_equal(z[:,1].reshape(-1),y[1]))

输出:

True
<class 'numpy.ndarray'> (50005149, 2)
True
True

我对上面的解释是 nonzero函数确实构造了一个 m通过 ndim大小的数组。
还有一个 np.argwhere(x)功能。不像 np.nonzero ，它将返回 m通过 ndim大小数组，而不是元组。原因表明 argwhere和 nonzero实际上是相同的，除了 nonzero以稍微不同的格式返回输出。令我惊讶的是， argwhere似乎实现如下(根据: https://numpy.org/doc/stable/reference/generated/numpy.argwhere.html ):

# nonzero does not behave well on 0d, so promote to 1d
    if np.ndim(a) == 0:
        a = shape_base.atleast_1d(a)
        # then remove the added dimension
        return argwhere(a)[:,:0]
    return transpose(nonzero(a))

因为 nonzero函数返回一个元组，看起来像 transpose argwhere 中的操作将无缘无故地需要额外的副本。使用计时器进行的快速实验表明确实会发生这种情况。
问题
谁能解释为什么 argwhere功能是这样实现的？例如，替代品

def faster_argwhere(a):
    # nonzero does not behave well on 0d, so promote to 1d
    if np.ndim(a) == 0:
        a = shape_base.atleast_1d(a)
        # then remove the added dimension
        a=a[:,:0]
    return np.nonzero(a)[0].base

已经看起来更好了，因为:

x = rg.integers(0,2,(10000,10000))
t_0 = time.perf_counter()
y = np.argwhere(x)
t_1 = time.perf_counter()
z = faster_argwhere(x)
t_2 = time.perf_counter()
print('elapsed time for argwhere:' + str(t_1-t_0) + ", and for other method:" + str(t_2-t_1))
print(np.array_equal(y,z))

产量:

elapsed time for argwhere:2.175326200000086, and for other method:1.7338391999999203
True

Answer 1

当我查看 nonzero前段时间，我看到它执行了两次传球。第一个 np.count_nonzero (在 c-api 版本中)确定返回大小，然后实际获取索引。没注意有没有分配n独立数组或只有一个 2d。但是现在你提到它，数组是一个常见的二维数组的 View :

In [464]: x = np.arange(12).reshape(3,4)  
In [468]: idx = np.nonzero(x%3==0)                                                                   
In [469]: idx                                                                                        
Out[469]: (array([0, 0, 1, 2]), array([0, 3, 2, 1]))
In [470]: idx[0].__array_interface__                                                                 
Out[470]: 
{'data': (54070800, False),
 'strides': (16,),
  ...}
In [471]: idx[1].__array_interface__                                                                 
Out[471]: 
{'data': (54070808, False),
 'strides': (16,),
 ....}

数据指针为 8更大，迈出 16 ，与此一致。
对于 3d 数组:

In [472]: x = np.arange(24).reshape(2,3,4)                                                           
In [473]: idx = np.nonzero(x%3==0)                                                                   
In [474]: idx[0].__array_interface__                                                                 
Out[474]: 
{'data': (54163904, False),
 'strides': (24,),
In [475]: idx[1].__array_interface__                                                                 
Out[475]: 
{'data': (54163912, False),
 'strides': (24,),
In [476]: idx[2].__array_interface__                                                                 
Out[476]: 
{'data': (54163920, False),
 'strides': (24,),

同样的事情 - 24 步等。
我不知道这些功能的历史。有 nonzero总是返回常见二维数组的 View ？最近的版本试图让我们不再使用 np.where而是 np.nonzero .我们还没有 np.arg_nonzero :)
如您所见， nonzero tuple 可以很好地用作索引。

x[np.nonzero(x)]

新手通常不理解这一点，而是认为他们需要一个元组列表。但是要使用那些他们必须迭代

 for tup in [(0,0),(1,0),...]: print(x[tup])

argwhere不太合适:

In [480]: x[idx]                                                                                     
Out[480]: array([ 0,  3,  6,  9, 12, 15, 18, 21])
In [481]: np.transpose(idx)                                                                          
Out[481]: 
array([[0, 0, 0],
       [0, 0, 3],
        ....
       [1, 2, 1]])
In [482]: [x[tuple(i)] for i in np.transpose(idx)]                                                   
Out[482]: [0, 3, 6, 9, 12, 15, 18, 21]

我不知道如果没有 tuple 是否可行是否在早期版本中；它现在引发错误。
所以是的，有可能制作 argwhere更高效，通过访问 nonzero直接打基础。但是有什么好处呢？

In [487]: idx[0].base                                                                                
Out[487]: 
array([[0, 0, 0],
       [0, 0, 3],
         ...
       [1, 2, 1]])

In [492]: timeit x[idx]                                                                              
2.46 µs ± 8.17 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [493]: timeit [x[tuple(i)] for i in idx[0].base]                                                  
21.7 µs ± 836 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [494]: timeit [x[tuple(i)] for i in np.transpose(idx)]                                            
33.2 µs ± 94.3 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

最终 why此类问题只能通过代码中的评论或开发者论坛中的讨论或 github 问题来回答。