r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import numpy as np

2from ase.atoms import Atoms

5class Quaternions(Atoms):

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)))

18 def set_shapes(self, shapes):

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

21 def set_quaternions(self, quaternions):

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

24 def get_shapes(self):

25 return self.get_array('shapes')

27 def get_quaternions(self):

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

31class Quaternion:

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

34 assert len(qin) == 4

35 self.q = np.array(qin)

37 def __str__(self):

38 return self.q.__str__()

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])

48 def conjugate(self):

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

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]

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

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)])

72 def rotation_matrix(self):

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

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

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]])

91 def axis_angle(self):

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

93 by this Quaternion"""

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

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

105 return n, theta

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."""

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])

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])

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))

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

150 def arc_distance(self, other):

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

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

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

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]

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

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)])

178 @staticmethod

179 def from_matrix(matrix):

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

181 m = np.array(matrix)

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

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

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

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

221 @staticmethod

222 def from_axis_angle(n, theta):

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

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]))

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."""

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

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

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))

245 return q_c * q_b * q_a