Add BSplineSE3 class. (#128)

Author: myeatman-bdai

spatialmath/__init__.py

@@ -15,6 +15,7 @@
 )
 from spatialmath.quaternion import Quaternion, UnitQuaternion
 from spatialmath.DualQuaternion import DualQuaternion, UnitDualQuaternion
+from spatialmath.spline import BSplineSE3
 
 # from spatialmath.Plucker import *
 # from spatialmath import base as smb
@@ -43,6 +44,7 @@
     "LineSegment2",
     "Polygon2",
     "Ellipse",
+    "BSplineSE3",
 ]
 
 try:

spatialmath/base/animate.py

@@ -212,6 +212,7 @@ def update(frame, animation):
                 # assume it is an SO(3) or SE(3)
                 T = frame
             # ensure result is SE(3)
+            
             if T.shape == (3, 3):
                 T = smb.r2t(T)
 
@@ -308,7 +309,7 @@ def __init__(self, anim: Animate, h, xs, ys, zs):
             self.anim = anim
 
         def draw(self, T):
-            p = T @ self.p
+            p = T.A @ self.p
             self.h.set_data(p[0, :], p[1, :])
             self.h.set_3d_properties(p[2, :])
 
@@ -365,7 +366,8 @@ def __init__(self, anim, h):
             self.anim = anim
 
         def draw(self, T):
-            p = T @ self.p
+            # import ipdb; ipdb.set_trace()
+            p = T.A @ self.p
 
             # reshape it
             p = p[0:3, :].T.reshape(3, 2, 3)
@@ -419,7 +421,7 @@ def __init__(self, anim, h, x, y, z):
             self.anim = anim
 
         def draw(self, T):
-            p = T @ self.p
+            p = T.A @ self.p
             # x2, y2, _ = proj3d.proj_transform(
             #   p[0], p[1], p[2], self.anim.ax.get_proj())
             # self.h.set_position((x2, y2))

spatialmath/spline.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2024 Boston Dynamics AI Institute LLC.
+# MIT Licence, see details in top-level file: LICENCE
+
+"""
+Classes for parameterizing a trajectory in SE3 with B-splines. 
+
+Copies parts of the API from scipy's B-spline class.
+"""
+
+from typing import Any, Dict, List, Optional
+from scipy.interpolate import BSpline
+from spatialmath import SE3
+import numpy as np
+import matplotlib.pyplot as plt
+from spatialmath.base.transforms3d import tranimate, trplot
+
+
+class BSplineSE3:
+    """A class to parameterize a trajectory in SE3 with a 6-dimensional B-spline.
+
+    The SE3 control poses are converted to se3 twists (the lie algebra) and a B-spline
+    is created for each dimension of the twist, using the corresponding element of the twists
+    as the control point for the spline.
+
+    For detailed information about B-splines, please see this wikipedia article.
+    https://en.wikipedia.org/wiki/Non-uniform_rational_B-spline
+    """
+
+    def __init__(
+        self,
+        control_poses: List[SE3],
+        degree: int = 3,
+        knots: Optional[List[float]] = None,
+    ) -> None:
+        """Construct BSplineSE3 object. The default arguments generate a cubic B-spline
+        with uniformly spaced knots.
+
+        - control_poses: list of SE3 objects that govern the shape of the spline.
+        - degree: int that controls degree of the polynomial that governs any given point on the spline.
+        - knots: list of floats that govern which control points are active during evaluating the spline
+        at a given t input. If none, they are automatically, uniformly generated based on number of control poses and
+        degree of spline.
+        """
+
+        self.control_poses = control_poses
+
+        # a matrix where each row is a control pose as a twist
+        # (so each column is a vector of control points for that dim of the twist)
+        self.control_pose_matrix = np.vstack(
+            [np.array(element.twist()) for element in control_poses]
+        )
+
+        self.degree = degree
+
+        if knots is None:
+            knots = np.linspace(0, 1, len(control_poses) - degree + 1, endpoint=True)
+            knots = np.append(
+                [0.0] * degree, knots
+            )  # ensures the curve starts on the first control pose
+            knots = np.append(
+                knots, [1] * degree
+            )  # ensures the curve ends on the last control pose
+        self.knots = knots
+
+        self.splines = [
+            BSpline(knots, self.control_pose_matrix[:, i], degree)
+            for i in range(0, 6)  # twists are length 6
+        ]
+
+    def __call__(self, t: float) -> SE3:
+        """Returns pose of spline at t.
+
+        t: Normalized time value [0,1] to evaluate the spline at.
+        """
+        twist = np.hstack([spline(t) for spline in self.splines])
+        return SE3.Exp(twist)
+
+    def visualize(
+        self,
+        num_samples: int,
+        length: float = 1.0,
+        repeat: bool = False,
+        ax: Optional[plt.Axes] = None,
+        kwargs_trplot: Dict[str, Any] = {"color": "green"},
+        kwargs_tranimate: Dict[str, Any] = {"wait": True},
+        kwargs_plot: Dict[str, Any] = {},
+    ) -> None:
+        """Displays an animation of the trajectory with the control poses."""
+        out_poses = [self(t) for t in np.linspace(0, 1, num_samples)]
+        x = [pose.x for pose in out_poses]
+        y = [pose.y for pose in out_poses]
+        z = [pose.z for pose in out_poses]
+
+        if ax is None:
+            fig = plt.figure(figsize=(10, 10))
+            ax = fig.add_subplot(projection="3d")
+
+        trplot(
+            [np.array(self.control_poses)], ax=ax, length=length, **kwargs_trplot
+        )  # plot control points
+        ax.plot(x, y, z, **kwargs_plot)  # plot x,y,z trajectory
+
+        tranimate(
+            out_poses, repeat=repeat, length=length, **kwargs_tranimate
+        )  # animate pose along trajectory

tests/test_spline.py

@@ -0,0 +1,32 @@
+import numpy.testing as nt
+import numpy as np
+import matplotlib.pyplot as plt
+import unittest
+import sys
+import pytest
+
+from spatialmath import BSplineSE3, SE3
+
+
+class TestBSplineSE3(unittest.TestCase):
+    control_poses = [
+        SE3.Trans([e, 2 * np.cos(e / 2 * np.pi), 2 * np.sin(e / 2 * np.pi)])
+        * SE3.Ry(e / 8 * np.pi)
+        for e in range(0, 8)
+    ]
+
+    @classmethod
+    def tearDownClass(cls):
+        plt.close("all")
+
+    def test_constructor(self):
+        BSplineSE3(self.control_poses)
+
+    def test_evaluation(self):
+        spline = BSplineSE3(self.control_poses)
+        nt.assert_almost_equal(spline(0).A, self.control_poses[0].A)
+        nt.assert_almost_equal(spline(1).A, self.control_poses[-1].A)
+
+    def test_visualize(self):
+        spline = BSplineSE3(self.control_poses)
+        spline.visualize(num_samples=100, repeat=False)