From 0abeafcb1d1673dd9f4bff6a593176b52877b6c8 Mon Sep 17 00:00:00 2001
From: The visu3d Authors <no-reply@google.com>
Date: Fri, 26 Apr 2024 04:18:57 -0700
Subject: [PATCH] Add rot_to_rad function to visu3d.math.rotation_utils

PiperOrigin-RevId: 628363278
---
 CHANGELOG.md                       |  3 ++-
 visu3d/math/__init__.py            |  1 +
 visu3d/math/rotation_utils.py      | 34 ++++++++++++++++++++++++++++++
 visu3d/math/rotation_utils_test.py | 10 +++++++++
 4 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b9c1eda..da8b719 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,7 +9,8 @@ Changelog follow https://keepachangelog.com/ format.
 ## [Unreleased]
 
 *  Drop Python 3.10 support
-*   Add `hash()` support (e.g. after `etree.spec_like`).
+*  Add `hash()` support (e.g. after `etree.spec_like`).
+*  Added: `visu3d.math.rotation_utils.rot_to_rad`.
 
 ## [1.5.4] - 2023-11-23
 
diff --git a/visu3d/math/__init__.py b/visu3d/math/__init__.py
index d6b7c70..09e1ea8 100644
--- a/visu3d/math/__init__.py
+++ b/visu3d/math/__init__.py
@@ -24,6 +24,7 @@
 from visu3d.math.rotation_utils import is_rot
 from visu3d.math.rotation_utils import RAD2DEG
 from visu3d.math.rotation_utils import rot_to_euler
+from visu3d.math.rotation_utils import rot_to_rad
 from visu3d.math.rotation_utils import rot_x
 from visu3d.math.rotation_utils import rot_y
 from visu3d.math.rotation_utils import rot_z
diff --git a/visu3d/math/rotation_utils.py b/visu3d/math/rotation_utils.py
index fd39347..9b3b164 100644
--- a/visu3d/math/rotation_utils.py
+++ b/visu3d/math/rotation_utils.py
@@ -204,6 +204,40 @@ def rot_to_euler(
   return theta_x, theta_y, theta_z
 
 
+@enp.check_and_normalize_arrays
+def rot_to_rad(
+    rot: FloatArray['*B 3 3'],
+    *,
+    xnp: enp.NpModule = ...,
+) -> FloatArray['*B']:
+  """Compute the absolute angle of rotation in radians of a 3x3 rotation matrix.
+
+  With a change of coordinate frame, any rotation matrix can be expressed as:
+  [[cos(t), -sin(t), 0], [sin(t), cos(t), 0], [0, 0, 1]], which gives
+  tr(A) = 2 * cos(t) + 1. This expression still holds for the original matrix as
+  trace(SAS^{-1})=trace(A).
+
+  ```python
+  rads = acos(0.5 * (trace(R) - 1))
+  ```
+
+  See:
+  https://en.wikipedia.org/wiki/Rotation_matrix#Determining_the_angle
+
+  Args:
+    rot: Rotation matrix
+    xnp: Np module used (jnp, tnp, np,...)
+
+  Returns:
+    rads: The absolute angle in radians represented by rot
+  """
+  trace = xnp.trace(rot, axis1=-2, axis2=-1)
+  # Because tr(A) = 2 * cos(t) + 1, it should lie in the interval [-1, 3]
+  # In case this doesn't happen due to numerical errors, we apply clipping.
+  trace = xnp.clip(trace, -1, 3)
+  return xnp.arccos(0.5 * (trace - 1))
+
+
 @enp.check_and_normalize_arrays
 def is_orth(
     rot: FloatArray['3 3'],
diff --git a/visu3d/math/rotation_utils_test.py b/visu3d/math/rotation_utils_test.py
index ef036bf..5b2839a 100644
--- a/visu3d/math/rotation_utils_test.py
+++ b/visu3d/math/rotation_utils_test.py
@@ -25,6 +25,16 @@
 enable_tf_np_mode = enp.testing.set_tnp
 
 
+@enp.testing.parametrize_xnp(skip=['torch'])
+def test_is_rot_to_rad(xnp: enp.NpModule):
+  rot = xnp.asarray([
+      [-1.0, 0.0, 0.0],
+      [0.0, -1.0, 0.0],
+      [0.0, 0.0, 1.0],
+  ])
+  np.testing.assert_allclose(v3d.math.rot_to_rad(rot), np.pi, atol=1e-6)
+
+
 @enp.testing.parametrize_xnp()
 def test_is_rotation_matrix(xnp: enp.NpModule):
   assert v3d.math.is_rot(xnp.eye(3))