Fix '*_quat_to_xyz' helpers broken by PR#1478. (#1484)

Author: duburcqa

genesis/utils/geom.py

@@ -149,9 +149,7 @@ def ti_quat_to_xyz(quat):
 
         siny_cosp = q_wz - q_xy
         cosy_cosp = 1.0 - (q_yy + q_zz)
-        hypot_min = ti.min(siny_cosp, cosy_cosp)
-        hypot_max = ti.max(siny_cosp, cosy_cosp)
-        cosp = ti.select(hypot_max > 0.0, hypot_max * ti.sqrt(1.0 + (hypot_min / hypot_max) ** 2), 0.0)
+        cosp = ti.sqrt(cosy_cosp**2 + siny_cosp**2)
 
         roll = ti.atan2(q_wx - q_yz, 1.0 - (q_xx + q_yy))
         pitch = ti.atan2(q_xz + q_wy, cosp)
@@ -652,10 +650,7 @@ def _np_quat_to_xyz(quat, rpy=False, out=None):
         siny_cosp = q_wz - q_xy
     cosr_cosp = 1.0 - (q_xx + q_yy)
     cosy_cosp = 1.0 - (q_yy + q_zz)
-
-    hypot_min = np.minimum(siny_cosp, cosy_cosp)
-    hypot_max = np.maximum(siny_cosp, cosy_cosp)
-    cosp = np.where(hypot_max > 0.0, hypot_max * np.sqrt(1.0 + (hypot_min / hypot_max) ** 2), 0.0)
+    cosp = np.sqrt(cosy_cosp**2 + siny_cosp**2)
 
     out_[..., 0] = np.arctan2(sinr_cosp, cosr_cosp)
     out_[..., 1] = np.arctan2(sinp, cosp)
@@ -678,7 +673,9 @@ def _tc_quat_to_xyz(quat, rpy=False, out=None):
     q_yy, q_yz = torch.unbind(q_y * q_vec_s[..., 1:], -1)
     q_zz = q_z[..., 0] * q_vec_s[..., 2]
 
-    # Compute some intermediary quantities
+    # Compute some intermediary quantities.
+    # Numerical robustness of 'cos(pitch)' could be improved using 'hypot' implementation from Eigen:
+    # https://gitlab.com/libeigen/eigen/-/blob/master/Eigen/src/Core/MathFunctionsImpl.h#L149
     if rpy:
         sinp = q_wy - q_xz
         sinr_cosp = q_wx + q_yz
@@ -689,14 +686,7 @@ def _tc_quat_to_xyz(quat, rpy=False, out=None):
         siny_cosp = q_wz - q_xy
     cosr_cosp = 1.0 - (q_xx + q_yy)
     cosy_cosp = 1.0 - (q_yy + q_zz)
-
-    # Use numerical robust formula for computing cos(pitch)
-    # See Eigen implementation for reference:
-    # https://gitlab.com/libeigen/eigen/-/blob/master/Eigen/src/Geometry/EulerAngles.h#L45
-    # https://gitlab.com/libeigen/eigen/-/blob/master/Eigen/src/Core/MathFunctionsImpl.h#L149
-    hypot_min = torch.minimum(siny_cosp, cosy_cosp)
-    hypot_max = torch.maximum(siny_cosp, cosy_cosp)
-    cosp = torch.where(hypot_max > 0.0, hypot_max * torch.sqrt(1.0 + (hypot_min / hypot_max) ** 2), 0.0)
+    cosp = torch.sqrt(cosy_cosp**2 + siny_cosp**2)
 
     # Roll (x-axis rotation)
     out[..., 0] = torch.atan2(sinr_cosp, cosr_cosp)

tests/test_utils.py

@@ -206,14 +206,14 @@ def test_utils_geom_taichi_inverse(batch_shape):
         ti_value_in_args, ti_transform_args, ti_value_out_args, ti_value_inv_out_args = [], [], [], []
         for i, shape_arg in enumerate(map(tuple, (*shapes_in, *shapes_value_args, *shapes_value_args))):
             if shape_arg in ((4, 4), (3, 3)):
-                R = gu.rotvec_to_R(np.random.rand(*batch_shape, 3).astype(gs.np_float))
+                R = gu.rotvec_to_R(np.random.randn(*batch_shape, 3).clip(-1.0, 1.0).astype(gs.np_float))
                 if shape_arg == (4, 4):
-                    trans = np.random.rand(*batch_shape, 3).astype(gs.np_float)
+                    trans = np.random.randn(*batch_shape, 3).astype(gs.np_float)
                     np_arg = gu.trans_R_to_T(trans, R)
                 else:
                     np_arg = R
             else:
-                np_arg = np.random.rand(*batch_shape, *shape_arg).astype(gs.np_float)
+                np_arg = np.random.randn(*batch_shape, *shape_arg).clip(-1.0, 1.0).astype(gs.np_float)
 
             ti_type = ti.Vector if len(shape_arg) == 1 else ti.Matrix
             ti_arg = ti_type.field(*shape_arg, dtype=gs.ti_float, shape=batch_shape)
@@ -251,7 +251,7 @@ def test_utils_geom_taichi_identity(batch_shape):
         for shape_arg in (*shape_args, shape_args[0]):
             ti_type = ti.Vector if len(shape_arg) == 1 else ti.Matrix
             ti_arg = ti_type.field(*shape_arg, dtype=gs.ti_float, shape=batch_shape)
-            ti_arg.from_numpy(np.random.rand(*batch_shape, *shape_arg).astype(gs.np_float))
+            ti_arg.from_numpy(np.random.randn(*batch_shape, *shape_arg).clip(-1.0, 1.0).astype(gs.np_float))
             ti_args.append(ti_arg)
 
         num_funcs = len(ti_funcs)
@@ -273,9 +273,9 @@ def test_utils_geom_tensor_identity(batch_shape):
         np_args, tc_args = [], []
         for shape_arg in (*shape_args, shape_args[0]):
             if tuple(shape_arg) == (3, 3):
-                np_arg = gu.rotvec_to_R(np.random.rand(*batch_shape, 3).astype(gs.np_float))
+                np_arg = gu.rotvec_to_R(np.random.randn(*batch_shape, 3).clip(-1.0, 1.0).astype(gs.np_float))
             else:
-                np_arg = np.random.rand(*batch_shape, *shape_arg).astype(gs.np_float)
+                np_arg = np.random.randn(*batch_shape, *shape_arg).clip(-1.0, 1.0).astype(gs.np_float)
             tc_arg = torch.as_tensor(np_arg, dtype=gs.tc_float, device=gs.device)
             np_args.append(np_arg)
             tc_args.append(tc_arg)