呼叫函数

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转换为相同形状的三维数组，但f（）中的np.sum（）将始终返回一个标量。

Answer 1

1. 通过np.ix_()

2. 呼叫np.sum()与axis=-1构建从x轴，y轴，ARR数据的开放网格。

下面是代码：

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)