diff --git a/spatialmath/__init__.py b/spatialmath/__init__.py index 63287dcf..e6ef1f77 100644 --- a/spatialmath/__init__.py +++ b/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: diff --git a/spatialmath/base/animate.py b/spatialmath/base/animate.py index 8efa3a4e..3876a2ea 100755 --- a/spatialmath/base/animate.py +++ b/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)) diff --git a/spatialmath/spline.py b/spatialmath/spline.py new file mode 100644 index 00000000..f81dcec3 --- /dev/null +++ b/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 diff --git a/tests/test_spline.py b/tests/test_spline.py new file mode 100644 index 00000000..f518fcfb --- /dev/null +++ b/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)