执行矩阵乘法来计算python中的旋转

时间:2018-02-18 02:20:12

标签: python numpy matrix rotation matrix-multiplication

我有一堆图像坐标存储在大小为Nx2xHxW的矩阵中。 “N”表示图像的数量,“H”和“W”是高度和宽度的缩写,并且该矩阵的每个nx1xHxW实例存储图像n的x和y坐标。

将旋转应用于2x1向量V是通过创建2x2旋转矩阵R:

来完成的
[[np.cos(theta), -np.sin(theta)],
 [np.sin(theta), np.sin(theta)]]

并执行矩阵乘法R * V,得到一个新的2x1向量。

现在我有一个大小为Nx2x2的旋转矩阵R,我想将旋转计算应用于M的每个实例。一个简单的方法是:

B = 8
N = 2
H = 100
W = 200
M = np.random.rand(B, N, H, W)
R = np.random.rand(B, N, N)

for i in range(B):
    tmp1 = R[i, 0, 0] * M[i, 0, :, :] + R[i, 0, 1] * M[i, 1, :, :]
    tmp2 = R[i, 1, 0] * M[i, 0, :, :] + R[i, 1, 1] * M[i, 1, :, :]
    M[i, 0, :, :] = tmp1
    M[i, 1, :, :] = tmp2

这显然非常低效,我想知道是否有更好的方法在python中执行此类任务。

谢谢!

1 个答案:

答案 0 :(得分:1)

一种方法是使用np.einsum

>>> Mpp = np.einsum('ijk,iklm->ijlm', R, M)

#verify
>>> for i in range(B):
...     tmp1 = R[i, 0, 0] * M[i, 0, :, :] + R[i, 0, 1] * M[i, 1, :, :]
...     tmp2 = R[i, 1, 0] * M[i, 0, :, :] + R[i, 1, 1] * M[i, 1, :, :]
...     M[i, 0, :, :] = tmp1
...     M[i, 1, :, :] = tmp2
... 
>>> 
>>> np.allclose(Mpp, M)
True