使用Numpy的最快方式 - 多维总和和产品

Question

我有以下维度的这些变量:

A   - (3,)
B   - (4,)
X_r - (3,K,N,nS)
X_u - (4,K,N,nS)
k - (K,)

我想计算(A.dot(X_r[:,:,n,s])*B.dot(X_u[:,:,n,s])).dot(k)每一个可能的n和s,我现在正在做的方式如下:

np.array([[(A.dot(X_r[:,:,n,s])*B.dot(X_u[:,:,n,s])).dot(k) for n in xrange(N)] for s in xrange(nS)]) #nSxN

但这是超级慢,我想知道是否有更好的方法,但我不确定.

然而,我正在做另一个计算,我相信它可以被优化:

np.sum(np.array([(X_r[:,:,n,s]*B.dot(X_u[:,:,n,s])).dot(k) for n in xrange(N)]),axis=0)

在这一个我创建一个numpy数组只是为了在一个轴上求和并丢弃后的数组.如果这是1-DI中的列表将使用reduce和优化它,我应该使用什么numpy数组？

Answer 1

使用很少的np.einsum调用 -

# Calculation of A.dot(X_r[:,:,n,s])
p1 = np.einsum('i,ijkl->jkl',A,X_r)

# Calculation of B.dot(X_u[:,:,n,s])
p2 = np.einsum('i,ijkl->jkl',B,X_u)

# Include .dot(k) part to get the final output
out = np.einsum('ijk,i->kj',p1*p2,k)

关于第二个例子，这解决了它：

p1 = np.einsum('i,ijkl->jkl',B,X_u)#OUT_DIM - (k,N,nS)
sol = np.einsum('ijkl,j->il',X_r*p1[None,:,:,:],k)#OUT_DIM (3,nS)