python - Numpy 矩阵与 2D 元素的乘法

Question

我有一个 a numpy ndarray 3x3 矩阵，如下所示

a =  ([[ uu, uv, uw],
       [ uv, vv, vw],
       [ uw, vw, ww]])

每个分量本身就是一个大小为 (N,M) 的二维数组，因此 a 矩阵具有 (3,3,N,M) 形状。

如何以 Python 方式执行 a*a 的矩阵乘法？使用 a@a 会引发以下错误(对于 N=1218 和 M=540):

ValueError: shapes (3,3,1218,540) and (3,3,1218,540) not aligned: 540 (dim 3) != 1218 (dim 2)

我希望能够执行此操作，就好像 a 的元素只是标量值，其中 a@a 不会抛出与其形状相关的错误，因为这是一个简单的 3x3 矩阵乘法。

谢谢。

Answer 1

假设您希望对最后两个轴上的每个元素执行矩阵乘法，我们可以使用 np.einsum -

np.einsum('ijkl,jmkl->imkl',a,a)

运行示例进行验证 -

In [43]: np.random.seed(0)

In [44]: a = np.random.rand(3,3,4,5)

In [45]: a[:,:,0,0].dot(a[:,:,0,0])
Out[45]: 
array([[0.71750146, 1.17057872, 1.11135764],
       [0.62938365, 0.86437796, 0.74541383],
       [1.04636618, 1.62011127, 1.35483565]])

In [46]: np.einsum('ijkl,jmkl->imkl',a,a)[:,:,0,0]
Out[46]: 
array([[0.71750146, 1.17057872, 1.11135764],
       [0.62938365, 0.86437796, 0.74541383],
       [1.04636618, 1.62011127, 1.35483565]])