compute mean of a set of rotations (#160)

petercorke · web-flow · commit 59873f14c92e · 2025-01-26T01:00:08.000-05:00
diff --git a/spatialmath/pose3d.py b/spatialmath/pose3d.py
@@ -843,22 +843,22 @@ def Exp(
 
     def UnitQuaternion(self) -> UnitQuaternion:
         """
-            SO3 as a unit quaternion instance
+        SO3 as a unit quaternion instance
 
-            :return: a unit quaternion representation
-            :rtype: UnitQuaternion instance
+        :return: a unit quaternion representation
+        :rtype: UnitQuaternion instance
 
-            ``R.UnitQuaternion()`` is an ``UnitQuaternion`` instance representing the same rotation
-            as the SO3 rotation ``R``.
+        ``R.UnitQuaternion()`` is an ``UnitQuaternion`` instance representing the same rotation
+        as the SO3 rotation ``R``.
 
-            Example:
+        Example:
 
-            .. runblock:: pycon
+        .. runblock:: pycon
 
-                >>> from spatialmath import SO3
-                >>> SO3.Rz(0.3).UnitQuaternion()
+            >>> from spatialmath import SO3
+            >>> SO3.Rz(0.3).UnitQuaternion()
 
-            """
+        """
         # Function level import to avoid circular dependencies
         from spatialmath import UnitQuaternion
 
@@ -931,6 +931,29 @@ def angdist(self, other: SO3, metric: int = 6) -> Union[float, ndarray]:
         else:
             return ad
 
+    def mean(self, tol: float = 20) -> SO3:
+        """Mean of a set of rotations
+
+        :param tol: iteration tolerance in units of eps, defaults to 20
+        :type tol: float, optional
+        :return: the mean rotation
+        :rtype: :class:`SO3` instance.
+
+        Computes the Karcher mean of the set of rotations within the SO(3) instance.
+
+        :references:
+            - `**Hartley, Trumpf** - "Rotation Averaging" - IJCV 2011 <https://users.cecs.anu.edu.au/~hartley/Papers/PDF/Hartley-Trumpf:Rotation-averaging:IJCV.pdf>`_, Algorithm 1, page 15.
+            - `Karcher mean <https://en.wikipedia.org/wiki/Karcher_mean>`_
+        """
+
+        eta = tol * np.finfo(float).eps
+        R_mean = self[0]  # initial guess
+        while True:
+            r = np.dstack((R_mean.inv() * self).log()).mean(axis=2)
+            if np.linalg.norm(r) < eta:
+                return R_mean
+            R_mean = R_mean @ self.Exp(r)  # update estimate and normalize
+
 
 # ============================== SE3 =====================================#
 
diff --git a/spatialmath/quaternion.py b/spatialmath/quaternion.py
@@ -45,10 +45,10 @@ def __init__(self, s: Any = None, v=None, check: Optional[bool] = True):
         r"""
         Construct a new quaternion
 
-        :param s: scalar
-        :type s: float
-        :param v: vector
-        :type v: 3-element array_like
+        :param s: scalar part
+        :type s: float or ndarray(N)
+        :param v: vector part
+        :type v: ndarray(3), ndarray(Nx3)
 
         - ``Quaternion()`` constructs a zero quaternion
         - ``Quaternion(s, v)`` construct a new quaternion from the scalar ``s``
@@ -78,7 +78,7 @@ def __init__(self, s: Any = None, v=None, check: Optional[bool] = True):
         super().__init__()
 
         if s is None and smb.isvector(v, 4):
-            v,s = (s,v)
+            v, s = (s, v)
 
         if v is None:
             # single argument
@@ -92,6 +92,11 @@ def __init__(self, s: Any = None, v=None, check: Optional[bool] = True):
             # Quaternion(s, v)
             self.data = [np.r_[s, smb.getvector(v)]]
 
+        elif (
+            smb.isvector(s) and smb.ismatrix(v, (None, 3)) and s.shape[0] == v.shape[0]
+        ):
+            # Quaternion(s, v) where s and v are arrays
+            self.data = [np.r_[_s, _v] for _s, _v in zip(s, v)]
         else:
             raise ValueError("bad argument to Quaternion constructor")
 
@@ -395,9 +400,23 @@ def log(self) -> Quaternion:
         :seealso: :meth:`Quaternion.exp` :meth:`Quaternion.log` :meth:`UnitQuaternion.angvec`
         """
         norm = self.norm()
-        s = math.log(norm)
-        v = math.acos(np.clip(self.s / norm, -1, 1)) * smb.unitvec(self.v)
-        return Quaternion(s=s, v=v)
+        s = np.log(norm)
+        if len(self) == 1:
+            if smb.iszerovec(self._A[1:4]):
+                v = np.zeros((3,))
+            else:
+                v = math.acos(np.clip(self._A[0] / norm, -1, 1)) * smb.unitvec(
+                    self._A[1:4]
+                )
+            return Quaternion(s=s, v=v)
+        else:
+            v = [
+                np.zeros((3,))
+                if smb.iszerovec(A[1:4])
+                else math.acos(np.clip(A[0] / n, -1, 1)) * smb.unitvec(A[1:4])
+                for A, n in zip(self._A, norm)
+            ]
+            return Quaternion(s=s, v=np.array(v))
 
     def exp(self, tol: float = 20) -> Quaternion:
         r"""
@@ -437,7 +456,11 @@ def exp(self, tol: float = 20) -> Quaternion:
         exp_s = math.exp(self.s)
         norm_v = smb.norm(self.v)
         s = exp_s * math.cos(norm_v)
-        v = exp_s * self.v / norm_v * math.sin(norm_v)
+        if smb.iszerovec(self.v, tol * _eps):
+            # result will be a unit quaternion
+            v = np.zeros((3,))
+        else:
+            v = exp_s * self.v / norm_v * math.sin(norm_v)
         if abs(self.s) < tol * _eps:
             # result will be a unit quaternion
             return UnitQuaternion(s=s, v=v)
@@ -1260,7 +1283,7 @@ def Eul(cls, *angles: List[float], unit: Optional[str] = "rad") -> UnitQuaternio
         Construct a new unit quaternion from Euler angles
 
         :param 𝚪: 3-vector of Euler angles
-        :type 𝚪: array_like
+        :type 𝚪: 3 floats, array_like(3) or ndarray(N,3)
         :param unit: angular units: 'rad' [default], or 'deg'
         :type unit: str
         :return: unit-quaternion
@@ -1286,20 +1309,23 @@ def Eul(cls, *angles: List[float], unit: Optional[str] = "rad") -> UnitQuaternio
         if len(angles) == 1:
             angles = angles[0]
 
-        return cls(smb.r2q(smb.eul2r(angles, unit=unit)), check=False)
+        if smb.isvector(angles, 3):
+            return cls(smb.r2q(smb.eul2r(angles, unit=unit)), check=False)
+        else:
+            return cls([smb.r2q(smb.eul2r(a, unit=unit)) for a in angles], check=False)
 
     @classmethod
     def RPY(
         cls,
-        *angles: List[float],
+        *angles,
         order: Optional[str] = "zyx",
         unit: Optional[str] = "rad",
     ) -> UnitQuaternion:
         r"""
         Construct a new unit quaternion from roll-pitch-yaw angles
 
         :param 𝚪: 3-vector of roll-pitch-yaw angles
-        :type 𝚪: array_like
+        :type 𝚪: 3 floats, array_like(3) or ndarray(N,3)
         :param unit: angular units: 'rad' [default], or 'deg'
         :type unit: str
         :param unit: rotation order: 'zyx' [default], 'xyz', or 'yxz'
@@ -1341,7 +1367,13 @@ def RPY(
         if len(angles) == 1:
             angles = angles[0]
 
-        return cls(smb.r2q(smb.rpy2r(angles, unit=unit, order=order)), check=False)
+        if smb.isvector(angles, 3):
+            return cls(smb.r2q(smb.rpy2r(angles, unit=unit, order=order)), check=False)
+        else:
+            return cls(
+                [smb.r2q(smb.rpy2r(a, unit=unit, order=order)) for a in angles],
+                check=False,
+            )
 
     @classmethod
     def OA(cls, o: ArrayLike3, a: ArrayLike3) -> UnitQuaternion:
@@ -1569,6 +1601,24 @@ def dotb(self, omega: ArrayLike3) -> R4:
         """
         return smb.qdotb(self._A, omega)
 
+    # def mean(self, tol: float = 20) -> SO3:
+    #     """Mean of a set of rotations
+
+    #     :param tol: iteration tolerance in units of eps, defaults to 20
+    #     :type tol: float, optional
+    #     :return: the mean rotation
+    #     :rtype: :class:`UnitQuaternion` instance.
+
+    #     Computes the Karcher mean of the set of rotations within the unit quaternion instance.
+
+    #     :references:
+    #         - `**Hartley, Trumpf** - "Rotation Averaging" - IJCV 2011 <https://users.cecs.anu.edu.au/~hartley/Papers/PDF/Hartley-Trumpf:Rotation-averaging:IJCV.pdf>`_
+    #         - `Karcher mean <https://en.wikipedia.org/wiki/Karcher_mean`_
+    #     """
+
+    #     R_mean = self.SO3().mean(tol=tol)
+    #     return R_mean.UnitQuaternion()
+
     def __mul__(
         left, right: UnitQuaternion
     ) -> UnitQuaternion:  # lgtm[py/not-named-self] pylint: disable=no-self-argument
diff --git a/tests/test_pose3d.py b/tests/test_pose3d.py
@@ -717,6 +717,37 @@ def test_functions_lie(self):
         nt.assert_equal(R, SO3.EulerVec(R.eulervec()))
         np.testing.assert_equal((R.inv() * R).eulervec(), np.zeros(3))
 
+        R = SO3()  # identity matrix case
+
+        # Check log and exponential map
+        nt.assert_equal(R, SO3.Exp(R.log()))
+        np.testing.assert_equal((R.inv() * R).log(), np.zeros([3, 3]))
+
+        # Check euler vector map
+        nt.assert_equal(R, SO3.EulerVec(R.eulervec()))
+        np.testing.assert_equal((R.inv() * R).eulervec(), np.zeros(3))
+
+    def test_mean(self):
+        rpy = np.ones((100, 1)) @ np.c_[0.1, 0.2, 0.3]
+        R = SO3.RPY(rpy)
+        self.assertEqual(len(R), 100)
+        m = R.mean()
+        self.assertIsInstance(m, SO3)
+        array_compare(m, R[0])
+
+        # range of angles, mean should be the middle one, index=25
+        R = SO3.Rz(np.linspace(start=0.3, stop=0.7, num=51))
+        m = R.mean()
+        self.assertIsInstance(m, SO3)
+        array_compare(m, R[25])
+
+        # now add noise
+        rng = np.random.default_rng(0)  # reproducible random numbers
+        rpy += rng.normal(scale=0.00001, size=(100, 3))
+        R = SO3.RPY(rpy)
+        m = R.mean()
+        array_compare(m, SO3.RPY(0.1, 0.2, 0.3))
+
 
 # ============================== SE3 =====================================#
 
diff --git a/tests/test_quaternion.py b/tests/test_quaternion.py
@@ -257,25 +257,79 @@ def test_staticconstructors(self):
                 UnitQuaternion.Rz(theta, "deg").R, rotz(theta, "deg")
             )
 
+    def test_constructor_RPY(self):
         # 3 angle
+        q = UnitQuaternion.RPY([0.1, 0.2, 0.3])
+        self.assertIsInstance(q, UnitQuaternion)
+        self.assertEqual(len(q), 1)
+        nt.assert_array_almost_equal(q.R, rpy2r(0.1, 0.2, 0.3))
+        q = UnitQuaternion.RPY(0.1, 0.2, 0.3)
+        self.assertIsInstance(q, UnitQuaternion)
+        self.assertEqual(len(q), 1)
+        nt.assert_array_almost_equal(q.R, rpy2r(0.1, 0.2, 0.3))
+        q = UnitQuaternion.RPY(np.r_[0.1, 0.2, 0.3])
+        self.assertIsInstance(q, UnitQuaternion)
+        self.assertEqual(len(q), 1)
+        nt.assert_array_almost_equal(q.R, rpy2r(0.1, 0.2, 0.3))
+
         nt.assert_array_almost_equal(
-            UnitQuaternion.RPY([0.1, 0.2, 0.3]).R, rpy2r(0.1, 0.2, 0.3)
+            UnitQuaternion.RPY([10, 20, 30], unit="deg").R,
+            rpy2r(10, 20, 30, unit="deg"),
         )
-
         nt.assert_array_almost_equal(
-            UnitQuaternion.Eul([0.1, 0.2, 0.3]).R, eul2r(0.1, 0.2, 0.3)
+            UnitQuaternion.RPY([0.1, 0.2, 0.3], order="xyz").R,
+            rpy2r(0.1, 0.2, 0.3, order="xyz"),
+        )
+
+        angles = np.array(
+            [[0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6]]
         )
+        q = UnitQuaternion.RPY(angles)
+        self.assertIsInstance(q, UnitQuaternion)
+        self.assertEqual(len(q), 4)
+        for i in range(4):
+            nt.assert_array_almost_equal(q[i].R, rpy2r(angles[i, :]))
+
+        q = UnitQuaternion.RPY(angles, order="xyz")
+        self.assertIsInstance(q, UnitQuaternion)
+        self.assertEqual(len(q), 4)
+        for i in range(4):
+            nt.assert_array_almost_equal(q[i].R, rpy2r(angles[i, :], order="xyz"))
 
+        angles *= 10
+        q = UnitQuaternion.RPY(angles, unit="deg")
+        self.assertIsInstance(q, UnitQuaternion)
+        self.assertEqual(len(q), 4)
+        for i in range(4):
+            nt.assert_array_almost_equal(q[i].R, rpy2r(angles[i, :], unit="deg"))
+
+    def test_constructor_Eul(self):
         nt.assert_array_almost_equal(
-            UnitQuaternion.RPY([10, 20, 30], unit="deg").R,
-            rpy2r(10, 20, 30, unit="deg"),
+            UnitQuaternion.Eul([0.1, 0.2, 0.3]).R, eul2r(0.1, 0.2, 0.3)
         )
 
         nt.assert_array_almost_equal(
             UnitQuaternion.Eul([10, 20, 30], unit="deg").R,
             eul2r(10, 20, 30, unit="deg"),
         )
 
+        angles = np.array(
+            [[0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6]]
+        )
+        q = UnitQuaternion.Eul(angles)
+        self.assertIsInstance(q, UnitQuaternion)
+        self.assertEqual(len(q), 4)
+        for i in range(4):
+            nt.assert_array_almost_equal(q[i].R, eul2r(angles[i, :]))
+
+        angles *= 10
+        q = UnitQuaternion.Eul(angles, unit="deg")
+        self.assertIsInstance(q, UnitQuaternion)
+        self.assertEqual(len(q), 4)
+        for i in range(4):
+            nt.assert_array_almost_equal(q[i].R, eul2r(angles[i, :], unit="deg"))
+
+    def test_constructor_AngVec(self):
         # (theta, v)
         th = 0.2
         v = unitvec([1, 2, 3])
@@ -286,6 +340,7 @@ def test_staticconstructors(self):
         )
         nt.assert_array_almost_equal(UnitQuaternion.AngVec(th, -v).R, angvec2r(th, -v))
 
+    def test_constructor_EulerVec(self):
         # (theta, v)
         th = 0.2
         v = unitvec([1, 2, 3])
@@ -830,6 +885,20 @@ def test_log(self):
         nt.assert_array_almost_equal(q1.log().exp(), q1)
         nt.assert_array_almost_equal(q2.log().exp(), q2)
 
+        q = Quaternion([q1, q2, q1, q2])
+        qlog = q.log()
+        nt.assert_array_almost_equal(qlog[0].exp(), q1)
+        nt.assert_array_almost_equal(qlog[1].exp(), q2)
+        nt.assert_array_almost_equal(qlog[2].exp(), q1)
+        nt.assert_array_almost_equal(qlog[3].exp(), q2)
+
+        q = UnitQuaternion()  # identity
+        qlog = q.log()
+        nt.assert_array_almost_equal(qlog.vec, np.zeros(4))
+        qq = qlog.exp()
+        self.assertIsInstance(qq, UnitQuaternion)
+        nt.assert_array_almost_equal(qq.vec, np.r_[1, 0, 0, 0])
+
     def test_concat(self):
         u = Quaternion()
         uu = Quaternion([u, u, u, u])
@@ -1018,6 +1087,27 @@ def test_miscellany(self):
         nt.assert_equal(q.inner(q), q.norm() ** 2)
         nt.assert_equal(q.inner(u), np.dot(q.vec, u.vec))
 
+    # def test_mean(self):
+    #     rpy = np.ones((100, 1)) @ np.c_[0.1, 0.2, 0.3]
+    #     q = UnitQuaternion.RPY(rpy)
+    #     self.assertEqual(len(q), 100)
+    #     m = q.mean()
+    #     self.assertIsInstance(m, UnitQuaternion)
+    #     nt.assert_array_almost_equal(m.vec, q[0].vec)
+
+    #     # range of angles, mean should be the middle one, index=25
+    #     q = UnitQuaternion.Rz(np.linspace(start=0.3, stop=0.7, num=51))
+    #     m = q.mean()
+    #     self.assertIsInstance(m, UnitQuaternion)
+    #     nt.assert_array_almost_equal(m.vec, q[25].vec)
+
+    #     # now add noise
+    #     rng = np.random.default_rng(0)  # reproducible random numbers
+    #     rpy += rng.normal(scale=0.1, size=(100, 3))
+    #     q = UnitQuaternion.RPY(rpy)
+    #     m = q.mean()
+    #     nt.assert_array_almost_equal(m.vec, q.RPY(0.1, 0.2, 0.3).vec)
+
 
 # ---------------------------------------------------------------------------------------#
 if __name__ == "__main__":
