以下是两种不同的选择.第一个使用ndarray.sum和NumPyinteger array indexing.第二个使用np.einsum.
def using_sum(T):
total = T.sum(axis=1).sum(axis=0)
m = np.arange(T.shape[0])
trace = T[m, m].sum(axis=0)
return total - trace
def using_einsum(T):
return np.einsum('mnijk->ijk', T) - np.einsum('nnijk->ijk', T)
np.einsum的第一个参数指定求和的下标.
“ mnijk-> ijk”表示T具有子句mnijk,求和后仅保留ijk下标.因此,对m和n下标进行求和.这使得
np.einsum(‘mnijk-> ijk’,T)[i,j,k]等于np.sum(T [:,:,i,j,k]),但以一个矢量化计算来计算整个数组.
类似地,“ nnijk-> ijk”告诉np.einsum T具有下标nnijk,并且再次只有ijk下标可以求和.因此,总和超过n.由于重复了n,因此对n的求和将计算出轨迹.
我喜欢np.einsum,因为它传达了计算的意图
简洁地.但自从开始以来,熟悉using_sum的工作方式也很好
它使用基本的NumPy操作.这是嵌套循环的一个很好的例子
通过使用对整个数组进行操作的NumPy方法可以避免这种情况.
这是一个perfplot,将orig与using_sum和using_einsum的性能作为n的函数进行比较,其中T的形状为(10,10,n,n,n):
import perfplot
import numpy as np
def orig(T):
_, _, nx, ny, nz = T.shape
r = np.zeros((nx, ny, nz))
for i in range(nx):
for j in range(ny):
for k in range(nz):
r[i,j,k] = np.sum(T[:,:,i,j,k])-np.trace(T[:,:,i,j,k])
return r
def using_einsum(T):
r = np.einsum('mnijk->ijk', T) - np.einsum('nnijk->ijk', T)
return r
def using_sum(T):
total = T.sum(axis=1).sum(axis=0)
m = np.arange(T.shape[0])
trace = T[m,m].sum(axis=0)
return total - trace
def make_T(n):
return np.random.random((10,10,n,n,n))
perfplot.show(
setup=make_T,
kernels=[orig, using_sum, using_einsum],
n_range=range(2, 80, 3),
xlabel='n')
perfplot.show还检查orig,using_sum和using_einsum返回的值是否相等.