| import torch |
| from einops import rearrange |
| from jaxtyping import Float |
| from torch import Tensor |
|
|
|
|
| @torch.no_grad() |
| def generate_wobble_transformation( |
| radius: Float[Tensor, "*#batch"], |
| t: Float[Tensor, " time_step"], |
| num_rotations: int = 1, |
| scale_radius_with_t: bool = True, |
| ) -> Float[Tensor, "*batch time_step 4 4"]: |
| |
| tf = torch.eye(4, dtype=torch.float32, device=t.device) |
| tf = tf.broadcast_to((*radius.shape, t.shape[0], 4, 4)).clone() |
| radius = radius[..., None] |
| if scale_radius_with_t: |
| radius = radius * t |
| tf[..., 0, 3] = torch.sin(2 * torch.pi * num_rotations * t) * radius |
| tf[..., 1, 3] = -torch.cos(2 * torch.pi * num_rotations * t) * radius |
| return tf |
|
|
|
|
| @torch.no_grad() |
| def generate_wobble( |
| extrinsics: Float[Tensor, "*#batch 4 4"], |
| radius: Float[Tensor, "*#batch"], |
| t: Float[Tensor, " time_step"], |
| ) -> Float[Tensor, "*batch time_step 4 4"]: |
| tf = generate_wobble_transformation(radius, t) |
| return rearrange(extrinsics, "... i j -> ... () i j") @ tf |
|
|
