ensure that angdiff returns a scalar for scalar args (#89)

Author: petercorke

spatialmath/base/__init__.py

@@ -216,6 +216,7 @@
     "isrot2",
     "trlog2",
     "trexp2",
+    "trnorm2",
     "tr2jac2",
     "trinterp2",
     "trprint2",

spatialmath/base/transforms2d.py

@@ -28,6 +28,7 @@
 import spatialmath.base as smb
 from spatialmath.base.types import *
 from spatialmath.base.transformsNd import rt2tr
+from spatialmath.base.vectors import unitvec
 
 _eps = np.finfo(np.float64).eps
 
@@ -679,6 +680,72 @@ def trexp2(
         raise ValueError(" First argument must be SO(2), 1-vector, SE(2) or 3-vector")
 
 
+@overload  # pragma: no cover
+def trnorm2(R: SO2Array) -> SO2Array:
+    ...
+
+
+def trnorm2(T: SE2Array) -> SE2Array:
+    r"""
+    Normalize an SO(2) or SE(2) matrix
+
+    :param T: SE(2) or SO(2) matrix
+    :type T: ndarray(3,3) or ndarray(2,2)
+    :return: normalized SE(2) or SO(2) matrix
+    :rtype: ndarray(3,3) or ndarray(2,2)
+    :raises ValueError: bad arguments
+
+    - ``trnorm(R)`` is guaranteed to be a proper orthogonal matrix rotation
+      matrix (2,2) which is *close* to the input matrix R (2,2).
+    - ``trnorm(T)`` as above but the rotational submatrix of the homogeneous
+      transformation T (3,3) is normalised while the translational part is
+      unchanged.
+
+    The steps in normalization are:
+
+    #. If :math:`\mathbf{R} = [a, b]`
+    #. Form unit vectors :math:`\hat{b}
+    #. Form the orthogonal planar vector :math:`\hat{a} = [\hat{b}_y  -\hat{b}_x]`
+    #. Form the normalized SO(2) matrix :math:`\mathbf{R} = [\hat{a}, \hat{b}]`
+
+    .. runblock:: pycon
+
+        >>> from spatialmath.base import trnorm, troty
+        >>> from numpy import linalg
+        >>> T = trot2(45, 'deg', t=[3, 4])
+        >>> linalg.det(T[:2,:2]) - 1 # is a valid SO(3)
+        >>> T = T @ T @ T @ T @ T @ T @ T @ T @ T @ T @ T @ T @ T
+        >>> linalg.det(T[:2,:2]) - 1  # not quite a valid SE(2) anymore
+        >>> T = trnorm2(T)
+        >>> linalg.det(T[:2,:2]) - 1  # once more a valid SE(2)
+
+    .. note::
+
+        - Only the direction of a-vector (the z-axis) is unchanged.
+        - Used to prevent finite word length arithmetic causing transforms to
+          become 'unnormalized', ie. determinant :math:`\ne 1`.
+    """
+
+    if not ishom2(T) and not isrot2(T):
+        raise ValueError("expecting SO(2) or SE(2)")
+
+    a = T[:, 0]
+    b = T[:, 1]
+
+    b = unitvec(b)
+    # fmt: off
+    R = np.array([
+        [ b[1], b[0]], 
+        [-b[0], b[1]]
+    ])
+    # fmt: on
+
+    if ishom2(T):
+        return rt2tr(cast(SO2Array, R), T[:2, 2])
+    else:
+        return R
+
+
 @overload  # pragma: no cover
 def tradjoint2(T: SO2Array) -> R1x1:
     ...

spatialmath/base/transforms3d.py

@@ -1533,13 +1533,17 @@ def trexp(S, theta=None, check=True):
         raise ValueError(" First argument must be SO(3), 3-vector, SE(3) or 6-vector")
 
 
+@overload  # pragma: no cover
+def trnorm(R: SO3Array) -> SO3Array:
+    ...
+
+
 def trnorm(T: SE3Array) -> SE3Array:
     r"""
     Normalize an SO(3) or SE(3) matrix
 
-    :param R: SE(3) or SO(3) matrix
-    :type R: ndarray(4,4) or ndarray(3,3)
-    :param T1: second SE(3) matrix
+    :param T: SE(3) or SO(3) matrix
+    :type T: ndarray(4,4) or ndarray(3,3)
     :return: normalized SE(3) or SO(3) matrix
     :rtype: ndarray(4,4) or ndarray(3,3)
     :raises ValueError: bad arguments
@@ -1565,9 +1569,9 @@ def trnorm(T: SE3Array) -> SE3Array:
         >>> T = troty(45, 'deg', t=[3, 4, 5])
         >>> linalg.det(T[:3,:3]) - 1 # is a valid SO(3)
         >>> T = T @ T @ T @ T @ T @ T @ T @ T @ T @ T @ T @ T @ T
-        >>> linalg.det(T[:3,:3]) - 1  # not quite a valid SO(3) anymore
+        >>> linalg.det(T[:3,:3]) - 1  # not quite a valid SE(3) anymore
         >>> T = trnorm(T)
-        >>> linalg.det(T[:3,:3]) - 1  # once more a valid SO(3)
+        >>> linalg.det(T[:3,:3]) - 1  # once more a valid SE(3)
 
     .. note::
 

spatialmath/base/vectors.py

@@ -647,7 +647,11 @@ def angdiff(a, b=None):
         b = getvector(b)
         a = a - b  # cannot use -= here, numpy wont broadcast
 
-    return np.mod(a + math.pi, 2 * math.pi) - math.pi
+    y = np.mod(a + math.pi, 2 * math.pi) - math.pi
+    if isinstance(y, np.ndarray) and len(y) == 1:
+        return float(y)
+    else:
+        return y
 
 
 def angle_std(theta: ArrayLike) -> float:

tests/base/test_transforms2d.py

@@ -94,6 +94,20 @@ def test_trexp2(self):
         T = transl2(1, 2) @ trot2(0.5)
         nt.assert_array_almost_equal(trexp2(logm(T)), T)
 
+    def test_trnorm2(self):
+        R = rot2(0.4)
+        R = np.round(R, 3)  # approx SO(2)
+        R = trnorm2(R)
+        self.assertTrue(isrot2(R, check=True))
+
+        R = rot2(0.4)
+        R = np.round(R, 3)  # approx SO(2)
+        T = rt2tr(R, [3, 4])
+
+        T = trnorm2(T)
+        self.assertTrue(ishom2(T, check=True))
+        nt.assert_almost_equal(T[:2, 2], [3, 4])
+
     def test_transl2(self):
         nt.assert_array_almost_equal(
             transl2(1, 2), np.array([[1, 0, 1], [0, 1, 2], [0, 0, 1]])

tests/base/test_transforms3d.py

@@ -430,6 +430,20 @@ def test_tr2rpy(self):
         a = rpy2tr(ang, order=seq)
         nt.assert_array_almost_equal(rpy2tr(tr2rpy(a, order=seq), order=seq), a)
 
+    def test_trnorm(self):
+        R = rpy2r(0.2, 0.3, 0.4)
+        R = np.round(R, 3)  # approx SO(3)
+        R = trnorm(R)
+        self.assertTrue(isrot(R, check=True))
+
+        R = rpy2r(0.2, 0.3, 0.4)
+        R = np.round(R, 3)  # approx SO(3)
+        T = rt2tr(R, [3, 4, 5])
+
+        T = trnorm(T)
+        self.assertTrue(ishom(T, check=True))
+        nt.assert_almost_equal(T[:3, 3], [3, 4, 5])
+
     def test_tr2eul(self):
         eul = np.r_[0.1, 0.2, 0.3]
         R = eul2r(eul)

tests/base/test_vectors.py

@@ -226,16 +226,23 @@ def test_iszero(self):
 
     def test_angdiff(self):
         self.assertEqual(angdiff(0, 0), 0)
+        self.assertIsInstance(angdiff(0, 0), float)
         self.assertEqual(angdiff(pi, 0), -pi)
         self.assertEqual(angdiff(-pi, pi), 0)
 
-        nt.assert_array_almost_equal(angdiff([0, -pi, pi], 0), [0, -pi, -pi])
+        x = angdiff([0, -pi, pi], 0)
+        nt.assert_array_almost_equal(x, [0, -pi, -pi])
+        self.assertIsInstance(x, np.ndarray)
         nt.assert_array_almost_equal(angdiff([0, -pi, pi], pi), [-pi, 0, 0])
 
-        nt.assert_array_almost_equal(angdiff(0, [0, -pi, pi]), [0, -pi, -pi])
+        x = angdiff(0, [0, -pi, pi])
+        nt.assert_array_almost_equal(x, [0, -pi, -pi])
+        self.assertIsInstance(x, np.ndarray)
         nt.assert_array_almost_equal(angdiff(pi, [0, -pi, pi]), [-pi, 0, 0])
 
-        nt.assert_array_almost_equal(angdiff([1, 2, 3], [1, 2, 3]), [0, 0, 0])
+        x = angdiff([1, 2, 3], [1, 2, 3])
+        nt.assert_array_almost_equal(x, [0, 0, 0])
+        self.assertIsInstance(x, np.ndarray)
 
     def test_wrap(self):
         self.assertAlmostEqual(wrap_0_2pi(0), 0)