-
Notifications
You must be signed in to change notification settings - Fork 0
/
quaternion.py
129 lines (109 loc) · 4.74 KB
/
quaternion.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import numpy as np
import numbers
class kuaternion:
"""
A simple class implementing basic quaternion arithmetic.
"""
def __init__(self, w_or_q, x=None, y=None, z=None):
"""
Initializes a Quaternion object
:param w_or_q: A scalar representing the real part of the quaternion, another Quaternion object or a
four-element array containing the quaternion values
:param x: The first imaginary part if w_or_q is a scalar
:param y: The second imaginary part if w_or_q is a scalar
:param z: The third imaginary part if w_or_q is a scalar
"""
self._q = np.array([1, 0, 0, 0])
if x is not None and y is not None and z is not None:
w = w_or_q
q = np.array([w, x, y, z])
elif isinstance(w_or_q, kuaternion):
q = np.array(w_or_q.q)
else:
q = np.array(w_or_q)
if len(q) != 4:
raise ValueError("Expecting a 4-element array or w x y z as parameters")
self._set_q(q)
# Quaternion specific interfaces
def conj(self):
"""
Returns the conjugate of the quaternion
:rtype : Quaternion
:return: the conjugate of the quaternion
"""
return kuaternion(self._q[0], -self._q[1], -self._q[2], -self._q[3])
def to_angle_axis(self):
"""
Returns the quaternion's rotation represented by an Euler angle and axis.
If the quaternion is the identity quaternion (1, 0, 0, 0), a rotation along the x axis with angle 0 is returned.
:return: rad, x, y, z
"""
if self[0] == 1 and self[1] == 0 and self[2] == 0 and self[3] == 0:
return 0, 1, 0, 0
rad = np.arccos(self[0]) * 2
imaginary_factor = np.sin(rad / 2)
if abs(imaginary_factor) < 1e-8:
return 0, 1, 0, 0
x = self._q[1] / imaginary_factor
y = self._q[2] / imaginary_factor
z = self._q[3] / imaginary_factor
return rad, x, y, z
@staticmethod
def from_angle_axis(rad, x, y, z):
s = np.sin(rad / 2)
return kuaternion(np.cos(rad / 2), x*s, y*s, z*s)
def to_euler_angles(self):
pitch = np.arcsin(2 * self[1] * self[2] + 2 * self[0] * self[3])
if np.abs(self[1] * self[2] + self[3] * self[0] - 0.5) < 1e-8:
roll = 0
yaw = 2 * np.arctan2(self[1], self[0])
elif np.abs(self[1] * self[2] + self[3] * self[0] + 0.5) < 1e-8:
roll = -2 * np.arctan2(self[1], self[0])
yaw = 0
else:
roll = np.arctan2(2 * self[0] * self[1] - 2 * self[2] * self[3], 1 - 2 * self[1] ** 2 - 2 * self[3] ** 2)
yaw = np.arctan2(2 * self[0] * self[2] - 2 * self[1] * self[3], 1 - 2 * self[2] ** 2 - 2 * self[3] ** 2)
return roll, pitch, yaw
def to_euler123(self):
roll = np.arctan2(-2*(self[2]*self[3] - self[0]*self[1]), self[0]**2 - self[1]**2 - self[2]**2 + self[3]**2)
pitch = np.arcsin(2*(self[1]*self[3] + self[0]*self[1]))
yaw = np.arctan2(-2*(self[1]*self[2] - self[0]*self[3]), self[0]**2 + self[1]**2 - self[2]**2 - self[3]**2)
return roll, pitch, yaw
def __mul__(self, other):
"""
multiply the given quaternion with another quaternion or a scalar
:param other: a Quaternion object or a number
:return:
"""
if isinstance(other, kuaternion):
w = self._q[0]*other._q[0] - self._q[1]*other._q[1] - self._q[2]*other._q[2] - self._q[3]*other._q[3]
x = self._q[0]*other._q[1] + self._q[1]*other._q[0] + self._q[2]*other._q[3] - self._q[3]*other._q[2]
y = self._q[0]*other._q[2] - self._q[1]*other._q[3] + self._q[2]*other._q[0] + self._q[3]*other._q[1]
z = self._q[0]*other._q[3] + self._q[1]*other._q[2] - self._q[2]*other._q[1] + self._q[3]*other._q[0]
return kuaternion(w, x, y, z)
elif isinstance(other, numbers.Number):
q = self._q * other
return kuaternion(q)
def __add__(self, other):
"""
add two quaternions element-wise or add a scalar to each element of the quaternion
:param other:
:return:
"""
if not isinstance(other, kuaternion):
if len(other) != 4:
raise TypeError("Quaternions must be added to other quaternions or a 4-element array")
q = self.q + other
else:
q = self.q + other.q
return kuaternion(q)
# Implementing other interfaces to ease working with the class
def _set_q(self, q):
self._q = q
def _get_q(self):
return self._q
q = property(_get_q, _set_q)
def __getitem__(self, item):
return self._q[item]
def __array__(self):
return self._q