Python-Numpy : 3D matrix * 2D vector fast calculation

Question

大家好，

这是我想做的。我有两个数组:

• rotation_matrices 包含 50 个二维旋转矩阵。每个旋转矩阵的形状为 (2,2)。因此，rotation_matrices 的形状为 (2,2,50)。
• vectors 包含 50 个二维向量。因此，它的形状为 (2,50)。

我想要(如果它存在)一个单行 numpy 操作，它给我包含旋转向量的 (2,50) 数组，我们称它为 rotated_vectors。我的意思是，rotated_vectors 的第 k 个元素包含第 k 个旋转矩阵与第 k 个向量的乘积。

目前，我想出了以下循环:

for ind,elt in enumerate(np.arange(nb_of_vectors)):
            rotated_vector[ind] = np.dot( rotation_matrices[:,:,ind], vectors[:,ind] )

我认为还有改进的余地。如果您有任何建议，欢迎您。

感谢您的宝贵时间。

猎豹

Answer 1

您的坐标轴顺序不正常。首先，您需要将矩阵轴放在最后:

rotation_matrices = np.rollaxis(rotation_matrices, -1)  # shape (50, 2, 2)
vectors = np.rollaxis(vectors, -1)                      # shape (50, 2)

这将使您现有的循环更具可读性:

for ind in np.arange(nb_of_vectors):
    rotated_vector[ind] = np.dot(rotation_matrices[ind], vectors[ind])

但是，您可以使用矩阵乘法运算符(或 python < 3.5 中的 np.matmul)

rotated_vectors = (a @ vectors[...,None])[...,0]
# rotated_vectors = np.matmul(a, vectors[...,None])[...,0]

[...,None] 将向量数组 (shape (n,) 转换为列矩阵数组 (shape (n, 1)), 尾部的 [...,0] 将列矩阵转换回向量