python - 使用 numpy.select 避免 div/0 错误

Question

我正在尝试使用 numpy.select 有条件地为数组赋值。但是，使用 select 需要计算所有可能的赋值而不仅仅是相关赋值，这可能会导致 div/0 错误，例如:

import numpy as np
def testfunc(z):
    conditionlist = [z < 0, z == 0, z > 0]
    choicelist = [1 / z, 0, 1 + z]
    return np.select(conditionlist, choicelist)

if __name__ == "__main__":

     print testfunc(np.array([0]))

此代码将失败并出现 div/0 错误，尽管 1/z where z = 0 实际上永远不需要分配给返回的数组。

如何在不遇到 div/0 错误的情况下有条件地为 numpy 数组赋值？循环是唯一的选择吗？

Answer 1

np.select([z < 0, z == 0, z > 0], [1 / (z + (z == 0)), 0, 1 + z])

z == 0 给出一个 bool 数组。将此添加到 z 会得到一个没有零元素的数组，该数组等于 z 用于 np.select 将使用的索引。