Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/scipy/spatial/transform/_rotation_spline.py : 10%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import numpy as np
2from scipy.linalg import solve_banded
3from .rotation import Rotation
6def _create_skew_matrix(x):
7 """Create skew-symmetric matrices corresponding to vectors.
9 Parameters
10 ----------
11 x : ndarray, shape (n, 3)
12 Set of vectors.
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
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)
33def _angular_rate_to_rotvec_dot_matrix(rotvecs):
34 """Compute matrices to transform angular rates to rot. vector derivatives.
36 The matrices depend on the current attitude represented as a rotation
37 vector.
39 Parameters
40 ----------
41 rotvecs : ndarray, shape (n, 3)
42 Set of rotation vectors.
44 Returns
45 -------
46 ndarray, shape (n, 3, 3)
47 """
48 norm = np.linalg.norm(rotvecs, axis=1)
49 k = np.empty_like(norm)
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
58 skew = _create_skew_matrix(rotvecs)
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)
65 return result
68def _rotvec_dot_to_angular_rate_matrix(rotvecs):
69 """Compute matrices to transform rot. vector derivatives to angular rates.
71 The matrices depend on the current attitude represented as a rotation
72 vector.
74 Parameters
75 ----------
76 rotvecs : ndarray, shape (n, 3)
77 Set of rotation vectors.
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)
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
92 mask = ~mask
93 nm = norm[mask]
94 k1[mask] = 0.5 - nm ** 2 / 24
95 k2[mask] = 1 / 6 - nm ** 2 / 120
97 skew = _create_skew_matrix(rotvecs)
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)
104 return result
107def _angular_acceleration_nonlinear_term(rotvecs, rotvecs_dot):
108 """Compute the non-linear term in angular acceleration.
110 The angular acceleration contains a quadratic term with respect to
111 the derivative of the rotation vector. This function computes that.
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.
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)
130 k1 = np.empty_like(norm)
131 k2 = np.empty_like(norm)
132 k3 = np.empty_like(norm)
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
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
146 dp = dp[:, None]
147 k1 = k1[:, None]
148 k2 = k2[:, None]
149 k3 = k3[:, None]
151 return dp * (k1 * cp + k2 * ccp) + k3 * dccp
154def _compute_angular_rate(rotvecs, rotvecs_dot):
155 """Compute angular rates given rotation vectors and its derivatives.
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.
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)
172def _compute_angular_acceleration(rotvecs, rotvecs_dot, rotvecs_dot_dot):
173 """Compute angular acceleration given rotation vector and its derivatives.
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.
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))
192def _create_block_3_diagonal_matrix(A, B, d):
193 """Create a 3-diagonal block matrix as banded.
195 The matrix has the following structure:
197 DB...
198 ADB..
199 .ADB.
200 ..ADB
201 ...AD
203 The blocks A, B and D are 3-by-3 matrices. The D matrices has the form
204 d * I.
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.
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))
223 A_i = np.empty_like(A, dtype=int)
224 A_i[:] = ind[:, None]
225 A_i += 3 * (1 + ind_blocks[:, None, None])
227 A_j = np.empty_like(A, dtype=int)
228 A_j[:] = ind
229 A_j += 3 * ind_blocks[:, None, None]
231 B_i = np.empty_like(B, dtype=int)
232 B_i[:] = ind[:, None]
233 B_i += 3 * ind_blocks[:, None, None]
235 B_j = np.empty_like(B, dtype=int)
236 B_j[:] = ind
237 B_j += 3 * (1 + ind_blocks[:, None, None])
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)))
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
251class RotationSpline(object):
252 """Interpolate rotations with continuous angular rate and acceleration.
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.
259 Refer to [1]_ for math and implementation details.
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.
269 Methods
270 -------
271 __call__
273 References
274 ----------
275 .. [1] `Smooth Attitude Interpolation
276 <https://github.com/scipy/scipy/files/2932755/attitude_interpolation.pdf>`_
278 Examples
279 --------
280 >>> from scipy.spatial.transform import Rotation, RotationSpline
282 Define the sequence of times and rotations from the Euler angles:
284 >>> times = [0, 10, 20, 40]
285 >>> angles = [[-10, 20, 30], [0, 15, 40], [-30, 45, 30], [20, 45, 90]]
286 >>> rotations = Rotation.from_euler('XYZ', angles, degrees=True)
288 Create the interpolator object:
290 >>> spline = RotationSpline(times, rotations)
292 Interpolate the Euler angles, angular rate and acceleration:
294 >>> angular_rate = np.rad2deg(spline(times, 1))
295 >>> angular_acceleration = np.rad2deg(spline(times, 2))
296 >>> times_plot = np.linspace(times[0], times[-1], 100)
297 >>> angles_plot = spline(times_plot).as_euler('XYZ', degrees=True)
298 >>> angular_rate_plot = np.rad2deg(spline(times_plot, 1))
299 >>> angular_acceleration_plot = np.rad2deg(spline(times_plot, 2))
301 On this plot you see that Euler angles are continuous and smooth:
303 >>> import matplotlib.pyplot as plt
304 >>> plt.plot(times_plot, angles_plot)
305 >>> plt.plot(times, angles, 'x')
306 >>> plt.title("Euler angles")
307 >>> plt.show()
309 The angular rate is also smooth:
311 >>> plt.plot(times_plot, angular_rate_plot)
312 >>> plt.plot(times, angular_rate, 'x')
313 >>> plt.title("Angular rate")
314 >>> plt.show()
316 The angular acceleration is continuous, but not smooth. Also note that
317 the angular acceleration is not a piecewise-linear function, because
318 it is different from the second derivative of the rotation vector (which
319 is a piecewise-linear function as in the cubic spline).
321 >>> plt.plot(times_plot, angular_acceleration_plot)
322 >>> plt.plot(times, angular_acceleration, 'x')
323 >>> plt.title("Angular acceleration")
324 >>> plt.show()
325 """
326 # Parameters for the solver for angular rate.
327 MAX_ITER = 10
328 TOL = 1e-9
330 def _solve_for_angular_rates(self, dt, angular_rates, rotvecs):
331 angular_rate_first = angular_rates[0].copy()
333 A = _angular_rate_to_rotvec_dot_matrix(rotvecs)
334 A_inv = _rotvec_dot_to_angular_rate_matrix(rotvecs)
335 M = _create_block_3_diagonal_matrix(
336 2 * A_inv[1:-1] / dt[1:-1, None, None],
337 2 * A[1:-1] / dt[1:-1, None, None],
338 4 * (1 / dt[:-1] + 1 / dt[1:]))
340 b0 = 6 * (rotvecs[:-1] * dt[:-1, None] ** -2 +
341 rotvecs[1:] * dt[1:, None] ** -2)
342 b0[0] -= 2 / dt[0] * A_inv[0].dot(angular_rate_first)
343 b0[-1] -= 2 / dt[-1] * A[-1].dot(angular_rates[-1])
345 for iteration in range(self.MAX_ITER):
346 rotvecs_dot = _matrix_vector_product_of_stacks(A, angular_rates)
347 delta_beta = _angular_acceleration_nonlinear_term(
348 rotvecs[:-1], rotvecs_dot[:-1])
349 b = b0 - delta_beta
350 angular_rates_new = solve_banded((5, 5), M, b.ravel())
351 angular_rates_new = angular_rates_new.reshape((-1, 3))
353 delta = np.abs(angular_rates_new - angular_rates[:-1])
354 angular_rates[:-1] = angular_rates_new
355 if np.all(delta < self.TOL * (1 + np.abs(angular_rates_new))):
356 break
358 rotvecs_dot = _matrix_vector_product_of_stacks(A, angular_rates)
359 angular_rates = np.vstack((angular_rate_first, angular_rates[:-1]))
361 return angular_rates, rotvecs_dot
363 def __init__(self, times, rotations):
364 from scipy.interpolate import PPoly
366 if len(rotations) == 1:
367 raise ValueError("`rotations` must contain at least 2 rotations.")
369 times = np.asarray(times, dtype=float)
370 if times.ndim != 1:
371 raise ValueError("`times` must be 1-dimensional.")
373 if len(times) != len(rotations):
374 raise ValueError("Expected number of rotations to be equal to "
375 "number of timestamps given, got {} rotations "
376 "and {} timestamps."
377 .format(len(rotations), len(times)))
379 dt = np.diff(times)
380 if np.any(dt <= 0):
381 raise ValueError("Values in `times` must be in a strictly "
382 "increasing order.")
384 rotvecs = (rotations[:-1].inv() * rotations[1:]).as_rotvec()
385 angular_rates = rotvecs / dt[:, None]
387 if len(rotations) == 2:
388 rotvecs_dot = angular_rates
389 else:
390 angular_rates, rotvecs_dot = self._solve_for_angular_rates(
391 dt, angular_rates, rotvecs)
393 dt = dt[:, None]
394 coeff = np.empty((4, len(times) - 1, 3))
395 coeff[0] = (-2 * rotvecs + dt * angular_rates
396 + dt * rotvecs_dot) / dt ** 3
397 coeff[1] = (3 * rotvecs - 2 * dt * angular_rates
398 - dt * rotvecs_dot) / dt ** 2
399 coeff[2] = angular_rates
400 coeff[3] = 0
402 self.times = times
403 self.rotations = rotations
404 self.interpolator = PPoly(coeff, times)
406 def __call__(self, times, order=0):
407 """Compute interpolated values.
409 Parameters
410 ----------
411 times : float or array_like
412 Times of interest.
413 order : {0, 1, 2}, optional
414 Order of differentiation:
416 * 0 (default) : return Rotation
417 * 1 : return the angular rate in rad/sec
418 * 2 : return the angular acceleration in rad/sec/sec
420 Returns
421 -------
422 Interpolated Rotation, angular rate or acceleration.
423 """
424 if order not in [0, 1, 2]:
425 raise ValueError("`order` must be 0, 1 or 2.")
427 times = np.asarray(times, dtype=float)
428 if times.ndim > 1:
429 raise ValueError("`times` must be at most 1-dimensional.")
431 singe_time = times.ndim == 0
432 times = np.atleast_1d(times)
434 rotvecs = self.interpolator(times)
435 if order == 0:
436 index = np.searchsorted(self.times, times, side='right')
437 index -= 1
438 index[index < 0] = 0
439 n_segments = len(self.times) - 1
440 index[index > n_segments - 1] = n_segments - 1
441 result = self.rotations[index] * Rotation.from_rotvec(rotvecs)
442 elif order == 1:
443 rotvecs_dot = self.interpolator(times, 1)
444 result = _compute_angular_rate(rotvecs, rotvecs_dot)
445 elif order == 2:
446 rotvecs_dot = self.interpolator(times, 1)
447 rotvecs_dot_dot = self.interpolator(times, 2)
448 result = _compute_angular_acceleration(rotvecs, rotvecs_dot,
449 rotvecs_dot_dot)
450 else:
451 assert False
453 if singe_time:
454 result = result[0]
456 return result