Aluode/PerceptionLabPortable
0
1import numpy as np
2from scipy.linalg import solve_banded
3from ._rotation import Rotation
4
5
6def _create_skew_matrix(x):
7 """Create skew-symmetric matrices corresponding to vectors.
8
9 Parameters
10 ----------
11 x : ndarray, shape (n, 3)
12 Set of vectors.
13
14 Returns
15 -------
16 ndarray, shape (n, 3, 3)
17 """
18 result = np.zeros((len(x), 3, 3))
19 result[:, 0, 1] = -x[:, 2]
20 result[:, 0, 2] = x[:, 1]
21 result[:, 1, 0] = x[:, 2]
22 result[:, 1, 2] = -x[:, 0]
23 result[:, 2, 0] = -x[:, 1]
24 result[:, 2, 1] = x[:, 0]
25 return result
26
27
28def _matrix_vector_product_of_stacks(A, b):
29 """Compute the product of stack of matrices and vectors."""
30 return np.einsum("ijk,ik->ij", A, b)
31
32
33def _angular_rate_to_rotvec_dot_matrix(rotvecs):
34 """Compute matrices to transform angular rates to rot. vector derivatives.
35
36 The matrices depend on the current attitude represented as a rotation
37 vector.
38
39 Parameters
40 ----------
41 rotvecs : ndarray, shape (n, 3)
42 Set of rotation vectors.
43
44 Returns
45 -------
46 ndarray, shape (n, 3, 3)
47 """
48 norm = np.linalg.norm(rotvecs, axis=1)
49 k = np.empty_like(norm)
50
51 mask = norm > 1e-4
52 nm = norm[mask]
53 k[mask] = (1 - 0.5 * nm / np.tan(0.5 * nm)) / nm**2
54 mask = ~mask
55 nm = norm[mask]
56 k[mask] = 1/12 + 1/720 * nm**2
57
58 skew = _create_skew_matrix(rotvecs)
59
60 result = np.empty((len(rotvecs), 3, 3))
61 result[:] = np.identity(3)
62 result[:] += 0.5 * skew
63 result[:] += k[:, None, None] * np.matmul(skew, skew)
64
65 return result
66
67
68def _rotvec_dot_to_angular_rate_matrix(rotvecs):
69 """Compute matrices to transform rot. vector derivatives to angular rates.
70
71 The matrices depend on the current attitude represented as a rotation
72 vector.
73
74 Parameters
75 ----------
76 rotvecs : ndarray, shape (n, 3)
77 Set of rotation vectors.
78
79 Returns
80 -------
81 ndarray, shape (n, 3, 3)
82 """
83 norm = np.linalg.norm(rotvecs, axis=1)
84 k1 = np.empty_like(norm)
85 k2 = np.empty_like(norm)
86
87 mask = norm > 1e-4
88 nm = norm[mask]
89 k1[mask] = (1 - np.cos(nm)) / nm ** 2
90 k2[mask] = (nm - np.sin(nm)) / nm ** 3
91
92 mask = ~mask
93 nm = norm[mask]
94 k1[mask] = 0.5 - nm ** 2 / 24
95 k2[mask] = 1 / 6 - nm ** 2 / 120
96
97 skew = _create_skew_matrix(rotvecs)
98
99 result = np.empty((len(rotvecs), 3, 3))
100 result[:] = np.identity(3)
101 result[:] -= k1[:, None, None] * skew
102 result[:] += k2[:, None, None] * np.matmul(skew, skew)
103
104 return result
105
106
107def _angular_acceleration_nonlinear_term(rotvecs, rotvecs_dot):
108 """Compute the non-linear term in angular acceleration.
109
110 The angular acceleration contains a quadratic term with respect to
111 the derivative of the rotation vector. This function computes that.
112
113 Parameters
114 ----------
115 rotvecs : ndarray, shape (n, 3)
116 Set of rotation vectors.
117 rotvecs_dot : ndarray, shape (n, 3)
118 Set of rotation vector derivatives.
119
120 Returns
121 -------
122 ndarray, shape (n, 3)
123 """
124 norm = np.linalg.norm(rotvecs, axis=1)
125 dp = np.sum(rotvecs * rotvecs_dot, axis=1)
126 cp = np.cross(rotvecs, rotvecs_dot)
127 ccp = np.cross(rotvecs, cp)
128 dccp = np.cross(rotvecs_dot, cp)
129
130 k1 = np.empty_like(norm)
131 k2 = np.empty_like(norm)
132 k3 = np.empty_like(norm)
133
134 mask = norm > 1e-4
135 nm = norm[mask]
136 k1[mask] = (-nm * np.sin(nm) - 2 * (np.cos(nm) - 1)) / nm ** 4
137 k2[mask] = (-2 * nm + 3 * np.sin(nm) - nm * np.cos(nm)) / nm ** 5
138 k3[mask] = (nm - np.sin(nm)) / nm ** 3
139
140 mask = ~mask
141 nm = norm[mask]
142 k1[mask] = 1/12 - nm ** 2 / 180
143 k2[mask] = -1/60 + nm ** 2 / 12604
144 k3[mask] = 1/6 - nm ** 2 / 120
145
146 dp = dp[:, None]
147 k1 = k1[:, None]
148 k2 = k2[:, None]
149 k3 = k3[:, None]
150
151 return dp * (k1 * cp + k2 * ccp) + k3 * dccp
152
153
154def _compute_angular_rate(rotvecs, rotvecs_dot):
155 """Compute angular rates given rotation vectors and its derivatives.
156
157 Parameters
158 ----------
159 rotvecs : ndarray, shape (n, 3)
160 Set of rotation vectors.
161 rotvecs_dot : ndarray, shape (n, 3)
162 Set of rotation vector derivatives.
163
164 Returns
165 -------
166 ndarray, shape (n, 3)
167 """
168 return _matrix_vector_product_of_stacks(
169 _rotvec_dot_to_angular_rate_matrix(rotvecs), rotvecs_dot)
170
171
172def _compute_angular_acceleration(rotvecs, rotvecs_dot, rotvecs_dot_dot):
173 """Compute angular acceleration given rotation vector and its derivatives.
174
175 Parameters
176 ----------
177 rotvecs : ndarray, shape (n, 3)
178 Set of rotation vectors.
179 rotvecs_dot : ndarray, shape (n, 3)
180 Set of rotation vector derivatives.
181 rotvecs_dot_dot : ndarray, shape (n, 3)
182 Set of rotation vector second derivatives.
183
184 Returns
185 -------
186 ndarray, shape (n, 3)
187 """
188 return (_compute_angular_rate(rotvecs, rotvecs_dot_dot) +
189 _angular_acceleration_nonlinear_term(rotvecs, rotvecs_dot))
190
191
192def _create_block_3_diagonal_matrix(A, B, d):
193 """Create a 3-diagonal block matrix as banded.
194
195 The matrix has the following structure:
196
197 DB...
198 ADB..
199 .ADB.
200 ..ADB
201 ...AD
202
203 The blocks A, B and D are 3-by-3 matrices. The D matrices has the form
204 d * I.
205
206 Parameters
207 ----------
208 A : ndarray, shape (n, 3, 3)
209 Stack of A blocks.
210 B : ndarray, shape (n, 3, 3)
211 Stack of B blocks.
212 d : ndarray, shape (n + 1,)
213 Values for diagonal blocks.
214
215 Returns
216 -------
217 ndarray, shape (11, 3 * (n + 1))
218 Matrix in the banded form as used by `scipy.linalg.solve_banded`.
219 """
220 ind = np.arange(3)
221 ind_blocks = np.arange(len(A))
222
223 A_i = np.empty_like(A, dtype=int)
224 A_i[:] = ind[:, None]
225 A_i += 3 * (1 + ind_blocks[:, None, None])
226
227 A_j = np.empty_like(A, dtype=int)
228 A_j[:] = ind
229 A_j += 3 * ind_blocks[:, None, None]
230
231 B_i = np.empty_like(B, dtype=int)
232 B_i[:] = ind[:, None]
233 B_i += 3 * ind_blocks[:, None, None]
234
235 B_j = np.empty_like(B, dtype=int)
236 B_j[:] = ind
237 B_j += 3 * (1 + ind_blocks[:, None, None])
238
239 diag_i = diag_j = np.arange(3 * len(d))
240 i = np.hstack((A_i.ravel(), B_i.ravel(), diag_i))
241 j = np.hstack((A_j.ravel(), B_j.ravel(), diag_j))
242 values = np.hstack((A.ravel(), B.ravel(), np.repeat(d, 3)))
243
244 u = 5
245 l = 5
246 result = np.zeros((u + l + 1, 3 * len(d)))
247 result[u + i - j, j] = values
248 return result
249
250
251class RotationSpline:
252 """Interpolate rotations with continuous angular rate and acceleration.
253
254 The rotation vectors between each consecutive orientation are cubic
255 functions of time and it is guaranteed that angular rate and acceleration
256 are continuous. Such interpolation are analogous to cubic spline
257 interpolation.
258
259 Refer to [1]_ for math and implementation details.
260
261 Parameters
262 ----------
263 times : array_like, shape (N,)
264 Times of the known rotations. At least 2 times must be specified.
265 rotations : `Rotation` instance
266 Rotations to perform the interpolation between. Must contain N
267 rotations.
268
269 Methods
270 -------
271 __call__
272
273 References
274 ----------
275 .. [1] `Smooth Attitude Interpolation
276 <https://github.com/scipy/scipy/files/2932755/attitude_interpolation.pdf>`_
277
278 Examples
279 --------
280 >>> from scipy.spatial.transform import Rotation, RotationSpline
281 >>> import numpy as np
282
283 Define the sequence of times and rotations from the Euler angles:
284
285 >>> times = [0, 10, 20, 40]
286 >>> angles = [[-10, 20, 30], [0, 15, 40], [-30, 45, 30], [20, 45, 90]]
287 >>> rotations = Rotation.from_euler('XYZ', angles, degrees=True)
288
289 Create the interpolator object:
290
291 >>> spline = RotationSpline(times, rotations)
292
293 Interpolate the Euler angles, angular rate and acceleration:
294
295 >>> angular_rate = np.rad2deg(spline(times, 1))
296 >>> angular_acceleration = np.rad2deg(spline(times, 2))
297 >>> times_plot = np.linspace(times[0], times[-1], 100)
298 >>> angles_plot = spline(times_plot).as_euler('XYZ', degrees=True)
299 >>> angular_rate_plot = np.rad2deg(spline(times_plot, 1))
300 >>> angular_acceleration_plot = np.rad2deg(spline(times_plot, 2))
301
302 On this plot you see that Euler angles are continuous and smooth:
303
304 >>> import matplotlib.pyplot as plt
305 >>> plt.plot(times_plot, angles_plot)
306 >>> plt.plot(times, angles, 'x')
307 >>> plt.title("Euler angles")
308 >>> plt.show()
309
310 The angular rate is also smooth:
311
312 >>> plt.plot(times_plot, angular_rate_plot)
313 >>> plt.plot(times, angular_rate, 'x')
314 >>> plt.title("Angular rate")
315 >>> plt.show()
316
317 The angular acceleration is continuous, but not smooth. Also note that
318 the angular acceleration is not a piecewise-linear function, because
319 it is different from the second derivative of the rotation vector (which
320 is a piecewise-linear function as in the cubic spline).
321
322 >>> plt.plot(times_plot, angular_acceleration_plot)
323 >>> plt.plot(times, angular_acceleration, 'x')
324 >>> plt.title("Angular acceleration")
325 >>> plt.show()
326 """
327 # Parameters for the solver for angular rate.
328 MAX_ITER = 10
329 TOL = 1e-9
330
331 def _solve_for_angular_rates(self, dt, angular_rates, rotvecs):
332 angular_rate_first = angular_rates[0].copy()
333
334 A = _angular_rate_to_rotvec_dot_matrix(rotvecs)
335 A_inv = _rotvec_dot_to_angular_rate_matrix(rotvecs)
336 M = _create_block_3_diagonal_matrix(
337 2 * A_inv[1:-1] / dt[1:-1, None, None],
338 2 * A[1:-1] / dt[1:-1, None, None],
339 4 * (1 / dt[:-1] + 1 / dt[1:]))
340
341 b0 = 6 * (rotvecs[:-1] * dt[:-1, None] ** -2 +
342 rotvecs[1:] * dt[1:, None] ** -2)
343 b0[0] -= 2 / dt[0] * A_inv[0].dot(angular_rate_first)
344 b0[-1] -= 2 / dt[-1] * A[-1].dot(angular_rates[-1])
345
346 for iteration in range(self.MAX_ITER):
347 rotvecs_dot = _matrix_vector_product_of_stacks(A, angular_rates)
348 delta_beta = _angular_acceleration_nonlinear_term(
349 rotvecs[:-1], rotvecs_dot[:-1])
350 b = b0 - delta_beta
351 angular_rates_new = solve_banded((5, 5), M, b.ravel())
352 angular_rates_new = angular_rates_new.reshape((-1, 3))
353
354 delta = np.abs(angular_rates_new - angular_rates[:-1])
355 angular_rates[:-1] = angular_rates_new
356 if np.all(delta < self.TOL * (1 + np.abs(angular_rates_new))):
357 break
358
359 rotvecs_dot = _matrix_vector_product_of_stacks(A, angular_rates)
360 angular_rates = np.vstack((angular_rate_first, angular_rates[:-1]))
361
362 return angular_rates, rotvecs_dot
363
364 def __init__(self, times, rotations):
365 from scipy.interpolate import PPoly
366
367 if rotations.single:
368 raise ValueError("`rotations` must be a sequence of rotations.")
369
370 if len(rotations) == 1:
371 raise ValueError("`rotations` must contain at least 2 rotations.")
372
373 times = np.asarray(times, dtype=float)
374 if times.ndim != 1:
375 raise ValueError("`times` must be 1-dimensional.")
376
377 if len(times) != len(rotations):
378 raise ValueError("Expected number of rotations to be equal to "
379 "number of timestamps given, "
380 f"got {len(rotations)} rotations "
381 f"and {len(times)} timestamps.")
382
383 dt = np.diff(times)
384 if np.any(dt <= 0):
385 raise ValueError("Values in `times` must be in a strictly "
386 "increasing order.")
387
388 rotvecs = (rotations[:-1].inv() * rotations[1:]).as_rotvec()
389 angular_rates = rotvecs / dt[:, None]
390
391 if len(rotations) == 2:
392 rotvecs_dot = angular_rates
393 else:
394 angular_rates, rotvecs_dot = self._solve_for_angular_rates(
395 dt, angular_rates, rotvecs)
396
397 dt = dt[:, None]
398 coeff = np.empty((4, len(times) - 1, 3))
399 coeff[0] = (-2 * rotvecs + dt * angular_rates
400 + dt * rotvecs_dot) / dt ** 3
401 coeff[1] = (3 * rotvecs - 2 * dt * angular_rates
402 - dt * rotvecs_dot) / dt ** 2
403 coeff[2] = angular_rates
404 coeff[3] = 0
405
406 self.times = times
407 self.rotations = rotations
408 self.interpolator = PPoly(coeff, times)
409
410 def __call__(self, times, order=0):
411 """Compute interpolated values.
412
413 Parameters
414 ----------
415 times : float or array_like
416 Times of interest.
417 order : {0, 1, 2}, optional
418 Order of differentiation:
419
420 * 0 (default) : return Rotation
421 * 1 : return the angular rate in rad/sec
422 * 2 : return the angular acceleration in rad/sec/sec
423
424 Returns
425 -------
426 Interpolated Rotation, angular rate or acceleration.
427 """
428 if order not in [0, 1, 2]:
429 raise ValueError("`order` must be 0, 1 or 2.")
430
431 times = np.asarray(times, dtype=float)
432 if times.ndim > 1:
433 raise ValueError("`times` must be at most 1-dimensional.")
434
435 singe_time = times.ndim == 0
436 times = np.atleast_1d(times)
437
438 rotvecs = self.interpolator(times)
439 if order == 0:
440 index = np.searchsorted(self.times, times, side='right')
441 index -= 1
442 index[index < 0] = 0
443 n_segments = len(self.times) - 1
444 index[index > n_segments - 1] = n_segments - 1
445 result = self.rotations[index] * Rotation.from_rotvec(rotvecs)
446 elif order == 1:
447 rotvecs_dot = self.interpolator(times, 1)
448 result = _compute_angular_rate(rotvecs, rotvecs_dot)
449 elif order == 2:
450 rotvecs_dot = self.interpolator(times, 1)
451 rotvecs_dot_dot = self.interpolator(times, 2)
452 result = _compute_angular_acceleration(rotvecs, rotvecs_dot,
453 rotvecs_dot_dot)
454 else:
455 assert False
456
457 if singe_time:
458 result = result[0]
459
460 return result
461