python - 在 numpy 或 scipy 中有效计算 3d 旋转矩阵列表

Question

我有一个 N 单位归一化 3D 向量 p 的列表，存储在形状为 (N, 3) 的 numpy ndarray 中。我还有另一个这样的列表，q。我想计算一个形状为 (N, 3, 3) 的 ndarray U ，存储将 p 中的每个点旋转到相应点 q 的旋转矩阵强>。

旋转矩阵U的列表应满足:

np.all(np.einsum('ijk,ik->ij', U, p) ==  q)

在逐点的基础上，问题简化为能够计算绕某个轴旋转某个角度的旋转矩阵。解决单点情况的代码如下:

def rotation_matrix(angle, direction):
    direction = np.atleast_1d(direction).astype('f4')
    sina = np.sin(angle)
    cosa = np.cos(angle)
    direction = direction/np.sqrt(np.sum(direction*direction))

    R = np.diag([cosa, cosa, cosa])
    R += np.outer(direction, direction) * (1.0 - cosa)
    direction *= sina
    R += np.array(((0.0, -direction[2], direction[1]),
                (direction[2], 0.0, -direction[0]),
                (-direction[1], direction[0],  0.0)))
    return R

我需要的是一个与上述函数完全相同的函数，但它不接受单个角度和单个方向，而是接受形状 (npts, ) 的 angles 数组和 directions 形状数组 (npts, 3)。下面的代码仅部分完成 - 问题是 np.diag 和 np.outer 都不接受 axis 参数

def rotation_matrices(angles, directions):
    directions = np.atleast_2d(directions)
    angles = np.atleast_1d(angles)
    npts = directions.shape[0]
    directions = directions/np.sqrt(np.sum(directions*directions, axis=1)).reshape((npts, 1))

    sina = np.sin(angles)
    cosa = np.cos(angles)

    #  Lines below require extension to 2d case - np.diag and np.outer do not support axis arguments
    R = np.diag([cosa, cosa, cosa])
    R += np.outer(directions, directions) * (1.0 - cosa)
    directions *= sina
    R += np.array(((0.0, -directions[2], directions[1]),
                (directions[2], 0.0, -directions[0]),
                (-directions[1], directions[0],  0.0)))
    return R

numpy 或 scipy 是否有一个紧凑的向量化函数，可以避免使用 for 循环来计算适当的旋转矩阵？问题是 np.diag 和 np.outer 都不接受 axis 作为参数。我的应用程序的 N 非常大，为 1e7 或更大，因此出于性能原因，需要保持所有相关轴对齐的矢量化函数。

Answer 1

暂时先放在这里，稍后会解释。使用 @jaime 的回答 here 中的 levi-cevita 符号以及罗德里格斯公式的矩阵形式here以及一些基于k = (a x b)/sin(theta)

的代数

def rotmatx(p, q):
    eijk = np.zeros((3, 3, 3))
    eijk[0, 1, 2] = eijk[1, 2, 0] = eijk[2, 0, 1] = 1
    eijk[0, 2, 1] = eijk[2, 1, 0] = eijk[1, 0, 2] = -1
    d = (p * q).sum(-1)[:, None, None]
    c = (p.dot(eijk) @ q[..., None]).squeeze()   # cross product (optimized)
    cx = c.dot(eijk)
    return np.eye(3) + cx + cx @ cx / (1 + d)

编辑:该死。问题改变了。

def rotation_matrices(angles, directions):
    eijk = np.zeros((3, 3, 3))
    eijk[0, 1, 2] = eijk[1, 2, 0] = eijk[2, 0, 1] = 1
    eijk[0, 2, 1] = eijk[2, 1, 0] = eijk[1, 0, 2] = -1
    theta = angles[:, None, None]
    K = directions.dot(eijk)
    return np.eye(3) + K * np.sin(theta) + K @ K * (1 - np.cos(theta))