numpy中的分段函數具有多個參數

Question

我試圖定義一個函數（帳篷圖），如下所示：

def f(r, x):
    return np.piecewise([r, x], [x < 0.5, x >= 0.5], [lambda r, x: 2*r*x, lambda r, x: 2*r*(1-x)])

r，x將是numpy數組：

no_r = 10001
r = np.linspace(0, 4, no_r)
x = np.random.rand(no_r)

我希望結果是一個匹配r和x形狀的numpy數組，該數組使用具有相同索引的數組r和x的每對元素來計算。 例如，如果r = [0，1，2，3]和x = [0.1，0.7，0.3，1]，則結果應為[0，0.6，1.2，0]。 發生錯誤：“布爾索引與維度0上的索引數組不匹配；維度為2，但相應的布爾維度為10001”那么我應該怎么做才能達到預期的目的？

Answer 1

您想要得到的結果可以通過np.select完成，例如：

def f(r, x):
    return np.select([x < 0.5,x >= 0.5], [2*r*x, 2*r*(1-x)])

然后用

r = np.array([0, 1, 2, 3])
x = np.array([0.1, 0.7, 0.3, 1])

print (f(r,x))
[0.  0.6 1.2 0. ]

編輯：在這種情況下，只有兩個條件是互斥的，您也可以使用np.where ：

def f(r,x):
    return np.where(x<0.5,2*r*x, 2*r*(1-x))

將給出相同的結果。