向量化python函數

Question

我有一個名爲old_func的舊函數，它輸入兩個位置參數x和y。該函數的輸入是這樣寫的使用元組作爲輸入：

def old_func(position): 

    x, y = position 
    return x**2 + y**2 

我現在想調用該函數在值的網格的快速簡便的方法：

xx = numpy.linspace(0, 1, 100) 
yy = numpy.linspace(0, 1, 100) 
X, Y = numpy.meshgrid(xx, yy) 

array_positions = (X,Y) 
old_fun(array_positions) 

意圖x中的每個操作都在X上完成，y也是如此。我嘗試使用numpy.vectorize向量化該功能，但這不起作用。我寧願不更改函數來接受NumPy數組，因爲這會花費太長時間。

Answer 1

下面的代碼是訣竅，併爲您節省創建array_positions的麻煩。

---

首先，使用ravel弄平X和Y成形狀(10000,)的NumPy的陣列。

X_flattened = X.ravel() 
Y_flattened = Y.ravel() 

然後，使用apply_along_axis實現自定義函數迭代地向下這些扁平陣列的長度。

float_array = np.apply_along_axis(old_func, 0, (X_flattened, Y_flattened)) 

最後，將輸出數組重新整形爲所需的形狀(100, 100)。

np.reshape(float_array, (100, 100)) 

Answer 2

自己的代碼應該（和不）工作的優良：

def old_fun(position): 

    x, y = position 
    z = x**2 + y**2 
    return z 

xx = numpy.linspace(0,1,100) 
yy = numpy.linspace(0,1,100) 
X,Y = numpy.meshgrid(xx, yy) 

array_positions = (X,Y) 
Z = old_fun(array_positions) 

Z.shape現在(100, 100)。

一般來說，numpy數組可以與任何標準運算符一起使用，如+和**。請注意，儘管您的old_fun將元組作爲輸入，並且此元組需要由兩個值組成，但只要此類型支持數學運算符，則這兩個值的類型可以是任何類型。標準的Python標量和numpy數組都支持這些，所以代碼工作正常。

的注意事項有關@ JuanManuel的回答是：雖然工作完全正常，以及，apply_along_axis是專門並不意味着在這個意義上矢量化，人們通常要使用它，即獲得良好的表現出來numpy的的。 apply_along_axis將導致Python循環緩慢，而不是像正確矢量化代碼那樣的快速C循環。你自己的代碼使用適當的矢量化，所以用它來代替。