Created internal quaternion conversion function for SO3 (bdaiinstitute#119)

bokorn-bdaii · web-flow · commit 3a8cb036b547 · 2024-05-13T08:53:23.000-04:00
diff --git a/spatialmath/pose3d.py b/spatialmath/pose3d.py
@@ -34,6 +34,9 @@
 
 from spatialmath.twist import Twist3
 
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from spatialmath.quaternion import UnitQuaternion
 
 # ============================== SO3 =====================================#
 
@@ -834,6 +837,29 @@ def Exp(
         else:
             return cls(smb.trexp(cast(R3, S), check=check), check=False)
 
+    def UnitQuaternion(self) -> UnitQuaternion:
+        """
+            SO3 as a unit quaternion instance
+
+            :return: a unit quaternion representation
+            :rtype: UnitQuaternion instance
+
+            ``R.UnitQuaternion()`` is an ``UnitQuaternion`` instance representing the same rotation
+            as the SO3 rotation ``R``.
+
+            Example:
+
+            .. runblock:: pycon
+
+                >>> from spatialmath import SO3
+                >>> SO3.Rz(0.3).UnitQuaternion()
+
+            """
+        # Function level import to avoid circular dependencies
+        from spatialmath import UnitQuaternion
+
+        return UnitQuaternion(smb.r2q(self.R), check=False)
+
     def angdist(self, other: SO3, metric: int = 6) -> Union[float, ndarray]:
         r"""
         Angular distance metric between rotations
diff --git a/tests/test_pose3d.py b/tests/test_pose3d.py
@@ -8,7 +8,7 @@
 we will assume that the primitives rotx,trotx, etc. all work
 """
 from math import pi
-from spatialmath import SE3, SO3, SE2
+from spatialmath import SE3, SO3, SE2, UnitQuaternion
 import numpy as np
 from spatialmath.base import *
 from spatialmath.baseposematrix import BasePoseMatrix
@@ -233,6 +233,20 @@ def test_constructor_TwoVec(self):
         # x axis should equal normalized x vector
         nt.assert_almost_equal(R.R[:, 0], v3 / np.linalg.norm(v3), 5)
 
+    def test_conversion(self):
+        R = SO3.AngleAxis(0.7, [1,2,3])
+        q = UnitQuaternion([11,7,3,-6])
+
+        R_from_q = SO3(q.R)
+        q_from_R = UnitQuaternion(R)
+
+        nt.assert_array_almost_equal(R.UnitQuaternion(), q_from_R)
+        nt.assert_array_almost_equal(R.UnitQuaternion().SO3(), R)
+
+        nt.assert_array_almost_equal(q.SO3(), R_from_q)
+        nt.assert_array_almost_equal(q.SO3().UnitQuaternion(), q)
+
+
     def test_shape(self):
         a = SO3()
         self.assertEqual(a._A.shape, a.shape)
