OneScience-Group/SurfDock

1import math2 3import torch4 5 6def quaternion_to_matrix(quaternions):7    """8    From https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html9    Convert rotations given as quaternions to rotation matrices.10 11    Args:12        quaternions: quaternions with real part first,13            as tensor of shape (..., 4).14 15    Returns:16        Rotation matrices as tensor of shape (..., 3, 3).17    """18    r, i, j, k = torch.unbind(quaternions, -1)19    two_s = 2.0 / (quaternions * quaternions).sum(-1)20 21    o = torch.stack(22        (23            1 - two_s * (j * j + k * k),24            two_s * (i * j - k * r),25            two_s * (i * k + j * r),26            two_s * (i * j + k * r),27            1 - two_s * (i * i + k * k),28            two_s * (j * k - i * r),29            two_s * (i * k - j * r),30            two_s * (j * k + i * r),31            1 - two_s * (i * i + j * j),32        ),33        -1,34    )35    return o.reshape(quaternions.shape[:-1] + (3, 3))36 37 38def axis_angle_to_quaternion(axis_angle):39    """40    From https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html41    Convert rotations given as axis/angle to quaternions.42 43    Args:44        axis_angle: Rotations given as a vector in axis angle form,45            as a tensor of shape (..., 3), where the magnitude is46            the angle turned anticlockwise in radians around the47            vector's direction.48 49    Returns:50        quaternions with real part first, as tensor of shape (..., 4).51    """52    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)53    half_angles = 0.5 * angles54    eps = 1e-655    small_angles = angles.abs() < eps56    sin_half_angles_over_angles = torch.empty_like(angles)57    sin_half_angles_over_angles[~small_angles] = (58            torch.sin(half_angles[~small_angles]) / angles[~small_angles]59    )60    # for x small, sin(x/2) is about x/2 - (x/2)^3/661    # so sin(x/2)/x is about 1/2 - (x*x)/4862    sin_half_angles_over_angles[small_angles] = (63            0.5 - (angles[small_angles] * angles[small_angles]) / 4864    )65    quaternions = torch.cat(66        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-167    )68    return quaternions69 70 71def axis_angle_to_matrix(axis_angle):72    """73    From https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html74    Convert rotations given as axis/angle to rotation matrices.75 76    Args:77        axis_angle: Rotations given as a vector in axis angle form,78            as a tensor of shape (..., 3), where the magnitude is79            the angle turned anticlockwise in radians around the80            vector's direction.81 82    Returns:83        Rotation matrices as tensor of shape (..., 3, 3).84    """85    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))86 87 88def rigid_transform_Kabsch_3D_torch(A, B):89    # R = 3x3 rotation matrix, t = 3x1 column vector90    # This already takes residue identity into account.91 92    assert A.shape[1] == B.shape[1]93    num_rows, num_cols = A.shape94    if num_rows != 3:95        raise Exception(f"matrix A is not 3xN, it is {num_rows}x{num_cols}")96    num_rows, num_cols = B.shape97    if num_rows != 3:98        raise Exception(f"matrix B is not 3xN, it is {num_rows}x{num_cols}")99 100 101    # find mean column wise: 3 x 1102    centroid_A = torch.mean(A, axis=1, keepdims=True)103    centroid_B = torch.mean(B, axis=1, keepdims=True)104 105    # subtract mean106    Am = A - centroid_A107    Bm = B - centroid_B108 109    H = Am @ Bm.T110 111    # find rotation112    U, S, Vt = torch.linalg.svd(H)113 114    R = Vt.T @ U.T115    # special reflection case116    if torch.linalg.det(R) < 0:117        # print("det(R) < R, reflection detected!, correcting for it ...")118        SS = torch.diag(torch.tensor([1.,1.,-1.], device=A.device))119        R = (Vt.T @ SS) @ U.T120    assert math.fabs(torch.linalg.det(R) - 1) < 3e-3  # note I had to change this error bound to be higher121 122    t = -R @ centroid_A + centroid_B123    return R, t124