Calculates the natural cubic spline approximation to the batch of controls given. Also calculates its derivative. Example: t = torch.linspace(0, 1, 7) # (2, 1) are batch dimensions. 7 is the time dimension (of the same length as t). 3 is the channel dimension. x = torch.
| 232 | return t, a, b, c, d |
| 233 | |
| 234 | class NaturalCubicSpline: |
| 235 | """Calculates the natural cubic spline approximation to the batch of controls given. Also calculates its derivative. |
| 236 | |
| 237 | Example: |
| 238 | t = torch.linspace(0, 1, 7) |
| 239 | # (2, 1) are batch dimensions. 7 is the time dimension (of the same length as t). 3 is the channel dimension. |
| 240 | x = torch.rand(2, 1, 7, 3) |
| 241 | coeffs = natural_cubic_spline_coeffs(t, x) |
| 242 | |
| 243 | # ...at this point you can save the coeffs, put them through PyTorch's Datasets and DataLoaders, etc... |
| 244 | |
| 245 | spline = NaturalCubicSpline(coeffs) |
| 246 | |
| 247 | point = torch.tensor(0.4) |
| 248 | # will be a tensor of shape (2, 1, 3), corresponding to batch and channel dimensions |
| 249 | out = spline.derivative(point) |
| 250 | |
| 251 | point = torch.tensor([0.4, 0.5]) |
| 252 | # will be a tensor of shape (2, 1, 2, 3), corresponding to batch, time and channel dimensions |
| 253 | out = spline.derivative(point) |
| 254 | """ |
| 255 | |
| 256 | def __init__(self, coeffs, **kwargs): |
| 257 | """ |
| 258 | Arguments: |
| 259 | coeffs: As returned by `natural_cubic_spline_coeffs`. |
| 260 | """ |
| 261 | super(NaturalCubicSpline, self).__init__(**kwargs) |
| 262 | |
| 263 | t, a, b, c, d = coeffs |
| 264 | |
| 265 | self._t = t |
| 266 | self._a = a |
| 267 | self._b = b |
| 268 | self._c = c |
| 269 | self._d = d |
| 270 | |
| 271 | def evaluate(self, t): |
| 272 | maxlen = self._b.size(-2) - 1 |
| 273 | inners = torch.zeros((t.shape[0], t.shape[1], 3)).to(t.device) |
| 274 | for i_b in range(self._t.shape[0]): |
| 275 | index = torch.bucketize(t.detach()[i_b], self._t[i_b]) - 1 |
| 276 | index = index.clamp(0, maxlen) # clamp because t may go outside of [t[0], t[-1]]; this is fine |
| 277 | # will never access the last element of self._t; this is correct behaviour |
| 278 | fractional_part = t[i_b] - self._t[i_b][index] |
| 279 | fractional_part = fractional_part.unsqueeze(-1) |
| 280 | inner = self._c[i_b, index, :] + self._d[i_b, index, :] * fractional_part |
| 281 | inner = self._b[i_b, index, :] + inner * fractional_part |
| 282 | inner = self._a[i_b, index, :] + inner * fractional_part |
| 283 | inners[i_b] = inner |
| 284 | return inners |
| 285 | |
| 286 | def derivative(self, t, order=1): |
| 287 | fractional_part, index = self._interpret_t(t) |
| 288 | fractional_part = fractional_part.unsqueeze(-1) |
| 289 | if order == 1: |
| 290 | inner = 2 * self._c[..., index, :] + 3 * self._d[..., index, :] * fractional_part |
| 291 | deriv = self._b[..., index, :] + inner * fractional_part |