diff --git a/spatialmath/__init__.py b/spatialmath/__init__.py index e6ef1f77..5a329f58 100644 --- a/spatialmath/__init__.py +++ b/spatialmath/__init__.py @@ -15,7 +15,7 @@ ) from spatialmath.quaternion import Quaternion, UnitQuaternion from spatialmath.DualQuaternion import DualQuaternion, UnitDualQuaternion -from spatialmath.spline import BSplineSE3 +from spatialmath.spline import CubicBSplineSE3, FitCubicBSplineSE3 # from spatialmath.Plucker import * # from spatialmath import base as smb @@ -44,7 +44,8 @@ "LineSegment2", "Polygon2", "Ellipse", - "BSplineSE3", + "CubicBSplineSE3", + "FitCubicBSplineSE3" ] try: diff --git a/spatialmath/pose3d.py b/spatialmath/pose3d.py index 4c2dae06..0e0c760e 100644 --- a/spatialmath/pose3d.py +++ b/spatialmath/pose3d.py @@ -1041,6 +1041,13 @@ def shape(self) -> Tuple[int, int]: """ return (4, 4) + @SO3.R.setter + def R(self, r: SO3Array) -> None: + if len(self) > 1: + raise ValueError("can only assign rotation to length 1 object") + so3 = SO3(r) + self.A[:3, :3] = so3.R + @property def t(self) -> R3: """ diff --git a/spatialmath/spline.py b/spatialmath/spline.py index f81dcec3..00823c1d 100644 --- a/spatialmath/spline.py +++ b/spatialmath/spline.py @@ -9,76 +9,77 @@ from typing import Any, Dict, List, Optional from scipy.interpolate import BSpline -from spatialmath import SE3 +from spatialmath import SE3, Twist3, SO3 import numpy as np import matplotlib.pyplot as plt from spatialmath.base.transforms3d import tranimate, trplot +from scipy.optimize import minimize +from scipy.interpolate import splrep -class BSplineSE3: - """A class to parameterize a trajectory in SE3 with a 6-dimensional B-spline. +class CubicBSplineSE3: + """A class to parameterize a trajectory in SE3 with a cubic 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. + The position and orientation are calculated disjointly. The position uses the + "classic" B-spline formulation. The orientation calculation is based on + interpolation between the identity element and the control SO3 element, them applying + each interpolated SO3 element as a group action. For detailed information about B-splines, please see this wikipedia article. https://en.wikipedia.org/wiki/Non-uniform_rational_B-spline """ + degree = 3 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. + """Construct a CubicBSplineSE3 object with a open uniform knot vector. - 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: + self.knots = self.knots_from_num_control_poses(len(control_poses)) + + def __call__(self, time:float): """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) + """ + + spline_no_coeff = BSpline.design_matrix([time], self.knots, self.degree) #the B in sum(alpha_i * B_i) = S(t) + rows,cols = spline_no_coeff.nonzero() + + current_so3 = SO3() + for row,col in zip(rows,cols): + control_so3: SO3 = SO3(self.control_poses[col]) + current_so3 = control_so3.interp1(spline_no_coeff[row,col]) * current_so3 + + xyz = np.array([0,0,0]) + for row,col in zip(rows,cols): + control_point = self.control_poses[col].t + xyz = xyz + control_point*spline_no_coeff[row,col] + + return SE3.Rt(t = xyz, R = current_so3) + + @classmethod + def knots_from_num_control_poses(self, num_control_poses: int): + """ Return open uniform knots vector based on number of control poses. """ + # use open uniform knot vector + knots = np.linspace(0, 1, num_control_poses-2, endpoint=True) + knots = np.append( + [0] * CubicBSplineSE3.degree, knots + ) # ensures the curve starts on the first control pose + knots = np.append( + knots, [1] * CubicBSplineSE3.degree + ) # ensures the curve ends on the last control pose + return knots def visualize( self, num_samples: int, - length: float = 1.0, + length: float = 0.2, repeat: bool = False, ax: Optional[plt.Axes] = None, kwargs_trplot: Dict[str, Any] = {"color": "green"}, @@ -103,3 +104,133 @@ def visualize( tranimate( out_poses, repeat=repeat, length=length, **kwargs_tranimate ) # animate pose along trajectory + + +class FitCubicBSplineSE3: + + def __init__(self, pose_data: List[SE3], timestamps, num_control_points: int = 6, method = "slsqp") -> None: + """ + Timestamps should be normalized to [0,1] and sorted. + Outputs control poses, but the optimization is done on the [pos,axis-angle], flattened into a 1d vector. + """ + self.pose_data = pose_data + self.timestamps = timestamps + self.num_control_pose = num_control_points + self.method = method + + # initialize control poses with data points + sample_points = np.linspace(self.timestamps[0], self.timestamps[-1], self.num_control_pose) + closest_timestamp_indices = np.searchsorted(self.timestamps, sample_points) + for i, index in enumerate(closest_timestamp_indices): + if index < 0: + closest_timestamp_indices[i] = 0 + if index >= len(timestamps): + closest_timestamp_indices[i] = len(timestamps) - 1 + + self.spline = CubicBSplineSE3(control_poses=[SE3.CopyFrom(self.pose_data[index]) for index in closest_timestamp_indices]) + + def objective_function_xyz(self, xyz_flat: Optional[np.ndarray] = None): + """L-infinity norm of euclidean distance between data points and spline""" + + # data massage + if xyz_flat is not None: + self._assign_xyz_to_control_poses(xyz_flat) + + # objective + error_vector = self.euclidean_distance() + return np.linalg.norm(error_vector, ord=np.inf) + + def objective_function_so3(self, so3_twists_flat: Optional[np.ndarray] = None): + """L-infinity norm of angular distance between data points and spline""" + + # data massage + if so3_twists_flat is not None: + self._assign_so3_twist_to_control_poses(so3_twists_flat) + + # objective + error_vector = self.ang_distance() + return np.linalg.norm(error_vector, ord=np.inf) + + def fit(self, disp: bool = False): + """ Find the spline control points that minimize the distance from the spline to the data points. + """ + so3_result = self.fit_so3(disp) + pos_result = self.fit_xyz(disp) + + return pos_result, so3_result + + def fit_xyz(self, disp: bool = False): + """ Solve fitting problem for x,y,z coordinates. + """ + xyz_flat = np.concatenate([pose.t for pose in self.spline.control_poses]) + result = minimize(self.objective_function_xyz, xyz_flat, method=self.method, options = {"disp":disp}) + self._assign_xyz_to_control_poses(result.x) + + return result + + def fit_so3(self, disp: bool = False): + """ Solve fitting problem for SO3 coordinates. + """ + so3_twists_flat = np.concatenate([SO3(pose.R).log(twist=True) for pose in self.spline.control_poses]) + result = minimize(self.objective_function_so3, so3_twists_flat, method = self.method, options = {"disp":disp}) + self._assign_so3_twist_to_control_poses(result.x) + + return result + + def _assign_xyz_to_control_poses(self, xyz_flat: np.ndarray) -> None: + xyz_mat = xyz_flat.reshape(self.num_control_pose, 3) + for i, xyz in enumerate(xyz_mat): + self.spline.control_poses[i].t = xyz + + def _assign_so3_twist_to_control_poses(self, so3_twists_flat: np.ndarray) -> None: + so3_twists_mat = so3_twists_flat.reshape(self.num_control_pose, 3) + for i, so3_twist in enumerate(so3_twists_mat): + self.spline.control_poses[i].R = SO3.Exp(so3_twist) + + def ang_distance(self): + """ Returns vector of angular distance between spline and data points. + """ + return [pose.angdist(self.spline(timestamp)) for pose, timestamp in zip(self.pose_data, self.timestamps)] + + def euclidean_distance(self): + """ Returns vector of euclidean distance between spline and data points. + """ + return [np.linalg.norm(pose.t - self.spline(timestamp).t) for pose, timestamp in zip(self.pose_data, self.timestamps)] + + def visualize( + self, + num_samples: int, + length: float = 0.2, + repeat: bool = False, + 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 and data points.""" + out_poses = [self.spline(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] + + + fig = plt.figure(figsize=(10, 10)) + + ax = fig.add_subplot(1, 2, 1) + ax.plot(self.timestamps, self.ang_distance(), "--x") + ax.plot(self.timestamps, self.euclidean_distance(), "--x") + ax.legend(["angular distance", "euclidiean distance"]) + ax.set_xlabel("Normalized time") + + ax = fig.add_subplot(1, 2, 2, projection="3d") + + trplot( + [np.array(self.spline.control_poses)], ax=ax, length=length*1.2, color= "green" + ) # plot control points + ax.plot(x, y, z, **kwargs_plot) # plot x,y,z trajectory + trplot( + [np.array(self.pose_data)], ax=ax, length=length, color= "cornflowerblue", + ) # plot data points + + 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 index f518fcfb..edd71ec2 100644 --- a/tests/test_spline.py +++ b/tests/test_spline.py @@ -2,10 +2,7 @@ import numpy as np import matplotlib.pyplot as plt import unittest -import sys -import pytest - -from spatialmath import BSplineSE3, SE3 +from spatialmath import FitCubicBSplineSE3, CubicBSplineSE3, SE3, SO3 class TestBSplineSE3(unittest.TestCase): @@ -20,13 +17,44 @@ def tearDownClass(cls): plt.close("all") def test_constructor(self): - BSplineSE3(self.control_poses) + CubicBSplineSE3(self.control_poses) def test_evaluation(self): - spline = BSplineSE3(self.control_poses) + spline = CubicBSplineSE3(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 = CubicBSplineSE3(self.control_poses) spline.visualize(num_samples=100, repeat=False) + +class TestFitBSplineSE3(unittest.TestCase): + + num_data_points = 16 + num_samples = 100 + num_control_points = 6 + + timestamps = np.linspace(0, 1, num_data_points) + trajectory = [ + SE3.Rt(t = [t*4, 4*np.sin(t * 2*np.pi* 0.5), 4*np.cos(t * 2*np.pi * 0.5)], + R= SO3.Rx( t*2*np.pi* 0.5)) + for t in timestamps + ] + + @classmethod + def tearDownClass(cls): + plt.close("all") + + def test_constructor(self): + pass + + def test_evaluation_and_visualization(self): + fit_se3_spline = FitCubicBSplineSE3(self.trajectory, self.timestamps, num_control_points=self.num_control_points) + + result = fit_se3_spline.fit(disp=True) + + assert len(result) == 2 + assert fit_se3_spline.objective_function_xyz() < 0.01 + assert fit_se3_spline.objective_function_so3() < 0.01 + + fit_se3_spline.visualize(num_samples=self.num_samples, repeat=False, length=0.4, kwargs_tranimate={"wait": True, "interval" : 10}) \ No newline at end of file