Question

我有两个 numpy 数组 A 和 B，它们的维度都是 [2,2,n]，其中 n 是一个非常大的数字。我想在前两个维度中矩阵乘以A和B得到C，即C=AB，其中C的维度为[2,2,n] .

最简单的方法是使用 for 循环，即

for i in range(n):
    C[:,:,i] = np.matmul(A[:,:,i],B[:,:,i])

然而，这是低效的，因为 n 非常大。使用 numpy 执行此操作的最有效方法是什么？

Answer 1

您可以执行以下操作：

new_array = np.einsum('ijk,jlk->ilk', A, B)

Answer 2

你想要的是Numpy中默认的数组乘法

In [22]: a = np.arange(8).reshape((2,2,2))+1 ; a[:,:,0], a[:,:,1]
Out[22]: 
(array([[1, 3],
        [5, 7]]),
 array([[2, 4],
        [6, 8]]))

In [23]: aa = a*a ; aa[:,:,0], aa[:,:,1]
Out[23]: 
(array([[ 1,  9],
        [25, 49]]),
 array([[ 4, 16],
        [36, 64]]))

请注意，我强调了 array，因为 Numpy 的数组看起来像矩阵，但实际上是 Numpy's ndarrays。

后脚本

我猜您真正想要的是矩阵形状为(n,2,2)的数组，这样您就可以使用单个索引处理单个2×2 矩阵，例如，

In [27]: n = 3
    ...: a = np.arange(n*2*2)+1 ; a_22n, a_n22 = a.reshape((2,2,n)), a.reshape((n,2,2))
    ...: print(a_22n[0])
    ...: print(a_n22[0])
[[1 2 3]
 [4 5 6]]
[[1 2]
 [3 4]]

Post Post Scriptum

重新语义正确：

In [13]: import numpy as np
    ...: n = 3
    ...: a = np.arange(2*2*n).reshape((2,2,n))+1
    ...: p = lambda t,a,n:print(t,*(a[:,:,i]for i in range(n)),sep=',\n')
    ...: p('Original array', a, n)
    ...: p('Using `einsum("ijk,jlk->ilk", ...)`', np.einsum('ijk,jlk->ilk', a, a), n)
    ...: p('Using standard multiplication', a*a, n)
Original array,
[[ 1  4]
 [ 7 10]],
[[ 2  5]
 [ 8 11]],
[[ 3  6]
 [ 9 12]]
Using `einsum("ijk,jlk->ilk", ...)`,
[[ 29  44]
 [ 77 128]],
[[ 44  65]
 [104 161]],
[[ 63  90]
 [135 198]]
Using standard multiplication,
[[  1  16]
 [ 49 100]],
[[  4  25]
 [ 64 121]],
[[  9  36]
 [ 81 144]]

逐元素numpy矩阵乘法

2 个答案: