使用 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（其中 z = 0）实际上不需要分配给返回的数组。

如何有条件地将值分配给 numpy 数组而不遇到 div/0 错误？循环是唯一的选择吗？

Answer 1

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

z == 0给出一个布尔值数组. 将其添加到z给出一个不带零元素的数组，该数组等于将使用的z索引np.select。