如何使用numpy einsum函数对3D数组执行矩阵乘法以获取2D乘积矩阵？

问题描述

如何不使用循环结构而是使用numpy einsum函数来执行以下代码？我希望乘积矩阵是2D矩阵而不是3D。简单地做 “ D = np.einsum（'ijk，ijk-> jk'，A，B） D + = np.einsum（'ijk，ijk-> jk'，B，C）“给出不同的结果。我应该引入任何中间临时数组还是使用einsum函数的东西？

import numpy as np
A = np.array( [[[1,2,3,0],[ 4,1,1]],[[2,3],[1,4]]] )

B = np.array( [[[0,1],[0,5,0]],[[0,2],[3,1]]] )
C = np.array( [[[0,1]]] )
X = np.zeros([2,4])
for i in range(2):
  for j in range(2):
    for k in range(4):

        X[j,k] = A[i,j,k]*B[i,k]
        X[j,k] += B[i,k]*C[i,k]
D = np.einsum('ijk,ijk->jk',A,B)
D += np.einsum('ijk,B,C)

        

解决方法

在我的评论之后，只需一步即可替换i循环

In [64]:   X = np.zeros([2,4])
    ...:   i=1
    ...:   for j in range(2):
    ...:     for k in range(4):
    ...: 
    ...:         X[j,k] = A[i,j,k]*B[i,k]
    ...:         X[j,k] += B[i,k]*C[i,k]
In [65]: X
Out[65]: 
array([[ 0.,0.,4.,6.],[12.,3.,8.,5.]])

这是循环产生的值。您已经舍弃了i=0值。

但是您不需要循环：

In [68]: A[1,:,:]*B[1,:]+B[1,:]*C[1,:]
Out[68]: 
array([[ 0,4,6],[12,3,8,5]])

In [69]: A*B+B*C
Out[69]: 
array([[[ 0,15,1],[ 0,10,0]],[[ 0,5]]])

与einsum相同：

In [71]: np.einsum('ijk,ijk->ijk',A,B)+np.einsum('ijk,B,C)
Out[71]: 
array([[[ 0,5]]])

，如果您想对i求和：

In [72]: np.einsum('ijk,ijk->jk',C)
Out[72]: 
array([[ 0,19,7],11,18,5]])
In [73]: (A*B+B*C).sum(axis=0)
Out[73]: 
array([[ 0,5]])