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和Bwith的扩展维度np.newaxis/None来引入broadcasting矢量化解决方案，如下所示 -

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

解释： 基本上在和的柱状版本np.outer(np.transpose(A[i]), B[i])之间进行元素乘法。您将针对 中的相应行对 中的所有行重复此操作。请注意，这似乎没有产生任何影响A[i]B[i]ABnp.transpose()np.outer它会处理预期的元素乘法。

我会用矢量化语言描述这些步骤，从而实现，就像这样 -

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

最后，扩展版本将是 : A[...,None]& B[:,None,:]，将其相乘即可得到所需的输出。