Coverage for /builds/ase/ase/ase/quaternions.py: 84.00%
125 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3import numpy as np
6class Quaternion:
8 def __init__(self, qin=[1, 0, 0, 0]):
9 assert len(qin) == 4
10 self.q = np.array(qin)
12 def __str__(self):
13 return self.q.__str__()
15 def __mul__(self, other):
16 sw, sx, sy, sz = self.q
17 ow, ox, oy, oz = other.q
18 return Quaternion([sw * ow - sx * ox - sy * oy - sz * oz,
19 sw * ox + sx * ow + sy * oz - sz * oy,
20 sw * oy + sy * ow + sz * ox - sx * oz,
21 sw * oz + sz * ow + sx * oy - sy * ox])
23 def conjugate(self):
24 return Quaternion(-self.q * np.array([-1., 1., 1., 1.]))
26 def rotate(self, vector):
27 """Apply the rotation matrix to a vector."""
28 qw, qx, qy, qz = self.q[0], self.q[1], self.q[2], self.q[3]
29 x, y, z = vector[0], vector[1], vector[2]
31 ww = qw * qw
32 xx = qx * qx
33 yy = qy * qy
34 zz = qz * qz
35 wx = qw * qx
36 wy = qw * qy
37 wz = qw * qz
38 xy = qx * qy
39 xz = qx * qz
40 yz = qy * qz
42 return np.array(
43 [(ww + xx - yy - zz) * x + 2 * ((xy - wz) * y + (xz + wy) * z),
44 (ww - xx + yy - zz) * y + 2 * ((xy + wz) * x + (yz - wx) * z),
45 (ww - xx - yy + zz) * z + 2 * ((xz - wy) * x + (yz + wx) * y)])
47 def rotation_matrix(self):
49 qw, qx, qy, qz = self.q[0], self.q[1], self.q[2], self.q[3]
51 ww = qw * qw
52 xx = qx * qx
53 yy = qy * qy
54 zz = qz * qz
55 wx = qw * qx
56 wy = qw * qy
57 wz = qw * qz
58 xy = qx * qy
59 xz = qx * qz
60 yz = qy * qz
62 return np.array([[ww + xx - yy - zz, 2 * (xy - wz), 2 * (xz + wy)],
63 [2 * (xy + wz), ww - xx + yy - zz, 2 * (yz - wx)],
64 [2 * (xz - wy), 2 * (yz + wx), ww - xx - yy + zz]])
66 def axis_angle(self):
67 """Returns axis and angle (in radians) for the rotation described
68 by this Quaternion"""
70 sinth_2 = np.linalg.norm(self.q[1:])
72 if sinth_2 == 0:
73 # The angle is zero
74 theta = 0.0
75 n = np.array([0, 0, 1])
76 else:
77 theta = np.arctan2(sinth_2, self.q[0]) * 2
78 n = self.q[1:] / sinth_2
80 return n, theta
82 def euler_angles(self, mode='zyz'):
83 """Return three Euler angles describing the rotation, in radians.
84 Mode can be zyz or zxz. Default is zyz."""
85 # https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0276302
86 if mode == 'zyz':
87 a, b, c, d = self.q[0], self.q[3], self.q[2], -self.q[1]
88 elif mode == 'zxz':
89 a, b, c, d = self.q[0], self.q[3], self.q[1], self.q[2]
90 else:
91 raise ValueError(f'Invalid Euler angles mode {mode}')
93 beta = 2 * np.arccos(
94 np.sqrt((a**2 + b**2) / (a**2 + b**2 + c**2 + d**2))
95 )
96 gap = np.arctan2(b, a) # (gamma + alpha) / 2
97 gam = np.arctan2(d, c) # (gamma - alpha) / 2
98 if np.isclose(beta, 0):
99 # gam is meaningless here
100 alpha = 0
101 gamma = 2 * gap - alpha
102 elif np.isclose(beta, np.pi):
103 # gap is meaningless here
104 alpha = 0
105 gamma = 2 * gam + alpha
106 else:
107 alpha = gap - gam
108 gamma = gap + gam
110 return np.array([alpha, beta, gamma])
112 def arc_distance(self, other):
113 """Gives a metric of the distance between two quaternions,
114 expressed as 1-|q1.q2|"""
116 return 1.0 - np.abs(np.dot(self.q, other.q))
118 @staticmethod
119 def rotate_byq(q, vector):
120 """Apply the rotation matrix to a vector."""
121 qw, qx, qy, qz = q[0], q[1], q[2], q[3]
122 x, y, z = vector[0], vector[1], vector[2]
124 ww = qw * qw
125 xx = qx * qx
126 yy = qy * qy
127 zz = qz * qz
128 wx = qw * qx
129 wy = qw * qy
130 wz = qw * qz
131 xy = qx * qy
132 xz = qx * qz
133 yz = qy * qz
135 return np.array(
136 [(ww + xx - yy - zz) * x + 2 * ((xy - wz) * y + (xz + wy) * z),
137 (ww - xx + yy - zz) * y + 2 * ((xy + wz) * x + (yz - wx) * z),
138 (ww - xx - yy + zz) * z + 2 * ((xz - wy) * x + (yz + wx) * y)])
140 @staticmethod
141 def from_matrix(matrix):
142 """Build quaternion from rotation matrix."""
143 m = np.array(matrix)
144 assert m.shape == (3, 3)
146 # Now we need to find out the whole quaternion
147 # This method takes into account the possibility of qw being nearly
148 # zero, so it picks the stablest solution
150 if m[2, 2] < 0:
151 if (m[0, 0] > m[1, 1]):
152 # Use x-form
153 qx = np.sqrt(1 + m[0, 0] - m[1, 1] - m[2, 2]) / 2.0
154 fac = 1.0 / (4 * qx)
155 qw = (m[2, 1] - m[1, 2]) * fac
156 qy = (m[0, 1] + m[1, 0]) * fac
157 qz = (m[0, 2] + m[2, 0]) * fac
158 else:
159 # Use y-form
160 qy = np.sqrt(1 - m[0, 0] + m[1, 1] - m[2, 2]) / 2.0
161 fac = 1.0 / (4 * qy)
162 qw = (m[0, 2] - m[2, 0]) * fac
163 qx = (m[0, 1] + m[1, 0]) * fac
164 qz = (m[1, 2] + m[2, 1]) * fac
165 else:
166 if (m[0, 0] < -m[1, 1]):
167 # Use z-form
168 qz = np.sqrt(1 - m[0, 0] - m[1, 1] + m[2, 2]) / 2.0
169 fac = 1.0 / (4 * qz)
170 qw = (m[1, 0] - m[0, 1]) * fac
171 qx = (m[2, 0] + m[0, 2]) * fac
172 qy = (m[1, 2] + m[2, 1]) * fac
173 else:
174 # Use w-form
175 qw = np.sqrt(1 + m[0, 0] + m[1, 1] + m[2, 2]) / 2.0
176 fac = 1.0 / (4 * qw)
177 qx = (m[2, 1] - m[1, 2]) * fac
178 qy = (m[0, 2] - m[2, 0]) * fac
179 qz = (m[1, 0] - m[0, 1]) * fac
181 return Quaternion(np.array([qw, qx, qy, qz]))
183 @staticmethod
184 def from_axis_angle(n, theta):
185 """Build quaternion from axis (n, vector of 3 components) and angle
186 (theta, in radianses)."""
188 n = np.array(n, float) / np.linalg.norm(n)
189 return Quaternion(np.concatenate([[np.cos(theta / 2.0)],
190 np.sin(theta / 2.0) * n]))
192 @staticmethod
193 def from_euler_angles(a, b, c, mode='zyz'):
194 """Build quaternion from Euler angles, given in radians. Default
195 mode is ZYZ, but it can be set to ZXZ as well."""
197 q_a = Quaternion.from_axis_angle([0, 0, 1], a)
198 q_c = Quaternion.from_axis_angle([0, 0, 1], c)
200 if mode == 'zyz':
201 q_b = Quaternion.from_axis_angle([0, 1, 0], b)
202 elif mode == 'zxz':
203 q_b = Quaternion.from_axis_angle([1, 0, 0], b)
204 else:
205 raise ValueError(f'Invalid Euler angles mode {mode}')
207 return q_c * q_b * q_a