import numpy as np class Quaternions: """ Quaternions is a wrapper around a numpy ndarray that allows it to act as if it were an narray of a quater data type. Therefore addition, subtraction, multiplication, division, negation, absolute, are all defined in terms of quater operations such as quater multiplication. This allows for much neater code and many routines which conceptually do the same thing to be written in the same way for point data and for rotation data. The Quaternions class has been desgined such that it should support broadcasting and slicing in all of the usual ways. """ def __init__(self, qs): if isinstance(qs, np.ndarray): if len(qs.shape) == 1: qs = np.array([qs]) self.qs = qs return if isinstance(qs, Quaternions): self.qs = qs return raise TypeError('Quaternions must be constructed from iterable, numpy array, or Quaternions, not %s' % type(qs)) def __str__(self): return "Quaternions(" + str(self.qs) + ")" def __repr__(self): return "Quaternions(" + repr(self.qs) + ")" """ Helper Methods for Broadcasting and Data extraction """ @classmethod def _broadcast(cls, sqs, oqs, scalar=False): if isinstance(oqs, float): return sqs, oqs * np.ones(sqs.shape[:-1]) ss = np.array(sqs.shape) if not scalar else np.array(sqs.shape[:-1]) os = np.array(oqs.shape) if len(ss) != len(os): raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape)) if np.all(ss == os): return sqs, oqs if not np.all((ss == os) | (os == np.ones(len(os))) | (ss == np.ones(len(ss)))): raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape)) sqsn, oqsn = sqs.copy(), oqs.copy() for a in np.where(ss == 1)[0]: sqsn = sqsn.repeat(os[a], axis=a) for a in np.where(os == 1)[0]: oqsn = oqsn.repeat(ss[a], axis=a) return sqsn, oqsn """ Adding Quaterions is just Defined as Multiplication """ def __add__(self, other): return self * other def __sub__(self, other): return self / other """ Quaterion Multiplication """ def __mul__(self, other): """ Quaternion multiplication has three main methods. When multiplying a Quaternions array by Quaternions normal quater multiplication is performed. When multiplying a Quaternions array by a vector array of the same shape, where the last axis is 3, it is assumed to be a Quaternion by 3D-Vector multiplication and the 3D-Vectors are rotated in space by the Quaternions. When multipplying a Quaternions array by a scalar or vector of different shape it is assumed to be a Quaternions by Scalars multiplication and the Quaternions are scaled using Slerp and the identity quaternions. """ """ If Quaternions type do Quaternions * Quaternions """ if isinstance(other, Quaternions): sqs, oqs = Quaternions._broadcast(self.qs, other.qs) q0 = sqs[..., 0]; q1 = sqs[..., 1]; q2 = sqs[..., 2]; q3 = sqs[..., 3]; r0 = oqs[..., 0]; r1 = oqs[..., 1]; r2 = oqs[..., 2]; r3 = oqs[..., 3]; qs = np.empty(sqs.shape) qs[..., 0] = r0 * q0 - r1 * q1 - r2 * q2 - r3 * q3 qs[..., 1] = r0 * q1 + r1 * q0 - r2 * q3 + r3 * q2 qs[..., 2] = r0 * q2 + r1 * q3 + r2 * q0 - r3 * q1 qs[..., 3] = r0 * q3 - r1 * q2 + r2 * q1 + r3 * q0 return Quaternions(qs) """ If array type do Quaternions * Vectors """ if isinstance(other, np.ndarray) and other.shape[-1] == 3: vs = Quaternions(np.concatenate([np.zeros(other.shape[:-1] + (1,)), other], axis=-1)) return (self * (vs * -self)).imaginaries """ If float do Quaternions * Scalars """ if isinstance(other, np.ndarray) or isinstance(other, float): return Quaternions.slerp(Quaternions.id_like(self), self, other) raise TypeError('Cannot multiply/add Quaternions with type %s' % str(type(other))) def __div__(self, other): """ When a Quaternion type is supplied, division is defined as multiplication by the inverse of that Quaternion. When a scalar or vector is supplied it is defined as multiplicaion of one over the supplied value. Essentially a scaling. """ if isinstance(other, Quaternions): return self * (-other) if isinstance(other, np.ndarray): return self * (1.0 / other) if isinstance(other, float): return self * (1.0 / other) raise TypeError('Cannot divide/subtract Quaternions with type %s' + str(type(other))) def __eq__(self, other): return self.qs == other.qs def __ne__(self, other): return self.qs != other.qs def __neg__(self): """ Invert Quaternions """ return Quaternions(self.qs * np.array([[1, -1, -1, -1]])) def __abs__(self): """ Unify Quaternions To Single Pole """ qabs = self.normalized().copy() top = np.sum((qabs.qs) * np.array([1, 0, 0, 0]), axis=-1) bot = np.sum((-qabs.qs) * np.array([1, 0, 0, 0]), axis=-1) qabs.qs[top < bot] = -qabs.qs[top < bot] return qabs def __iter__(self): return iter(self.qs) def __len__(self): return len(self.qs) def __getitem__(self, k): return Quaternions(self.qs[k]) def __setitem__(self, k, v): self.qs[k] = v.qs @property def lengths(self): return np.sum(self.qs ** 2.0, axis=-1) ** 0.5 @property def reals(self): return self.qs[..., 0] @property def imaginaries(self): return self.qs[..., 1:4] @property def shape(self): return self.qs.shape[:-1] def repeat(self, n, **kwargs): return Quaternions(self.qs.repeat(n, **kwargs)) def normalized(self): return Quaternions(self.qs / self.lengths[..., np.newaxis]) def log(self): norm = abs(self.normalized()) imgs = norm.imaginaries lens = np.sqrt(np.sum(imgs ** 2, axis=-1)) lens = np.arctan2(lens, norm.reals) / (lens + 1e-10) return imgs * lens[..., np.newaxis] def constrained(self, axis): rl = self.reals im = np.sum(axis * self.imaginaries, axis=-1) t1 = -2 * np.arctan2(rl, im) + np.pi t2 = -2 * np.arctan2(rl, im) - np.pi top = Quaternions.exp(axis[np.newaxis] * (t1[:, np.newaxis] / 2.0)) bot = Quaternions.exp(axis[np.newaxis] * (t2[:, np.newaxis] / 2.0)) img = self.dot(top) > self.dot(bot) ret = top.copy() ret[img] = top[img] ret[~img] = bot[~img] return ret def constrained_x(self): return self.constrained(np.array([1, 0, 0])) def constrained_y(self): return self.constrained(np.array([0, 1, 0])) def constrained_z(self): return self.constrained(np.array([0, 0, 1])) def dot(self, q): return np.sum(self.qs * q.qs, axis=-1) def copy(self): return Quaternions(np.copy(self.qs)) def reshape(self, s): self.qs.reshape(s) return self def interpolate(self, ws): return Quaternions.exp(np.average(abs(self).log, axis=0, weights=ws)) def euler(self, order='xyz'): # fix the wrong convert, this should convert to world euler by default. q = self.normalized().qs q0 = q[..., 0] q1 = q[..., 1] q2 = q[..., 2] q3 = q[..., 3] es = np.zeros(self.shape + (3,)) if order == 'xyz': es[..., 0] = np.arctan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2)) es[..., 1] = np.arcsin((2 * (q0 * q2 - q3 * q1)).clip(-1, 1)) es[..., 2] = np.arctan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3)) elif order == 'yzx': es[..., 0] = np.arctan2(2 * (q1 * q0 - q2 * q3), -q1 * q1 + q2 * q2 - q3 * q3 + q0 * q0) es[..., 1] = np.arctan2(2 * (q2 * q0 - q1 * q3), q1 * q1 - q2 * q2 - q3 * q3 + q0 * q0) es[..., 2] = np.arcsin((2 * (q1 * q2 + q3 * q0)).clip(-1, 1)) else: raise NotImplementedError('Cannot convert from ordering %s' % order) """ # These conversion don't appear to work correctly for Maya. # http://bediyap.com/programming/convert-quaternion-to-euler-rotations/ if order == 'xyz': es[fa + (0,)] = np.arctan2(2 * (q0 * q3 - q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3) es[fa + (1,)] = np.arcsin((2 * (q1 * q3 + q0 * q2)).clip(-1,1)) es[fa + (2,)] = np.arctan2(2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3) elif order == 'yzx': es[fa + (0,)] = np.arctan2(2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3) es[fa + (1,)] = np.arcsin((2 * (q1 * q2 + q0 * q3)).clip(-1,1)) es[fa + (2,)] = np.arctan2(2 * (q0 * q2 - q1 * q3), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3) elif order == 'zxy': es[fa + (0,)] = np.arctan2(2 * (q0 * q2 - q1 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3) es[fa + (1,)] = np.arcsin((2 * (q0 * q1 + q2 * q3)).clip(-1,1)) es[fa + (2,)] = np.arctan2(2 * (q0 * q3 - q1 * q2), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3) elif order == 'xzy': es[fa + (0,)] = np.arctan2(2 * (q0 * q2 + q1 * q3), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3) es[fa + (1,)] = np.arcsin((2 * (q0 * q3 - q1 * q2)).clip(-1,1)) es[fa + (2,)] = np.arctan2(2 * (q0 * q1 + q2 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3) elif order == 'yxz': es[fa + (0,)] = np.arctan2(2 * (q1 * q2 + q0 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3) es[fa + (1,)] = np.arcsin((2 * (q0 * q1 - q2 * q3)).clip(-1,1)) es[fa + (2,)] = np.arctan2(2 * (q1 * q3 + q0 * q2), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3) elif order == 'zyx': es[fa + (0,)] = np.arctan2(2 * (q0 * q1 + q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3) es[fa + (1,)] = np.arcsin((2 * (q0 * q2 - q1 * q3)).clip(-1,1)) es[fa + (2,)] = np.arctan2(2 * (q0 * q3 + q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3) else: raise KeyError('Unknown ordering %s' % order) """ # https://github.com/ehsan/ogre/blob/master/OgreMain/src/OgreMatrix3.cpp # Use this class and convert from matrix return es def average(self): if len(self.shape) == 1: import numpy.core.umath_tests as ut system = ut.matrix_multiply(self.qs[:, :, np.newaxis], self.qs[:, np.newaxis, :]).sum(axis=0) w, v = np.linalg.eigh(system) qiT_dot_qref = (self.qs[:, :, np.newaxis] * v[np.newaxis, :, :]).sum(axis=1) return Quaternions(v[:, np.argmin((1. - qiT_dot_qref ** 2).sum(axis=0))]) else: raise NotImplementedError('Cannot average multi-dimensionsal Quaternions') def angle_axis(self): norm = self.normalized() s = np.sqrt(1 - (norm.reals ** 2.0)) s[s == 0] = 0.001 angles = 2.0 * np.arccos(norm.reals) axis = norm.imaginaries / s[..., np.newaxis] return angles, axis def transforms(self): qw = self.qs[..., 0] qx = self.qs[..., 1] qy = self.qs[..., 2] qz = self.qs[..., 3] x2 = qx + qx; y2 = qy + qy; z2 = qz + qz; xx = qx * x2; yy = qy * y2; wx = qw * x2; xy = qx * y2; yz = qy * z2; wy = qw * y2; xz = qx * z2; zz = qz * z2; wz = qw * z2; m = np.empty(self.shape + (3, 3)) m[..., 0, 0] = 1.0 - (yy + zz) m[..., 0, 1] = xy - wz m[..., 0, 2] = xz + wy m[..., 1, 0] = xy + wz m[..., 1, 1] = 1.0 - (xx + zz) m[..., 1, 2] = yz - wx m[..., 2, 0] = xz - wy m[..., 2, 1] = yz + wx m[..., 2, 2] = 1.0 - (xx + yy) return m def ravel(self): return self.qs.ravel() @classmethod def id(cls, n): if isinstance(n, tuple): qs = np.zeros(n + (4,)) qs[..., 0] = 1.0 return Quaternions(qs) if isinstance(n, int): qs = np.zeros((n, 4)) qs[:, 0] = 1.0 return Quaternions(qs) raise TypeError('Cannot Construct Quaternion from %s type' % str(type(n))) @classmethod def id_like(cls, a): qs = np.zeros(a.shape + (4,)) qs[..., 0] = 1.0 return Quaternions(qs) @classmethod def exp(cls, ws): ts = np.sum(ws ** 2.0, axis=-1) ** 0.5 ts[ts == 0] = 0.001 ls = np.sin(ts) / ts qs = np.empty(ws.shape[:-1] + (4,)) qs[..., 0] = np.cos(ts) qs[..., 1] = ws[..., 0] * ls qs[..., 2] = ws[..., 1] * ls qs[..., 3] = ws[..., 2] * ls return Quaternions(qs).normalized() @classmethod def slerp(cls, q0s, q1s, a): fst, snd = cls._broadcast(q0s.qs, q1s.qs) fst, a = cls._broadcast(fst, a, scalar=True) snd, a = cls._broadcast(snd, a, scalar=True) len = np.sum(fst * snd, axis=-1) neg = len < 0.0 len[neg] = -len[neg] snd[neg] = -snd[neg] amount0 = np.zeros(a.shape) amount1 = np.zeros(a.shape) linear = (1.0 - len) < 0.01 omegas = np.arccos(len[~linear]) sinoms = np.sin(omegas) amount0[linear] = 1.0 - a[linear] amount1[linear] = a[linear] amount0[~linear] = np.sin((1.0 - a[~linear]) * omegas) / sinoms amount1[~linear] = np.sin(a[~linear] * omegas) / sinoms return Quaternions( amount0[..., np.newaxis] * fst + amount1[..., np.newaxis] * snd) @classmethod def between(cls, v0s, v1s): a = np.cross(v0s, v1s) w = np.sqrt((v0s ** 2).sum(axis=-1) * (v1s ** 2).sum(axis=-1)) + (v0s * v1s).sum(axis=-1) return Quaternions(np.concatenate([w[..., np.newaxis], a], axis=-1)).normalized() @classmethod def from_angle_axis(cls, angles, axis): axis = axis / (np.sqrt(np.sum(axis ** 2, axis=-1)) + 1e-10)[..., np.newaxis] sines = np.sin(angles / 2.0)[..., np.newaxis] cosines = np.cos(angles / 2.0)[..., np.newaxis] return Quaternions(np.concatenate([cosines, axis * sines], axis=-1)) @classmethod def from_euler(cls, es, order='xyz', world=False): axis = { 'x': np.array([1, 0, 0]), 'y': np.array([0, 1, 0]), 'z': np.array([0, 0, 1]), } q0s = Quaternions.from_angle_axis(es[..., 0], axis[order[0]]) q1s = Quaternions.from_angle_axis(es[..., 1], axis[order[1]]) q2s = Quaternions.from_angle_axis(es[..., 2], axis[order[2]]) return (q2s * (q1s * q0s)) if world else (q0s * (q1s * q2s)) @classmethod def from_transforms(cls, ts): d0, d1, d2 = ts[..., 0, 0], ts[..., 1, 1], ts[..., 2, 2] q0 = (d0 + d1 + d2 + 1.0) / 4.0 q1 = (d0 - d1 - d2 + 1.0) / 4.0 q2 = (-d0 + d1 - d2 + 1.0) / 4.0 q3 = (-d0 - d1 + d2 + 1.0) / 4.0 q0 = np.sqrt(q0.clip(0, None)) q1 = np.sqrt(q1.clip(0, None)) q2 = np.sqrt(q2.clip(0, None)) q3 = np.sqrt(q3.clip(0, None)) c0 = (q0 >= q1) & (q0 >= q2) & (q0 >= q3) c1 = (q1 >= q0) & (q1 >= q2) & (q1 >= q3) c2 = (q2 >= q0) & (q2 >= q1) & (q2 >= q3) c3 = (q3 >= q0) & (q3 >= q1) & (q3 >= q2) q1[c0] *= np.sign(ts[c0, 2, 1] - ts[c0, 1, 2]) q2[c0] *= np.sign(ts[c0, 0, 2] - ts[c0, 2, 0]) q3[c0] *= np.sign(ts[c0, 1, 0] - ts[c0, 0, 1]) q0[c1] *= np.sign(ts[c1, 2, 1] - ts[c1, 1, 2]) q2[c1] *= np.sign(ts[c1, 1, 0] + ts[c1, 0, 1]) q3[c1] *= np.sign(ts[c1, 0, 2] + ts[c1, 2, 0]) q0[c2] *= np.sign(ts[c2, 0, 2] - ts[c2, 2, 0]) q1[c2] *= np.sign(ts[c2, 1, 0] + ts[c2, 0, 1]) q3[c2] *= np.sign(ts[c2, 2, 1] + ts[c2, 1, 2]) q0[c3] *= np.sign(ts[c3, 1, 0] - ts[c3, 0, 1]) q1[c3] *= np.sign(ts[c3, 2, 0] + ts[c3, 0, 2]) q2[c3] *= np.sign(ts[c3, 2, 1] + ts[c3, 1, 2]) qs = np.empty(ts.shape[:-2] + (4,)) qs[..., 0] = q0 qs[..., 1] = q1 qs[..., 2] = q2 qs[..., 3] = q3 return cls(qs)