Skip to content

Add BSplineSE3 class. #128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Jul 9, 2024
Merged
2 changes: 2 additions & 0 deletions spatialmath/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -43,6 +44,7 @@
"LineSegment2",
"Polygon2",
"Ellipse",
"BSplineSE3",
]

try:
Expand Down
8 changes: 5 additions & 3 deletions spatialmath/base/animate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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, :])

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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))
Expand Down
105 changes: 105 additions & 0 deletions spatialmath/spline.py
Original file line number Diff line number Diff line change
@@ -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
32 changes: 32 additions & 0 deletions tests/test_spline.py
Original file line number Diff line number Diff line change
@@ -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)
Loading