[英]Call function with internal sum over 2D array
假设我有一个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()将始终返回标量。
通过np.ix_()
从xaxis,yaxis,arr数据构造一个开放式网格
用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)
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.