Coverage for src/pyroboplan/trajectory/polynomial.py: 75%
152 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-03-02 22:03 -0500
1"""
2Implementations of polynomial trajectories.
4Some good resources:
5 * Video with derivation: https://robotacademy.net.au/lesson/1d-polynomial-trajectory
6 * Chapter 9 of this book: https://hades.mech.northwestern.edu/images/7/7f/MR.pdf
7"""
9from abc import ABC
10import numpy as np
11import warnings
14class PolynomialTrajectoryBase(ABC):
15 """Abstract base class for polynomial trajectories."""
17 def evaluate(self, t):
18 """
19 Evaluates the trajectory at a specific time.
21 Parameters
22 ----------
23 t : float
24 The time, in seconds, at which to evaluate the trajectory.
26 Returns
27 -------
28 tuple(array-like or float, array-like or float, array-like or float, array-like or float)
29 The tuple of (q, qd, qdd, qddd) trajectory values at the specified time.
30 If the trajectory is single-DOF, these will be scalars, otherwise they are arrays.
31 """
32 q = np.zeros(self.num_dims) # Position
33 qd = np.zeros(self.num_dims) # Velocity
34 qdd = np.zeros(self.num_dims) # Acceleration
35 qddd = np.zeros(self.num_dims) # Jerk
37 # Handle edge cases.
38 if t < self.segment_times[0][0]:
39 warnings.warn("Cannot evaluate trajectory before its start time.")
40 return None
41 elif t > self.segment_times[-1][-1]:
42 warnings.warn("Cannot evaluate trajectory after its end time.")
43 return None
45 segment_idx = 0
46 evaluated = False
47 while not evaluated:
48 t_segment_start = self.segment_times[segment_idx][0]
49 t_segment_final = self.segment_times[segment_idx][-1]
50 if t <= t_segment_final:
51 for dim in range(self.num_dims):
52 coeffs = self.coeffs[dim][segment_idx]
53 dt = t - t_segment_start
54 q[dim] = np.polyval(coeffs, dt)
55 qd[dim] = np.polyval(np.polyder(coeffs, 1), dt)
56 qdd[dim] = np.polyval(np.polyder(coeffs, 2), dt)
57 qddd[dim] = np.polyval(np.polyder(coeffs, 3), dt)
58 evaluated = True
59 else:
60 segment_idx += 1
62 # If the trajectory is single-DOF, return the values as scalars.
63 if len(q) == 1:
64 q = q[0]
65 qd = qd[0]
66 qdd = qdd[0]
67 qddd = qddd[0]
68 return q, qd, qdd, qddd
70 def generate(self, dt):
71 """
72 Generates a full trajectory at a specified sample time.
74 Parameters
75 ----------
76 dt : float
77 The sample time, in seconds.
79 Returns
80 -------
81 tuple(array-like, array-like, array-like, array-like, array-like)
82 The tuple of (t, q, qd, qdd, qddd) trajectory values generated at the sample time.
83 The time vector is 1D, and the others are 2D (time and dimension).
84 """
86 # Initialize the data structure
87 t_start = self.segment_times[0][0]
88 t_final = self.segment_times[-1][-1]
89 t_vec = np.arange(t_start, t_final, dt)
90 if t_vec[-1] != t_final: 90 ↛ 92line 90 didn't jump to line 92 because the condition on line 90 was always true
91 t_vec = np.append(t_vec, t_final)
92 num_pts = len(t_vec)
94 q = np.zeros([self.num_dims, num_pts]) # Position
95 qd = np.zeros([self.num_dims, num_pts]) # Velocity
96 qdd = np.zeros([self.num_dims, num_pts]) # Acceleration
97 qddd = np.zeros([self.num_dims, num_pts]) # Jerk
99 t_idx = 0
100 segment_idx = 0
102 while t_idx < num_pts:
103 t = t_vec[t_idx]
104 t_segment_start = self.segment_times[segment_idx][0]
105 t_segment_final = self.segment_times[segment_idx][-1]
106 if t <= t_segment_final:
107 for dim in range(self.num_dims):
108 coeffs = self.coeffs[dim][segment_idx]
109 dt = t - t_segment_start
110 q[dim, t_idx] = np.polyval(coeffs, dt)
111 qd[dim, t_idx] = np.polyval(np.polyder(coeffs, 1), dt)
112 qdd[dim, t_idx] = np.polyval(np.polyder(coeffs, 2), dt)
113 qddd[dim, t_idx] = np.polyval(np.polyder(coeffs, 3), dt)
114 t_idx += 1
115 else:
116 segment_idx += 1
118 return t_vec, q, qd, qdd, qddd
120 def visualize(
121 self,
122 dt=0.01,
123 joint_names=None,
124 show_position=True,
125 show_velocity=False,
126 show_acceleration=False,
127 show_jerk=False,
128 ):
129 """
130 Visualizes a generated trajectory with one figure window per dimension.
132 Parameters
133 ----------
134 dt : float, optional
135 The sample time at which to generate the trajectory by evaluating the polynomial.
136 joint_names : list[str], optional
137 The joint names to use for the figure titles.
138 If unset, uses the text "Dimension 1", "Dimension 2", etc.
139 show_position : bool, optional
140 If true, shows the trajectory position values.
141 show_velocity : bool, optional
142 If true, shows the trajectory velocity values.
143 show_acceleration : bool, optional
144 If true, shows the trajectory acceleration values.
145 show_jerk : bool, optional
146 If true, shows the trajectory jerk values.
147 """
148 import matplotlib.pyplot as plt
150 t_vec, q, qd, qdd, qddd = self.generate(dt)
151 vertical_line_scale_factor = 1.2
153 for dim in range(self.num_dims):
154 if joint_names is not None:
155 title = joint_names[dim]
156 else:
157 title = f"Dimension {dim + 1}"
158 plt.figure(title)
159 plt.cla()
161 # Positions, velocities, and accelerations
162 legend = []
163 min_pos = min_vel = min_accel = min_jerk = np.inf
164 max_pos = max_vel = max_accel = max_jerk = -np.inf
165 if show_position:
166 plt.plot(t_vec, q[dim, :])
167 min_pos = np.min(q[dim, :])
168 max_pos = np.max(q[dim, :])
169 legend.append("Position")
170 if show_velocity:
171 plt.plot(t_vec, qd[dim, :])
172 min_vel = np.min(qd[dim, :])
173 max_vel = np.max(qd[dim, :])
174 legend.append("Velocity")
175 if show_acceleration:
176 plt.plot(t_vec, qdd[dim, :])
177 min_accel = np.min(qdd[dim, :])
178 max_accel = np.max(qdd[dim, :])
179 legend.append("Acceleration")
180 if show_jerk:
181 plt.plot(t_vec, qddd[dim, :])
182 min_jerk = np.min(qddd[dim, :])
183 max_jerk = np.max(qddd[dim, :])
184 legend.append("Jerk")
185 plt.legend(legend)
187 # Times
188 min_val = vertical_line_scale_factor * min(
189 [min_pos, min_vel, min_accel, min_jerk]
190 )
191 max_val = vertical_line_scale_factor * max(
192 [max_pos, max_vel, max_accel, max_jerk]
193 )
194 plt.vlines(
195 [t[1] for t in self.segment_times],
196 min_val,
197 max_val,
198 colors="k",
199 linestyles=":",
200 linewidth=1.0,
201 )
203 plt.show()
206class CubicPolynomialTrajectory(PolynomialTrajectoryBase):
207 """Describes a cubic (3rd-order) polynomial trajectory."""
209 def __init__(self, t, q, qd=0.0):
210 """
211 Creates a cubic (3rd-order) polynomial trajectory.
213 Parameters
214 ----------
215 t : array-like
216 An array of length N corresponding to the time breakpoints.
217 q : array-like
218 A M-by-N array of M dimensions and N waypoints corresponding to the position breakpoints.
219 qd : array-like, optional
220 A M-by-N array of M dimensions and N waypoints corresponding to the velocity breakpoints.
221 It can also be a scalar or a 1D array to copy to all dimensions/breakpoints.
222 Defaults to zero endpoint velocities.
223 """
224 # Validate inputs.
225 num_segments = len(t) - 1
227 if len(q.shape) == 1:
228 q = q[np.newaxis, :]
229 if q.shape[1] != num_segments + 1: 229 ↛ 230line 229 didn't jump to line 230 because the condition on line 229 was never true
230 raise ValueError(
231 "Position breakpoints must have the same number of columns as the time breakpoints."
232 )
233 self.num_dims = q.shape[0]
235 if isinstance(qd, float):
236 qd = qd * np.ones_like(q)
237 if len(qd.shape) == 1:
238 qd = qd[np.newaxis, :]
240 # Fit cubic trajectory coefficients to each segment.
241 self.segment_times = []
242 self.coeffs = [[] for _ in range(self.num_dims)]
243 for idx in range(num_segments):
244 t_breakpoints = [t[idx], t[idx + 1]]
245 self.segment_times.append(t_breakpoints)
247 for dim in range(self.num_dims):
248 # These are in the form:
249 # q[0], q[T], qd[0], qd[T]
250 endpoints = np.array(
251 [
252 q[dim, idx],
253 q[dim, idx + 1],
254 qd[dim, idx],
255 qd[dim, idx + 1],
256 ]
257 )
258 T = t_breakpoints[1] - t_breakpoints[0]
259 M = np.array(
260 [
261 [0, 0, 0, 1],
262 [T**3, T**2, T, 1],
263 [0, 0, 1, 0],
264 [3 * T**2, 2 * T, 1, 0],
265 ]
266 )
267 coeffs = np.linalg.solve(M, endpoints)
268 self.coeffs[dim].append(coeffs)
271class QuinticPolynomialTrajectory(PolynomialTrajectoryBase):
272 """Describes a quintic (5th-order) polynomial trajectory."""
274 def __init__(self, t, q, qd=0.0, qdd=0.0):
275 """
276 Creates a quintic (5th-order) polynomial trajectory.
278 Parameters
279 ----------
280 t : array-like
281 An array of length N corresponding to the time breakpoints.
282 q : array-like
283 A M-by-N array of M dimensions and N waypoints corresponding to the position breakpoints.
284 qd : array-like, optional
285 A M-by-N array of M dimensions and N waypoints corresponding to the velocity breakpoints.
286 It can also be a scalar or a 1D array to copy to all dimensions/breakpoints.
287 Defaults to zero endpoint velocities.
288 qdd : array-like, optional
289 A M-by-N array of M dimensions and N waypoints corresponding to the acceleration breakpoints.
290 It can also be a scalar or a 1D array to copy to all dimensions/breakpoints.
291 Defaults to zero endpoint accelerations.
292 """
293 # Validate inputs.
294 num_segments = len(t) - 1
296 if len(q.shape) == 1:
297 q = q[np.newaxis, :]
298 if q.shape[1] != num_segments + 1: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true
299 raise ValueError(
300 "Position breakpoints must have the same number of columns as the time breakpoints."
301 )
302 self.num_dims = q.shape[0]
304 if isinstance(qd, float):
305 qd = qd * np.ones_like(q)
306 if len(qd.shape) == 1:
307 qd = qd[np.newaxis, :]
309 if isinstance(qdd, float):
310 qdd = qdd * np.ones_like(q)
311 if len(qdd.shape) == 1:
312 qdd = qdd[np.newaxis, :]
314 # Fit quintic trajectory coefficients to each segment.
315 self.segment_times = []
316 self.coeffs = [[] for _ in range(self.num_dims)]
317 for idx in range(num_segments):
318 t_breakpoints = [t[idx], t[idx + 1]]
319 self.segment_times.append(t_breakpoints)
321 for dim in range(self.num_dims):
322 # These are in the form:
323 # q[0], q[T], qd[0], qd[T], qdd[0], qdd[T]
324 endpoints = np.array(
325 [
326 q[dim, idx],
327 q[dim, idx + 1],
328 qd[dim, idx],
329 qd[dim, idx + 1],
330 qdd[dim, idx],
331 qdd[dim, idx + 1],
332 ]
333 )
334 T = t_breakpoints[1] - t_breakpoints[0]
335 M = np.array(
336 [
337 [0, 0, 0, 0, 0, 1],
338 [T**5, T**4, T**3, T**2, T, 1],
339 [0, 0, 0, 0, 1, 0],
340 [5 * T**4, 4 * T**3, 3 * T**2, 2 * T, 1, 0],
341 [0, 0, 0, 2, 0, 0],
342 [20 * T**3, 12 * T**2, 6 * T, 2, 0, 0],
343 ]
344 )
345 coeffs = np.linalg.solve(M, endpoints)
346 self.coeffs[dim].append(coeffs)