1import numpy as np 

2from ase.atoms import Atoms 

3 

4 

5class Quaternions(Atoms): 

6 

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

8 quaternions = None 

9 if 'quaternions' in kwargs: 

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

11 del kwargs['quaternions'] 

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

13 if quaternions is not None: 

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

15 # set default shapes 

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

17 

18 def set_shapes(self, shapes): 

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

20 

21 def set_quaternions(self, quaternions): 

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

23 

24 def get_shapes(self): 

25 return self.get_array('shapes') 

26 

27 def get_quaternions(self): 

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

29 

30 

31class Quaternion: 

32 

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

34 assert(len(qin) == 4) 

35 self.q = np.array(qin) 

36 

37 def __str__(self): 

38 return self.q.__str__() 

39 

40 def __mul__(self, other): 

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

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

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

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

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

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

47 

48 def conjugate(self): 

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

50 

51 def rotate(self, vector): 

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

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

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

55 

56 ww = qw * qw 

57 xx = qx * qx 

58 yy = qy * qy 

59 zz = qz * qz 

60 wx = qw * qx 

61 wy = qw * qy 

62 wz = qw * qz 

63 xy = qx * qy 

64 xz = qx * qz 

65 yz = qy * qz 

66 

67 return np.array( 

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

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

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

71 

72 def rotation_matrix(self): 

73 

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

75 

76 ww = qw * qw 

77 xx = qx * qx 

78 yy = qy * qy 

79 zz = qz * qz 

80 wx = qw * qx 

81 wy = qw * qy 

82 wz = qw * qz 

83 xy = qx * qy 

84 xz = qx * qz 

85 yz = qy * qz 

86 

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

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

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

90 

91 def axis_angle(self): 

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

93 by this Quaternion""" 

94 

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

96 

97 if sinth_2 == 0: 

98 # The angle is zero 

99 theta = 0.0 

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

101 else: 

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

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

104 

105 return n, theta 

106 

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

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

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

110 

111 if mode == 'zyz': 

112 # These are (a+c)/2 and (a-c)/2 respectively 

113 apc = np.arctan2(self.q[3], self.q[0]) 

114 amc = np.arctan2(self.q[1], self.q[2]) 

115 

116 a, c = (apc+amc), (apc-amc) 

117 cos_amc = np.cos(amc) 

118 if cos_amc != 0: 

119 sinb2 = self.q[2]/cos_amc 

120 else: 

121 sinb2 = self.q[1]/np.sin(amc) 

122 cos_apc = np.cos(apc) 

123 if cos_apc != 0: 

124 cosb2 = self.q[0]/cos_apc 

125 else: 

126 cosb2 = self.q[3]/np.sin(apc) 

127 b = np.arctan2(sinb2, cosb2)*2 

128 elif mode == 'zxz': 

129 # These are (a+c)/2 and (a-c)/2 respectively 

130 apc = np.arctan2(self.q[3], self.q[0]) 

131 amc = np.arctan2(-self.q[2], self.q[1]) 

132 

133 a, c = (apc+amc), (apc-amc) 

134 cos_amc = np.cos(amc) 

135 if cos_amc != 0: 

136 sinb2 = self.q[1]/cos_amc 

137 else: 

138 sinb2 = self.q[2]/np.sin(amc) 

139 cos_apc = np.cos(apc) 

140 if cos_apc != 0: 

141 cosb2 = self.q[0]/cos_apc 

142 else: 

143 cosb2 = self.q[3]/np.sin(apc) 

144 b = np.arctan2(sinb2, cosb2)*2 

145 else: 

146 raise ValueError('Invalid Euler angles mode {0}'.format(mode)) 

147 

148 return np.array([a, b, c]) 

149 

150 def arc_distance(self, other): 

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

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

153 

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

155 

156 @staticmethod 

157 def rotate_byq(q, vector): 

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

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

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

161 

162 ww = qw * qw 

163 xx = qx * qx 

164 yy = qy * qy 

165 zz = qz * qz 

166 wx = qw * qx 

167 wy = qw * qy 

168 wz = qw * qz 

169 xy = qx * qy 

170 xz = qx * qz 

171 yz = qy * qz 

172 

173 return np.array( 

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

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

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

177 

178 @staticmethod 

179 def from_matrix(matrix): 

180 """Build quaternion from rotation matrix.""" 

181 m = np.array(matrix) 

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

183 

184 # Now we need to find out the whole quaternion 

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

186 # zero, so it picks the stablest solution 

187 

188 if m[2, 2] < 0: 

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

190 # Use x-form 

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

192 fac = 1.0 / (4 * qx) 

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

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

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

196 else: 

197 # Use y-form 

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

199 fac = 1.0 / (4 * qy) 

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

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

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

203 else: 

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

205 # Use z-form 

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

207 fac = 1.0 / (4 * qz) 

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

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

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

211 else: 

212 # Use w-form 

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

214 fac = 1.0 / (4 * qw) 

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

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

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

218 

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

220 

221 @staticmethod 

222 def from_axis_angle(n, theta): 

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

224 (theta, in radianses).""" 

225 

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

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

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

229 

230 @staticmethod 

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

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

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

234 

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

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

237 

238 if mode == 'zyz': 

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

240 elif mode == 'zxz': 

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

242 else: 

243 raise ValueError('Invalid Euler angles mode {0}'.format(mode)) 

244 

245 return q_c*q_b*q_a