Coverage for src/robotics_numpy/transforms/rotations.py: 70%

156 statements  

« prev     ^ index     » next       coverage.py v7.9.2, created at 2025-07-14 16:02 +0200

1""" 

2Rotation transformations for robotics-numpy 

3 

4This module provides efficient rotation matrix operations using only NumPy. 

5All functions are designed to work with both single rotations and batches. 

6 

7Conventions: 

8- Rotation matrices are 3x3 NumPy arrays 

9- Euler angles follow ZYX convention (yaw-pitch-roll) 

10- Quaternions are [w, x, y, z] (scalar first) 

11- All angles are in radians 

12""" 

13 

14from typing import Tuple, Union 

15 

16import numpy as np 

17 

18# Type aliases 

19ArrayLike = Union[float, int, np.ndarray] 

20RotationMatrix = np.ndarray # 3x3 rotation matrix 

21Quaternion = np.ndarray # [w, x, y, z] quaternion 

22 

23 

24def rotx(angle: ArrayLike) -> RotationMatrix: 

25 """ 

26 Create rotation matrix about X-axis. 

27 

28 Args: 

29 angle: Rotation angle in radians 

30 

31 Returns: 

32 3x3 rotation matrix 

33 

34 Examples: 

35 >>> R = rotx(np.pi/2) # 90 degree rotation about X 

36 >>> R = rotx([0, np.pi/4, np.pi/2]) # Batch operation 

37 """ 

38 angle = np.asarray(angle) 

39 c = np.cos(angle) 

40 s = np.sin(angle) 

41 

42 if angle.ndim == 0: 

43 # Single rotation 

44 return np.array([ 

45 [1, 0, 0], 

46 [0, c, -s], 

47 [0, s, c] 

48 ]) 

49 else: 

50 # Batch rotations 

51 batch_size = angle.shape[0] 

52 R = np.zeros((batch_size, 3, 3)) 

53 R[:, 0, 0] = 1 

54 R[:, 1, 1] = c 

55 R[:, 1, 2] = -s 

56 R[:, 2, 1] = s 

57 R[:, 2, 2] = c 

58 return R 

59 

60 

61def roty(angle: ArrayLike) -> RotationMatrix: 

62 """ 

63 Create rotation matrix about Y-axis. 

64 

65 Args: 

66 angle: Rotation angle in radians 

67 

68 Returns: 

69 3x3 rotation matrix 

70 """ 

71 angle = np.asarray(angle) 

72 c = np.cos(angle) 

73 s = np.sin(angle) 

74 

75 if angle.ndim == 0: 

76 return np.array([ 

77 [c, 0, s], 

78 [0, 1, 0], 

79 [-s, 0, c] 

80 ]) 

81 else: 

82 batch_size = angle.shape[0] 

83 R = np.zeros((batch_size, 3, 3)) 

84 R[:, 0, 0] = c 

85 R[:, 0, 2] = s 

86 R[:, 1, 1] = 1 

87 R[:, 2, 0] = -s 

88 R[:, 2, 2] = c 

89 return R 

90 

91 

92def rotz(angle: ArrayLike) -> RotationMatrix: 

93 """ 

94 Create rotation matrix about Z-axis. 

95 

96 Args: 

97 angle: Rotation angle in radians 

98 

99 Returns: 

100 3x3 rotation matrix 

101 """ 

102 angle = np.asarray(angle) 

103 c = np.cos(angle) 

104 s = np.sin(angle) 

105 

106 if angle.ndim == 0: 

107 return np.array([ 

108 [c, -s, 0], 

109 [s, c, 0], 

110 [0, 0, 1] 

111 ]) 

112 else: 

113 batch_size = angle.shape[0] 

114 R = np.zeros((batch_size, 3, 3)) 

115 R[:, 0, 0] = c 

116 R[:, 0, 1] = -s 

117 R[:, 1, 0] = s 

118 R[:, 1, 1] = c 

119 R[:, 2, 2] = 1 

120 return R 

121 

122 

123def rpy2r(roll: float, pitch: float, yaw: float) -> RotationMatrix: 

124 """ 

125 Convert roll-pitch-yaw angles to rotation matrix. 

126 

127 Applies rotations in order: Rz(yaw) * Ry(pitch) * Rx(roll) 

128 

129 Args: 

130 roll: Rotation about X-axis (radians) 

131 pitch: Rotation about Y-axis (radians) 

132 yaw: Rotation about Z-axis (radians) 

133 

134 Returns: 

135 3x3 rotation matrix 

136 

137 Examples: 

138 >>> R = rpy2r(0.1, 0.2, 0.3) # Small rotations 

139 >>> R = rpy2r(0, 0, np.pi/2) # 90 deg yaw 

140 """ 

141 return rotz(yaw) @ roty(pitch) @ rotx(roll) 

142 

143 

144def r2rpy(R: RotationMatrix) -> Tuple[float, float, float]: 

145 """ 

146 Convert rotation matrix to roll-pitch-yaw angles. 

147 

148 Args: 

149 R: 3x3 rotation matrix 

150 

151 Returns: 

152 Tuple of (roll, pitch, yaw) in radians 

153 

154 Raises: 

155 ValueError: If R is not a valid rotation matrix 

156 

157 Examples: 

158 >>> R = rpy2r(0.1, 0.2, 0.3) 

159 >>> roll, pitch, yaw = r2rpy(R) 

160 """ 

161 if not is_rotation_matrix(R): 

162 raise ValueError("Input is not a valid rotation matrix") 

163 

164 # Extract angles using ZYX convention 

165 pitch = np.arcsin(-R[2, 0]) 

166 

167 # Check for gimbal lock 

168 if np.abs(np.cos(pitch)) < 1e-6: 

169 # Gimbal lock case 

170 roll = 0.0 

171 if pitch > 0: 

172 yaw = np.arctan2(R[0, 1], R[1, 1]) 

173 else: 

174 yaw = np.arctan2(-R[0, 1], R[1, 1]) 

175 else: 

176 roll = np.arctan2(R[2, 1], R[2, 2]) 

177 yaw = np.arctan2(R[1, 0], R[0, 0]) 

178 

179 return roll, pitch, yaw 

180 

181 

182def eul2r(phi: float, theta: float, psi: float, 

183 convention: str = 'ZYZ') -> RotationMatrix: 

184 """ 

185 Convert Euler angles to rotation matrix. 

186 

187 Args: 

188 phi: First rotation angle (radians) 

189 theta: Second rotation angle (radians) 

190 psi: Third rotation angle (radians) 

191 convention: Euler angle convention ('ZYZ', 'ZYX') 

192 

193 Returns: 

194 3x3 rotation matrix 

195 

196 Examples: 

197 >>> R = eul2r(0.1, 0.2, 0.3, 'ZYZ') 

198 """ 

199 if convention == 'ZYZ': 

200 return rotz(phi) @ roty(theta) @ rotz(psi) 

201 elif convention == 'ZYX': 

202 return rotz(phi) @ roty(theta) @ rotx(psi) 

203 else: 

204 raise ValueError(f"Unsupported convention: {convention}") 

205 

206 

207def r2eul(R: RotationMatrix, convention: str = 'ZYZ') -> Tuple[float, float, float]: 

208 """ 

209 Convert rotation matrix to Euler angles. 

210 

211 Args: 

212 R: 3x3 rotation matrix 

213 convention: Euler angle convention ('ZYZ', 'ZYX') 

214 

215 Returns: 

216 Tuple of three Euler angles in radians 

217 """ 

218 if not is_rotation_matrix(R): 

219 raise ValueError("Input is not a valid rotation matrix") 

220 

221 if convention == 'ZYZ': 

222 theta = np.arccos(np.clip(R[2, 2], -1, 1)) 

223 

224 if np.abs(np.sin(theta)) < 1e-6: 

225 # Singularity case 

226 phi = 0.0 

227 if theta < 1e-6: # theta ≈ 0 

228 psi = np.arctan2(R[1, 0], R[0, 0]) 

229 else: # theta ≈ π 

230 psi = np.arctan2(-R[1, 0], R[0, 0]) 

231 else: 

232 phi = np.arctan2(R[1, 2], R[0, 2]) 

233 psi = np.arctan2(R[2, 1], -R[2, 0]) 

234 

235 return phi, theta, psi 

236 

237 elif convention == 'ZYX': 

238 return r2rpy(R) 

239 

240 else: 

241 raise ValueError(f"Unsupported convention: {convention}") 

242 

243 

244def quat2r(q: Quaternion) -> RotationMatrix: 

245 """ 

246 Convert quaternion to rotation matrix. 

247 

248 Args: 

249 q: Quaternion [w, x, y, z] (scalar first) 

250 

251 Returns: 

252 3x3 rotation matrix 

253 

254 Examples: 

255 >>> q = [1, 0, 0, 0] # Identity quaternion 

256 >>> R = quat2r(q) 

257 """ 

258 q = np.asarray(q, dtype=float) 

259 if q.shape != (4,): 

260 raise ValueError("Quaternion must be a 4-element array") 

261 

262 # Normalize quaternion 

263 q = quat_normalize(q) 

264 w, x, y, z = q 

265 

266 # Convert to rotation matrix 

267 R = np.array([ 

268 [1 - 2*(y*y + z*z), 2*(x*y - w*z), 2*(x*z + w*y)], 

269 [2*(x*y + w*z), 1 - 2*(x*x + z*z), 2*(y*z - w*x)], 

270 [2*(x*z - w*y), 2*(y*z + w*x), 1 - 2*(x*x + y*y)] 

271 ]) 

272 

273 return R 

274 

275 

276def r2quat(R: RotationMatrix) -> Quaternion: 

277 """ 

278 Convert rotation matrix to quaternion. 

279 

280 Args: 

281 R: 3x3 rotation matrix 

282 

283 Returns: 

284 Quaternion [w, x, y, z] 

285 

286 Examples: 

287 >>> R = np.eye(3) # Identity matrix 

288 >>> q = r2quat(R) # Returns [1, 0, 0, 0] 

289 """ 

290 if not is_rotation_matrix(R): 

291 raise ValueError("Input is not a valid rotation matrix") 

292 

293 # Shepperd's method for numerical stability 

294 trace = np.trace(R) 

295 

296 if trace > 0: 

297 s = np.sqrt(trace + 1.0) * 2 # s = 4 * w 

298 w = 0.25 * s 

299 x = (R[2, 1] - R[1, 2]) / s 

300 y = (R[0, 2] - R[2, 0]) / s 

301 z = (R[1, 0] - R[0, 1]) / s 

302 elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]: 

303 s = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2 # s = 4 * x 

304 w = (R[2, 1] - R[1, 2]) / s 

305 x = 0.25 * s 

306 y = (R[0, 1] + R[1, 0]) / s 

307 z = (R[0, 2] + R[2, 0]) / s 

308 elif R[1, 1] > R[2, 2]: 

309 s = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2 # s = 4 * y 

310 w = (R[0, 2] - R[2, 0]) / s 

311 x = (R[0, 1] + R[1, 0]) / s 

312 y = 0.25 * s 

313 z = (R[1, 2] + R[2, 1]) / s 

314 else: 

315 s = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2 # s = 4 * z 

316 w = (R[1, 0] - R[0, 1]) / s 

317 x = (R[0, 2] + R[2, 0]) / s 

318 y = (R[1, 2] + R[2, 1]) / s 

319 z = 0.25 * s 

320 

321 return np.array([w, x, y, z]) 

322 

323 

324def quat_multiply(q1: Quaternion, q2: Quaternion) -> Quaternion: 

325 """ 

326 Multiply two quaternions. 

327 

328 Args: 

329 q1: First quaternion [w, x, y, z] 

330 q2: Second quaternion [w, x, y, z] 

331 

332 Returns: 

333 Product quaternion q1 * q2 

334 """ 

335 q1 = np.asarray(q1) 

336 q2 = np.asarray(q2) 

337 

338 w1, x1, y1, z1 = q1 

339 w2, x2, y2, z2 = q2 

340 

341 w = w1*w2 - x1*x2 - y1*y2 - z1*z2 

342 x = w1*x2 + x1*w2 + y1*z2 - z1*y2 

343 y = w1*y2 - x1*z2 + y1*w2 + z1*x2 

344 z = w1*z2 + x1*y2 - y1*x2 + z1*w2 

345 

346 return np.array([w, x, y, z]) 

347 

348 

349def quat_inverse(q: Quaternion) -> Quaternion: 

350 """ 

351 Compute quaternion inverse (conjugate for unit quaternions). 

352 

353 Args: 

354 q: Quaternion [w, x, y, z] 

355 

356 Returns: 

357 Inverse quaternion 

358 """ 

359 q = np.asarray(q) 

360 return np.array([q[0], -q[1], -q[2], -q[3]]) 

361 

362 

363def quat_normalize(q: Quaternion) -> Quaternion: 

364 """ 

365 Normalize quaternion to unit length. 

366 

367 Args: 

368 q: Quaternion [w, x, y, z] 

369 

370 Returns: 

371 Normalized quaternion 

372 """ 

373 q = np.asarray(q, dtype=float) 

374 norm = np.linalg.norm(q) 

375 if norm < 1e-8: 

376 raise ValueError("Cannot normalize zero quaternion") 

377 return q / norm 

378 

379 

380def is_rotation_matrix(R: np.ndarray, tol: float = 1e-6) -> bool: 

381 """ 

382 Check if matrix is a valid rotation matrix. 

383 

384 Args: 

385 R: Matrix to check 

386 tol: Tolerance for checks 

387 

388 Returns: 

389 True if R is a valid rotation matrix 

390 

391 Examples: 

392 >>> R = np.eye(3) 

393 >>> is_rotation_matrix(R) # True 

394 >>> is_rotation_matrix(np.ones((3, 3))) # False 

395 """ 

396 if R.shape != (3, 3): 

397 return False 

398 

399 # Check if orthogonal: R @ R.T ≈ I 

400 should_be_identity = R @ R.T 

401 identity = np.eye(3) 

402 if not np.allclose(should_be_identity, identity, atol=tol): 

403 return False 

404 

405 # Check if determinant is 1 (proper rotation) 

406 if not np.isclose(np.linalg.det(R), 1.0, atol=tol): 

407 return False 

408 

409 return True 

410 

411 

412# Convenience functions for common rotations 

413def rx(angle: ArrayLike) -> RotationMatrix: 

414 """Alias for rotx.""" 

415 return rotx(angle) 

416 

417 

418def ry(angle: ArrayLike) -> RotationMatrix: 

419 """Alias for roty.""" 

420 return roty(angle) 

421 

422 

423def rz(angle: ArrayLike) -> RotationMatrix: 

424 """Alias for rotz.""" 

425 return rotz(angle) 

426 

427 

428# Constants 

429IDENTITY_QUAT = np.array([1.0, 0.0, 0.0, 0.0]) 

430IDENTITY_ROT = np.eye(3)