Make pose (+) point composition return a BasePose subtype (#50)

JeffLIrion · Nov 6, 2023 · ffde16f · ffde16f
1 parent 33d7f06
commit ffde16f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 9 deletions.
diff --git a/graphslam/pose/se2.py b/graphslam/pose/se2.py
@@ -171,7 +171,7 @@ def __add__(self, other):
 
         Returns
         -------
-        PoseSE2, np.ndarray
+        PoseSE2, PoseR2
             The result of pose composition
 
         """
@@ -185,9 +185,8 @@ def __add__(self, other):
         if isinstance(other, PoseR2) or (isinstance(other, np.ndarray) and len(other) == 2):
             # pose (+) point
             # fmt: off
-            return np.array([self[0] + other[0] * np.cos(self[2]) - other[1] * np.sin(self[2]),
-                             self[1] + other[0] * np.sin(self[2]) + other[1] * np.cos(self[2])],
-                            dtype=np.float64)
+            return PoseR2([self[0] + other[0] * np.cos(self[2]) - other[1] * np.sin(self[2]),
+                           self[1] + other[0] * np.sin(self[2]) + other[1] * np.cos(self[2])])
             # fmt: on
 
         raise NotImplementedError

diff --git a/graphslam/pose/se3.py b/graphslam/pose/se3.py
@@ -159,7 +159,7 @@ def __add__(self, other):
 
         Returns
         -------
-        PoseSE3, np.ndarray
+        PoseSE3, PoseR3
             The result of pose composition
 
         """
@@ -198,10 +198,9 @@ def __add__(self, other):
         if isinstance(other, PoseR3) or (isinstance(other, np.ndarray) and len(other) == 3):
             # pose (+) point
             # fmt: off
-            return np.array([self[0] + other[0] + 2. * (-(self[4]**2 + self[5]**2) * other[0] + (self[3] * self[4] - self[5] * self[6]) * other[1] + (self[4] * self[6] + self[3] * self[5]) * other[2]),
-                             self[1] + other[1] + 2. * ((self[5] * self[6] + self[3] * self[4]) * other[0] - (self[3]**2 + self[5]**2) * other[1] + (self[4] * self[5] - self[3] * self[6]) * other[2]),
-                             self[2] + other[2] + 2. * ((self[3] * self[5] - self[4] * self[6]) * other[0] + (self[3] * self[6] + self[4] * self[5]) * other[1] - (self[3]**2 + self[4]**2) * other[2])],
-                            dtype=np.float64)
+            return PoseR3([self[0] + other[0] + 2. * (-(self[4]**2 + self[5]**2) * other[0] + (self[3] * self[4] - self[5] * self[6]) * other[1] + (self[4] * self[6] + self[3] * self[5]) * other[2]),
+                           self[1] + other[1] + 2. * ((self[5] * self[6] + self[3] * self[4]) * other[0] - (self[3]**2 + self[5]**2) * other[1] + (self[4] * self[5] - self[3] * self[6]) * other[2]),
+                           self[2] + other[2] + 2. * ((self[3] * self[5] - self[4] * self[6]) * other[0] + (self[3] * self[6] + self[4] * self[5]) * other[1] - (self[3]**2 + self[4]**2) * other[2])])
             # fmt: on
 
         raise NotImplementedError