Quaternion.py

# construct a random SO(3) matrix

import numpy as np

# We construct a random element in SO(3)
def rand_so3():
    A = np.random.randn(3,3)
    # normalize the first column
    A[:,0]=A[:,0]/np.linalg.norm(A[:,0], 2)
    # make the 2nd column orthogonal to first column
    A[:,1]=A[:,1] - np.dot(A[:,0], A[:,1])*A[:,0]
    # normalize the second column
    A[:,1]=A[:,1]/np.linalg.norm(A[:,1], 2)
    # The third column is just the cross product of the first two columns => det = 1
    A[:,2]=np.cross(A[:,0],A[:,1])
    return A


def fromSO3_fast(A):
    q_r = 0.5*np.sqrt(1 + np.trace(A))
    q_i = (A[2,1]- A[1,2])/(4*q_r)
    q_j = (A[0,2]- A[2,0])/(4*q_r)
    q_k = (A[1,0]- A[0,1])/(4*q_r)
    return Quaternion(np.array([q_r, q_i, q_j, q_k]))

    
class Quaternion(object):
    def __init__(self, q):
        assert len(q) == 4
        self.__q = q
    
    @property
    def conjugate(self):
        return Quaternion(np.array([self.__q[0], -self.__q[1], -self.__q[2], -self.__q[3]]))
    
    @property
    def versor(self):
        return Quaternion(self.__q/self.norm)
    
    @property
    def tuple(self):
        return tuple(self.__q)
    
    @property
    def norm(self):
        return np.linalg.norm(self.__q, 2)
    
    def __repr__(self):
        return "Q{0}".format(self.__q)
    
    def __mul__(self, other):
        w1, x1, y1, z1 = self.tuple
        w2, x2, y2, z2 = other.tuple
        
        w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
        x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
        y = w1 * y2 + y1 * w2 + z1 * x2 - x1 * z2
        z = w1 * z2 + z1 * w2 + x1 * y2 - y1 * x2
        return Quaternion(np.array([w, x, y, z]))
    
    def rotate(self, x):
        # construct the pure quaternion
        z = Quaternion(np.array([0, x[0], x[1], x[2]]))
        return (self.versor*z*self.conjugate.versor).vector
    
    @property
    def so3(self):    
        return np.apply_along_axis(self.rotate, 0, np.eye(3))
        
    @property
    def vector(self):
        return self.__q[1:]
    
    @property
    def real(self):
        return self.__q[0]