Python - Numpy：3D 矩阵 * 2D 向量快速计算答案

【问题标题】：Python - Numpy : 3D matrix * 2D vector fast calculationPython - Numpy：3D 矩阵 * 2D 向量快速计算
【发布时间】：2016-09-02 19:25:18
【问题描述】：

大家好，

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

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] )

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

感谢您的宝贵时间。

捷豹

【问题讨论】：

标签： python arrays performance numpy matrix

【解决方案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 np.matmul）

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

[...,None] 将向量数组（形状 (n,) 转换为列矩阵数组（形状 (n, 1)），而结尾的 [...,0] 将列矩阵转换回向量

【讨论】：

感谢您的回答。不幸的是，我使用的 numpy 版本不知道函数“matmul”。我用其他一些功能尝试了你的建议，但它不适用于我。也许是版本问题。
是的，matmul 很新

【解决方案2】：

对于需要沿一个或多个轴对齐的此类缩减，可以使用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 中进行迭代基本上是很糟糕！

【讨论】：

感谢您的回答。它显着改进了我的代码。我尽量避免循环。我希望我会习惯这种编码方式=)
@Jagaral 希望你能做到！祝你好运。
@Jagaral 不要忘记接受其中一个答案。在此处阅读有关接受答案的更多信息：meta.stackexchange.com/questions/5234/…

【解决方案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()

在我的机器上比您的原始代码快很多 (14x)，但比 einsum 慢 (2.6x)

【讨论】：