Rigid-motion least-squares algnment of two sets of corresponding points.

rigid_motion_alignment.py

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.utils import check_array
from sklearn.utils.validation import check_is_fitted, column_or_1d
import numpy as np


def rigid_motion_alignment(P, Q, sample_weight=None):
    if sample_weight is None:
        sample_weight = np.ones(P.shape[1])
    p_centroid = np.sum(sample_weight * P, axis=1) / np.sum(sample_weight)
    q_centroid = np.sum(sample_weight * Q, axis=1) / np.sum(sample_weight)
    X = P - p_centroid.reshape(-1, 1)
    Y = Q - q_centroid.reshape(-1, 1)
    S = np.dot(X * sample_weight, Y.T)
    U, s, V = np.linalg.svd(S, full_matrices=False)
    V = V.T  # np.linalg.svd returns vh, which is V.T
    s = np.ones(s.shape)
    s[-1] = np.linalg.det(np.dot(V, U.T))
    R = np.dot(V * s, U.T)
    t = q_centroid - np.dot(R, p_centroid)
    return R, t


class RigidMotionAlignment(BaseEstimator, TransformerMixin):
    """Rigid motion alignment of two matrices
    Calculate and apply a rigid transformation that optimally aligns two sets of corresponding points in the least
    squares sense (https://igl.ethz.ch/projects/ARAP/svd_rot.pdf)

    Attributes
    ----------
    R_: The rotation matrix
    t_: The translation matrix
    """
    def __init__(self):
        self.R_ = None
        self.t_ = None

    def fit(self, X, Y, sample_weight=None):
        """Find the rotation and translation matrices.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
            The data to be aligned.
        Y : array-like, shape = [n_samples, n_features]
            Data to which X should be aligned.
        """
        X = check_array(X)
        Y = check_array(Y)
        if X.shape != Y.shape:
            raise ValueError('X and y must be sets of corresponding points with the same dimensions.')
        if sample_weight is not None:
            sample_weight = column_or_1d(sample_weight)
        self.R_, self.t_ = rigid_motion_alignment(np.transpose(X), np.transpose(Y), sample_weight)
        return self

    def transform(self, X):
        """Rotate X using the rotation and translation matrices found by fit()

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
            The data to be aligned
        
        Returns
        -------
        X_rot, a rigidly transformed X which is closest to the Y from fit() in the least-squares sense.
        """
        X = check_array(X)
        check_is_fitted(self, 'R_')
        return np.dot(self.R_, np.transpose(X)).T + self.t_

See this paper explaining the least-squares approach to rigid motion alignment.