| | 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 |
| |
|
