a/b，但如果b==0则返回0

我想要一个除法方法，它返回a / b, 0if bis 0，nan如果至少有一个aor bis nan，并且适用于标量和数组输入，即，

import numpy as np


def div(a, b):
    out = a / b
    # out[b == 0] = 0  TypeError: 'float' object does not support item assignment
    # out = np.nan_to_num(out, nan=0.0)
    return out


assert div(2, 1) == 2
assert np.all(div(np.array([2, 2]), np.array([1, 1])) == np.array([2, 2]))
assert np.isnan(div(np.nan, np.nan))
assert div(np.array(0), np.array(0)) == 0

np.nan_to_num几乎可以解决问题，但我很想控制条件（b == 0vs isnan）。

回答

np.where(b==0, 0, a/b)