估计两个二维点云之间的旋转

Question

我有两个元素数量相同的二维点云。对于这些元素，我知道它们的对应关系，即对于 PC1 中的每个点，我都知道 PC2 中的对应元素，反之亦然。

我现在想估计这两个点云之间的旋转。也就是说，我想找到角度alpha，我必须围绕原点旋转PC1中的所有点，以使PC1和PC2中对应点之间的距离最小。

我可以使用 scipy 的线性优化器解决这个问题（见下文）；但是，这种优化位于我的代码关键路径的循环内，是当前的瓶颈。

import numpy as np
from scipy.optimize import minimize_scalar
from math import sin, cos

# generate some data for demonstration purposes
# points in each point cloud are ordered by correspondence
num_points = 10

distance = np.random.rand(num_points) * 3
radii = np.random.rand(num_points) * 2*np.pi
pc1 = distance[:, None] * np.stack([np.cos(radii), np.sin(radii)], axis=1)

distance = np.random.rand(num_points) * 3
radii = np.random.rand(num_points) * 2*np.pi
pc2 = distance[:, None] * np.stack([np.cos(radii), np.sin(radii)], axis=1)

# solve using scipy
def score(alpha):
    rot_matrix = np.array([
        [cos(alpha), -sin(alpha)],
        [sin(alpha), cos(alpha)]
    ])
    pc1_rotated = (rot_matrix @ pc1.T).T

    sum_of_squares = np.sum((pc2 - pc1_rotated)**2, axis=1)
    mse = np.mean(sum_of_squares)

    return mse

# simple solution via scipy
result = minimize_scalar(
            score,
            bounds=(0, 2*np.pi),
            method="bounded",
            options={"maxiter": 1000},
        )

if result.success:
    print(f"Best angle: {result.x}")
else:
    raise RuntimeError(f"IK failed. Reason: {result.message}")

对于这个问题有更快（可能是分析）的解决方案吗？

Answer 1

由于minimize_scalar 仅使用无导数方法，优化运行时间很大程度上取决于评估目标函数score 所需的时间。因此，我建议尽可能加速此功能。

让我们为您的函数和优化器计时作为基准参考：

In [68]: %timeit score(0.5)
20.2 µs ± 203 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [69]: %timeit result = minimize_scalar(score,bounds=(0, 2*np.pi),method="bounded",options={"maxiter": 1000})
415 µs ± 7.27 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

• 首先，请注意(rot_matrix @ pc1.T).T 与pc1 @ rot_matrix.T 相同，即我们只需要转置一个矩阵而不是两个。

• 接下来，请注意 -sin(alpha) = cos(alpha + 5*pi/2) 和 sin(alpha) = cos(alpha + 3*pi/2)。这意味着我们只需要调用一次np.cos 来创建rot_matrix，而不是调用四次math.sin 或math.cos。

• 最后，您可以通过np.einsum 更有效地计算mse。

考虑所有点，函数可以如下所示：

k1 = 5*np.pi/2
k2 = 3*np.pi/2

def score2(alpha):
    rot_matrixT = np.cos((alpha, alpha+k2, alpha + k1, alpha)).reshape(2,2)
    pc1_rotated = pc1 @ rot_matrixT
    diff = pc2 - pc1_rotated
    return np.einsum('ij,ij->', diff, diff) / num_points

再次给函数计时

In [70]: %timeit score(0.5)
9.26 µs ± 84.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

因此，优化器要快得多：

In [71]: %timeit result = minimize_scalar(score, bounds=(0, 2*np.pi), method="bounded", options={"maxiter": 1000})
279 µs ± 1.79 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

如果还不够快，你可以通过Numba及时编译你的函数：

In [60]: from numba import njit

In [61]: @njit
    ...: def score3(alpha):
    ...:     rot_matrix = np.array([
    ...:         [cos(alpha), -sin(alpha)],
    ...:         [sin(alpha), cos(alpha)]
    ...:     ])
    ...:     pc1_rotated = (rot_matrix @ pc1.T).T
    ...:     sum_of_squares = np.sum((pc2 - pc1_rotated)**2, axis=1)
    ...:     mse = np.mean(sum_of_squares)
    ...:     return mse

In [62]: %timeit score3(0.5)
2.97 µs ± 47.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

或使用Cython 重写它。为了完整起见，这里有一个快速的 Cython 实现：

In [45]: %%cython -c=-O3 -c=-march=native -c=-Wno-deprecated-declarations -c=-Wno-#warnings
    ...:
    ...: from libc.math cimport cos, sin
    ...: cimport numpy as np
    ...: import numpy as np
    ...: from cython cimport wraparound, boundscheck
    ...:
    ...: @wraparound(False)
    ...: @boundscheck(False)
    ...: cpdef double score4(double alpha, double[:, ::1] pc1, double[:, ::1] pc2):
    ...:     cdef int i
    ...:     cdef int N = pc1.shape[0]
    ...:     cdef double diff1 = 0.0
    ...:     cdef double diff2 = 0.0
    ...:     cdef double   mse = 0.0
    ...:     cdef double  rmT00 = cos(alpha)
    ...:     cdef double  rmT01 = sin(alpha)
    ...:     cdef double  rmT10 = -rmT01
    ...:     cdef double  rmT11 = rmT00
    ...:
    ...:     for i in range(N):
    ...:         diff1 = pc2[i,0] - (pc1[i,0]*rmT00 + pc1[i,1]*rmT10)
    ...:         diff2 = pc2[i,1] - (pc1[i,0]*rmT01 + pc1[i,1]*rmT11)
    ...:         mse  += diff1*diff1 + diff2*diff2
    ...:     return mse / N

产生

In [48]: %timeit score4(0.5, pc1, pc2)
1.05 µs ± 15.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

最后但并非最不重要的一点是，您可以写下问题的一阶必要条件，并检查它是否可以解析解决。否则，您可以尝试对生成的非线性方程进行数值求解。