avoid division by zero in numpy.where()

user_353337 0
Viewed 877

I have two numpy arrays a, b of the same shape, b has a few zeros. I would like to set an output array to a / b where b is not zero, and a otherwise. The following works, but yields a warning because a / b is computed everywhere first.

import numpy

a = numpy.random.rand(4, 5)
b = numpy.random.rand(4, 5)
b[b < 0.3] = 0.0

A = numpy.where(b > 0.0, a / b, a)

/tmp/l.py:7: RuntimeWarning: divide by zero encountered in true_divide
  A = numpy.where(b > 0.0, a / b, a)

Filtering the division with a mask doesn't perserve the shape, so this doesn't work:

import numpy

a = numpy.random.rand(4, 5)
b = numpy.random.rand(4, 5)
b[b < 0.3] = 0.0

mask = b > 0.0
A = numpy.where(mask, a[mask] / b[mask], a)

ValueError: operands could not be broadcast together with shapes (4,5) (14,) (4,5)

Any hints on how to avoid the warning?

numpy python
2 Answers

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

enter image description here

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


b = perfplot.bench(
    setup=setup,
    kernels=[copy_slash, copy_divide, slash_where, numexpr_eval],
    n_range=[2 ** k for k in range(24)],
    xlabel="n"
)
b.save("out.png")
user_32938810

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

user_26472790
Related