python - 使用二维数组的内部求和调用函数

Question

假设我有一个 python 函数 f()，它接受 2 个标量和 1 个类似数组的参数:

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)