cleanUrl: "kabsch-algorithm-in-python"
description: "Describes how we can implement Kabsch algorithm in Python"

Kabsch algorithm

$N\times3$ point coordinate matrices $P$ and $Q$가 있을 때, Kabsch 알고리즘은 아래의 3개 step으로 진행된다.

1. Translation

   Move the whole structures so that the centroid is at the origin

2. Cross-covariance matrix (3x3 matrix)

   $$ H=P^TQ $$

   # in python, using einsum:
   H = einsum('ni,nj->ij', P, Q)
   

3. Computing rotation matrix $R$ with singular value decomposition

   $$ U \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

Optimal translation vector의 유도 과정 - 직관

왜 optimal translation vector가 centroid vector의 차 $c_q - c_p$가 아닐까? 직관적인 설명은 아래와 같다.

우리는 rotation → translation 순으로 rigid (Euclidean) transformation이 일어난다고 정의하므로 (참고), 먼저 수행되는 rotation에 의해서 P의 centroid도 원점을 중심으로 함께 회전하게 된다. 따라서 이렇게 이미 회전한 상태의 centroid $R_{\text{opt}}c_p$와 $c_q$의 차이 $R_{\text{opt}}c_p - c_q$ 를 optimal translation vector로 계산해야 하는 것이다.

Optimal translation vector의 유도 과정 - 수식

Alignment의 궁극적인 목적은 아래와 같이 최적의 rotation matrix $R_{\text{opt}}$와 translation vector $t_{\text{opt}}$를 구하는 것으로, 편의상 P의 한 점 $x_p$ $x_p$$x_q$를 이용하여 아래와 같이 나타낼 수 있다.

$$ \begin{equation}R_{\text{opt}}x_p+t_{\text{opt}}\approx x_q\end{equation} $$

이 때, $R_{\text{opt}}$ 는 Kabsch 알고리즘에서 구할 수 있지만, 이 알고리즘이 올바르게 수행되기 위해서는 P와 Q를 먼저 origin-centered point cloud (즉, 각각의 centroid $c_p$, $c_q$가 빼진 벡터들)가 되도록 평행이동시켜야 함에 주의해야 한다. 결과적으로 Kabsch 알고리즘을 통해 구한 $R_{\text{opt}}$ 를 포함하는 관계식을 아래와 같이 나타낼 수 있다.

$$ \begin{equation}R_{\text{opt}}(x_p - c_p) \approx (x_q - c_q)\end{equation} $$

(2)의 식을 (1)의 꼴로 정리하면

$$ \begin{equation}R_{\text{opt}}x_p - R_{\text{opt}}c_p + c_q \approx x_q \end{equation} $$

이고, 두 식의 비교를 통해

$$ \begin{equation} t_{\text{opt}} = -R_{\text{opt}}c_p + c_q \end{equation} $$