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

1""" 

2SE(3) and SO(3) classes for robotics-numpy 

3 

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 

7 

8These classes provide intuitive APIs while maintaining high performance through NumPy. 

9""" 

10 

11from typing import List, Optional, Tuple, Union 

12 

13import numpy as np 

14 

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) 

37 

38# Type aliases 

39ArrayLike = Union[float, int, np.ndarray, List[float]] 

40Vector3 = Union[np.ndarray, List[float]] 

41Matrix3x3 = np.ndarray 

42Matrix4x4 = np.ndarray 

43 

44 

45class SO3: 

46 """ 

47 Special Orthogonal group SO(3) - 3D rotations. 

48 

49 Represents rotations in 3D space using rotation matrices. 

50 Provides convenient methods for creation, composition, and conversion. 

51 

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

58 

59 def __init__(self, matrix: Optional[Matrix3x3] = None): 

60 """ 

61 Initialize SO(3) rotation. 

62 

63 Args: 

64 matrix: 3x3 rotation matrix (default: identity) 

65 

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

78 

79 @property 

80 def matrix(self) -> Matrix3x3: 

81 """Get the rotation matrix.""" 

82 return self._matrix.copy() 

83 

84 @property 

85 def R(self) -> Matrix3x3: 

86 """Alias for matrix property.""" 

87 return self.matrix 

88 

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

94 

95 @classmethod 

96 def Ry(cls, angle: float) -> 'SO3': 

97 """Create rotation about Y-axis.""" 

98 return cls(roty(angle)) 

99 

100 @classmethod 

101 def Rz(cls, angle: float) -> 'SO3': 

102 """Create rotation about Z-axis.""" 

103 return cls(rotz(angle)) 

104 

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

109 

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

114 

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

119 

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) 

132 

133 # Conversion methods 

134 def rpy(self) -> Tuple[float, float, float]: 

135 """Extract roll-pitch-yaw angles.""" 

136 return r2rpy(self._matrix) 

137 

138 def eul(self, convention: str = 'ZYZ') -> Tuple[float, float, float]: 

139 """Extract Euler angles.""" 

140 return r2eul(self._matrix, convention) 

141 

142 def quaternion(self) -> np.ndarray: 

143 """Extract quaternion [w, x, y, z].""" 

144 return r2quat(self._matrix) 

145 

146 # Operations 

147 def inv(self) -> 'SO3': 

148 """Compute inverse (transpose for rotation matrices).""" 

149 return SO3(self._matrix.T) 

150 

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) 

156 

157 def __rmul__(self, other): 

158 """Right multiplication (not supported for SO3).""" 

159 raise TypeError("Right multiplication not supported") 

160 

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

165 

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

175 

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 

182 

183 def __repr__(self) -> str: 

184 """String representation.""" 

185 return f"SO3(\n{self._matrix})" 

186 

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" 

191 

192 

193class SE3: 

194 """ 

195 Special Euclidean group SE(3) - rigid body transformations. 

196 

197 Represents rigid body transformations (rotation + translation) in 3D space 

198 using 4x4 homogeneous transformation matrices. 

199 

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

206 

207 def __init__(self, matrix: Optional[Matrix4x4] = None): 

208 """ 

209 Initialize SE(3) transformation. 

210 

211 Args: 

212 matrix: 4x4 homogeneous transformation matrix (default: identity) 

213 

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

226 

227 @property 

228 def matrix(self) -> Matrix4x4: 

229 """Get the transformation matrix.""" 

230 return self._matrix.copy() 

231 

232 @property 

233 def T(self) -> Matrix4x4: 

234 """Alias for matrix property.""" 

235 return self.matrix 

236 

237 @property 

238 def R(self) -> Matrix3x3: 

239 """Get rotation part as 3x3 matrix.""" 

240 return extract_rotation(self._matrix) 

241 

242 @property 

243 def t(self) -> np.ndarray: 

244 """Get translation part as 3-element vector.""" 

245 return extract_translation(self._matrix) 

246 

247 @property 

248 def rotation(self) -> SO3: 

249 """Get rotation part as SO3 object.""" 

250 return SO3(self.R) 

251 

252 @property 

253 def translation(self) -> np.ndarray: 

254 """Alias for t property.""" 

255 return self.t 

256 

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

262 

263 @classmethod 

264 def Rx(cls, angle: float) -> 'SE3': 

265 """Create rotation about X-axis.""" 

266 return cls(rotmat(rotx(angle))) 

267 

268 @classmethod 

269 def Ry(cls, angle: float) -> 'SE3': 

270 """Create rotation about Y-axis.""" 

271 return cls(rotmat(roty(angle))) 

272 

273 @classmethod 

274 def Rz(cls, angle: float) -> 'SE3': 

275 """Create rotation about Z-axis.""" 

276 return cls(rotmat(rotz(angle))) 

277 

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

284 

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

291 

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

298 

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) 

305 

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) 

310 

311 def eul(self, convention: str = 'ZYZ') -> Tuple[float, float, float]: 

312 """Extract Euler angles of rotation part.""" 

313 return r2eul(self.R, convention) 

314 

315 # Operations 

316 def inv(self) -> 'SE3': 

317 """Compute inverse transformation.""" 

318 return SE3(homogeneous_inverse(self._matrix)) 

319 

320 def __mul__(self, other: Union['SE3', Vector3, np.ndarray]) -> Union['SE3', np.ndarray]: 

321 """ 

322 Multiply with SE3 (composition) or transform points. 

323 

324 Args: 

325 other: SE3 transformation or 3D point(s) 

326 

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

344 

345 def __rmul__(self, other): 

346 """Right multiplication (not supported for SE3).""" 

347 raise TypeError("Right multiplication not supported") 

348 

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

353 

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

363 

364 def translate(self, t: Vector3) -> 'SE3': 

365 """Apply additional translation.""" 

366 return SE3.Trans(t) * self 

367 

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 

373 

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

378 

379 # Euclidean distance between translations 

380 dt = np.linalg.norm(self.t - other.t) 

381 

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

386 

387 return dt + dR # Simple combination 

388 

389 def __repr__(self) -> str: 

390 """String representation.""" 

391 return f"SE3(\n{self._matrix})" 

392 

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

399 

400 

401# Convenience functions 

402def SE3_identity() -> SE3: 

403 """Create identity SE(3) transformation.""" 

404 return SE3() 

405 

406def SO3_identity() -> SO3: 

407 """Create identity SO(3) rotation.""" 

408 return SO3()