Repository navigation
Expand file tree
/
Copy pathspline.py
More file actions
105 lines (85 loc) · 3.81 KB
/
Copy pathspline.py
File metadata and controls
105 lines (85 loc) · 3.81 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
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