James Gregson's Website

Nov 04, 2019

Useful Transformation Functions in Python

These are collections of functions that I end up re-implementing frequently.

from typing import *
import numpy as np

def closest_rotation( M: np.ndarray ):
    U,s,Vt = np.linalg.svd( M[:3,:3] )
    if np.linalg.det(U@Vt) < 0:
        Vt[2] *= -1.0
    return U@Vt

def mat2quat( M: np.ndarray ):
    if np.abs(np.linalg.det(M[:3,:3])-1.0) > 1e-5:
        raise ValueError('Matrix determinant is not 1')
    if np.abs( np.linalg.norm( M[:3,:3].T@M[:3,:3] - np.eye(3)) ) > 1e-5:
        raise ValueError('Matrix is not orthogonal')
    w = np.sqrt( 1.0 + M[0,0]+M[1,1]+M[2,2])/2.0
    x = (M[2,1]-M[1,2])/(4*w)
    y = (M[0,2]-M[2,0])/(4*w)
    z = (M[1,0]-M[0,1])/(4*w)
    return np.array((x,y,z,w),float)

def quat2mat( q: np.ndarray ):
    qx,qy,qz,qw = q/np.linalg.norm(q)
    return np.array((
        (1 - 2*qy*qy - 2*qz*qz,     2*qx*qy - 2*qz*qw,     2*qx*qz + 2*qy*qw),
        (    2*qx*qy + 2*qz*qw, 1 - 2*qx*qx - 2*qz*qz,     2*qy*qz - 2*qx*qw),
        (    2*qx*qz - 2*qy*qw,         2*qy*qz + 2*qx*qw, 1 - 2*qx*qx - 2*qy*qy),
    ))

if __name__ == '__main__':
    for i in range( 1000 ):
        q = np.random.standard_normal(4)
        q /= np.linalg.norm(q)
        A = quat2mat( q )
        s = mat2quat( A )
        if np.linalg.norm( s-q ) > 1e-6 and np.linalg.norm( s+q ) > 1e-6:
            raise ValueError('uh oh.')

        R = closest_rotation( A+np.random.standard_normal((3,3))*1e-6 )
        if np.linalg.norm(A-R) > 1e-5:
            print( np.linalg.norm( A-R ) )
            raise ValueError('uh oh.')