avoid division by zero in numpy.where()
Simply initialize output array with the fallback values (condition-not-satisfying values) or array and then mask to select the condition-satisfying values to assign -
out = a.copy()
out[mask] /= b[mask]
If you are looking for performance, we can use a modified b
for the division -
out = a / np.where(mask, b, 1)
Going further, super-charge it with numexpr
for this specific case of positive values in b
(>=0) -
import numexpr as ne
out = ne.evaluate('a / (1 - mask + b)')
Benchmarking
Code to reproduce the plot:
import perfplot
import numpy
import numexpr
numpy.random.seed(0)
def setup(n):
a = numpy.random.rand(n)
b = numpy.random.rand(n)
b[b < 0.3] = 0.0
mask = b > 0
return a, b, mask
def copy_slash(data):
a, b, mask = data
out = a.copy()
out[mask] /= b[mask]
return out
def copy_divide(data):
a, b, mask = data
out = a.copy()
return numpy.divide(a, b, out=out, where=mask)
def slash_where(data):
a, b, mask = data
return a / numpy.where(mask, b, 1.0)
def numexpr_eval(data):
a, b, mask = data
return numexpr.evaluate('a / (1 - mask + b)')
perfplot.save(
"out.png",
setup=setup,
kernels=[copy_slash, copy_divide, slash_where, numexpr_eval],
n_range=[2 ** k for k in range(22)],
xlabel="n"
)
A slight variation on Divakar's answer is to use the where
and out
arguments of Numpy's divide function
out = a.copy()
np.divide(a, b, out=out, where=mask)
For big arrays, this seems to be twice as fast:
In [1]: import numpy as np
In [2]: a = np.random.rand(1000, 1000)
...: b = np.random.rand(1000, 1000)
...: b[b < 0.3] = 0.0
In [3]: def f(a, b):
...: mask = b > 0
...: out = a.copy()
...: out[mask] = a[mask] / b[mask]
...: return out
...:
In [4]: def g(a, b):
...: mask = b > 0
...: out = a.copy()
...: np.divide(a, b, out=out, where=mask)
...: return out
...:
In [5]: (f(a, b) == g(a, b)).all() # sanity check
Out[5]: True
In [6]: timeit f(a,b)
26.7 ms ± 52.6 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [7]: timeit g(a,b)
12.2 ms ± 36 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
The reason why this is faster is likely since this avoids making a temporary array for the right-hand-side, and since the 'masking' is done internally to the divide
function, instead of by the indexing of a[mask]
, b[mask]
and out[mask]
.