add support for single x_y_z_quaternion

theseus/geometry/se3.py

@@ -32,8 +32,6 @@ def __init__(
         super().__init__(data=data, name=name, dtype=dtype)
         if x_y_z_quaternion is not None:
             if x_y_z_quaternion is not None:
-                if x_y_z_quaternion.ndim == 1:
-                    x_y_z_quaternion = x_y_z_quaternion.unsqueeze(0)
                 self.update_from_x_y_z_quaternion(x_y_z_quaternion=x_y_z_quaternion)
 
     @staticmethod
@@ -64,6 +62,25 @@ def _SE3_matrix_check(matrix: torch.Tensor):
             raise ValueError("SE(3) can only be 3x4 matrices.")
         SO3._SO3_matrix_check(matrix.data[:, :3, :3])
 
+    @staticmethod
+    def x_y_z_unit_quaternion_to_SE3(x_y_z_quaternion: torch.Tensor) -> "SE3":
+        if x_y_z_quaternion.ndim == 1:
+            x_y_z_quaternion = x_y_z_quaternion.unsqueeze(0)
+
+        if x_y_z_quaternion.ndim != 2 and x_y_z_quaternion.shape[1] != 7:
+            raise ValueError("x_y_z_quaternion can only be 7-D vectors.")
+
+        ret = SE3()
+
+        batch_size = x_y_z_quaternion.shape[0]
+        ret.data = torch.empty(batch_size, 3, 4).to(
+            device=x_y_z_quaternion.device, dtype=x_y_z_quaternion.dtype
+        )
+        ret[:, :3, :3] = SO3.unit_quaternion_to_SO3(x_y_z_quaternion[:, 3:]).data
+        ret[:, :3, 3] = x_y_z_quaternion[:, :3]
+
+        return ret
+
     @staticmethod
     def _hat_matrix_check(matrix: torch.Tensor):
         if matrix.ndim != 3 or matrix.shape[1:] != (4, 4):
@@ -81,15 +98,8 @@ def _hat_matrix_check(matrix: torch.Tensor):
 
     @staticmethod
     def exp_map(tangent_vector: torch.Tensor) -> LieGroup:
-        if tangent_vector.ndim != 2 or tangent_vector.shape[1] != 6:
-            raise ValueError("The tangent vectors can only be 6-D vectors.")
-        ret = SE3(dtype=tangent_vector.dtype)
-
-        # TODO: Not Implemented yet
         raise NotImplementedError
 
-        return ret
-
     def _log_map_impl(self) -> torch.Tensor:
         raise NotImplementedError
 
@@ -109,15 +119,7 @@ def to_matrix(self) -> torch.Tensor:
         return ret
 
     def update_from_x_y_z_quaternion(self, x_y_z_quaternion: torch.Tensor):
-        if x_y_z_quaternion.ndim != 2 and x_y_z_quaternion.shape[1] != 7:
-            raise ValueError("x_y_z_quaternion can only be 7-D vectors.")
-
-        batch_size = x_y_z_quaternion.shape[0]
-        self.data = torch.empty(batch_size, 3, 4).to(
-            device=x_y_z_quaternion.device, dtype=x_y_z_quaternion.dtype
-        )
-        self[:, :3, :3] = SO3.unit_quaternion_to_matrix(x_y_z_quaternion[:, 3:])
-        self[:, :3, 3] = x_y_z_quaternion[:, :3]
+        self.update(SE3.x_y_z_unit_quaternion_to_SE3(x_y_z_quaternion))
 
     def update_from_rot_and_trans(self, rotation: SO3, translation: Point3):
         if rotation.shape[0] != translation.shape[0]:
@@ -166,7 +168,7 @@ def transform_to(
     ) -> Point3:
         raise NotImplementedError
 
-    def untransform_to(
+    def transform_from(
         self,
         point: Union[Point3, torch.Tensor],
         jacobians: Optional[List[torch.Tensor]] = None,

theseus/geometry/so3.py

@@ -31,16 +31,14 @@ def __init__(
             self._SO3_matrix_check(data)
         super().__init__(data=data, name=name, dtype=dtype)
         if quaternion is not None:
-            if quaternion.ndim == 1:
-                quaternion = quaternion.unsqueeze(0)
             self.update_from_unit_quaternion(quaternion)
 
     @staticmethod
     def _init_data() -> torch.Tensor:  # type: ignore
         return torch.eye(3, 3).view(1, 3, 3)
 
     def update_from_unit_quaternion(self, quaternion: torch.Tensor):
-        self.update(self.unit_quaternion_to_matrix(quaternion))
+        self.update(self.unit_quaternion_to_SO3(quaternion))
 
     def dof(self) -> int:
         return 3
@@ -244,8 +242,11 @@ def _rotate_shape_check(self, point: Union[Point3, torch.Tensor]):
             )
 
     @staticmethod
-    def unit_quaternion_to_matrix(quaternion: torch.torch.Tensor):
+    def unit_quaternion_to_SO3(quaternion: torch.torch.Tensor):
+        if quaternion.ndim == 1:
+            quaternion = quaternion.unsqueeze(0)
         SO3._unit_quaternion_check(quaternion)
+
         q0 = quaternion[:, 0]
         q1 = quaternion[:, 1]
         q2 = quaternion[:, 2]
@@ -260,7 +261,9 @@ def unit_quaternion_to_matrix(quaternion: torch.torch.Tensor):
         q22 = q2 * q2
         q23 = q2 * q3
         q33 = q3 * q3
-        ret = torch.zeros(quaternion.shape[0], 3, 3).to(
+
+        ret = SO3()
+        ret.data = torch.zeros(quaternion.shape[0], 3, 3).to(
             dtype=quaternion.dtype, device=quaternion.device
         )
         ret[:, 0, 0] = 2 * (q00 + q11) - 1