使用布尔语句的Numpy向量化函数赋值

Question

我想使用快速方式分配一个具有布尔值的函数。 这是一个简单的例子。 我希望对任意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：具有多个元素的数组的真值是不明确的。 使用a.any（）或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)

Answer 2

也可以使用np.select - 在这种情况下，它比必要的更加冗长，但可扩展到许多条件

def func(x):
    condlist =   [
                  x < a, 
                  (x >= a) & (x <= b), 
                  x > b
                 ]
    choicelist = [
                  a**3, 
                  x**3, 
                  x**2
                 ]
    return np.select(condlist, choicelist)

要么

def func(x):
    condlist =   [
                  x < a,  
                  x > b
                 ]
    choicelist = [
                  a**3, 
                  x**2
                 ]
    return np.select(condlist, choicelist, default = x**3)

Answer 3

如果你想（或有）保留原始函数，你可以使用numpy的vectorize函数来创建接受numpy数组作为输入的“矢量化”版本。

请注意，此功能仅为方便起见而没有提供性能改进，因为它在内部仅实现for循环。 所以没有“真正的”矢量化！

import numpy as np

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

vfunc = np.vectorize(func)  # this is it!

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

print(arr_func)

上面的代码工作并产生以下输出：

[ 0.125  0.125  0.125  0.125  0.125  0.125  0.36   0.49   0.64   0.81 ]