(t, x)
| 297 | |
| 298 | |
| 299 | def cubic_spline_coeffs(t, x): |
| 300 | # x should be a tensor of shape (..., length) |
| 301 | # Will return the b, two_c, three_d coefficients of the derivative of the cubic spline interpolating the path. |
| 302 | |
| 303 | length = x.size(-1) |
| 304 | |
| 305 | if length < 2: |
| 306 | # In practice this should always already be caught in __init__. |
| 307 | raise ValueError("Must have a time dimension of size at least 2.") |
| 308 | elif length == 2: |
| 309 | a = x[..., :1] |
| 310 | b = (x[..., 1:] - x[..., :1]) / (t[..., 1:] - t[..., :1]) |
| 311 | two_c = torch.zeros(*x.shape[:-1], 1, dtype=x.dtype, device=x.device) |
| 312 | three_d = torch.zeros(*x.shape[:-1], 1, dtype=x.dtype, device=x.device) |
| 313 | else: |
| 314 | # Set up some intermediate values |
| 315 | time_diffs = t[..., 1:] - t[..., :-1] |
| 316 | time_diffs_reciprocal = time_diffs.reciprocal() |
| 317 | time_diffs_reciprocal_squared = time_diffs_reciprocal ** 2 |
| 318 | three_path_diffs = 3 * (x[..., 1:] - x[..., :-1]) |
| 319 | six_path_diffs = 2 * three_path_diffs |
| 320 | path_diffs_scaled = three_path_diffs * time_diffs_reciprocal_squared[:, None, :] |
| 321 | |
| 322 | # Solve a tridiagonal linear system to find the derivatives at the knots |
| 323 | system_diagonal = torch.empty((x.shape[0], length), dtype=x.dtype, device=x.device) |
| 324 | system_diagonal[..., :-1] = time_diffs_reciprocal |
| 325 | system_diagonal[..., -1] = 0 |
| 326 | system_diagonal[..., 1:] += time_diffs_reciprocal |
| 327 | system_diagonal *= 2 |
| 328 | system_rhs = torch.empty_like(x) |
| 329 | system_rhs[..., :-1] = path_diffs_scaled |
| 330 | system_rhs[..., -1] = 0 |
| 331 | system_rhs[..., 1:] += path_diffs_scaled |
| 332 | knot_derivatives = tridiagonal_solve(system_rhs, time_diffs_reciprocal, |
| 333 | system_diagonal, time_diffs_reciprocal) |
| 334 | |
| 335 | # Do some algebra to find the coefficients of the spline |
| 336 | time_diffs_reciprocal = time_diffs_reciprocal[:, None, :] |
| 337 | time_diffs_reciprocal_squared = time_diffs_reciprocal_squared[:, None, :] |
| 338 | a = x[..., :-1] |
| 339 | b = knot_derivatives[..., :-1] |
| 340 | two_c = (six_path_diffs * time_diffs_reciprocal |
| 341 | - 4 * knot_derivatives[..., :-1] |
| 342 | - 2 * knot_derivatives[..., 1:]) * time_diffs_reciprocal |
| 343 | three_d = (-six_path_diffs * time_diffs_reciprocal |
| 344 | + 3 * (knot_derivatives[..., :-1] |
| 345 | + knot_derivatives[..., 1:])) * time_diffs_reciprocal_squared |
| 346 | |
| 347 | return a, b, two_c, three_d |
| 348 | |
| 349 | |
| 350 |
no test coverage detected