python - 从两个矩阵乘积的条目填充 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]

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

Answer 1

您可以使用 numpy.einsum :

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

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

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