Numpy "Where" function can not avoid evaluate Sqrt(negative)

There is a much better way of doing this. Let's take a look at what your code is doing to see why.

np.where accepts three arrays as inputs. Arrays do not support lazy evaluation.

d = np.where(c >= 0, np.sqrt(c), c)

This line is therefore equivalent to doing

a = (c >= 0)
b = np.sqrt(c)
d = np.where(a, b, c)

Notice that the inputs are computed immediately, before where ever gets called.

Luckily, you don't need to use where at all. Instead, just use a boolean mask:

mask = (c >= 0)
d = np.empty_like(c)
d[mask] = np.sqrt(c[mask])
d[~mask] = c[~mask]

If you expect a lot of negatives, you can copy all the elements instead of just the negative ones:

d = c.copy()
d[mask] = np.sqrt(c[mask])

An even better solution might be to use masked arrays:

d = np.ma.masked_array(c, c < 0)
d = np.ma.sqrt(d)

To access the whole data array, with the masked portion unaltered, use d.data.


np.sqrt is a ufunc and accepts a where parameter. It can be used as a mask in this case:

In [61]: c = np.arange(10)-5.0
In [62]: d = c.copy()
In [63]: np.sqrt(c, where=c>=0, out=d);
In [64]: d
Out[64]: 
array([-5.        , -4.        , -3.        , -2.        , -1.        ,
        0.        ,  1.        ,  1.41421356,  1.73205081,  2.        ])

In contrast to the np.where case, this does not evaluate the function at the ~where elements.

Tags:

Python

Numpy