计算数组的期望，速度改进的建议

我有一个 (N0, N1, N2, N3) 矩阵 V 和一个 (N1, N1) 矩阵 M。N1 通常约为 30-50，N0xN1xN2xN3 约为 1 000 000。我想要一个新的 Matrix EV，其中 i0 , i1, i2, i3 条目由下式给出

np.sum(V[i0, :, i2, i3] * M[i1, :])    

我当前实现这一目标的代码是：

  V_exp = np.tile(V[:, :, :, :, None], (1, 1, 1, 1, N1))
  M_exp = np.tile(M.T[None, :, None, None, :], (N0, 1, N2, N3, 1))
  EV = np.sum(V_exp * M_exp, axis = 1)
  EV = np.rollaxis(EV, 3, 1)

我必须多次执行此操作，这是我的代码的绝对瓶颈。我想知道我的代码是否有提高速度的潜力。我感谢建议！

一次调用np.einsum http://docs.scipy.org/doc/numpy-1.10.0/reference/generated/numpy.einsum.html扩展后会一次性执行所有这些操作V to a 5D形状与np.newaxis/None http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#numpy.newaxis，就像这样-

EV  = np.einsum('ijklm,mj->imkl',V[...,None],M)

因此，我们避免使用任何中间数组来实现内存高效的解决方案。

解释

(1) 涉及扩展维度的起始代码：

V_exp = np.tile(V[:, :, :, :, None], (1, 1, 1, 1, N1))
M_exp = np.tile(M.T[None, :, None, None, :], (N0, 1, N2, N3, 1))
output1 = V_exp * M_exp

In np.einsum的方式，它会被翻译为：

np.einsum('ijklm,mj->ijklm',V[...,None],M)

请注意，我们已经使用了mj而不是通常的jm for M对应于M.T[None, :, None, None, :].

(2) 接下来，我们有：

EV = np.sum(V_exp * M_exp, axis = 1)

因此，我们总结一下axis = 1，所以 einsum 调用需要输出字符串说明符被改变自->ijklm to ->iklm.

（3）最后：

EV = np.rollaxis(EV, 3, 1)

对于本次移植rollaxis，我们只需要push up axis=3 to axis=1的位置和下推每一个axes=1,2向右一个位置。就这样输出字符串说明符会改变自->iklm to ->imkl给我们：

np.einsum('ijklm,mj->imkl',V[...,None],M)