simple library for homogeneous transformation using quaternions

transforms.py

# Standard Library
from math import pi as M_PI
from math import cos, sin, sqrt

class Vector(object):
    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z

    def __repr__(self):
        return "[x: %s, y: %s, z: %s]" % (self.x, self.y, self.z)

    def __iter__(self):
        for i in (self.x, self.y, self.z):
            yield i

    def __getitem__(self, i):
        return tuple(self)[i]

    def __mul__(self, s):
        return self.__class__(self.x*s, self.y*s, self.z*s)

    def __rmul__(self, s):
        return self.__mul__(s)

    def __add__(self, v):
        return self.__class__(
            self.x + v.x,
            self.y + v.y,
            self.z + v.z)

    def __neg__(self):
        return self.__class__(-self.x, -self.y, -self.z)

    def norm2(self):
        return sum(map(lambda i: i**2, self))

    def norm(self):
        return sqrt(self.norm2())

    def normalized(self):
        return 1./ self.norm() * self

class Quaternion(object):
    @classmethod
    def from_axisangle(cls, v, theta):
        vn = v.normalized()
        theta *= 0.5
        return cls(
            w=cos(theta),
            x=vn.x * sin(theta),
            y=vn.y * sin(theta),
            z=vn.z * sin(theta))

    @classmethod
    def from_rpy(cls, r, p, y):
        cr = cos(r * 0.5)
        sr = sin(r * 0.5)
        cp = cos(p * 0.5)
        sp = sin(p * 0.5)
        cy = cos(y * 0.5)
        sy = sin(y * 0.5)

        return cls(
            w=cy * cr * cp + sy * sr * sp,
            x=cy * sr * cp - sy * cr * sp,
            y=cy * cr * sp + sy * sr * cp,
            z=sy * cr * cp - cy * sr * sp)

    @classmethod
    def _conjugate(cls, q):
        return cls(q.w, -q.x, -q.y, -q.z)

    @classmethod
    def _mul(cls, q0, q1):
        return cls(
            w=q0.w*q1.w - q0.x*q1.x  - q0.y*q1.y  - q0.z*q1.z,
            x=q0.w*q1.x + q0.x*q1.w  + q0.y*q1.z  - q0.z*q1.y,
            y=q0.w*q1.y - q0.x*q1.z  + q0.y*q1.w  + q0.z*q1.x,
            z=q0.w*q1.z + q0.x*q1.y  - q0.y*q1.x  + q0.z*q1.w)

    @classmethod
    def _rot(cls, q, v):
        q2 = cls(w=0, x=v.x, y=v.y, z=v.z)
        # q * v * q^-1
        return Vector(*cls._mul(cls._mul(q, q2), cls._conjugate(q))[1:])

    def __init__(self, w, x, y, z):
        self.w = w
        self.x = x
        self.y = y
        self.z = z

    def __repr__(self):
        return "[w: %s, x: %s, y: %s, z: %s]" % (self.w, self.x, self.y, self.z)

    def __iter__(self):
        for i in (self.w, self.x, self.y, self.z):
            yield i

    def __getitem__(self, i):
        return tuple(self)[i]

    def __mul__(self, other):
        if isinstance(other, type(self)):
            return self.__class__._mul(self, other)
        elif isinstance(other, Vector):
            return self.__class__._rot(self, other)
        elif isinstance(other, float):
            return self.__class__(*map(lambda i: i*other, self))

    def __div__(self, s):
        return self.__mul__(1./s)

    def norm2(self):
        return sum(map(lambda i: i**2, self))

    def norm(self):
        return sqrt(self.norm2())

    def conjugate(self):
        return self.__class__._conjugate(self)

    def inv(self):
        return self.__class__._conjugate(self) / self.norm2()

class Transform(object):
    @staticmethod
    def _transform(T, v):
        return T.q*v + T.v

    @classmethod
    def _chain(cls, T1, T2):
        return cls(q=T1.q*T2.q, v=cls._transform(T1, T2.v))

    def __init__(self, q, v):
        assert isinstance(q, Quaternion)
        assert isinstance(v, Vector)
        self.q = q
        self.v = v

    def __mul__(self, other):
        if isinstance(other, type(self)):
            return self.__class__._chain(self, other)
        elif isinstance(other, Vector):
            return self.__class__._transform(self, other)

    def inv(self):
        return self.__class__(
            q=self.q.inv(),
            v=-(self.q.inv()*self.v))


if __name__ == '__main__':
    print("### Testing! ###")

    v1 = Vector(1,2,3)
    v2 = v1 * 2
    v3 = Vector(1,1,1)
    assert v1[0] == 1
    assert v1.x == 1
    assert v2[1] == 4
    assert v2.y == 4
    q1 = Quaternion.from_axisangle(v3, M_PI/4.)
    rpy1 = (-M_PI/2., 0, -M_PI/2)
    rpy2 = (1.5707963267949, 3.2899025714389E-17, 2.57436064669165)

    for i in q1:
        print(i)

    print(list(q1))
    print(tuple(q1))
    print(q1)


    import tf.transformations as tft
    import numpy as np

    fix_order = lambda q: [q[3], q[0], q[1], q[2]]
    print("\n### Quaternion from RPY ###")
    print(Quaternion.from_rpy(*rpy1))
    print(fix_order(tft.quaternion_from_euler(*rpy1)))
    print("")
    print(Quaternion.from_rpy(*rpy2))
    print(fix_order(tft.quaternion_from_euler(*rpy2)))

    print("\n### Quaternion Multiplication ###")
    print(Quaternion.from_rpy(*rpy1)*Quaternion.from_rpy(*rpy2))
    print(fix_order(
        tft.quaternion_from_matrix(
            tft.euler_matrix(*rpy1).dot(tft.euler_matrix(*rpy2)))))

    print("\n### Rotation Transforms ###")
    print(Quaternion.from_rpy(*rpy1)*v1)
    print(tft.euler_matrix(*rpy1).dot(list(v1)+[1.]))

    print("")
    print(Quaternion.from_rpy(*rpy2)*v2)
    print(tft.euler_matrix(*rpy2).dot(list(v2)+[1.]))

    print("\n### Transforms ###")
    print(Transform(Quaternion.from_rpy(*rpy1), v1) * v2)
    print(tft.compose_matrix(angles=rpy1, translate=v1).dot(list(v2)+[1.]))

    print("\n### Transform Chains ###")
    T1 = Transform(Quaternion.from_rpy(*rpy1), v1)
    T2 = Transform(Quaternion.from_rpy(*rpy2), v2)
    print(T1 * T2 * v3)
    T1 = tft.compose_matrix(angles=rpy1, translate=v1)
    T2 = tft.compose_matrix(angles=rpy2, translate=v2)
    print(T1.dot(T2).dot(list(v3)+[1.]))

    print("\n### Transform Inverse ###")
    T1 = Transform(Quaternion.from_rpy(*rpy1), v1)
    T2 = Transform(Quaternion.from_rpy(*rpy2), v2).inv()
    print(T1 * T2 * v3)
    T1 = tft.compose_matrix(angles=rpy1, translate=v1)
    T2 = np.linalg.inv(tft.compose_matrix(angles=rpy2, translate=v2))
    print(T1.dot(T2).dot(list(v3)+[1.]))