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

Question

如何在不使用循环结构的情况下执行以下代码，而是使用 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)

        

Answer 1

跟进我的评论，仅用 1 步替换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]])