Inefficiency of Numpy Argwhere

user_3281410 0
Viewed 428

I am learning about methods like argwhere and nonzero in NumPy. It appears that thenumpy.nonzero(x) the function returns a tuple of one-dimensional ndarray objects so that the output of this function could be used for indexing.

I have not ready the C source code of nonzero because I don't know how to find it. However, I imagine that the nonzero function will construct an m by ndim ndarray object (for some m that depends on the input) that will hold the indices of nonzero elements of a. To verify that this guess is correct, I tried:

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

which outputs:

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

My interpretation of the above is that the nonzero function does construct a single m by ndim sized array.

There is also a np.argwhere(x) function. Unlike np.nonzero, it will return an m by ndim size array, not a tuple. Reason suggests that argwhere and nonzero would be practically the same, except that nonzero returns the output in a slightly different format. To my surprise, argwhere seems to be implemented as follows (according to: 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))

Because the nonzero function returns a tuple, it seems like the transpose operation in argwhere is going to require an extra copy for no reason. A quick experiment with a timer suggested that this copy does happen.

Question Can someone explain why the argwhere function is implemented this way? For example, the alternative

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

already seems better, since:

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

yields:

elapsed time for argwhere:2.175326200000086, and for other method:1.7338391999999203
True
numpy python
1 Answers

When I looked at nonzero some time ago, I saw that it performed two passes. First a np.count_nonzero (in the c-api version) to determine the return size, and a second to actually get the indices. I didn't pay attention as to whether it allocated n independent arrays or just one 2d. But now that you mention it, the arrays are views of a common 2d array:

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,),
 ....}

The data pointer is 8 larger, and strides 16, consistent with that.

For a 3d array:

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

The same thing - 24 strides etc.

I don't know the history of these functions. Has nonzero always returned views of a common 2d array? Recent versions have tried to ween us away from using np.where instead np.nonzero. We don't yet have a np.arg_nonzero :)

As you note, the nonzero tuple works well as index.

x[np.nonzero(x)]

Novices often don't appreciated this, and instead think they need a list of tuples. But to use those they have to iterate

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

The argwhere isn't quite right for that:

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]

I don't know if that worked without the tuple in earlier versions or not; it raises an error now.

So yes it might be possible to make argwhere more efficient, by accessing the nonzero base directly. But what's the benefit?

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)

Ultimately why questions like this can only be answered by comments in the code or discussions in the developer's forums or the github issues.

user_9019250
Related