Python-Numpy：3D矩阵* 2D矢量快速计算

Question

大家好，

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

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

. 我想要（如果存在）单行numpy操作，该操作为我提供了（2,50）包含旋转矢量的数组，我们称其为 。 contains the product of the k-ith rotation matrix with the k-ith vector. 我的意思是， 的第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

对于需要沿一个或多个轴对齐的缩小，可以使用np.einsum

rotated_vector = np.einsum('ijk,jk->ki',rotation_matrices,vectors)

请注意，输出的形状为(N,2) ，其中N是向量的数量。 如果相反，您正在寻找形状为(2,N)的输出，并且需要原始代码为： rotated_vector[:,ind] = np.dot(...) ，只需编辑输出字符串表示ik而不是ki 。

运行时测试-

In [24]: def org_app(rotation_matrices,vectors):
    ...:     nb_of_vectors = vectors.shape[1]
    ...:     r = np.zeros((nb_of_vectors,2))
    ...:     for ind,elt in enumerate(np.arange(nb_of_vectors)):
    ...:         r[ind] = np.dot( rotation_matrices[:,:,ind], vectors[:,ind] )
    ...:     return r
    ...: 

In [25]: # Input arrays
    ...: rotation_matrices = np.random.rand(2,2,50)
    ...: vectors = np.random.rand(2,50)
    ...: 

In [26]: out1 = org_app(rotation_matrices,vectors)

In [27]: out2 = np.einsum('ijk,jk->ki',rotation_matrices,vectors)

In [28]: np.allclose(out1,out2) # Verify results
Out[28]: True

In [29]: %timeit org_app(rotation_matrices,vectors)
10000 loops, best of 3: 196 µs per loop

In [30]: %timeit np.einsum('ijk,jk->ki',rotation_matrices,vectors)
100000 loops, best of 3: 5.12 µs per loop

这再次证明了为什么在NumPy中进行迭代基本上是不好的 ！

Answer 2

您的轴处于异常状态。 首先，您需要将矩阵轴放在最后：

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]将向量数组（形状(n,)转换为列矩阵（形状(n, 1) ），尾随[...,0]将列矩阵转换回向量

Answer 3

这是一个明确的公式版本

result = np.array([vectors[0,:]*rotation_matrices[0,0,:] +
                   vectors[1,:]*rotation_matrices[0,1,:],
                   vectors[0,:]*rotation_matrices[1,0,:] +
                   vectors[1,:]*rotation_matrices[1,1,:]]).transpose()

多（14倍），比比你原来的代码，但速度较慢（2.6倍），更快的einsum我的机器上