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我的機器上