kimodo/geometry.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Rotation and representation conversions: axis-angle, quaternion, matrix, 6D continuous."""

import torch
import torch.nn.functional as F


def angle_to_Y_rotation_matrix(angle: torch.Tensor) -> torch.Tensor:
    """Build a rotation matrix around the Y axis from a scalar angle (radians).

    Shape: angle.shape + (3, 3).
    """
    cos, sin = torch.cos(angle), torch.sin(angle)
    one, zero = torch.ones_like(angle), torch.zeros_like(angle)
    mat = torch.stack((cos, zero, sin, zero, one, zero, -sin, zero, cos), -1)
    mat = mat.reshape(angle.shape + (3, 3))
    return mat


def matrix_to_cont6d(matrix: torch.Tensor) -> torch.Tensor:
    """Convert rotation matrix to 6D continuous representation (first two columns).

    Shape: (..., 3, 3) -> (..., 6).
    """
    cont_6d = torch.concat([matrix[..., 0], matrix[..., 1]], dim=-1)
    return cont_6d


def cont6d_to_matrix(cont6d: torch.Tensor) -> torch.Tensor:
    """Convert 6D continuous representation to rotation matrix (Gram–Schmidt on two columns).

    Last dim must be 6.
    """
    assert cont6d.shape[-1] == 6, "The last dimension must be 6"
    x_raw = cont6d[..., 0:3]
    y_raw = cont6d[..., 3:6]

    x = x_raw / torch.norm(x_raw, dim=-1, keepdim=True)
    z = torch.cross(x, y_raw, dim=-1)
    z = z / torch.norm(z, dim=-1, keepdim=True)

    y = torch.cross(z, x, dim=-1)

    x = x[..., None]
    y = y[..., None]
    z = z[..., None]

    mat = torch.cat([x, y, z], dim=-1)
    return mat


def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor:
    """Convert axis-angle to rotation matrix.

    Args:
        axis_angle: (..., 3) axis-angle vectors (angle = norm, axis = normalized)
    Returns:
        rotmat: (..., 3, 3) rotation matrices
    """
    eps = 1e-6
    angle = torch.norm(axis_angle, dim=-1, keepdim=True)  # (..., 1)
    axis = axis_angle / (angle + eps)

    x, y, z = axis.unbind(-1)

    zero = torch.zeros_like(x)
    K = torch.stack([zero, -z, y, z, zero, -x, -y, x, zero], dim=-1).reshape(*axis.shape[:-1], 3, 3)

    eye = torch.eye(3, device=axis.device, dtype=axis.dtype)
    eye = eye.expand(*axis.shape[:-1], 3, 3)

    sin = torch.sin(angle)[..., None]
    cos = torch.cos(angle)[..., None]

    R = eye + sin * K + (1 - cos) * (K @ K)
    return R


def matrix_to_axis_angle(R: torch.Tensor) -> torch.Tensor:
    """Convert rotation matrix to axis-angle via quaternions (more numerically stable).

    Args:
        R: (..., 3, 3) rotation matrices
    Returns:
        axis_angle: (..., 3)
    """
    # Go through quaternions for numerical stability
    quat = matrix_to_quaternion(R)  # (..., 4) with (w, x, y, z)
    return quaternion_to_axis_angle(quat)


def quaternion_to_axis_angle(quat: torch.Tensor) -> torch.Tensor:
    """Convert quaternion to axis-angle representation.

    Args:
        quat: (..., 4) quaternions with real part first (w, x, y, z)
    Returns:
        axis_angle: (..., 3)
    """
    eps = 1e-6

    # Ensure canonical form to avoid sign ambiguity.
    # Primary: prefer w > 0. When w ≈ 0 (angle ≈ π), prefer first nonzero xyz > 0.
    w = quat[..., 0:1]
    xyz = quat[..., 1:]

    # Find first significant component of xyz for tie-breaking when w ≈ 0
    first_significant = xyz[..., 0:1]  # use x component as tie-breaker

    # Flip if: w < 0, OR (w ≈ 0 AND first xyz component < 0)
    should_flip = (w < -eps) | ((w.abs() <= eps) & (first_significant < 0))
    quat = torch.where(should_flip, -quat, quat)

    w = quat[..., 0]
    xyz = quat[..., 1:]

    # sin(angle/2) = ||xyz||
    sin_half_angle = xyz.norm(dim=-1)

    # angle = 2 * atan2(sin(angle/2), cos(angle/2))
    # This is more stable than 2 * acos(w) near angle=0
    angle = 2.0 * torch.atan2(sin_half_angle, w)

    # axis = xyz / sin(angle/2), but handle small angles
    # For small angles: axis-angle ≈ 2 * xyz (since sin(x) ≈ x for small x)
    small_angle = sin_half_angle.abs() < eps

    # Safe division
    scale = torch.where(
        small_angle,
        2.0 * torch.ones_like(angle),  # small angle: axis_angle ≈ 2 * xyz
        angle / sin_half_angle.clamp(min=eps),
    )

    return xyz * scale.unsqueeze(-1)


def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
    """Returns torch.sqrt(torch.max(0, x)) subgradient is zero where x is 0."""
    return torch.sqrt(x * (x > 0).to(x.dtype))


def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
    """Convert rotations given as rotation matrices to quaternions.

    Args:
        matrix: Rotation matrices as tensor of shape (..., 3, 3).
    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")

    batch_dim = matrix.shape[:-2]
    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(matrix.reshape(batch_dim + (9,)), dim=-1)

    q_abs = _sqrt_positive_part(
        torch.stack(
            [
                1.0 + m00 + m11 + m22,
                1.0 + m00 - m11 - m22,
                1.0 - m00 + m11 - m22,
                1.0 - m00 - m11 + m22,
            ],
            dim=-1,
        )
    )

    quat_by_rijk = torch.stack(
        [
            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
        ],
        dim=-2,
    )

    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))

    return (
        (F.one_hot(q_abs.argmax(dim=-1), num_classes=4)[..., None] * quat_candidates)
        .sum(dim=-2)
        .reshape(batch_dim + (4,))
    )


def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
    """Convert rotations given as quaternions to rotation matrices.

    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    r, i, j, k = torch.unbind(quaternions, -1)
    two_s = 2.0 / (quaternions * quaternions).sum(-1)

    o = torch.stack(
        (
            1 - two_s * (j * j + k * k),
            two_s * (i * j - k * r),
            two_s * (i * k + j * r),
            two_s * (i * j + k * r),
            1 - two_s * (i * i + k * k),
            two_s * (j * k - i * r),
            two_s * (i * k - j * r),
            two_s * (j * k + i * r),
            1 - two_s * (i * i + j * j),
        ),
        -1,
    )
    return o.reshape(quaternions.shape[:-1] + (3, 3))