refactor SO3 and remove functorch context for log_map_impl() and to_q…

…uaternion()

theseus/geometry/so3.py

@@ -281,44 +281,23 @@ def _log_map_impl(
         )
         ret = sine_axis * scale.view(-1, 1)
 
-        if _FunctorchContext.get_context():
-            # # theta ~ pi
-            ddiag = torch.diagonal(self.tensor, dim1=1, dim2=2)
-            # Find the index of major coloumns and diagonals
-            major = torch.logical_and(
-                ddiag[:, 1] > ddiag[:, 0], ddiag[:, 1] > ddiag[:, 2]
-            ) + 2 * torch.logical_and(
-                ddiag[:, 2] > ddiag[:, 0], ddiag[:, 2] > ddiag[:, 1]
-            )
-            aux = torch.ones(self.shape[0], dtype=torch.bool)
-            sel_rows = 0.5 * (self[aux, major] + self[aux, :, major])
-            sel_rows[aux, major] -= cosine
-            axis = sel_rows / torch.where(
-                near_zero,
-                non_zero,
-                sel_rows.norm(dim=1),
-            ).view(-1, 1)
-            sign_tmp = sine_axis[aux, major].sign()
-            sign = torch.where(sign_tmp != 0, sign_tmp, torch.ones_like(sign_tmp))
-            ret = torch.where(
-                near_pi.view(-1, 1), axis * (theta * sign).view(-1, 1), ret
-            )
-        else:
-            if near_pi.any():
-                ddiag = torch.diagonal(self[near_pi], dim1=1, dim2=2)
-                # Find the index of major coloumns and diagonals
-                major = torch.logical_and(
-                    ddiag[:, 1] > ddiag[:, 0], ddiag[:, 1] > ddiag[:, 2]
-                ) + 2 * torch.logical_and(
-                    ddiag[:, 2] > ddiag[:, 0], ddiag[:, 2] > ddiag[:, 1]
-                )
-                sel_rows = 0.5 * (self[near_pi, major] + self[near_pi, :, major])
-                aux = torch.ones(sel_rows.shape[0], dtype=torch.bool)
-                sel_rows[aux, major] -= cosine[near_pi]
-                axis = sel_rows / sel_rows.norm(dim=1, keepdim=True)
-                sign_tmp = sine_axis[near_pi, major].sign()
-                sign = torch.where(sign_tmp != 0, sign_tmp, torch.ones_like(sign_tmp))
-                ret[near_pi] = axis * (theta[near_pi] * sign).view(-1, 1)
+        # # theta ~ pi
+        ddiag = torch.diagonal(self.tensor, dim1=1, dim2=2)
+        # Find the index of major coloumns and diagonals
+        major = torch.logical_and(
+            ddiag[:, 1] > ddiag[:, 0], ddiag[:, 1] > ddiag[:, 2]
+        ) + 2 * torch.logical_and(ddiag[:, 2] > ddiag[:, 0], ddiag[:, 2] > ddiag[:, 1])
+        aux = torch.ones(self.shape[0], dtype=torch.bool)
+        sel_rows = 0.5 * (self[aux, major] + self[aux, :, major])
+        sel_rows[aux, major] -= cosine
+        axis = sel_rows / torch.where(
+            near_zero,
+            non_zero,
+            sel_rows.norm(dim=1),
+        ).view(-1, 1)
+        sign_tmp = sine_axis[aux, major].sign()
+        sign = torch.where(sign_tmp != 0, sign_tmp, torch.ones_like(sign_tmp))
+        ret = torch.where(near_pi.view(-1, 1), axis * (theta * sign).view(-1, 1), ret)
 
         if jacobians is not None:
             SO3._check_jacobians_list(jacobians)
@@ -388,58 +367,29 @@ def to_quaternion(self) -> torch.Tensor:
         ret[:, 0] = w
         ret[:, 1:] = 0.5 * sine_axis / torch.where(near_pi, non_zero, w).view(-1, 1)
 
-        if _FunctorchContext.get_context():
-            # theta ~ pi
-            ddiag = torch.diagonal(self.tensor, dim1=1, dim2=2)
-            # Find the index of major coloumns and diagonals
-            major = torch.logical_and(
-                ddiag[:, 1] > ddiag[:, 0], ddiag[:, 1] > ddiag[:, 2]
-            ) + 2 * torch.logical_and(
-                ddiag[:, 2] > ddiag[:, 0], ddiag[:, 2] > ddiag[:, 1]
-            )
-            aux = torch.ones(self.shape[0], dtype=torch.bool)
-            sel_rows = 0.5 * (self[aux, major] + self[aux, :, major])
-            cosine_near_pi = 0.5 * (self[:, 0, 0] + self[:, 1, 1] + self[:, 2, 2] - 1)
-            sel_rows[aux, major] -= cosine_near_pi
-            axis = (
-                sel_rows
-                / torch.where(
-                    near_zero.view(-1, 1),
-                    non_zero.view(-1, 1),
-                    sel_rows.norm(dim=1, keepdim=True),
-                )
-                * sine_axis[aux, major].sign().view(-1, 1)
-            )
-            sine_half_theta = (
-                (0.5 * (1 - cosine_near_pi)).clamp(0, 1).sqrt().view(-1, 1)
-            )
-            ret[:, 1:] = torch.where(
-                near_pi.view(-1, 1), axis * sine_half_theta, ret[:, 1:]
-            )
-        else:
-            # theta ~ pi
-            ddiag = torch.diagonal(self[near_pi], dim1=1, dim2=2)
-            # Find the index of major coloumns and diagonals
-            major = torch.logical_and(
-                ddiag[:, 1] > ddiag[:, 0], ddiag[:, 1] > ddiag[:, 2]
-            ) + 2 * torch.logical_and(
-                ddiag[:, 2] > ddiag[:, 0], ddiag[:, 2] > ddiag[:, 1]
+        # theta ~ pi
+        ddiag = torch.diagonal(self.tensor, dim1=1, dim2=2)
+        # Find the index of major coloumns and diagonals
+        major = torch.logical_and(
+            ddiag[:, 1] > ddiag[:, 0], ddiag[:, 1] > ddiag[:, 2]
+        ) + 2 * torch.logical_and(ddiag[:, 2] > ddiag[:, 0], ddiag[:, 2] > ddiag[:, 1])
+        aux = torch.ones(self.shape[0], dtype=torch.bool)
+        sel_rows = 0.5 * (self[aux, major] + self[aux, :, major])
+        cosine_near_pi = 0.5 * (self[:, 0, 0] + self[:, 1, 1] + self[:, 2, 2] - 1)
+        sel_rows[aux, major] -= cosine_near_pi
+        axis = (
+            sel_rows
+            / torch.where(
+                near_zero.view(-1, 1),
+                non_zero.view(-1, 1),
+                sel_rows.norm(dim=1, keepdim=True),
             )
-            sel_rows = 0.5 * (self[near_pi, major] + self[near_pi, :, major])
-            aux = torch.ones(sel_rows.shape[0], dtype=torch.bool)
-            cosine_near_pi = 0.5 * (
-                self[near_pi, 0, 0] + self[near_pi, 1, 1] + self[near_pi, 2, 2] - 1
-            )
-            sel_rows[aux, major] -= cosine_near_pi
-            axis = (
-                sel_rows
-                / sel_rows.norm(dim=1, keepdim=True)
-                * sine_axis[near_pi, major].sign().view(-1, 1)
-            )
-            sine_half_theta = (
-                (0.5 * (1 - cosine_near_pi)).clamp(0, 1).sqrt().view(-1, 1)
-            )
-            ret[near_pi, 1:] = axis * sine_half_theta
+            * sine_axis[aux, major].sign().view(-1, 1)
+        )
+        sine_half_theta = (0.5 * (1 - cosine_near_pi)).clamp(0, 1).sqrt().view(-1, 1)
+        ret[:, 1:] = torch.where(
+            near_pi.view(-1, 1), axis * sine_half_theta, ret[:, 1:]
+        )
 
         return ret
 
@@ -579,6 +529,180 @@ def unrotate(
 
         return ret
 
+    def _deprecated_log_map_impl(
+        self, jacobians: Optional[List[torch.Tensor]] = None
+    ) -> torch.Tensor:
+        sine_axis = self.tensor.new_zeros(self.shape[0], 3)
+        sine_axis[:, 0] = 0.5 * (self[:, 2, 1] - self[:, 1, 2])
+        sine_axis[:, 1] = 0.5 * (self[:, 0, 2] - self[:, 2, 0])
+        sine_axis[:, 2] = 0.5 * (self[:, 1, 0] - self[:, 0, 1])
+        cosine = 0.5 * (self[:, 0, 0] + self[:, 1, 1] + self[:, 2, 2] - 1)
+        sine = sine_axis.norm(dim=1)
+        theta = torch.atan2(sine, cosine)
+
+        near_zero = theta < self._NEAR_ZERO_EPS
+
+        near_pi = 1 + cosine <= self._NEAR_PI_EPS
+        # theta != pi
+        near_zero_or_near_pi = torch.logical_or(near_zero, near_pi)
+        # Compute the approximation of theta / sin(theta) when theta is near to 0
+        non_zero = torch.ones(1, dtype=self.dtype, device=self.device)
+        sine_nz = torch.where(near_zero_or_near_pi, non_zero, sine)
+        scale = torch.where(
+            near_zero_or_near_pi,
+            1 + sine**2 / 6,
+            theta / sine_nz,
+        )
+        ret = sine_axis * scale.view(-1, 1)
+
+        if _FunctorchContext.get_context():
+            # # theta ~ pi
+            ddiag = torch.diagonal(self.tensor, dim1=1, dim2=2)
+            # Find the index of major coloumns and diagonals
+            major = torch.logical_and(
+                ddiag[:, 1] > ddiag[:, 0], ddiag[:, 1] > ddiag[:, 2]
+            ) + 2 * torch.logical_and(
+                ddiag[:, 2] > ddiag[:, 0], ddiag[:, 2] > ddiag[:, 1]
+            )
+            aux = torch.ones(self.shape[0], dtype=torch.bool)
+            sel_rows = 0.5 * (self[aux, major] + self[aux, :, major])
+            sel_rows[aux, major] -= cosine
+            axis = sel_rows / torch.where(
+                near_zero,
+                non_zero,
+                sel_rows.norm(dim=1),
+            ).view(-1, 1)
+            sign_tmp = sine_axis[aux, major].sign()
+            sign = torch.where(sign_tmp != 0, sign_tmp, torch.ones_like(sign_tmp))
+            ret = torch.where(
+                near_pi.view(-1, 1), axis * (theta * sign).view(-1, 1), ret
+            )
+        else:
+            if near_pi.any():
+                ddiag = torch.diagonal(self[near_pi], dim1=1, dim2=2)
+                # Find the index of major coloumns and diagonals
+                major = torch.logical_and(
+                    ddiag[:, 1] > ddiag[:, 0], ddiag[:, 1] > ddiag[:, 2]
+                ) + 2 * torch.logical_and(
+                    ddiag[:, 2] > ddiag[:, 0], ddiag[:, 2] > ddiag[:, 1]
+                )
+                sel_rows = 0.5 * (self[near_pi, major] + self[near_pi, :, major])
+                aux = torch.ones(sel_rows.shape[0], dtype=torch.bool)
+                sel_rows[aux, major] -= cosine[near_pi]
+                axis = sel_rows / sel_rows.norm(dim=1, keepdim=True)
+                sign_tmp = sine_axis[near_pi, major].sign()
+                sign = torch.where(sign_tmp != 0, sign_tmp, torch.ones_like(sign_tmp))
+                ret[near_pi] = axis * (theta[near_pi] * sign).view(-1, 1)
+
+        if jacobians is not None:
+            SO3._check_jacobians_list(jacobians)
+
+            theta2 = theta**2
+            sine_theta = sine * theta
+            two_cosine_minus_two = 2 * cosine - 2
+            two_cosine_minus_two_nz = torch.where(
+                near_zero, non_zero, two_cosine_minus_two
+            )
+            theta2_nz = torch.where(near_zero, non_zero, theta2)
+
+            a = torch.where(
+                near_zero, 1 - theta2 / 12, -sine_theta / two_cosine_minus_two_nz
+            )
+            b = torch.where(
+                near_zero,
+                1.0 / 12 + theta2 / 720,
+                (sine_theta + two_cosine_minus_two)
+                / (theta2_nz * two_cosine_minus_two_nz),
+            )
+
+            jac = (b.view(-1, 1) * ret).view(-1, 3, 1) * ret.view(-1, 1, 3)
+
+            half_ret = 0.5 * ret
+            jac[:, 0, 1] -= half_ret[:, 2]
+            jac[:, 1, 0] += half_ret[:, 2]
+            jac[:, 0, 2] += half_ret[:, 1]
+            jac[:, 2, 0] -= half_ret[:, 1]
+            jac[:, 1, 2] -= half_ret[:, 0]
+            jac[:, 2, 1] += half_ret[:, 0]
+
+            diag_jac = torch.diagonal(jac, dim1=1, dim2=2)
+            diag_jac += a.view(-1, 1)
+
+            jacobians.append(jac)
+
+        return ret
+
+    def _deprecated_to_quaternion(self) -> torch.Tensor:
+        sine_axis = self.tensor.new_zeros(self.shape[0], 3)
+        sine_axis[:, 0] = 0.5 * (self[:, 2, 1] - self[:, 1, 2])
+        sine_axis[:, 1] = 0.5 * (self[:, 0, 2] - self[:, 2, 0])
+        sine_axis[:, 2] = 0.5 * (self[:, 1, 0] - self[:, 0, 1])
+        w = 0.5 * (1 + self[:, 0, 0] + self[:, 1, 1] + self[:, 2, 2]).clamp(0, 4).sqrt()
+
+        near_zero = w > 1 - self._NEAR_ZERO_EPS
+        near_pi = w <= self._NEAR_PI_EPS
+        non_zero = self.tensor.new_ones([1])
+
+        ret = self.tensor.new_zeros(self.shape[0], 4)
+        # theta != pi
+        ret[:, 0] = w
+        ret[:, 1:] = 0.5 * sine_axis / torch.where(near_pi, non_zero, w).view(-1, 1)
+
+        if _FunctorchContext.get_context():
+            # theta ~ pi
+            ddiag = torch.diagonal(self.tensor, dim1=1, dim2=2)
+            # Find the index of major coloumns and diagonals
+            major = torch.logical_and(
+                ddiag[:, 1] > ddiag[:, 0], ddiag[:, 1] > ddiag[:, 2]
+            ) + 2 * torch.logical_and(
+                ddiag[:, 2] > ddiag[:, 0], ddiag[:, 2] > ddiag[:, 1]
+            )
+            aux = torch.ones(self.shape[0], dtype=torch.bool)
+            sel_rows = 0.5 * (self[aux, major] + self[aux, :, major])
+            cosine_near_pi = 0.5 * (self[:, 0, 0] + self[:, 1, 1] + self[:, 2, 2] - 1)
+            sel_rows[aux, major] -= cosine_near_pi
+            axis = (
+                sel_rows
+                / torch.where(
+                    near_zero.view(-1, 1),
+                    non_zero.view(-1, 1),
+                    sel_rows.norm(dim=1, keepdim=True),
+                )
+                * sine_axis[aux, major].sign().view(-1, 1)
+            )
+            sine_half_theta = (
+                (0.5 * (1 - cosine_near_pi)).clamp(0, 1).sqrt().view(-1, 1)
+            )
+            ret[:, 1:] = torch.where(
+                near_pi.view(-1, 1), axis * sine_half_theta, ret[:, 1:]
+            )
+        else:
+            # theta ~ pi
+            ddiag = torch.diagonal(self[near_pi], dim1=1, dim2=2)
+            # Find the index of major coloumns and diagonals
+            major = torch.logical_and(
+                ddiag[:, 1] > ddiag[:, 0], ddiag[:, 1] > ddiag[:, 2]
+            ) + 2 * torch.logical_and(
+                ddiag[:, 2] > ddiag[:, 0], ddiag[:, 2] > ddiag[:, 1]
+            )
+            sel_rows = 0.5 * (self[near_pi, major] + self[near_pi, :, major])
+            aux = torch.ones(sel_rows.shape[0], dtype=torch.bool)
+            cosine_near_pi = 0.5 * (
+                self[near_pi, 0, 0] + self[near_pi, 1, 1] + self[near_pi, 2, 2] - 1
+            )
+            sel_rows[aux, major] -= cosine_near_pi
+            axis = (
+                sel_rows
+                / sel_rows.norm(dim=1, keepdim=True)
+                * sine_axis[near_pi, major].sign().view(-1, 1)
+            )
+            sine_half_theta = (
+                (0.5 * (1 - cosine_near_pi)).clamp(0, 1).sqrt().view(-1, 1)
+            )
+            ret[near_pi, 1:] = axis * sine_half_theta
+
+        return ret
+
 
 rand_so3 = SO3.rand
 randn_so3 = SO3.randn