有效计算 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 and the matrix form of the Rodrigues formula here 中的 levi-cevita 符号和一些基于 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))

Answer 2

正在删除 Nx3x3 矩阵批量旋转的另一个解决方案。其中 3x3 分量表示

中的矢量分量

[[11, 12, 13],
 [21, 22, 23],
 [31, 32, 33]]

现在 np.einsum 的矩阵旋转是：

data = np.random.uniform(size=(500, 3, 3))
rotmat = np.random.uniform(size=(3, 3))

data_rot = np.einsum('ij,...jk,lk->...il', rotmat, data, rotmat)

这相当于

for data_mat in data:
    np.dot(np.dot(rotmat, data_mat), rotmat.T)

np.dot 循环的加速约为 250 倍。

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

efficiently calculate list of 3d rotation matrices in numpy or scipy

python

numpy

matrix

scientific-computing

scipy