了解numpy.where的运行时以及等效的替代方法

Question

According to http://docs.scipy.org/doc/numpy/reference/generated/numpy.where.html , if x and y are given and input arrays are 1-D, where is equivalent to [xv if c else yv for (c,xv, yv) in zip(x!=0, 1/x, x)] . When doing runtime benchmarks, however, they have significantly different speeds:

x = np.array(range(-500, 500))

%timeit np.where(x != 0, 1/x, x)
10000 loops, best of 3: 23.9 µs per loop

%timeit [xv if c else yv for (c,xv, yv) in zip(x!=0, 1/x, x)]
1000 loops, best of 3: 232 µs per loop

Is there a way I can rewrite the second form so that it has a similar runtime to the first? The reason I ask is because I'd like to use a slightly modified version of the second case to avoid division by zero errors:

[1 / xv if c else xv for (c,xv) in zip(x!=0, x)]

Another question: the first case returns a numpy array while the second case returns a list. Is the most efficient way to have the second case return an array is to first make a list and then convert the list to an array?

np.array([xv if c else yv for (c,xv, yv) in zip(x!=0, 1/x, x)])

Thanks!

Answer 1

You just asked about 'delaying' the 'where':

numpy.where : how to delay evaluating parameters?

and someone else just asked about divide by zero:

Replace all elements of a matrix by their inverses

When people say that where is similar to the list comprehension, they attempt to describe the action, not the actual implementation.

np.where called with just one argument is the same as np.nonzero . This quickly (in compiled code) loops through the argument, and collects the indices of all non-zero values.

np.where when called with 3 arguments, returns a new array, collecting values from the 2 and 3rd arguments based on the nonzero values. But it's important to realize that those arguments must be other arrays. They are not functions that it evaluates element by element.

So the where is more like:

m1 = 1/xv
m2 = xv
[v1 if c else v2 for (c, v1, v2) in zip(x!=0, m1, m2)]

It's easy to run this iteration in compiled code because it just involves 3 arrays of matching size (matching via broadcasting).

np.array([...]) is a reasonable way of converting a list (or list comprehension) into an array. It may be a little slower than some alternatives because np.array is a powerful general purpose function. np.fromiter([], dtype) may be faster in some cases, because it isn't as general (you have to specify dtype, and it it only works with 1d).

There are 2 time proven strategies for getting more speed in element-by-element calculations:

• use packages like numba and cython to rewrite the problem as c code

• rework your calculations to use existing numpy methods. The use of masking to avoid divide by zero is a good example of this.

=====================

np.ma.where , the version for masked arrays is written in Python. Its code might be instructive. Note in particular this piece:

# Construct an empty array and fill it
d = np.empty(fc.shape, dtype=ndtype).view(MaskedArray)
np.copyto(d._data, xv.astype(ndtype), where=fc)
np.copyto(d._data, yv.astype(ndtype), where=notfc)

It makes a target, and then selectively copies values from the 2 inputs arrays, based on the condition array.

Answer 2

You can avoid division by zero while maintaining performance by using advanced indexing:

x = np.arange(-500, 500)

result = np.empty(x.shape, dtype=float) # set the dtype to whatever is appropriate
nonzero = x != 0
result[nonzero] = 1/x[nonzero]
result[~nonzero] = 0

Answer 3

If you for some reason want to bypass an error with numpy it might be worth looking into the errstate context:

x  = np.array(range(-500, 500))

with np.errstate(divide='ignore'): #ignore zero-division error
    x = 1/x
x[x!=x] = 0 #convert inf and NaN's to 0

Answer 4

Consider changing the array in place by using np.put() :

In [56]: x = np.linspace(-1, 1, 5)

In [57]: x
Out[57]: array([-1. , -0.5,  0. ,  0.5,  1. ])

In [58]: indices = np.argwhere(x != 0)

In [59]: indices
Out[59]:
array([[0],
       [1],
       [3],
       [4]], dtype=int64)

In [60]: np.put(x, indices, 1/x[indices])

In [61]: x
Out[61]: array([-1., -2.,  0.,  2.,  1.])

The approach above does not create a new array, which could be very convenient if x is a large array.