Broadcast scalar integration steps over twist sequences

Author: antonsoo

spatialmath/twist.py

@@ -1014,7 +1014,7 @@ def SE3(self, theta=1, unit="rad"):
 
         if len(theta) == 1:
             # theta is a scalar
-            return SE3(smb.trexp(self.S * theta))
+            return SE3([smb.trexp(S * theta[0]) for S in self.data])
         else:
             # theta is a vector
             if len(self) == 1:
@@ -1072,6 +1072,8 @@ def exp(self, theta=1, unit="rad"):
 
         if len(self) == 1:
             return SE3([smb.trexp(self.S * t) for t in theta], check=False)
+        elif len(theta) == 1:
+            return SE3([smb.trexp(s * theta[0]) for s in self.data], check=False)
         elif len(self) == len(theta):
             return SE3([smb.trexp(s * t) for s, t in zip(self.S, theta)], check=False)
         else:
@@ -1530,15 +1532,19 @@ def SE2(self, theta=1, unit="rad"):
         """
         from spatialmath.pose2d import SE2
 
-        if unit != "rad" and self.isprismatic:
+        if unit != "rad" and any(s.isprismatic for s in self):
             print("Twist3.exp: using degree mode for a prismatic twist")
 
         theta = smb.getunit(theta, unit)
 
         if len(theta) == 1:
-            return SE2(smb.trexp2(self.S * theta))
-        else:
+            return SE2([smb.trexp2(S * theta[0]) for S in self.data])
+        elif len(self) == 1:
             return SE2([smb.trexp2(self.S * t) for t in theta])
+        elif len(self) == len(theta):
+            return SE2([smb.trexp2(S * t) for S, t in zip(self.data, theta)])
+        else:
+            raise ValueError("length of twist and theta not consistent")
 
     def skewa(self):
         """
@@ -1605,6 +1611,8 @@ def exp(self, theta=1, unit="rad"):
 
         if len(self) == 1:
             return SE2([smb.trexp2(self.S * t) for t in theta], check=False)
+        elif len(theta) == 1:
+            return SE2([smb.trexp2(s * theta[0]) for s in self.data], check=False)
         elif len(self) == len(theta):
             return SE2([smb.trexp2(s * t) for s, t in zip(self.S, theta)], check=False)
         else:

tests/test_twist_sequence.py

@@ -0,0 +1,49 @@
+import numpy as np
+import pytest
+from scipy.linalg import expm
+
+from spatialmath import SE2, SE3, Twist2, Twist3
+
+
+@pytest.mark.parametrize("twist,pose", [(Twist2, SE2), (Twist3, SE3)])
+@pytest.mark.parametrize("method", ["exp", "pose"])
+@pytest.mark.parametrize("theta", [None, 0, -0.4, [0.2, -0.7]])
+def test_twist_sequence_exponential(twist, pose, method, theta):
+    """A shared step and individual steps both integrate every twist."""
+    vectors = (
+        [[1, -2, 0.3], [2, 1, -0.4]]
+        if twist is Twist2
+        else [[1, -2, 3, 0.2, 0.3, -0.4], [2, 1, -1, -0.1, 0.5, 0.2]]
+    )
+    sequence = twist([np.array(v) for v in vectors])
+    convert = getattr(sequence, "exp" if method == "exp" else pose.__name__)
+    result = convert() if theta is None else convert(theta)
+    steps = np.broadcast_to(1 if theta is None else theta, (2,))
+    assert isinstance(result, pose)
+    assert len(result) == 2
+    for value, source, step in zip(result, sequence, steps):
+        np.testing.assert_allclose(value.A, expm(source.skewa() * step), atol=1e-14)
+
+
+@pytest.mark.parametrize("twist,pose", [(Twist2, SE2), (Twist3, SE3)])
+@pytest.mark.parametrize("method", ["exp", "pose"])
+def test_twist_sequence_degrees_and_singleton(twist, pose, method):
+    """Degrees, singleton expansion and mismatched lengths retain their rules."""
+    source = twist([1, 2, 0.3] if twist is Twist2 else [1, 2, 3, 0.2, 0.3, 0.4])
+    name = "exp" if method == "exp" else pose.__name__
+    angles = [15, -30]
+    expanded = getattr(source, name)(angles, unit="deg")
+    assert len(expanded) == 2
+    for value, angle in zip(expanded, angles):
+        np.testing.assert_allclose(
+            value.A, expm(source.skewa() * np.deg2rad(angle)), atol=1e-14
+        )
+    sequence = twist([source, source])
+    shared = getattr(sequence, name)(30, unit="deg")
+    assert len(shared) == 2
+    for value in shared:
+        np.testing.assert_allclose(
+            value.A, expm(source.skewa() * np.deg2rad(30)), atol=1e-14
+        )
+    with pytest.raises(ValueError):
+        getattr(sequence, name)([0.1, 0.2, 0.3])