[英]Call function with internal sum over 2D array
Suppose I have a python function f() that accepts 2 scalar and 1 array_like parameters: 假设我有一个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)
What I want to do is to evaluate f() for different values of a and b while keeping the same 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)
Is there a way to vectorize this? 有办法向量化吗? I'm thinking of converting XX, YY and arr into 3D arrays of the same shape but the np.sum() in f() will always return a scalar.
我正在考虑将XX,YY和arr转换为相同形状的3D数组,但f()中的np.sum()将始终返回标量。
Construct an open mesh from xaxis, yaxis, arr data by np.ix_()
通过
np.ix_()
从xaxis,yaxis,arr数据构造一个开放式网格
Call np.sum()
with axis=-1
. 用
axis=-1
调用np.sum()
。
Here is the code: 这是代码:
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.