避免在 numpy.where() 中被零除

Question

我有两个 numpy 数组a，b形状相同，b有几个零。我想设定的输出阵列a / b，其中b不为零，和a其它。以下有效，但会产生警告，因为a / b首先在任何地方计算。

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)

用 a 过滤除法mask不会保持形状，所以这不起作用：

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)

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)

关于如何避免警告的任何提示？

Answer 1

只需使用后备值（不满足条件的值）或数组初始化输出数组，然后掩码以选择要分配的满足条件的值 -

out = a.copy()
out[mask] /= b[mask]

如果您正在寻找性能，我们可以使用修改b后的除法 -

out = a / np.where(mask, b, 1)

更进一步，针对(>=0)numexpr中正值的这种特定情况对其进行超级充电 -b

import numexpr as ne
    
out = ne.evaluate('a / (1 - mask + b)')

标杆管理

重现情节的代码：

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