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

Question

标题可能没有我希望的那么准确，但这就是问题所在。基本上我从两个矩阵的乘积条目中填充一个 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]

事实证明，这个函数被调用的频率很高，看起来很time-consumming。我刚从 Python 开始，我怀疑通过避免这些循环可能有更有效的方法来编写这个函数，但我自己还没有弄清楚...

Answer 1

您可以使用 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)

Answer 2

除了在其他解决方案中列出的带有 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

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

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

Efficient way to fill up a 4d array from entries of a product of two matrices

python

arrays

numpy

vectorization