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
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-14 16:02 +0200
1"""
2Rotation transformations for robotics-numpy
4This module provides efficient rotation matrix operations using only NumPy.
5All functions are designed to work with both single rotations and batches.
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"""
14from typing import Tuple, Union
16import numpy as np
18# Type aliases
19ArrayLike = Union[float, int, np.ndarray]
20RotationMatrix = np.ndarray # 3x3 rotation matrix
21Quaternion = np.ndarray # [w, x, y, z] quaternion
24def rotx(angle: ArrayLike) -> RotationMatrix:
25 """
26 Create rotation matrix about X-axis.
28 Args:
29 angle: Rotation angle in radians
31 Returns:
32 3x3 rotation matrix
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)
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
61def roty(angle: ArrayLike) -> RotationMatrix:
62 """
63 Create rotation matrix about Y-axis.
65 Args:
66 angle: Rotation angle in radians
68 Returns:
69 3x3 rotation matrix
70 """
71 angle = np.asarray(angle)
72 c = np.cos(angle)
73 s = np.sin(angle)
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
92def rotz(angle: ArrayLike) -> RotationMatrix:
93 """
94 Create rotation matrix about Z-axis.
96 Args:
97 angle: Rotation angle in radians
99 Returns:
100 3x3 rotation matrix
101 """
102 angle = np.asarray(angle)
103 c = np.cos(angle)
104 s = np.sin(angle)
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
123def rpy2r(roll: float, pitch: float, yaw: float) -> RotationMatrix:
124 """
125 Convert roll-pitch-yaw angles to rotation matrix.
127 Applies rotations in order: Rz(yaw) * Ry(pitch) * Rx(roll)
129 Args:
130 roll: Rotation about X-axis (radians)
131 pitch: Rotation about Y-axis (radians)
132 yaw: Rotation about Z-axis (radians)
134 Returns:
135 3x3 rotation matrix
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)
144def r2rpy(R: RotationMatrix) -> Tuple[float, float, float]:
145 """
146 Convert rotation matrix to roll-pitch-yaw angles.
148 Args:
149 R: 3x3 rotation matrix
151 Returns:
152 Tuple of (roll, pitch, yaw) in radians
154 Raises:
155 ValueError: If R is not a valid rotation matrix
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")
164 # Extract angles using ZYX convention
165 pitch = np.arcsin(-R[2, 0])
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])
179 return roll, pitch, yaw
182def eul2r(phi: float, theta: float, psi: float,
183 convention: str = 'ZYZ') -> RotationMatrix:
184 """
185 Convert Euler angles to rotation matrix.
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')
193 Returns:
194 3x3 rotation matrix
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}")
207def r2eul(R: RotationMatrix, convention: str = 'ZYZ') -> Tuple[float, float, float]:
208 """
209 Convert rotation matrix to Euler angles.
211 Args:
212 R: 3x3 rotation matrix
213 convention: Euler angle convention ('ZYZ', 'ZYX')
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")
221 if convention == 'ZYZ':
222 theta = np.arccos(np.clip(R[2, 2], -1, 1))
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])
235 return phi, theta, psi
237 elif convention == 'ZYX':
238 return r2rpy(R)
240 else:
241 raise ValueError(f"Unsupported convention: {convention}")
244def quat2r(q: Quaternion) -> RotationMatrix:
245 """
246 Convert quaternion to rotation matrix.
248 Args:
249 q: Quaternion [w, x, y, z] (scalar first)
251 Returns:
252 3x3 rotation matrix
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")
262 # Normalize quaternion
263 q = quat_normalize(q)
264 w, x, y, z = q
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 ])
273 return R
276def r2quat(R: RotationMatrix) -> Quaternion:
277 """
278 Convert rotation matrix to quaternion.
280 Args:
281 R: 3x3 rotation matrix
283 Returns:
284 Quaternion [w, x, y, z]
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")
293 # Shepperd's method for numerical stability
294 trace = np.trace(R)
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
321 return np.array([w, x, y, z])
324def quat_multiply(q1: Quaternion, q2: Quaternion) -> Quaternion:
325 """
326 Multiply two quaternions.
328 Args:
329 q1: First quaternion [w, x, y, z]
330 q2: Second quaternion [w, x, y, z]
332 Returns:
333 Product quaternion q1 * q2
334 """
335 q1 = np.asarray(q1)
336 q2 = np.asarray(q2)
338 w1, x1, y1, z1 = q1
339 w2, x2, y2, z2 = q2
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
346 return np.array([w, x, y, z])
349def quat_inverse(q: Quaternion) -> Quaternion:
350 """
351 Compute quaternion inverse (conjugate for unit quaternions).
353 Args:
354 q: Quaternion [w, x, y, z]
356 Returns:
357 Inverse quaternion
358 """
359 q = np.asarray(q)
360 return np.array([q[0], -q[1], -q[2], -q[3]])
363def quat_normalize(q: Quaternion) -> Quaternion:
364 """
365 Normalize quaternion to unit length.
367 Args:
368 q: Quaternion [w, x, y, z]
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
380def is_rotation_matrix(R: np.ndarray, tol: float = 1e-6) -> bool:
381 """
382 Check if matrix is a valid rotation matrix.
384 Args:
385 R: Matrix to check
386 tol: Tolerance for checks
388 Returns:
389 True if R is a valid rotation matrix
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
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
405 # Check if determinant is 1 (proper rotation)
406 if not np.isclose(np.linalg.det(R), 1.0, atol=tol):
407 return False
409 return True
412# Convenience functions for common rotations
413def rx(angle: ArrayLike) -> RotationMatrix:
414 """Alias for rotx."""
415 return rotx(angle)
418def ry(angle: ArrayLike) -> RotationMatrix:
419 """Alias for roty."""
420 return roty(angle)
423def rz(angle: ArrayLike) -> RotationMatrix:
424 """Alias for rotz."""
425 return rotz(angle)
428# Constants
429IDENTITY_QUAT = np.array([1.0, 0.0, 0.0, 0.0])
430IDENTITY_ROT = np.eye(3)