Spaces:
Running on Zero
Running on Zero
File size: 2,723 Bytes
41ff959 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """Contains linear algebra related utility functions."""
from __future__ import annotations
import torch
from scipy.spatial.transform import Rotation
def rotation_matrices_from_quaternions(quaternions: torch.Tensor) -> torch.Tensor:
"""Convert batch of quaternions into rotation matrices.
Args:
quaternions: Quaternions in wxyz order with shape [..., 4].
Returns:
Rotation matrices with shape [..., 3, 3].
"""
device = quaternions.device
shape = quaternions.shape[:-1]
quaternions = quaternions / torch.linalg.norm(quaternions, dim=-1, keepdim=True)
real_part = quaternions[..., 0]
vector_part = quaternions[..., 1:]
vector_cross = get_cross_product_matrix(vector_part)
real_part = real_part[..., None, None]
matrix_outer = vector_part[..., :, None] * vector_part[..., None, :]
matrix_diag = real_part.square() * eyes(3, shape=shape, device=device)
matrix_cross_1 = 2 * real_part * vector_cross
matrix_cross_2 = vector_cross @ vector_cross
return matrix_outer + matrix_diag + matrix_cross_1 + matrix_cross_2
def quaternions_from_rotation_matrices(matrices: torch.Tensor) -> torch.Tensor:
"""Convert batch of rotation matrices to quaternions in wxyz order.
Args:
matrices: Rotation matrices with shape [..., 3, 3].
Returns:
Quaternions with shape [..., 4].
Note:
This operation is not differentiable and is performed on CPU.
"""
if matrices.shape[-2:] != (3, 3):
raise ValueError(f"matrices have invalid shape {matrices.shape}")
matrices_np = matrices.detach().cpu().numpy()
quaternions_np = Rotation.from_matrix(matrices_np.reshape(-1, 3, 3)).as_quat()
quaternions_np = quaternions_np[:, [3, 0, 1, 2]]
quaternions_np = quaternions_np.reshape(matrices_np.shape[:-2] + (4,))
return torch.as_tensor(quaternions_np, device=matrices.device, dtype=matrices.dtype)
def get_cross_product_matrix(vectors: torch.Tensor) -> torch.Tensor:
"""Generate cross product matrices for 3D vectors.
Args:
vectors: Input vectors with shape [..., 3].
Returns:
Cross-product matrices with shape [..., 3, 3].
"""
if vectors.shape[-1] != 3:
raise ValueError("Only 3-dimensional vectors are supported")
device = vectors.device
shape = vectors.shape[:-1]
unit_basis = eyes(3, shape=shape, device=device)
return torch.cross(vectors[..., :, None], unit_basis, dim=-2)
def eyes(
dim: int,
shape: tuple[int, ...],
device: torch.device | str | None = None,
) -> torch.Tensor:
"""Create a batch of identity matrices."""
return torch.eye(dim, device=device).broadcast_to(shape + (dim, dim)).clone()
|