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

时间:2017-12-03 21:37:48

标签: python numpy matrix scipy scientific-computing

我有一个N单位标准化的3D矢量 p 的列表,存储在具有形状(N,3)的numpy ndarray中。我有另一个这样的清单, q 。我想计算一个ndarray U 的形状(N,3,3)存储旋转矩阵,将 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

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

axis

numpy或scipy是否有一个紧凑的向量函数以一种避免使用for循环的方式计算适当的旋转矩阵?问题是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 np.diag都不接受np.outer作为参数。我的应用程序将具有N非常大,1e7或更大,因此出于性能原因需要保持所有相关轴对齐的矢量化函数。

2 个答案:

答案 0 :(得分:1)

暂时放在这里,稍后会解释。使用@ jaime回答here中的levi-cevita符号和Rodrigues公式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)
编辑:dang。问题改变了。

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))

答案 1 :(得分:1)

为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倍。