Coverage for src/robotics_numpy/transforms/pose.py: 66%
188 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"""
2SE(3) and SO(3) classes for robotics-numpy
4This module provides object-oriented interfaces for 3D transformations:
5- SE3: Special Euclidean group SE(3) for rigid body transformations
6- SO3: Special Orthogonal group SO(3) for rotations
8These classes provide intuitive APIs while maintaining high performance through NumPy.
9"""
11from typing import List, Optional, Tuple, Union
13import numpy as np
15from .homogeneous import (
16 extract_rotation,
17 extract_translation,
18 homogeneous_inverse,
19 is_homogeneous,
20 rotmat,
21 transform_point,
22 transform_points,
23 transl,
24)
25from .rotations import (
26 eul2r,
27 is_rotation_matrix,
28 quat2r,
29 r2eul,
30 r2quat,
31 r2rpy,
32 rotx,
33 roty,
34 rotz,
35 rpy2r,
36)
38# Type aliases
39ArrayLike = Union[float, int, np.ndarray, List[float]]
40Vector3 = Union[np.ndarray, List[float]]
41Matrix3x3 = np.ndarray
42Matrix4x4 = np.ndarray
45class SO3:
46 """
47 Special Orthogonal group SO(3) - 3D rotations.
49 Represents rotations in 3D space using rotation matrices.
50 Provides convenient methods for creation, composition, and conversion.
52 Examples:
53 >>> R1 = SO3.Rx(np.pi/2) # Rotation about X-axis
54 >>> R2 = SO3.RPY(0.1, 0.2, 0.3) # From roll-pitch-yaw
55 >>> R3 = R1 * R2 # Composition
56 >>> rpy = R3.rpy() # Extract RPY angles
57 """
59 def __init__(self, matrix: Optional[Matrix3x3] = None):
60 """
61 Initialize SO(3) rotation.
63 Args:
64 matrix: 3x3 rotation matrix (default: identity)
66 Examples:
67 >>> R1 = SO3() # Identity rotation
68 >>> R2 = SO3(np.eye(3)) # From matrix
69 >>> R3 = SO3.Rx(np.pi/4) # Factory method
70 """
71 if matrix is None:
72 self._matrix = np.eye(3)
73 else:
74 matrix = np.asarray(matrix)
75 if not is_rotation_matrix(matrix):
76 raise ValueError("Input is not a valid rotation matrix")
77 self._matrix = matrix.copy()
79 @property
80 def matrix(self) -> Matrix3x3:
81 """Get the rotation matrix."""
82 return self._matrix.copy()
84 @property
85 def R(self) -> Matrix3x3:
86 """Alias for matrix property."""
87 return self.matrix
89 # Factory methods for common rotations
90 @classmethod
91 def Rx(cls, angle: float) -> 'SO3':
92 """Create rotation about X-axis."""
93 return cls(rotx(angle))
95 @classmethod
96 def Ry(cls, angle: float) -> 'SO3':
97 """Create rotation about Y-axis."""
98 return cls(roty(angle))
100 @classmethod
101 def Rz(cls, angle: float) -> 'SO3':
102 """Create rotation about Z-axis."""
103 return cls(rotz(angle))
105 @classmethod
106 def RPY(cls, roll: float, pitch: float, yaw: float) -> 'SO3':
107 """Create rotation from roll-pitch-yaw angles."""
108 return cls(rpy2r(roll, pitch, yaw))
110 @classmethod
111 def Eul(cls, phi: float, theta: float, psi: float, convention: str = 'ZYZ') -> 'SO3':
112 """Create rotation from Euler angles."""
113 return cls(eul2r(phi, theta, psi, convention))
115 @classmethod
116 def Quaternion(cls, q: ArrayLike) -> 'SO3':
117 """Create rotation from quaternion [w, x, y, z]."""
118 return cls(quat2r(np.asarray(q)))
120 @classmethod
121 def Random(cls) -> 'SO3':
122 """Create random rotation matrix."""
123 # Generate random rotation using quaternion method
124 u = np.random.random(3)
125 q = np.array([
126 np.sqrt(1 - u[0]) * np.sin(2 * np.pi * u[1]),
127 np.sqrt(1 - u[0]) * np.cos(2 * np.pi * u[1]),
128 np.sqrt(u[0]) * np.sin(2 * np.pi * u[2]),
129 np.sqrt(u[0]) * np.cos(2 * np.pi * u[2])
130 ])
131 return cls.Quaternion(q)
133 # Conversion methods
134 def rpy(self) -> Tuple[float, float, float]:
135 """Extract roll-pitch-yaw angles."""
136 return r2rpy(self._matrix)
138 def eul(self, convention: str = 'ZYZ') -> Tuple[float, float, float]:
139 """Extract Euler angles."""
140 return r2eul(self._matrix, convention)
142 def quaternion(self) -> np.ndarray:
143 """Extract quaternion [w, x, y, z]."""
144 return r2quat(self._matrix)
146 # Operations
147 def inv(self) -> 'SO3':
148 """Compute inverse (transpose for rotation matrices)."""
149 return SO3(self._matrix.T)
151 def __mul__(self, other: 'SO3') -> 'SO3':
152 """Compose rotations: self * other."""
153 if not isinstance(other, SO3):
154 raise TypeError("Can only multiply SO3 with SO3")
155 return SO3(self._matrix @ other._matrix)
157 def __rmul__(self, other):
158 """Right multiplication (not supported for SO3)."""
159 raise TypeError("Right multiplication not supported")
161 def __pow__(self, n: int) -> 'SO3':
162 """Power operation (repeated multiplication)."""
163 if not isinstance(n, int):
164 raise TypeError("Exponent must be integer")
166 if n == 0:
167 return SO3() # Identity
168 elif n > 0:
169 result = SO3(self._matrix)
170 for _ in range(n - 1):
171 result = result * self
172 return result
173 else:
174 return (self.inv() ** (-n))
176 def rotate(self, vector: Vector3) -> np.ndarray:
177 """Rotate a 3D vector."""
178 v = np.asarray(vector)
179 if v.shape != (3,):
180 raise ValueError("Vector must be 3-element")
181 return self._matrix @ v
183 def __repr__(self) -> str:
184 """String representation."""
185 return f"SO3(\n{self._matrix})"
187 def __str__(self) -> str:
188 """Compact string representation."""
189 rpy = self.rpy()
190 return f"SO3: RPY=({rpy[0]:.3f}, {rpy[1]:.3f}, {rpy[2]:.3f}) rad"
193class SE3:
194 """
195 Special Euclidean group SE(3) - rigid body transformations.
197 Represents rigid body transformations (rotation + translation) in 3D space
198 using 4x4 homogeneous transformation matrices.
200 Examples:
201 >>> T1 = SE3.Trans(1, 2, 3) # Pure translation
202 >>> T2 = SE3.Rx(np.pi/2) # Pure rotation
203 >>> T3 = T1 * T2 # Composition
204 >>> p_new = T3 * [0, 0, 0] # Transform point
205 """
207 def __init__(self, matrix: Optional[Matrix4x4] = None):
208 """
209 Initialize SE(3) transformation.
211 Args:
212 matrix: 4x4 homogeneous transformation matrix (default: identity)
214 Examples:
215 >>> T1 = SE3() # Identity transformation
216 >>> T2 = SE3(np.eye(4)) # From matrix
217 >>> T3 = SE3.Trans(1, 2, 3) # Factory method
218 """
219 if matrix is None:
220 self._matrix = np.eye(4)
221 else:
222 matrix = np.asarray(matrix)
223 if not is_homogeneous(matrix):
224 raise ValueError("Input is not a valid SE(3) transformation matrix")
225 self._matrix = matrix.copy()
227 @property
228 def matrix(self) -> Matrix4x4:
229 """Get the transformation matrix."""
230 return self._matrix.copy()
232 @property
233 def T(self) -> Matrix4x4:
234 """Alias for matrix property."""
235 return self.matrix
237 @property
238 def R(self) -> Matrix3x3:
239 """Get rotation part as 3x3 matrix."""
240 return extract_rotation(self._matrix)
242 @property
243 def t(self) -> np.ndarray:
244 """Get translation part as 3-element vector."""
245 return extract_translation(self._matrix)
247 @property
248 def rotation(self) -> SO3:
249 """Get rotation part as SO3 object."""
250 return SO3(self.R)
252 @property
253 def translation(self) -> np.ndarray:
254 """Alias for t property."""
255 return self.t
257 # Factory methods
258 @classmethod
259 def Trans(cls, x: ArrayLike, y: Optional[float] = None, z: Optional[float] = None) -> 'SE3':
260 """Create pure translation transformation."""
261 return cls(transl(x, y, z))
263 @classmethod
264 def Rx(cls, angle: float) -> 'SE3':
265 """Create rotation about X-axis."""
266 return cls(rotmat(rotx(angle)))
268 @classmethod
269 def Ry(cls, angle: float) -> 'SE3':
270 """Create rotation about Y-axis."""
271 return cls(rotmat(roty(angle)))
273 @classmethod
274 def Rz(cls, angle: float) -> 'SE3':
275 """Create rotation about Z-axis."""
276 return cls(rotmat(rotz(angle)))
278 @classmethod
279 def RPY(cls, roll: float, pitch: float, yaw: float,
280 t: Optional[Vector3] = None) -> 'SE3':
281 """Create transformation from roll-pitch-yaw angles and optional translation."""
282 R = rpy2r(roll, pitch, yaw)
283 return cls(rotmat(R, t))
285 @classmethod
286 def Eul(cls, phi: float, theta: float, psi: float,
287 t: Optional[Vector3] = None, convention: str = 'ZYZ') -> 'SE3':
288 """Create transformation from Euler angles and optional translation."""
289 R = eul2r(phi, theta, psi, convention)
290 return cls(rotmat(R, t))
292 @classmethod
293 def Rt(cls, R: Union[Matrix3x3, SO3], t: Vector3) -> 'SE3':
294 """Create transformation from rotation and translation."""
295 if isinstance(R, SO3):
296 R = R.matrix
297 return cls(rotmat(R, t))
299 @classmethod
300 def Random(cls) -> 'SE3':
301 """Create random SE(3) transformation."""
302 R = SO3.Random()
303 t = np.random.random(3) * 10 - 5 # Random translation in [-5, 5]
304 return cls.Rt(R, t)
306 # Conversion methods
307 def rpy(self) -> Tuple[float, float, float]:
308 """Extract roll-pitch-yaw angles of rotation part."""
309 return r2rpy(self.R)
311 def eul(self, convention: str = 'ZYZ') -> Tuple[float, float, float]:
312 """Extract Euler angles of rotation part."""
313 return r2eul(self.R, convention)
315 # Operations
316 def inv(self) -> 'SE3':
317 """Compute inverse transformation."""
318 return SE3(homogeneous_inverse(self._matrix))
320 def __mul__(self, other: Union['SE3', Vector3, np.ndarray]) -> Union['SE3', np.ndarray]:
321 """
322 Multiply with SE3 (composition) or transform points.
324 Args:
325 other: SE3 transformation or 3D point(s)
327 Returns:
328 SE3 for composition, ndarray for point transformation
329 """
330 if isinstance(other, SE3):
331 # SE3 composition
332 return SE3(self._matrix @ other._matrix)
333 else:
334 # Point transformation
335 other = np.asarray(other)
336 if other.shape == (3,):
337 # Single point
338 return transform_point(self._matrix, other)
339 elif other.ndim == 2 and other.shape[1] == 3:
340 # Multiple points
341 return transform_points(self._matrix, other)
342 else:
343 raise ValueError("Can only transform 3D points or point arrays")
345 def __rmul__(self, other):
346 """Right multiplication (not supported for SE3)."""
347 raise TypeError("Right multiplication not supported")
349 def __pow__(self, n: int) -> 'SE3':
350 """Power operation (repeated multiplication)."""
351 if not isinstance(n, int):
352 raise TypeError("Exponent must be integer")
354 if n == 0:
355 return SE3() # Identity
356 elif n > 0:
357 result = SE3(self._matrix)
358 for _ in range(n - 1):
359 result = result * self
360 return result
361 else:
362 return (self.inv() ** (-n))
364 def translate(self, t: Vector3) -> 'SE3':
365 """Apply additional translation."""
366 return SE3.Trans(t) * self
368 def rotate(self, R: Union[Matrix3x3, SO3]) -> 'SE3':
369 """Apply additional rotation."""
370 if isinstance(R, SO3):
371 R = R.matrix
372 return SE3(rotmat(R)) * self
374 def distance(self, other: 'SE3') -> float:
375 """Compute distance to another SE3 transformation."""
376 if not isinstance(other, SE3):
377 raise TypeError("Can only compute distance to another SE3")
379 # Euclidean distance between translations
380 dt = np.linalg.norm(self.t - other.t)
382 # Angular distance between rotations
383 R_diff = self.R.T @ other.R
384 trace = np.trace(R_diff)
385 dR = np.arccos(np.clip((trace - 1) / 2, -1, 1))
387 return dt + dR # Simple combination
389 def __repr__(self) -> str:
390 """String representation."""
391 return f"SE3(\n{self._matrix})"
393 def __str__(self) -> str:
394 """Compact string representation."""
395 t = self.t
396 rpy = self.rpy()
397 return (f"SE3: t=({t[0]:.3f}, {t[1]:.3f}, {t[2]:.3f}), "
398 f"RPY=({rpy[0]:.3f}, {rpy[1]:.3f}, {rpy[2]:.3f})")
401# Convenience functions
402def SE3_identity() -> SE3:
403 """Create identity SE(3) transformation."""
404 return SE3()
406def SO3_identity() -> SO3:
407 """Create identity SO(3) rotation."""
408 return SO3()