Coverage for src\attipy\_transforms.py: 13%

1 import numpy as np

2 from numba import njit

3 from numpy.typing import NDArray

4 from ._vectorops import _normalize

5

6

7@njit # type: ignore[misc]

8 def _quaternion_from_matrix(A: np.ndarray) -> np.ndarray:

9 """

10 Convert a rotation matrix to a unit quaternion (see ref [1]_).

11

12 References

13 ----------

14 .. [1] https://en.wikipedia.org/wiki/Rotation_matrix#Quaternion

15 """

16

17 m00, m01, m02 = A[0]

18 m10, m11, m12 = A[1]

19 m20, m21, m22 = A[2]

20

21 trace = m00 + m11 + m22

22

23 if trace > 0.0:

24 s = 2.0 * np.sqrt(trace + 1.0)

25 w = 0.25 * s

26 x = (m21 - m12) / s

27 y = (m02 - m20) / s

28 z = (m10 - m01) / s

29 elif (m00 > m11) and (m00 > m22):

30 s = 2.0 * np.sqrt(1.0 + m00 - m11 - m22)

31 w = (m21 - m12) / s

32 x = 0.25 * s

33 y = (m01 + m10) / s

34 z = (m02 + m20) / s

35 elif m11 > m22:

36 s = 2.0 * np.sqrt(1.0 + m11 - m00 - m22)

37 w = (m02 - m20) / s

38 x = (m01 + m10) / s

39 y = 0.25 * s

40 z = (m12 + m21) / s

41 else:

42 s = 2.0 * np.sqrt(1.0 + m22 - m00 - m11)

43 w = (m10 - m01) / s

44 x = (m02 + m20) / s

45 y = (m12 + m21) / s

46 z = 0.25 * s

47

48 q = np.array([w, x, y, z])

49 return _normalize(q)

50

51

52@njit # type: ignore[misc]

53 def _rot_matrix_from_quaternion(q: NDArray[np.float64]) -> NDArray[np.float64]:

54 """

55 Compute the rotation matrix from a unit quaternion.

56

57 Parameters

58 ----------

59 q : numpy.ndarray, shape (4,)

60 Unit quaternion.

61

62 Returns

63 -------

64 rot : numpy.ndarray, shape (3, 3)

65 Rotation matrix.

66 """

67 q0, q1, q2, q3 = q

68

69 _2q1 = q1 + q1

70 _2q2 = q2 + q2

71 _2q3 = q3 + q3

72

73 _2q1q1 = q1 * _2q1

74 _2q1q2 = q1 * _2q2

75 _2q1q3 = q1 * _2q3

76 _2q2q2 = q2 * _2q2

77 _2q2q3 = q2 * _2q3

78 _2q3q3 = q3 * _2q3

79 _2q0q1 = q0 * _2q1

80 _2q0q2 = q0 * _2q2

81 _2q0q3 = q0 * _2q3

82

83 rot_00 = 1.0 - (_2q2q2 + _2q3q3)

84 rot_01 = _2q1q2 - _2q0q3

85 rot_02 = _2q1q3 + _2q0q2

86

87 rot_10 = _2q1q2 + _2q0q3

88 rot_11 = 1.0 - (_2q1q1 + _2q3q3)

89 rot_12 = _2q2q3 - _2q0q1

90

91 rot_20 = _2q1q3 - _2q0q2

92 rot_21 = _2q2q3 + _2q0q1

93 rot_22 = 1.0 - (_2q1q1 + _2q2q2)

94

95 rot = np.array(

96 [

97 [rot_00, rot_01, rot_02],

98 [rot_10, rot_11, rot_12],

99 [rot_20, rot_21, rot_22],

100 ]

101 )

102 return rot

103

104

105@njit # type: ignore[misc]

106 def _euler_zyx_from_quaternion(q: NDArray[np.float64]) -> NDArray[np.float64]:

107 """

108 Compute the Euler angles (ZYX convention) from a unit quaternion (see ref [1]_).

109

110 Parameters

111 ----------

112 q : numpy.ndarray, shape (4,)

113 Unit quaternion (representing transformation from-body-to-origin).

114

115 Returns

116 -------

117 numpy.ndarray, shape (3,)

118 Vector of Euler angles in radians (ZYX convention). Contains the following

119 three Euler angles in order:

120 - Roll (alpha): Rotation about the x-axis.

121 - Pitch (beta): Rotation about the y-axis.

122 - Yaw (gamma): Rotation about the z-axis.

123

124 References

125 ----------

126 .. [1] https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles

127 """

128 q_w, q_x, q_y, q_z = q

129

130 alpha = np.arctan2(2.0 * (q_y * q_z + q_x * q_w), 1.0 - 2.0 * (q_x**2 + q_y**2))

131 beta = -np.arcsin(2.0 * (q_x * q_z - q_y * q_w))

132 gamma = np.arctan2(2.0 * (q_x * q_y + q_z * q_w), 1.0 - 2.0 * (q_y**2 + q_z**2))

133

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

135

136

137@njit # type: ignore[misc]

138 def _rot_matrix_from_euler_zyx(euler: NDArray[np.float64]) -> NDArray[np.float64]:

139 """

140 Compute the rotation matrix (from-body-to-origin) from Euler angles.

141

142 Parameters

143 ----------

144 euler : numpy.ndarray, shape (3,)

145 Vector of Euler angles in radians (ZYX convention). Contains the following

146 three Euler angles in order:

147 - Roll (alpha): Rotation about the x-axis.

148 - Pitch (beta): Rotation about the y-axis.

149 - Yaw (gamma): Rotation about the z-axis.

150

151 Notes

152 -----

153 The Euler angles describe how to transition from the 'origin' frame to the 'body'

154 frame through three consecutive (passive, intrinsic) rotations in the ZYX order.

155 However, the returned rotation matrix represents the transformation of a vector

156 from the 'body' frame to the 'origin' frame.

157

158 Returns

159 -------

160 numpy.ndarray, shape (3, 3)

161 Rotation matrix.

162 """

163 alpha, beta, gamma = euler

164 cos_gamma = np.cos(gamma)

165 sin_gamma = np.sin(gamma)

166 cos_beta = np.cos(beta)

167 sin_beta = np.sin(beta)

168 cos_alpha = np.cos(alpha)

169 sin_alpha = np.sin(alpha)

170

171 rot_00 = cos_gamma * cos_beta

172 rot_01 = -sin_gamma * cos_alpha + cos_gamma * sin_beta * sin_alpha

173 rot_02 = sin_gamma * sin_alpha + cos_gamma * sin_beta * cos_alpha

174

175 rot_10 = sin_gamma * cos_beta

176 rot_11 = cos_gamma * cos_alpha + sin_gamma * sin_beta * sin_alpha

177 rot_12 = -cos_gamma * sin_alpha + sin_gamma * sin_beta * cos_alpha

178

179 rot_20 = -sin_beta

180 rot_21 = cos_beta * sin_alpha

181 rot_22 = cos_beta * cos_alpha

182

183 rot = np.array(

184 [[rot_00, rot_01, rot_02], [rot_10, rot_11, rot_12], [rot_20, rot_21, rot_22]]

185 )

186 return rot

187

188

189@njit # type: ignore[misc]

190 def _quaternion_from_euler_zyx(euler: NDArray[np.float64]) -> NDArray[np.float64]:

191 """

192 Compute the unit quaternion from Euler angles (see ref [1]_).

193

194 References

195 ----------

196 .. [1] https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles

197 """

198

199 alpha_half, beta_half, gamma_half = euler / 2.0

200 ca_half = np.cos(alpha_half)

201 sa_half = np.sin(alpha_half)

202 cb_half = np.cos(beta_half)

203 sb_half = np.sin(beta_half)

204 cg_half = np.cos(gamma_half)

205 sg_half = np.sin(gamma_half)

206

207 q_w = ca_half * cb_half * cg_half + sa_half * sb_half * sg_half

208 q_x = sa_half * cb_half * cg_half - ca_half * sb_half * sg_half

209 q_y = ca_half * sb_half * cg_half + sa_half * cb_half * sg_half

210 q_z = ca_half * cb_half * sg_half - sa_half * sb_half * cg_half

211

212 return np.array([q_w, q_x, q_y, q_z])