从两个矩阵的乘积的条目中填充 4d 数组的有效方法答案

【问题标题】：Efficient way to fill up a 4d array from entries of a product of two matrices从两个矩阵的乘积的条目中填充 4d 数组的有效方法
【发布时间】：2015-09-25 02:05:07
【问题描述】：

标题可能没有我希望的那么精确，但问题就在这里。基本上我正在从两个矩阵的乘积的条目中填充一个 4d numpy 数组。现在代码如下：

M = P.dot(U)
C_arr = np.zeros((b_size,b_size,N,N))
  for alpha in xrange(b_size):
      for beta in xrange(b_size):
          for i in xrange(N):
              for j in xrange(N):
                  C_arr[alpha,beta,i,j] = np.conjugate(M[i,alpha])*M[j,beta]

事实证明，这个函数被调用得非常频繁，而且看起来非常耗时。我刚开始使用 Python，我怀疑可以通过避免这些循环来更有效地编写此函数，但我自己无法弄清楚...

【问题讨论】：

标签： python arrays numpy vectorization

【解决方案1】：

除了在其他解决方案中列出的带有np.einsum 的简洁代码外，您还可以像这样使用带有np.outer 的外积-

np.outer(M.conj().ravel(),M.ravel()).reshape(N,b_size,N,b_size).transpose(1,3,0,2)

运行时测试 -

In [54]: # Create input and get shape parameters
    ...: M = np.random.rand(10,10)
    ...: N,b_size = M.shape
    ...: 

In [55]: %timeit np.einsum('ia,jb->abij', M.conj(), M)
10000 loops, best of 3: 26 µs per loop

In [56]: %timeit np.outer(M.conj().ravel(),M.ravel()).reshape(N,b_size,N,b_size).transpose(1,3,0,2)
10000 loops, best of 3: 55.6 µs per loop

In [57]: # Create input and get shape parameters
    ...: M = np.random.rand(40,40)
    ...: N,b_size = M.shape
    ...: 

In [58]: %timeit np.einsum('ia,jb->abij', M.conj(), M)
10 loops, best of 3: 31 ms per loop

In [59]: %timeit np.outer(M.conj().ravel(),M.ravel()).reshape(N,b_size,N,b_size).transpose(1,3,0,2)
10 loops, best of 3: 24.2 ms per loop

In [60]: # Create input and get shape parameters
    ...: M = np.random.rand(80,80)
    ...: N,b_size = M.shape
    ...: 

In [61]: %timeit np.einsum('ia,jb->abij', M.conj(), M)
1 loops, best of 3: 497 ms per loop

In [62]: %timeit np.outer(M.conj().ravel(),M.ravel()).reshape(N,b_size,N,b_size).transpose(1,3,0,2)
1 loops, best of 3: 399 ms per loop

因此，根据输入数组的形状，您可以选择任何一种方式。

【讨论】：

【解决方案2】：

你可以使用numpy.einsum:

C = np.einsum('ia,jb->abij', M.conj(), M)

或者，由于没有计算实际的总和（即，这是一种外积的形式），您可以在适当地重塑输入矩阵 M 后使用带有常规数组乘法的 numpy 广播：

nrows, ncols = M.shape
C = M.T.reshape(1, ncols, 1, nrows) * M.T.conj().reshape(ncols, 1, nrows, 1)

【讨论】：