具有2D數組內部和的調用函數

Question

假設我有一個python函數f（），它接受2個標量和1個array_like參數：

def f(a, b, arr):
    X = a * np.exp(-arr**2 / b)
    return np.sum(a * np.log(X) - arr)

我想做的是評估f（）的a和b的不同值，同時保持相同的arr：

XX, YY = np.meshgrid(A_axis, B_axis)
arr = np.arange(10)
ZZ = np.empty_like(XX)
for i in range(XX.shape[0]):
    for j in range(YY.shape[1]):
        ZZ[i,j] = f(XX[i,j], YY[i,j], arr)

有辦法向量化嗎？ 我正在考慮將XX，YY和arr轉換為相同形狀的3D數組，但f（）中的np.sum（）將始終返回標量。

Answer 1

1. 通過np.ix_()從xaxis，yaxis，arr數據構造一個開放式網格

2. 用axis=-1調用np.sum() 。

這是代碼：

import numpy as np

### original code
def f(a, b, arr):
    X = a * np.exp(-arr**2 / b)
    return np.sum(a * np.log(X) - arr)

A_axis = np.linspace(1, 5, 8)
B_axis = np.linspace(1, 2, 9)
XX, YY = np.meshgrid(A_axis, B_axis)
arr = np.arange(10)
ZZ = np.empty_like(XX)
for i in range(XX.shape[0]):
    for j in range(YY.shape[1]):
        ZZ[i,j] = f(XX[i,j], YY[i,j], arr)

### use broadcast        
def f(a, b, arr):
    X = a * np.exp(-arr**2 / b)
    return np.sum(a * np.log(X) - arr, axis=-1)
B, A, C = np.ix_(B_axis, A_axis, arr)
result = f(A, B, C)

print np.allclose(ZZ, result)