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

时间:2020-11-08 15:19:46

标签: python numpy numpy-einsum

如何不使用循环结构而是使用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, 2, 1, 1]],
                   [[2, 0, 0, 3],
                     [1, 0, 2, 4]]] )

B = np.array( [[[0, 2, 3, 1],
                 [0, 2, 5, 0]],
                  [[0, 1, 2, 2],
                    [3, 3, 2, 1]]] )
C = np.array( [[[0, 2, 2, 1],
                 [0, 2, 1, 0]],
                  [[0, 0, 2, 0],
                    [3, 1, 2, 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,j,k]
        X[j,k] += B[i,j,k]*C[i,j,k]
D = np.einsum('ijk,ijk->jk',A,B)
D += np.einsum('ijk,ijk->jk',B,C)

        

1 个答案:

答案 0 :(得分:0)

在我的评论之后,只需一步即可替换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,j,k]
    ...:         X[j,k] += B[i,j,k]*C[i,j,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,  0,  4,  6],
       [12,  3,  8,  5]])

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

       [[ 0,  0,  4,  6],
        [12,  3,  8,  5]]])

einsum相同:

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

       [[ 0,  0,  4,  6],
        [12,  3,  8,  5]]])

,如果您想对i求和:

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