roma

import torch import roma # Example: Convert a batch of rotation vectors to unit quaternions and then to rotation matrices batch_shape = (2, 3) # Example: 2 batches of 3 rotations each # Generate random rotation vectors (3D tensor) rotvec = torch.randn(batch_shape + (3,)) print(f"Rotation vector shape: {rotvec.shape}") # Convert rotation vectors to unit quaternions (XYZW convention) q = roma.rotvec_to_unitquat(rotvec) print(f"Unit quaternion shape: {q.shape}") # Convert unit quaternions to rotation matrices (3x3 tensor) R = roma.unitquat_to_rotmat(q) print(f"Rotation matrix shape: {R.shape}") # Direct conversion from rotation vector to rotation matrix R_direct = roma.rotvec_to_rotmat(rotvec) print(f"Direct Rotation matrix shape: {R_direct.shape}") assert torch.allclose(R, R_direct, atol=1e-6) # Example: Special Procrustes orthonormalization # Projects an arbitrary 3x3 matrix onto the closest rotation matrix random_matrix = torch.randn(batch_shape + (3, 3)) R_procrustes = roma.special_procrustes(random_matrix) print(f"Procrustes orthonormalized matrix shape: {R_procrustes.shape}")