矩阵向量运算的 Python 向量化

Question

我有一个形状为 (2,2,N) 的矩阵 A 和一个形状为 (2,N) 的矩阵 V

我想对以下内容进行矢量化：

F = np.zeros(N)
for k in xrange(N):
    F[k] = np.dot( A[:,:,k], V[:,k] ).sum()

任何方式都可以使用 tensordot 或任何其他 numpy 函数来完成，而无需显式循环？

Answer 1

与np.einsum -

F = np.einsum('ijk,jk->k',A,V)

我们可以通过将optimize 标志（检查文档）设置为True 来进一步优化它。