From c57a7271f85b9f8f6fee5f7558c5d20f5bc81ab4 Mon Sep 17 00:00:00 2001
From: Alexander Fabisch <afabisch@googlemail.com>
Date: Tue, 22 Oct 2024 18:33:15 +0200
Subject: [PATCH] Add norm_mrp

---
 doc/source/api.rst                   |  1 +
 pytransform3d/rotations/__init__.py  |  3 ++-
 pytransform3d/rotations/_mrp.py      | 25 ++++++++++++++++++++++++-
 pytransform3d/rotations/_mrp.pyi     |  3 +++
 pytransform3d/test/test_rotations.py | 22 ++++++++++++++++++++++
 5 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/doc/source/api.rst b/doc/source/api.rst
index 56948fc4..5b1f8846 100644
--- a/doc/source/api.rst
+++ b/doc/source/api.rst
@@ -148,6 +148,7 @@ Modified Rodrigues Parameters
 
    ~check_mrp
    ~mrp_near_singularity
+   ~norm_mrp
    ~concatenate_mrp
    ~mrp_from_axis_angle
    ~mrp_from_quaternion
diff --git a/pytransform3d/rotations/__init__.py b/pytransform3d/rotations/__init__.py
index 5f28a301..83e1e69c 100644
--- a/pytransform3d/rotations/__init__.py
+++ b/pytransform3d/rotations/__init__.py
@@ -93,7 +93,7 @@
     quaternion_double, quaternion_integrate, quaternion_gradient,
     concatenate_quaternions, q_conj, q_prod_vector, quaternion_diff,
     quaternion_dist, quaternion_from_euler)
-from ._mrp import mrp_near_singularity, concatenate_mrp
+from ._mrp import mrp_near_singularity, norm_mrp, concatenate_mrp
 from ._slerp import (slerp_weights, pick_closest_quaternion, quaternion_slerp,
                      axis_angle_slerp, rotor_slerp)
 from ._testing import (
@@ -220,6 +220,7 @@
     "quaternion_from_angle",
     "quaternion_from_euler",
     "mrp_near_singularity",
+    "norm_mrp",
     "concatenate_mrp",
     "cross_product_matrix",
     "mrp_from_quaternion",
diff --git a/pytransform3d/rotations/_mrp.py b/pytransform3d/rotations/_mrp.py
index f70cc982..be9b6195 100644
--- a/pytransform3d/rotations/_mrp.py
+++ b/pytransform3d/rotations/_mrp.py
@@ -1,9 +1,32 @@
 """Modified Rodrigues parameters."""
 import numpy as np
-from ._utils import check_mrp
+from ._utils import check_mrp, norm_angle
+from ._conversions import axis_angle_from_mrp, mrp_from_axis_angle
 from ._constants import two_pi
 
 
+def norm_mrp(mrp):
+    """Normalize angle of modified Rodrigues parameters to range [-pi, pi].
+
+    Normalization of modified Rodrigues parameters is required to avoid the
+    singularity at a rotation angle of 2 * pi.
+
+    Parameters
+    ----------
+    mrp : array-like, shape (3,)
+        Modified Rodrigues parameters.
+
+    Returns
+    -------
+    mrp : array, shape (3,)
+        Modified Rodrigues parameters with angle normalized to [-pi, pi].
+    """
+    mrp = check_mrp(mrp)
+    a = axis_angle_from_mrp(mrp)
+    a[3] = norm_angle(a[3])
+    return mrp_from_axis_angle(a)
+
+
 def mrp_near_singularity(mrp, tolerance=1e-6):
     """Check if modified Rodrigues parameters are close to singularity.
 
diff --git a/pytransform3d/rotations/_mrp.pyi b/pytransform3d/rotations/_mrp.pyi
index 2421ef38..4a99461f 100644
--- a/pytransform3d/rotations/_mrp.pyi
+++ b/pytransform3d/rotations/_mrp.pyi
@@ -2,6 +2,9 @@ import numpy as np
 import numpy.typing as npt
 
 
+def norm_mrp(mrp: npt.ArrayLike) -> np.ndarray: ...
+
+
 def mrp_near_singularity(mrp: npt.ArrayLike, tolerance: float = ...) -> bool: ...
 
 
diff --git a/pytransform3d/test/test_rotations.py b/pytransform3d/test/test_rotations.py
index 9766db44..21ec68ac 100644
--- a/pytransform3d/test/test_rotations.py
+++ b/pytransform3d/test/test_rotations.py
@@ -2404,6 +2404,28 @@ def test_axis_angle_from_mrp():
         [1.0, 0.0, 0.0, 0.0])
 
 
+def test_norm_mrp():
+    mrp_norm = pr.norm_mrp(
+        pr.mrp_from_axis_angle([1.0, 0.0, 0.0, 1.5 * np.pi]))
+    assert_array_almost_equal(
+        [-1.0, 0.0, 0.0, 0.5 * np.pi], pr.axis_angle_from_mrp(mrp_norm))
+
+    mrp_norm = pr.norm_mrp(
+        pr.mrp_from_axis_angle([1.0, 0.0, 0.0, -0.5 * np.pi]))
+    assert_array_almost_equal(
+        [-1.0, 0.0, 0.0, 0.5 * np.pi], pr.axis_angle_from_mrp(mrp_norm))
+
+    mrp_norm = pr.norm_mrp(
+        pr.mrp_from_axis_angle([1.0, 0.0, 0.0, 2.0 * np.pi]))
+    assert_array_almost_equal(
+        [1.0, 0.0, 0.0, 0.0], pr.axis_angle_from_mrp(mrp_norm))
+
+    mrp_norm = pr.norm_mrp(
+        pr.mrp_from_axis_angle([1.0, 0.0, 0.0, -2.0 * np.pi]))
+    assert_array_almost_equal(
+        [1.0, 0.0, 0.0, 0.0], pr.axis_angle_from_mrp(mrp_norm))
+
+
 def test_assert_euler_almost_equal():
     pr.assert_euler_equal(
         [0.2, 0.3, -0.5], [0.2 + np.pi, -0.3, -0.5 - np.pi], 0, 1, 0)