Source code for Quaternion

#!@PYTHON@
####!/usr/bin/env python
"""
Class :py:class:`Quaternion` works with quaternion rotations
============================================================

This software was developed in co-operation with Meng for analysis of data cxif5315.

See:
  - :py:class:`Graphics`
  - :py:class:`GlobalGraphics`
  - :py:class:`GlobalUtils`
  - :py:class:`NDArrGenerators`
  - :py:class:`Quaternion`
  - :py:class:`FiberAngles`
  - :py:class:`FiberIndexing`
  - :py:class:`PeakData`
  - :py:class:`PeakStore`
  - :py:class:`TDCheetahPeakRecord`
  - :py:class:`TDFileContainer`
  - :py:class:`TDGroup`
  - :py:class:`TDMatchRecord`
  - :py:class:`TDNodeRecord`
  - :py:class:`TDPeakRecord`
  - :py:class:`HBins`
  - :py:class:`HPolar`
  - :py:class:`HSpectrum`
  - :py:class:`RadialBkgd`
  - `Analysis of data for cxif5315 <https://confluence.slac.stanford.edu/display/PSDMInternal/Analysis+of+data+for+cxif5315>`_.

Created in April 2016 by Mikhail Dubrovin
"""
#------------------------------

import numpy as np
from math import sqrt, sin, cos, radians, degrees, fabs, atan2

ZERO_TOLERANCE = 1e-6

#------------------------------

[docs]def sin_cos(angle_deg) : """Returns sin and cos of angle_deg """ angle_rad = radians(angle_deg) s, c = sin(angle_rad), cos(angle_rad) if fabs(s) < ZERO_TOLERANCE : s = 0 if fabs(c) < ZERO_TOLERANCE : c = 0 return s, c
#------------------------------ #------------------------------ class Quaternion : def __init__(self, w=1, x=0, y=0, z=0) : self.w = w self.x = x self.y = y self.z = z def str_obj(self, cmt='Quaternion w,x,y,z: ', fmt='%7.3f') : pfmt = '%s%s %s %s %s' % (cmt, fmt, fmt, fmt, fmt) return pfmt % (self.w, self.x, self.y, self.z) def print_obj(self, cmt='Quaternion w,x,y,z: ', fmt='%7.3f') : print self.str_obj(cmt, fmt) #------------------------------ class Vector : def __init__(self, x, y, z) : self.u = self.x = x self.v = self.y = y self.w = self.z = z def str_obj(self, cmt='Vector: ', fmt='%6.3f') : pfmt = '%s%s %s %s' % (cmt, fmt, fmt, fmt) return pfmt % (self.u, self.v, self.w) def print_obj(self, cmt='Vector: ', fmt='%6.3f') : print self.str_obj(cmt, fmt) #------------------------------ class Matrix : def __init__(self, m00=1, m01=0, m02=0,\ m10=0, m11=1, m12=0,\ m20=0, m21=0, m22=1) : self.m00, self.m01, self.m02 = m00, m01, m02 self.m10, self.m11, self.m12 = m10, m11, m12 self.m20, self.m21, self.m22 = m20, m21, m22 def str_obj(self, cmt='Matrix:\n', fmt='%6.3f') : pfmt = '%s%s %s %s\n%s %s %s\n%s %s %s'%\ (cmt, fmt, fmt, fmt, fmt, fmt, fmt, fmt, fmt, fmt) return pfmt %\ (self.m00, self.m01, self.m02,\ self.m10, self.m11, self.m12,\ self.m20, self.m21, self.m22) def print_obj(self, cmt='Matrix:\n', fmt='%6.3f') : print self.str_obj(cmt, fmt) def product(self, A, B) : self.m00 = A.m00*B.m00 + A.m01*B.m10 + A.m02*B.m20 self.m01 = A.m00*B.m01 + A.m01*B.m11 + A.m02*B.m21 self.m02 = A.m00*B.m02 + A.m01*B.m12 + A.m02*B.m22 self.m10 = A.m10*B.m00 + A.m11*B.m10 + A.m12*B.m20 self.m11 = A.m10*B.m01 + A.m11*B.m11 + A.m12*B.m21 self.m12 = A.m10*B.m02 + A.m11*B.m12 + A.m12*B.m22 self.m20 = A.m20*B.m00 + A.m21*B.m10 + A.m22*B.m20 self.m21 = A.m20*B.m01 + A.m21*B.m11 + A.m22*B.m21 self.m22 = A.m20*B.m02 + A.m21*B.m12 + A.m22*B.m22 def as_list(self) : return [[self.m00, self.m01, self.m02],\ [self.m10, self.m11, self.m12],\ [self.m20, self.m21, self.m22]] def np_matrix(self, dtype=np.double) : return np.matrix(as_list, dtype, copy=True) def set_from_np_matrix(self, npm) : self.m00, self.m01, self.m02 = npm[0,0], npm[0,1], npm[0,2] self.m10, self.m11, self.m12 = npm[1,0], npm[1,1], npm[1,2] self.m20, self.m21, self.m22 = npm[2,0], npm[2,1], npm[2,2] def rotation_matrix_x(self, angle_deg) : s,c = sin_cos(angle_deg) self.m00, self.m01, self.m02 = 1, 0, 0 self.m10, self.m11, self.m12 = 0, c,-s self.m20, self.m21, self.m22 = 0, s, c def rotation_matrix_y(self, angle_deg) : s,c = sin_cos(angle_deg) self.m00, self.m01, self.m02 = c, 0, s self.m10, self.m11, self.m12 = 0, 1, 0 self.m20, self.m21, self.m22 =-s, 0, c def rotation_matrix_z(self, angle_deg) : s,c = sin_cos(angle_deg) self.m00, self.m01, self.m02 = c,-s, 0 self.m10, self.m11, self.m12 = s, c, 0 self.m20, self.m21, self.m22 = 0, 0, 1 def rotation_matrix(self, alpha, beta, gamma) : rotx = Matrix() roty = Matrix() rotz = Matrix() rotzy= Matrix() rotx.rotation_matrix_x(gamma) roty.rotation_matrix_y(beta) rotz.rotation_matrix_z(alpha) rotzy.product(rotz, roty) #rotx.print_obj('X-rotation matrix:') #roty.print_obj('Y-rotation matrix:') #rotz.print_obj('Z-rotation matrix:') #rotzy.print_obj('ZY-rotation matrix:') self.product(rotzy, rotx) def get_angles(self) : ang_x = atan2(self.m21, self.m22) ang_y = atan2(-self.m20, sqrt(self.m21*self.m21 + self.m22*self.m22)) ang_z = atan2(self.m10, self.m00) return degrees(ang_x), degrees(ang_y), degrees(ang_z) #------------------------------
[docs]def quaternion_modulus(q) : """q - quaternion If a quaternion represents a pure rotation, its modulus should be unity. Returns: the modulus of the given quaternion. """ return sqrt(q.w*q.w + q.x*q.x + q.y*q.y + q.z*q.z)
#------------------------------
[docs]def normalise_quaternion(q) : """q - quaternion Rescales the quaternion such that its modulus is unity. Returns: the normalised version of q """ mod = quaternion_modulus(q) w = q.w / mod x = q.x / mod y = q.y / mod z = q.z / mod return Quaternion(w, x, y, z)
#------------------------------
[docs]def random_quaternion() : """ Returns: a randomly generated, normalised, quaternion """ w, x, y, z = 2.0*np.random.random((4,)) - 1.0 q = Quaternion(w, x, y, z) return normalise_quaternion(q)
#------------------------------
[docs]def quaternion_valid(q, tol=0.001) : """q - quaternion Checks if the given quaternion is normalised. Returns: 1 if the quaternion is normalised, 0 if not. """ qmod = quaternion_modulus(q) if (qmod > 1+tol) or (qmod < 1-tol) : return 0 return 1
#------------------------------
[docs]def quat_rot(v, q) : """v - vector (in the form of a "struct rvec") q - quaternion Rotates a vector according to a quaternion. Returns: rotated vector vrot """ t01 = q.w*q.x t02 = q.w*q.y t03 = q.w*q.z t11 = q.x*q.x t12 = q.x*q.y t13 = q.x*q.z t22 = q.y*q.y t23 = q.y*q.z t33 = q.z*q.z u = (1.0 - 2.0 * (t22 + t33)) * v.u\ + (2.0 * (t12 + t03)) * v.v\ + (2.0 * (t13 - t02)) * v.w v = (2.0 * (t12 - t03)) * v.u\ + (1.0 - 2.0 * (t11 + t33)) * v.v\ + (2.0 * (t01 + t23)) * v.w w = (2.0 * (t02 + t13)) * v.u\ + (2.0 * (t23 - t01)) * v.v\ + (1.0 - 2.0 * (t11 + t22)) * v.w return Vector(u, v, w)
#------------------------------
[docs]def rotmatrix_from_quaternion(q) : """q - quaternion Evaluates rotation matrix from quaternion. Returns: rotation matrix as an object Matrix """ t01 = q.w*q.x t02 = q.w*q.y t03 = q.w*q.z t11 = q.x*q.x t12 = q.x*q.y t13 = q.x*q.z t22 = q.y*q.y t23 = q.y*q.z t33 = q.z*q.z return Matrix(1 - 2*(t22 + t33), 2*(t12 - t03), 2*(t13 + t02),\ 2*(t12 + t03), 1 - 2*(t11 + t33), 2*(t23 - t01),\ 2*(t13 - t02), 2*(t23 + t01), 1 - 2*(t11 + t22))
# Matrix from Thomas White #return Matrix(1 - 2*(t22 + t33), 2*(t12 + t03), 2*(t13 - t02),\ # 2*(t12 - t03), 1 - 2*(t11 + t33), 2*(t01 + t23),\ # 2*(t02 + t13), 2*(t23 - t01), 1 - 2*(t11 + t22)) #------------------------------
[docs]def quaternion_from_rotmatrix(m) : """m - 3-d rotation matrix, class Matrix Evaluates quaternion from rotation matrix. Implemented as https://en.wikipedia.org/wiki/Rotation_matrix Returns: normalised quaternion. """ Qxx, Qxy, Qxz = m.m00, m.m01, m.m02 Qyx, Qyy, Qyz = m.m10, m.m11, m.m12 Qzx, Qzy, Qzz = m.m20, m.m21, m.m22 t = Qxx+Qyy+Qzz r = sqrt(1+t) if fabs(r)<ZERO_TOLERANCE : r = ZERO_TOLERANCE s = 0.5/r w = 0.5*r x = (Qzy-Qyz)*s y = (Qxz-Qzx)*s z = (Qyx-Qxy)*s return Quaternion(w,x,y,z)
#------------------------------
[docs]def quaternion_for_angles(ax, ay, az) : """Returns: quaternion for three input rotation angles [deg]. """ #print 'Angles around x,y,z: %6.1f %6.1f %6.1f' % (ax, ay, az) m = Matrix() m.rotation_matrix(az, ay, ax) return quaternion_from_rotmatrix(m)
#------------------------------
[docs]def record_for_angles(ax, ay, az) : """Prints string like: Angles around x,y,z: 72.0 -3.5 176.4 quaternion: w,x,y,z: 0.007459 0.043148 0.586445 0.808804 """ print 'Angles around x,y,z: %6.1f %6.1f %6.1f' % (ax, ay, az), m = Matrix() m.rotation_matrix(az, ay, ax) q = quaternion_from_rotmatrix(m) q.print_obj(' quaternion:', fmt='%9.6f')
#------------------------------ #------------------------------ #----------- TEST ----------- #------------------------------ #------------------------------ def test_quaternion(tname) : v1 = Vector(1,0,0) v2 = Vector(0,1,0) v3 = Vector(0,0,1) v1.print_obj() v2.print_obj() v3.print_obj() q1 = Quaternion() q1.print_obj() m1 = Matrix() m1.print_obj() #------------------------------ def test_rotation_matrix(tname) : alpha, beta, gamma = 5, 5, 5 # angles degree m = Matrix() m.rotation_matrix(alpha, beta, gamma) m.print_obj('R3-rotation matrix:\n') #------------------------------ def test_quaternion_from_rotation_matrix(tname) : #alpha, beta, gamma = 5, 5, 5 # angles degree ax, ay, az =0, 90, 90 # angles degree vfmt = '%9.6f' print 'Inpurt angles around x,y,z: %.2f %.2f %.2f' % (ax, ay, az) m = Matrix() m.rotation_matrix(az, ay, ax) m.print_obj('R3-rotation matrix:\n', fmt=vfmt) q = quaternion_from_rotmatrix(m) q.print_obj('Associated with matrix quaternion:\n', fmt=vfmt) mq = rotmatrix_from_quaternion(q) mq.print_obj('R3-rotation matrix back from quaternion:\n', fmt=vfmt) axo, ayo, azo = mq.get_angles() print 'Output angles around x,y,z: %.2f %.2f %.2f' % (axo, ayo, azo) #------------------------------ def test_quaternion_table(tname) : ax, ay, az = 0, 0, 0 for ax in range(0, -180, -30) : record_for_angles(ax, ay, az) ax, ay, az = 0, 0, 0 for ay in range(0, 180, 30) : record_for_angles(ax, ay, az) ax, ay, az = 0, 0, 0 for az in range(0, 180, 30) : record_for_angles(ax, ay, az) #------------------------------ def test_quaternion_table_crystal(tname) : #ax, ay, az = 90, 0, 0 ax, ay, az = 90, -10, 0 #ax, ay, az = 90, -3.5, 0 #ax, ay, az = 72, -3.5, 0 #ax, ay, az = 108, 3.5, 0 #ax, ay, az = 102, 0, 0 #ax, ay, az = 90, 10, 0 #ax, ay, az = 108, 3.5, 0 #for az in range(0, 180, 30) : for az in range(0, 180, 1) : record_for_angles(ax, ay, az) #------------------------------ if __name__ == "__main__" : import sys; global sys tname = sys.argv[1] if len(sys.argv) > 1 else '0' print 50*'_', '\nTest %s:' % tname if tname == '0' : test_quaternion(tname) elif tname == '1' : test_rotation_matrix(tname) elif tname == '2' : test_quaternion_from_rotation_matrix(tname) elif tname == '3' : test_quaternion_table(tname) elif tname == '4' : test_quaternion_table_crystal(tname) else : sys.exit('Test %s is undefined' % tname) sys.exit('End of test %s' % tname) #------------------------------