python数据统计 矢量图_Python中具有不同维数的数组的矢量化计算

以下是两种不同的选择.第一个使用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返回的值是否相等.