2d 矩阵的 Numpy 矩阵乘法得到 3d 矩阵

Question

我有两个 numpy 数组，比如

A: = array([[0, 1],  
            [2, 3],  
            [4, 5]])   

B = array([[ 6,  7],  
           [ 8,  9],  
           [10, 11]])

对于A和B的每一行，分别说Ra和Rb，我想计算转置（Ra）*Rb。因此，对于 A 和 B 的给定值，我想要以下答案：

array([[[ 0,  0],  
        [ 6,  7]],  

       [[ 16,  18],  
        [ 24,  27]],  

       [[ 40,  44],  
        [ 50,  55]]])

为此，我编写了以下代码：

x = np.outer(np.transpose(A[0]), B[0])
for i in range(1,len(A)):
    x = np.append(x,np.outer(np.transpose(A[i]), B[i]),axis=0)

有没有更好的方法来完成这项任务。

Answer 1

您可以使用A 和B 的扩展尺寸和np.newaxis/None 来引入broadcasting 以获得像这样的矢量化解决方案-

A[...,None]*B[:,None,:]

解释： np.outer(np.transpose(A[i]), B[i]) 基本上在A[i] 和B[i] 的列 版本之间进行元素乘法。您正在对A 中的所有行重复此操作，而不是B 中的对应行。请注意，np.transpose() 似乎不会产生任何影响，因为np.outer 会处理预期的元素乘法。

我会用一种矢量化语言来描述这些步骤，然后像这样实现 -

1. 扩展A 和B 的尺寸以形成3D 形状，这样我们在这两个扩展版本中也保持axis=0 对齐并保持axis=0。因此，剩下的就是决定最后两个轴了。
2. 要引入元素乘法，推 A 在其原始 2D 版本中的 axis=1 到其 3D 版本中的 axis=1，从而在 axis=2 处创建一个单独的维度以进行扩展A 的版本。
3. 3D 版本的 A 的最后一个单例维度必须与 axis=1 中的元素对齐（在原始的 2D 版本的 B 中），才能让 broadcasting 发生。因此，B 的扩展版本会将其 2D 版本中的 axis=1 中的元素推送到其3D 版本中的axis=2，从而为axis=1 创建一个单一维度。

最后，扩展版本将是：A[...,None] 和 B[:,None,:]，乘以谁会给我们想要的输出。