add SE3.transform_from and SE3.transform_to (#80)

* add SE3 defintions

* add initialization method

* modify test_so3.check_SO3_to_quaternion to handle cases near pi

* add SE3.hat and SE.vee

* add SE3.adjoint

* add SE3.adjoint, vee and hat

* backup before switching to taoshaf.add_SE3

* add support for single x_y_z_quaternion

* modify SO3

* Remove repeated if in SE3 construction

* SE3.exp added but not tested

* add SE3.log_map()

* SE3.exp_map() and SE3.log_map() are tested

* add SE3.compose()

* simplify test_se3.py

* modify test_se3.py

* add SE3.transform_from and SE3.transform_to

* add SE3.transform_from and SE3.transform_to

* change err_msgs for SO3._rotate_shape_check and SE3._rotate_shape_check

* change error message for SO3._rotate_shape_check and SE3._transform_shape_check

* update err_msg for shape_check

theseus/geometry/se3.py

@@ -286,26 +286,90 @@ def vee(matrix: torch.Tensor) -> torch.Tensor:
         SE3._hat_matrix_check(matrix)
         return torch.cat((matrix[:, :3, 3], SO3.vee(matrix[:, :3, :3])), dim=1)
 
-    def _transform_shape_check(self, point: Union[Point3, torch.Tensor]):
-        raise NotImplementedError
-
     def _copy_impl(self, new_name: Optional[str] = None) -> "SE3":
         return SE3(data=self.data.clone(), name=new_name)
 
     # only added to avoid casting downstream
     def copy(self, new_name: Optional[str] = None) -> "SE3":
         return cast(SE3, super().copy(new_name=new_name))
 
+    def _transform_shape_check(self, point: Union[Point3, torch.Tensor]):
+        err_msg = (
+            f"SE3 can only transform vectors of shape [{self.shape[0]}, 3] or [1, 3], "
+            f"but the input has shape {point.shape}."
+        )
+
+        if isinstance(point, torch.Tensor):
+            if not point.ndim == 2 or point.shape[1] != 3:
+                raise ValueError(err_msg)
+        elif point.dof() != 3:
+            raise ValueError(err_msg)
+        if (
+            point.shape[0] != self.shape[0]
+            and point.shape[0] != 1
+            and self.shape[0] != 1
+        ):
+            raise ValueError(err_msg)
+
     def transform_to(
         self,
         point: Union[Point3, torch.Tensor],
         jacobians: Optional[List[torch.Tensor]] = None,
     ) -> Point3:
-        raise NotImplementedError
+        self._transform_shape_check(point)
+        batch_size = max(self.shape[0], point.shape[0])
+        if isinstance(point, torch.Tensor):
+            p = point.view(-1, 3, 1)
+        else:
+            p = point.data.view(-1, 3, 1)
+
+        ret = Point3(data=(self[:, :, :3] @ p).view(-1, 3))
+        ret += self[:, :, 3]
+
+        if jacobians is not None:
+            self._check_jacobians_list(jacobians)
+            # Right jacobians for SE(3) are computed
+            Jg = torch.zeros(batch_size, 3, 6, dtype=self.dtype, device=self.device)
+            Jg[:, :, :3] = self[:, :, :3]
+            Jg[:, :, 3:] = -self[:, :, :3] @ SO3.hat(p)
+            # Jacobians for point
+            Jpnt = Jg[:, :, :3]
+
+            jacobians.extend([Jg, Jpnt])
+
+        return ret
 
     def transform_from(
         self,
         point: Union[Point3, torch.Tensor],
         jacobians: Optional[List[torch.Tensor]] = None,
     ) -> Point3:
-        raise NotImplementedError
+        self._transform_shape_check(point)
+        batch_size = max(self.shape[0], point.shape[0])
+        if isinstance(point, torch.Tensor):
+            p = point.view(-1, 3, 1)
+        else:
+            p = point.data.view(-1, 3, 1)
+
+        temp = p - self[:, :, 3:]
+        ret = Point3(data=(self[:, :, :3].transpose(1, 2) @ temp).view(-1, 3))
+
+        if jacobians is not None:
+            self._check_jacobians_list(jacobians)
+            # Right jacobians for SE(3) are computed
+            Jg = torch.zeros(batch_size, 3, 6, dtype=self.dtype, device=self.device)
+            Jg[:, 0, 0] = -1
+            Jg[:, 1, 1] = -1
+            Jg[:, 2, 2] = -1
+            Jg[:, 0, 4] = -ret[:, 2]
+            Jg[:, 1, 3] = ret[:, 2]
+            Jg[:, 0, 5] = ret[:, 1]
+            Jg[:, 2, 3] = -ret[:, 1]
+            Jg[:, 1, 5] = -ret[:, 0]
+            Jg[:, 2, 4] = ret[:, 0]
+            # Jacobians for point
+            Jpnt = self[:, :, :3].transpose(1, 2).expand(batch_size, 3, 3)
+
+            jacobians.extend([Jg, Jpnt])
+
+        return ret

theseus/geometry/so3.py

@@ -226,7 +226,10 @@ def vee(matrix: torch.Tensor) -> torch.Tensor:
         return torch.stack((matrix[:, 2, 1], matrix[:, 0, 2], matrix[:, 1, 0]), dim=1)
 
     def _rotate_shape_check(self, point: Union[Point3, torch.Tensor]):
-        err_msg = "SO3 can only rotate 3-D vectors."
+        err_msg = (
+            f"SO3 can only rotate vectors of shape [{self.shape[0]}, 3] or [1, 3], "
+            f"but the input has shape {point.shape}."
+        )
         if isinstance(point, torch.Tensor):
             if not point.ndim == 2 or point.shape[1] != 3:
                 raise ValueError(err_msg)
@@ -237,9 +240,7 @@ def _rotate_shape_check(self, point: Union[Point3, torch.Tensor]):
             and point.shape[0] != 1
             and self.shape[0] != 1
         ):
-            raise ValueError(
-                "Input point batch size is not broadcastable with group batch size."
-            )
+            raise ValueError(err_msg)
 
     @staticmethod
     def unit_quaternion_to_SO3(quaternion: torch.torch.Tensor) -> "SO3":
@@ -299,7 +300,7 @@ def rotate(
         ret = Point3(data=(self.data @ p).view(-1, 3))
         if jacobians is not None:
             self._check_jacobians_list(jacobians)
-            # Jacobians for SO3: left-invariant jacobians are computed
+            # Right jacobians for SO(3) are computed
             Jrot = -self.data @ SO3.hat(p)
             # Jacobians for point
             Jpnt = self.to_matrix().expand(batch_size, 3, 3)
@@ -323,7 +324,7 @@ def unrotate(
         ret = Point3(data=(self.data.transpose(1, 2) @ p).view(-1, 3))
         if jacobians is not None:
             self._check_jacobians_list(jacobians)
-            # Jacobians for SO3: left-invariant jacobians are computed
+            # Left jacobians for SO3 are computed
             Jrot = torch.zeros(batch_size, 3, 3, dtype=self.dtype, device=self.device)
             Jrot[:, 0, 1] = -ret[:, 2]
             Jrot[:, 1, 0] = ret[:, 2]

theseus/geometry/tests/test_se3.py

@@ -9,17 +9,18 @@
 
 import theseus as th
 from theseus.constants import EPS
+from theseus.utils import numeric_jacobian
 
 from .common import check_adjoint, check_compose, check_exp_map, check_inverse
 
 
 def _create_random_se3(batch_size, rng):
-    tangent_vector_ang = torch.rand(batch_size, 3) - 0.5
+    tangent_vector_ang = torch.rand(batch_size, 3).double() - 0.5
     tangent_vector_ang /= tangent_vector_ang.norm(dim=1, keepdim=True)
-    tangent_vector_ang *= torch.rand(batch_size, 1) * 2 * np.pi - np.pi
-    tangent_vector_lin = torch.randn(batch_size, 3)
+    tangent_vector_ang *= torch.rand(batch_size, 1).double() * 2 * np.pi - np.pi
+    tangent_vector_lin = torch.randn(batch_size, 3).double()
     tangent_vector = torch.cat([tangent_vector_lin, tangent_vector_ang], dim=1)
-    return th.SE3.exp_map(tangent_vector.double())
+    return th.SE3.exp_map(tangent_vector)
 
 
 def check_SE3_log_map(tangent_vector, atol=EPS):
@@ -149,3 +150,42 @@ def test_adjoint():
         se3 = _create_random_se3(batch_size, rng)
         tangent = torch.randn(batch_size, 6).double()
         check_adjoint(se3, tangent)
+
+
+def test_transform_from_and_to():
+    rng = torch.Generator()
+    rng.manual_seed(0)
+    for _ in range(10):  # repeat a few times
+        for batch_size in [1, 20, 100]:
+            se3 = _create_random_se3(batch_size, rng)
+            point_tensor = torch.randn(batch_size, 3).double()
+            point_tensor_ext = torch.cat(
+                (point_tensor, torch.ones(batch_size, 1).double()), dim=1
+            )
+
+            jacobians_to = []
+            point_to = se3.transform_to(point_tensor, jacobians=jacobians_to)
+            expected_to = (se3.to_matrix() @ point_tensor_ext.unsqueeze(2))[:, :3]
+            jacobians_from = []
+            point_from = se3.transform_from(point_to, jacobians_from)
+
+            # Check the operation result
+            assert torch.allclose(expected_to.squeeze(2), point_to.data, atol=EPS)
+            assert torch.allclose(point_tensor, point_from.data, atol=EPS)
+
+            # Check the jacobians
+            expected_jac = numeric_jacobian(
+                lambda groups: groups[0].transform_to(groups[1]),
+                [se3, th.Point3(point_tensor)],
+                function_dim=3,
+            )
+            assert torch.allclose(jacobians_to[0], expected_jac[0])
+            assert torch.allclose(jacobians_to[1], expected_jac[1])
+            expected_jac = numeric_jacobian(
+                lambda groups: groups[0].transform_from(groups[1]),
+                [se3, point_to],
+                delta_mag=1e-5,
+                function_dim=3,
+            )
+            assert torch.allclose(jacobians_from[0], expected_jac[0])
+            assert torch.allclose(jacobians_from[1], expected_jac[1])