Merge PR #243: Accept a stacked (N, r, c) ndarray as a sequence of poses (SE3, SO3, SE2, SO2)

Author: taughz

spatialmath/baseposelist.py

@@ -154,6 +154,8 @@ def arghandler(
 
         #. None, an identity value is created
         #. a numpy.ndarray of the appropriate shape and value which is valid for the subclass
+        #. a stacked numpy.ndarray of shape (N, r, c), where (r, c) is the shape of one value
+           of the subclass, which is a sequence of N values
         #. a list whose elements all meet the criteria above
         #. an instance of the subclass
         #. a list whose elements are all singelton instances of the subclass
@@ -187,6 +189,20 @@ def arghandler(
         elif isinstance(arg, np.ndarray):
             # it's a numpy array
 
+            shape = tuple(self.shape)
+            if len(shape) == 2 and arg.ndim == 3:
+                # a 3D array passed to a matrix pose class is a stack
+                if arg.shape[1:] != shape:
+                    # a stack of matrices of the wrong size is rejected,
+                    # whatever the value of check
+                    return False
+                # a stacked (N, r, c) array is a sequence of N values
+                values = [self._import(x, check=check) for x in arg]
+                if any(x is None for x in values):
+                    return False
+                self.data = values
+                return True
+
             x = self._import(arg, check=check)
             if x is not None:
                 self.data = [x]

spatialmath/pose3d.py

@@ -1096,6 +1096,9 @@ def __init__(self, x=None, y=None, z=None, *, check=True):
         - ``SE3([X1, X2, ... XN])`` has ``N`` values
           given by the elements ``Xi`` each of which is an SE3 instance.
 
+        - ``SE3(A)`` where ``A`` is an ``(N, 4, 4)`` NumPy array has ``N`` values,
+          equivalent to ``SE3(list(A))``.
+
         :SymPy: supported
         """
         if y is None and z is None:
@@ -1115,7 +1118,10 @@ def __init__(self, x=None, y=None, z=None, *, check=True):
                 self.data = [smb.transl(T) for T in x]
 
             else:
-                raise ValueError("bad argument to constructor")
+                raise ValueError(
+                    "bad argument to constructor: expected an SE3, SO3 or SE2 instance, a 4x4 or "
+                    "(N, 4, 4) array, a list of those, a 3-vector or an (N, 3) array of translations"
+                )
 
         elif y is not None and z is not None:
             # SE3(x, y, z)

tests/test_pose3d.py

@@ -793,6 +793,26 @@ class TestSE3(unittest.TestCase):
     def tearDownClass(cls):
         plt.close("all")
 
+    def test_constructor_stacked_array(self):
+        # an (N, 4, 4) ndarray is a sequence of N poses, like a list of 4x4 arrays
+        poses = [SE3.Rx(0.3), SE3(1, 2, 3) * SE3.Ry(0.4), SE3.Rz(-0.2)]
+        stacked = np.array([T.A for T in poses])
+        for check in (True, False):
+            T = SE3(stacked, check=check)
+            self.assertEqual(len(T), 3)
+            for k in range(3):
+                nt.assert_array_almost_equal(T[k].A, poses[k].A)
+        self.assertEqual(len(SE3(np.empty((0, 4, 4)))), 0)
+        self.assertEqual(len(SO3(np.array([SO3.Rx(0.1).A, SO3.Ry(0.2).A]))), 2)
+
+        # stacked matrices that are not valid poses are still rejected by check=True
+        with self.assertRaises(ValueError):
+            SE3(2 * stacked)
+        # and a stack of matrices of the wrong size is rejected whatever the value of check
+        for check in (True, False):
+            with self.assertRaises(ValueError):
+                SE3(np.ones((3, 3, 3)), check=check)
+
     def test_constructor(self):
         # null constructor
         R = SE3()

tests/test_quaternion.py

@@ -954,6 +954,9 @@ def test_constructor(self):
         # tc.verifyError( @() Quaternion(2), 'SMTB:Quaternion:badarg')
         # tc.verifyError( @() Quaternion([1, 2, 3]), 'SMTB:Quaternion:badarg')
 
+        # an array that is not a 4-vector must not give an empty Quaternion
+        self.assertNotEqual(len(Quaternion(np.random.rand(3, 4))), 0)
+
     def test_string(self):
         u = Quaternion()
 

tests/test_twist.py

@@ -67,6 +67,12 @@ def test_conversion_SE3(self):
         T = SE3.Rx(0) * SE3(1, 2, 3)
         array_compare(Twist3(T).SE3(), T)
 
+        # from an SE(3) matrix, with and without checking
+        for check in (True, False):
+            tw = Twist3(T.A, check=check)
+            self.assertEqual(len(tw), 1)
+            array_compare(tw.SE3(), T)
+
     def test_conversion_se3(self):
         s = [1, 2, 3, 4, 5, 6]
         x = Twist3(s)
@@ -83,6 +89,12 @@ def test_conversion_se3(self):
             ),
         )
 
+        # from an se(3) matrix, with and without checking
+        for check in (True, False):
+            tw = Twist3(x.skewa(), check=check)
+            self.assertEqual(len(tw), 1)
+            array_compare(tw, x)
+
     def test_conversion_Plucker(self):
         pass
 
@@ -319,6 +331,12 @@ def test_conversion_SE2(self):
         self.assertIsInstance(tw.SE2(), SE2)
         self.assertEqual(len(tw.SE2()), 1)
 
+        # from an SE(2) matrix, with and without checking
+        for check in (True, False):
+            tw = Twist2(T.A, check=check)
+            self.assertEqual(len(tw), 1)
+            array_compare(tw.SE2(), T)
+
     def test_conversion_se2(self):
         s = [1, 2, 3]
         x = Twist2(s)
@@ -327,6 +345,12 @@ def test_conversion_se2(self):
             x.skewa(), np.array([[0.0, -3.0, 1.0], [3.0, 0.0, 2.0], [0.0, 0.0, 0.0]])
         )
 
+        # from an se(2) matrix, with and without checking
+        for check in (True, False):
+            tw = Twist2(x.skewa(), check=check)
+            self.assertEqual(len(tw), 1)
+            array_compare(tw, x)
+
     def test_list_constuctor(self):
         x = Twist2([1, 0, 0])