python - 带有 bool 语句的 Numpy 向量化函数赋值

Question

我想分配一个函数，其中有一个 bool 值评估，使用快速的方式。这是一个简单的例子。我希望针对任意 a 和 b 评估以下函数:

a = 0.5
b = 0.6
def func(x):
    x=max(x,a)
    if x>b:
        return x**2
    else:
        return x**3

然后我想以向量化的方式将函数值分配到一个数组中(为了速度):

xRange = np.arange(0, 1, 0.1)
arr_func = func(xRange)

但是我得到了错误:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

现在，我知道我可以在循环中分配值。但与向量化的等价物相比，这会很慢。我可以绕过这个异常并仍然以矢量化方式分配值吗？

Answer 1

如果我没看错你的代码，这个的矢量化版本将使用一对 np.where:

def func(x):
    x = np.where(a > x, a, x)
    return np.where(x > b, x**2, x**3)