如何不使用循环结构而是使用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)
答案 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]])