Search

Kabsch alignment algorithm in Python

Kabsch algorithm

N×3N\times3 point coordinate matrices PP and QQ가 있을 때, Kabsch 알고리즘은 아래의 3개 step으로 진행된다.
1.
Translation
Move the whole structures so that the centroid is at the origin
2.
Cross-covariance matrix (3x3 matrix)
H=PTQH=P^TQ
# in python, using einsum: H = einsum('ni,nj->ij', P, Q)
Python
복사
3.
Computing rotation matrix RR with singular value decomposition
UΣVT=SVD(H)d=sign(det⁡(VUT))R=V(10001000d)UTU \Sigma V^T=\text{SVD}(H) \\ d = \text{sign}(\det(VU^T)) \\ R = V \begin{pmatrix}1 & 0 & 0 \\ 0 & 1 & 0 \\ 0 & 0 & d\end{pmatrix} U^T

Python implementation

def kabsch(a: torch.Tensor, b: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Compute the optimal rotation matrix and translation vector that minimizes the RMSD between two sets of coordinates `a` and `b`, when applied to `a`. Args: a: A set of 3D coordinates. Shape: n, 3 b: A set of 3D coordinates. Shape: n, 3 Returns: rotation_matrix: Shape: 3, 3 translation_vector: Shape: 3 """ # compute centroids centroid_a = a.mean(dim=-2, keepdim=True) centroid_b = b.mean(dim=-2, keepdim=True) # compute centered coordinates a_centered = a - centroid_a b_centered = b - centroid_b # compute covariance matrix h = torch.einsum("ni,nj->ij", a_centered, b_centered) # compute optimal rotation matrix u, _, vt = torch.linalg.svd(h) v, ut = vt.transpose(-2, -1), u.transpose(-2, -1) d = torch.sign(torch.linalg.det(torch.einsum("ij,jk->ik", v, ut))) diag = torch.eye(3).expand_as(h).clone() diag[2, 2] = d rotation_matrix = torch.einsum("ij,jk,kl->...il", v, diag, ut) # compute optimal translation vector translation_vector = centroid_b.squeeze(-2) - torch.einsum( "ij,j->i", rotation_matrix, centroid_a.squeeze(-2) ) return rotation_matrix, translation_vector
Python
복사

Optimal translation vector의 유도 과정 - 직관

왜 optimal translation vector가 centroid vector의 차 cq−cpc_q - c_p가 아닐까? 직관적인 설명은 아래와 같다.
우리는 rotation → translation 순으로 rigid (Euclidean) transformation이 일어난다고 정의하므로 (참고), 먼저 수행되는 rotation에 의해서 P의 centroid도 원점을 중심으로 함께 회전하게 된다. 따라서 이렇게 이미 회전한 상태의 centroid RoptcpR_{\text{opt}}c_p와 cqc_q의 차이 Roptcp−cqR_{\text{opt}}c_p - c_q 를 optimal translation vector로 계산해야 하는 것이다.

Optimal translation vector의 유도 과정 - 수식

Alignment의 궁극적인 목적은 아래와 같이 최적의 rotation matrix RoptR_{\text{opt}}와 translation vector toptt_{\text{opt}}를 구하는 것으로, 편의상 P의 한 점 xpx_p xpx_pxqx_q를 이용하여 아래와 같이 나타낼 수 있다.
Roptxp+topt≈xq\begin{equation}R_{\text{opt}}x_p+t_{\text{opt}}\approx x_q\end{equation}
이 때, RoptR_{\text{opt}} 는 Kabsch 알고리즘에서 구할 수 있지만, 이 알고리즘이 올바르게 수행되기 위해서는 P와 Q를 먼저 origin-centered point cloud (즉, 각각의 centroid cpc_p, cqc_q가 빼진 벡터들)가 되도록 평행이동시켜야 함에 주의해야 한다. 결과적으로 Kabsch 알고리즘을 통해 구한 RoptR_{\text{opt}} 를 포함하는 관계식을 아래와 같이 나타낼 수 있다.
Ropt(xp−cp)≈(xq−cq)\begin{equation}R_{\text{opt}}(x_p - c_p) \approx (x_q - c_q)\end{equation}
(2)의 식을 (1)의 꼴로 정리하면
Roptxp−Roptcp+cq≈xq\begin{equation}R_{\text{opt}}x_p - R_{\text{opt}}c_p + c_q \approx x_q \end{equation}
이고, 두 식의 비교를 통해
topt=−Roptcp+cq\begin{equation} t_{\text{opt}} = -R_{\text{opt}}c_p + c_q \end{equation}
를 얻을 수 있다.

See also