quat.py

import numpy as np

def _fast_cross(a, b):
    o = np.empty(np.broadcast(a, b).shape)
    o[...,0] = a[...,1]*b[...,2] - a[...,2]*b[...,1]
    o[...,1] = a[...,2]*b[...,0] - a[...,0]*b[...,2]
    o[...,2] = a[...,0]*b[...,1] - a[...,1]*b[...,0]
    return o

def eye(shape, dtype=np.float32):
    return np.ones(list(shape) + [4], dtype=dtype) * np.asarray([1, 0, 0, 0], dtype=dtype)

def length(x):
    return np.sqrt(np.sum(x * x, axis=-1))

def normalize(x, eps=1e-8):
    return x / (length(x)[...,np.newaxis] + eps)

def abs(x):
    return np.where(x[...,0:1] > 0.0, x, -x)

def from_angle_axis(angle, axis):
    c = np.cos(angle / 2.0)[..., np.newaxis]
    s = np.sin(angle / 2.0)[..., np.newaxis]
    q = np.concatenate([c, s * axis], axis=-1)
    return q

def to_xform(x):
    
    out = np.empty(list(x.shape[:-1]) + [3, 3])
    out[...,0,0] = 1.0 - (x[...,2] * 2 * x[...,2] + x[...,3] * 2 * x[...,3])
    out[...,0,1] = x[...,1] * 2 * x[...,2] - x[...,0] * 2 * x[...,3]
    out[...,0,2] = x[...,1] * 2 * x[...,3] + x[...,0] * 2 * x[...,2]
    out[...,1,0] = x[...,1] * 2 * x[...,2] + x[...,0] * 2 * x[...,3]
    out[...,1,1] = 1.0 - (x[...,1] * 2 * x[...,1] + x[...,3] * 2 * x[...,3])
    out[...,1,2] = x[...,2] * 2 * x[...,3] - x[...,0] * 2 * x[...,1]
    out[...,2,0] = x[...,1] * 2 * x[...,3] - x[...,0] * 2 * x[...,2]
    out[...,2,1] = x[...,2] * 2 * x[...,3] + x[...,0] * 2 * x[...,1]
    out[...,2,2] = 1.0 - (x[...,1] * 2 * x[...,1] + x[...,2] * 2 * x[...,2])
    return out
    
def to_xform_xy(x):

    qw, qx, qy, qz = x[...,0:1], x[...,1:2], x[...,2:3], x[...,3:4]
    
    x2, y2, z2 = qx + qx, qy + qy, qz + qz
    xx, yy, wx = qx * x2, qy * y2, qw * x2
    xy, yz, wy = qx * y2, qy * z2, qw * y2
    xz, zz, wz = qx * z2, qz * z2, qw * z2
    
    return np.concatenate([
        np.concatenate([1.0 - (yy + zz), xy - wz], axis=-1)[...,np.newaxis,:],
        np.concatenate([xy + wz, 1.0 - (xx + zz)], axis=-1)[...,np.newaxis,:],
        np.concatenate([xz - wy, yz + wx], axis=-1)[...,np.newaxis,:],
    ], axis=-2)

def from_euler(e, order='zyx'):
    axis = {
        'x': np.asarray([1, 0, 0], dtype=np.float32),
        'y': np.asarray([0, 1, 0], dtype=np.float32),
        'z': np.asarray([0, 0, 1], dtype=np.float32)}

    q0 = from_angle_axis(e[..., 0], axis[order[0]])
    q1 = from_angle_axis(e[..., 1], axis[order[1]])
    q2 = from_angle_axis(e[..., 2], axis[order[2]])

    return mul(q0, mul(q1, q2))

def from_xform(ts):
    
    return normalize(
        np.where((ts[...,2,2] < 0.0)[...,np.newaxis],
            np.where((ts[...,0,0] >  ts[...,1,1])[...,np.newaxis],
                np.concatenate([
                    (ts[...,2,1]-ts[...,1,2])[...,np.newaxis], 
                    (1.0 + ts[...,0,0] - ts[...,1,1] - ts[...,2,2])[...,np.newaxis], 
                    (ts[...,1,0]+ts[...,0,1])[...,np.newaxis], 
                    (ts[...,0,2]+ts[...,2,0])[...,np.newaxis]], axis=-1),
                np.concatenate([
                    (ts[...,0,2]-ts[...,2,0])[...,np.newaxis], 
                    (ts[...,1,0]+ts[...,0,1])[...,np.newaxis], 
                    (1.0 - ts[...,0,0] + ts[...,1,1] - ts[...,2,2])[...,np.newaxis], 
                    (ts[...,2,1]+ts[...,1,2])[...,np.newaxis]], axis=-1)),
            np.where((ts[...,0,0] < -ts[...,1,1])[...,np.newaxis],
                np.concatenate([
                    (ts[...,1,0]-ts[...,0,1])[...,np.newaxis], 
                    (ts[...,0,2]+ts[...,2,0])[...,np.newaxis], 
                    (ts[...,2,1]+ts[...,1,2])[...,np.newaxis], 
                    (1.0 - ts[...,0,0] - ts[...,1,1] + ts[...,2,2])[...,np.newaxis]], axis=-1),
                np.concatenate([
                    (1.0 + ts[...,0,0] + ts[...,1,1] + ts[...,2,2])[...,np.newaxis], 
                    (ts[...,2,1]-ts[...,1,2])[...,np.newaxis], 
                    (ts[...,0,2]-ts[...,2,0])[...,np.newaxis], 
                    (ts[...,1,0]-ts[...,0,1])[...,np.newaxis]], axis=-1))))

    
def from_xform_xy(x):

    c2 = _fast_cross(x[...,0], x[...,1])
    c2 = c2 / np.sqrt(np.sum(np.square(c2), axis=-1))[...,np.newaxis]
    c1 = _fast_cross(c2, x[...,0])
    c1 = c1 / np.sqrt(np.sum(np.square(c1), axis=-1))[...,np.newaxis]
    c0 = x[...,0]
    
    return from_xform(np.concatenate([
        c0[...,np.newaxis], 
        c1[...,np.newaxis], 
        c2[...,np.newaxis]], axis=-1))

def inv(q):
    return np.asarray([1, -1, -1, -1], dtype=np.float32) * q

def mul(x, y):
    x0, x1, x2, x3 = x[...,0], x[...,1], x[...,2], x[...,3]
    y0, y1, y2, y3 = y[...,0], y[...,1], y[...,2], y[...,3]
    
    o = np.empty(np.broadcast(x, y).shape)
    o[...,0] = y0 * x0 - y1 * x1 - y2 * x2 - y3 * x3
    o[...,1] = y0 * x1 + y1 * x0 - y2 * x3 + y3 * x2
    o[...,2] = y0 * x2 + y1 * x3 + y2 * x0 - y3 * x1
    o[...,3] = y0 * x3 - y1 * x2 + y2 * x1 + y3 * x0
    return o

def inv_mul(x, y):
    return mul(inv(x), y)

def mul_inv(x, y):
    return mul(x, inv(y))

def mul_vec(q, x):
    t = 2.0 * _fast_cross(q[...,1:], x)
    return x + q[...,0][...,None] * t + _fast_cross(q[..., 1:], t)

def inv_mul_vec(q, x):
    return mul_vec(inv(q), x)

def unroll(x):
    y = x.copy()
    for i in range(1, len(x)):
        d0 = np.sum( y[i] * y[i-1], axis=-1)
        d1 = np.sum(-y[i] * y[i-1], axis=-1)
        y[i][d0 < d1] = -y[i][d0 < d1]
    return y

def between(x, y):
    return np.concatenate([
        np.sqrt(np.sum(x*x, axis=-1) * np.sum(y*y, axis=-1))[...,np.newaxis] + 
        np.sum(x * y, axis=-1)[...,np.newaxis], 
        _fast_cross(x, y)], axis=-1)
        
def log(x, eps=1e-5):
    length = np.sqrt(np.sum(np.square(x[...,1:]), axis=-1))[...,np.newaxis]
    halfangle = np.where(length < eps, np.ones_like(length), np.arctan2(length, x[...,0:1]) / length)
    return halfangle * x[...,1:]
    
def exp(x, eps=1e-5):
    halfangle = np.sqrt(np.sum(np.square(x), axis=-1))[...,np.newaxis]
    c = np.where(halfangle < eps, np.ones_like(halfangle), np.cos(halfangle))
    s = np.where(halfangle < eps, np.ones_like(halfangle), np.sinc(halfangle / np.pi))
    return np.concatenate([c, s * x], axis=-1)
    
def to_scaled_angle_axis(x, eps=1e-5):
    return 2.0 * log(x, eps)
    
def from_scaled_angle_axis(x, eps=1e-5):
    return exp(x / 2.0, eps)

def fk(lrot, lpos, parents):
    
    grot, gpos = lrot.copy(), lpos.copy()
    
    for i in range(1, len(parents)):
        p = parents[i]
        gpos[...,i,:] = mul_vec(grot[...,p,:], lpos[...,i,:]) + gpos[...,p,:]
        grot[...,i,:] = mul    (grot[...,p,:], lrot[...,i,:])
    
    return grot, gpos
    
def ik(grot, gpos, parents):
    lrot, lpos = grot.copy(), gpos.copy()
    lrot[...,1:,:] = mul(inv(grot[...,parents[1:],:]), grot[...,1:,:])
    lpos[...,1:,:] = mul_vec(inv(grot[...,parents[1:],:]), gpos[...,1:,:] - gpos[...,parents[1:],:])
    return lrot, lpos
    
def fk_vel(lrot, lpos, lvel, lang, parents):
    
    grot, gpos, gvel, gang = lrot.copy(), lpos.copy(), lvel.copy(), lang.copy()
    
    for i in range(1, len(parents)):
        p = parents[i]
        gpos[...,i,:] = mul_vec(grot[...,p,:], lpos[...,i,:]) + gpos[...,p,:]
        grot[...,i,:] = mul    (grot[...,p,:], lrot[...,i,:])
        gvel[...,i,:] = (mul_vec(grot[...,p,:], lvel[...,i,:]) + 
            _fast_cross(gang[...,p,:], mul_vec(grot[...,p,:], lpos[...,i,:])) +
            gvel[...,p,:])
        gang[...,i,:] = mul_vec(grot[...,p,:], lang[...,i,:]) + gang[...,p,:]
        
    return grot, gpos, gvel, gang
        
        
def to_euler(x, order='xyz'):
    
    q0 = x[...,0:1]
    q1 = x[...,1:2]
    q2 = x[...,2:3]
    q3 = x[...,3:4]
    
    if order == 'xyz':
    
        return np.concatenate([
            np.arctan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2)),
            np.arcsin((2 * (q0 * q2 - q3 * q1)).clip(-1,1)),
            np.arctan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))], axis=-1)
            
    elif order == 'yzx':
    
        return np.concatenate([
            np.arctan2(2 * (q1 * q0 - q2 * q3), -q1 * q1 + q2 * q2 - q3 * q3 + q0 * q0),
            np.arctan2(2 * (q2 * q0 - q1 * q3),  q1 * q1 - q2 * q2 - q3 * q3 + q0 * q0),
            np.arcsin((2 * (q1 * q2 + q3 * q0)).clip(-1,1))], axis=-1)
            
    else:
        raise NotImplementedError('Cannot convert from ordering %s' % order)