Coverage for /builds/kinetik161/ase/ase/quaternions.py: 77.78%

1import numpy as np 

2 

3from ase.atoms import Atoms 

4 

5 

6class Quaternions(Atoms): 

7 

8 def __init__(self, *args, **kwargs): 

9 quaternions = None 

10 if 'quaternions' in kwargs: 

11 quaternions = np.array(kwargs['quaternions']) 

12 del kwargs['quaternions'] 

13 Atoms.__init__(self, *args, **kwargs) 

14 if quaternions is not None: 

15 self.set_array('quaternions', quaternions, shape=(4,)) 

16 # set default shapes 

17 self.set_shapes(np.array([[3, 2, 1]] * len(self))) 

18 

19 def set_shapes(self, shapes): 

20 self.set_array('shapes', shapes, shape=(3,)) 

21 

22 def set_quaternions(self, quaternions): 

23 self.set_array('quaternions', quaternions, quaternion=(4,)) 

24 

25 def get_shapes(self): 

26 return self.get_array('shapes') 

27 

28 def get_quaternions(self): 

29 return self.get_array('quaternions').copy() 

30 

31 

32class Quaternion: 

33 

34 def __init__(self, qin=[1, 0, 0, 0]): 

35 assert len(qin) == 4 

36 self.q = np.array(qin) 

37 

38 def __str__(self): 

39 return self.q.__str__() 

40 

41 def __mul__(self, other): 

42 sw, sx, sy, sz = self.q 

43 ow, ox, oy, oz = other.q 

44 return Quaternion([sw * ow - sx * ox - sy * oy - sz * oz, 

45 sw * ox + sx * ow + sy * oz - sz * oy, 

46 sw * oy + sy * ow + sz * ox - sx * oz, 

47 sw * oz + sz * ow + sx * oy - sy * ox]) 

48 

49 def conjugate(self): 

50 return Quaternion(-self.q * np.array([-1., 1., 1., 1.])) 

51 

52 def rotate(self, vector): 

53 """Apply the rotation matrix to a vector.""" 

54 qw, qx, qy, qz = self.q[0], self.q[1], self.q[2], self.q[3] 

55 x, y, z = vector[0], vector[1], vector[2] 

56 

57 ww = qw * qw 

58 xx = qx * qx 

59 yy = qy * qy 

60 zz = qz * qz 

61 wx = qw * qx 

62 wy = qw * qy 

63 wz = qw * qz 

64 xy = qx * qy 

65 xz = qx * qz 

66 yz = qy * qz 

67 

68 return np.array( 

69 [(ww + xx - yy - zz) * x + 2 * ((xy - wz) * y + (xz + wy) * z), 

70 (ww - xx + yy - zz) * y + 2 * ((xy + wz) * x + (yz - wx) * z), 

71 (ww - xx - yy + zz) * z + 2 * ((xz - wy) * x + (yz + wx) * y)]) 

72 

73 def rotation_matrix(self): 

74 

75 qw, qx, qy, qz = self.q[0], self.q[1], self.q[2], self.q[3] 

76 

77 ww = qw * qw 

78 xx = qx * qx 

79 yy = qy * qy 

80 zz = qz * qz 

81 wx = qw * qx 

82 wy = qw * qy 

83 wz = qw * qz 

84 xy = qx * qy 

85 xz = qx * qz 

86 yz = qy * qz 

87 

88 return np.array([[ww + xx - yy - zz, 2 * (xy - wz), 2 * (xz + wy)], 

89 [2 * (xy + wz), ww - xx + yy - zz, 2 * (yz - wx)], 

90 [2 * (xz - wy), 2 * (yz + wx), ww - xx - yy + zz]]) 

91 

92 def axis_angle(self): 

93 """Returns axis and angle (in radians) for the rotation described 

94 by this Quaternion""" 

95 

96 sinth_2 = np.linalg.norm(self.q[1:]) 

97 

98 if sinth_2 == 0: 

99 # The angle is zero 

100 theta = 0.0 

101 n = np.array([0, 0, 1]) 

102 else: 

103 theta = np.arctan2(sinth_2, self.q[0]) * 2 

104 n = self.q[1:] / sinth_2 

105 

106 return n, theta 

107 

108 def euler_angles(self, mode='zyz'): 

109 """Return three Euler angles describing the rotation, in radians. 

110 Mode can be zyz or zxz. Default is zyz.""" 

111 # https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0276302 

112 if mode == 'zyz': 

113 a, b, c, d = self.q[0], self.q[3], self.q[2], -self.q[1] 

114 elif mode == 'zxz': 

115 a, b, c, d = self.q[0], self.q[3], self.q[1], self.q[2] 

116 else: 

117 raise ValueError(f'Invalid Euler angles mode {mode}') 

118 

119 beta = 2 * np.arccos( 

120 np.sqrt((a**2 + b**2) / (a**2 + b**2 + c**2 + d**2)) 

121 ) 

122 gap = np.arctan2(b, a) # (gamma + alpha) / 2 

123 gam = np.arctan2(d, c) # (gamma - alpha) / 2 

124 if np.isclose(beta, 0): 

125 # gam is meaningless here 

126 alpha = 0 

127 gamma = 2 * gap - alpha 

128 elif np.isclose(beta, np.pi): 

129 # gap is meaningless here 

130 alpha = 0 

131 gamma = 2 * gam + alpha 

132 else: 

133 alpha = gap - gam 

134 gamma = gap + gam 

135 

136 return np.array([alpha, beta, gamma]) 

137 

138 def arc_distance(self, other): 

139 """Gives a metric of the distance between two quaternions, 

140 expressed as 1-|q1.q2|""" 

141 

142 return 1.0 - np.abs(np.dot(self.q, other.q)) 

143 

144 @staticmethod 

145 def rotate_byq(q, vector): 

146 """Apply the rotation matrix to a vector.""" 

147 qw, qx, qy, qz = q[0], q[1], q[2], q[3] 

148 x, y, z = vector[0], vector[1], vector[2] 

149 

150 ww = qw * qw 

151 xx = qx * qx 

152 yy = qy * qy 

153 zz = qz * qz 

154 wx = qw * qx 

155 wy = qw * qy 

156 wz = qw * qz 

157 xy = qx * qy 

158 xz = qx * qz 

159 yz = qy * qz 

160 

161 return np.array( 

162 [(ww + xx - yy - zz) * x + 2 * ((xy - wz) * y + (xz + wy) * z), 

163 (ww - xx + yy - zz) * y + 2 * ((xy + wz) * x + (yz - wx) * z), 

164 (ww - xx - yy + zz) * z + 2 * ((xz - wy) * x + (yz + wx) * y)]) 

165 

166 @staticmethod 

167 def from_matrix(matrix): 

168 """Build quaternion from rotation matrix.""" 

169 m = np.array(matrix) 

170 assert m.shape == (3, 3) 

171 

172 # Now we need to find out the whole quaternion 

173 # This method takes into account the possibility of qw being nearly 

174 # zero, so it picks the stablest solution 

175 

176 if m[2, 2] < 0: 

177 if (m[0, 0] > m[1, 1]): 

178 # Use x-form 

179 qx = np.sqrt(1 + m[0, 0] - m[1, 1] - m[2, 2]) / 2.0 

180 fac = 1.0 / (4 * qx) 

181 qw = (m[2, 1] - m[1, 2]) * fac 

182 qy = (m[0, 1] + m[1, 0]) * fac 

183 qz = (m[0, 2] + m[2, 0]) * fac 

184 else: 

185 # Use y-form 

186 qy = np.sqrt(1 - m[0, 0] + m[1, 1] - m[2, 2]) / 2.0 

187 fac = 1.0 / (4 * qy) 

188 qw = (m[0, 2] - m[2, 0]) * fac 

189 qx = (m[0, 1] + m[1, 0]) * fac 

190 qz = (m[1, 2] + m[2, 1]) * fac 

191 else: 

192 if (m[0, 0] < -m[1, 1]): 

193 # Use z-form 

194 qz = np.sqrt(1 - m[0, 0] - m[1, 1] + m[2, 2]) / 2.0 

195 fac = 1.0 / (4 * qz) 

196 qw = (m[1, 0] - m[0, 1]) * fac 

197 qx = (m[2, 0] + m[0, 2]) * fac 

198 qy = (m[1, 2] + m[2, 1]) * fac 

199 else: 

200 # Use w-form 

201 qw = np.sqrt(1 + m[0, 0] + m[1, 1] + m[2, 2]) / 2.0 

202 fac = 1.0 / (4 * qw) 

203 qx = (m[2, 1] - m[1, 2]) * fac 

204 qy = (m[0, 2] - m[2, 0]) * fac 

205 qz = (m[1, 0] - m[0, 1]) * fac 

206 

207 return Quaternion(np.array([qw, qx, qy, qz])) 

208 

209 @staticmethod 

210 def from_axis_angle(n, theta): 

211 """Build quaternion from axis (n, vector of 3 components) and angle 

212 (theta, in radianses).""" 

213 

214 n = np.array(n, float) / np.linalg.norm(n) 

215 return Quaternion(np.concatenate([[np.cos(theta / 2.0)], 

216 np.sin(theta / 2.0) * n])) 

217 

218 @staticmethod 

219 def from_euler_angles(a, b, c, mode='zyz'): 

220 """Build quaternion from Euler angles, given in radians. Default 

221 mode is ZYZ, but it can be set to ZXZ as well.""" 

222 

223 q_a = Quaternion.from_axis_angle([0, 0, 1], a) 

224 q_c = Quaternion.from_axis_angle([0, 0, 1], c) 

225 

226 if mode == 'zyz': 

227 q_b = Quaternion.from_axis_angle([0, 1, 0], b) 

228 elif mode == 'zxz': 

229 q_b = Quaternion.from_axis_angle([1, 0, 0], b) 

230 else: 

231 raise ValueError(f'Invalid Euler angles mode {mode}') 

232 

233 return q_c * q_b * q_a