如何根据 numpy 数组中的值更改函数

Question

我正在尝试制作一个 function 的表面 plot，如下所示：

def model(param,x_1,x_2,x_3,x_4):
    est=param[0]+param[1]*(x_1+x_2*x_3+x_2**2*x_4)
    return est

重点是根据x_2的值，x_3=1 for x_2>=0 and x_4=1 for x_2<0 (else:0)。

当我尝试制作表面 plot 时，我很困惑如何制作网格，因为除了 x_1 和 x_2 之外还有 2 个变量。

为了计算 z 轴，我尝试将 function 修改为：

def function (param,x_1,x_2):
    if x_2>0:
      est=param[0]+param[1]*(x_1+x_2)
    else:
      est=param[0]+param[1]*(x_1+x_2**2)
    return est)

但是，它表示真值是不明确的。 我将其理解为 python 是否看到 x_2 中的所有值是否 >=0。

我也尝试过使用 np.sign()，但它并没有按照我想要的方式运行。

是否有任何方法可以根据数组中 elenemt 的值更改 function 和/或解决此问题而无需使用 for 循环手动计算 z 轴？

Answer 1

如果要检查所有大于 0 的值，请使用all ：

def function (param,x_1,x_2):
    if all(x_2>0):
      est=param[0]+param[1]*(x_1+x_2)
    else:
      est=param[0]+param[1]*(x_1+x_2**2)
    return est

但如果您想对每个值应用测试，请使用np.where ：

def function (param,x_1,x_2):
    return np.where(x_2 > 0,
                    param[0]+param[1]*(x_1+x_2),
                    param[0]+param[1]*(x_1+x_2**2))

Answer 2

我想你需要numpy.where ：

def function (param,x_1,x_2):
    return np.where(x_2>0, 
                    param[0]+param[1]*(x_1+x_2), 
                    param[0]+param[1]*(x_1+x_2**2))

它是如何工作的：

param = [10,8]
x_1 = np.array([1,2,3])
x_2 = np.array([0,4,10])

如果掩码中的True s 是使用来自param[0]+param[1]*(x_1+x_2)的值，否则来自param[0]+param[1]*(x_1+x_2**2) ：

print (x_2>0)
[False  True  True]

print (param[0]+param[1]*(x_1+x_2))
[ 18  58 114]

print (param[0]+param[1]*(x_1+x_2**2))
[ 18 154 834]

print (function(param,x_1,x_2))
[ 18  58 114]