pymomentum.quaternion

Quaternion Utilities

This module provides comprehensive utilities for working with quaternions in PyMomentum.

Quaternions are a mathematical representation of rotations in 3D space that offer several advantages over other rotation representations like Euler angles or rotation matrices:

No gimbal lock: Unlike Euler angles, quaternions don’t suffer from singularities
Compact representation: Only 4 components vs 9 for rotation matrices
Efficient composition: Quaternion multiplication is faster than matrix multiplication
Smooth interpolation: SLERP provides natural rotation interpolation

Quaternion Format

This module uses the (x, y, z, w) format where:

(x, y, z): Vector part representing the rotation axis scaled by sin(θ/2)
w: Scalar part representing cos(θ/2), where θ is the rotation angle

The identity quaternion is (0, 0, 0, 1), representing no rotation.

Core Operations

The module provides functions for:

Basic operations: multiply(), conjugate(), inverse(), normalize()
Conversions: from_axis_angle(), euler_xyz_to_quaternion(), from_rotation_matrix(), to_rotation_matrix()
Vector operations: rotate_vector(), from_two_vectors()
Interpolation: slerp(), blend()
Utilities: check(), split(), identity()

Example:

Basic quaternion operations:

import torch
from pymomentum import quaternion

# Create identity quaternion
q_identity = quaternion.identity()

# Create quaternion from axis-angle
axis_angle = torch.tensor([0.0, 0.0, 1.57])  # 90° rotation around Z
q_rot = quaternion.from_axis_angle(axis_angle)

# Rotate a vector
vector = torch.tensor([1.0, 0.0, 0.0])
rotated = quaternion.rotate_vector(q_rot, vector)

# Interpolate between quaternions
q_interp = quaternion.slerp(q_identity, q_rot, 0.5)

Note:

All functions expect quaternions as PyTorch tensors with the last dimension having size 4, following the (x, y, z, w) format. Most functions support batched operations for efficient processing of multiple quaternions.

pymomentum.quaternion.blend(quaternions: Tensor, weights_in: Tensor | None = None) → Tensor

Blend multiple quaternions together using the method described in https://stackoverflow.com/questions/12374087/average-of-multiple-quaternions and http://www.acsu.buffalo.edu/~johnc/ave_quat07.pdf.

Parameters:

quaternions – A tensor of shape (…, k, 4) representing the quaternions to blend.
weights_in – An optional tensor of shape (…, k) representing the weights for each quaternion. If not provided, all quaternions will be weighted equally.

Returns:

A tensor of shape (…, 4) representing the blended quaternion.

pymomentum.quaternion.check(q: Tensor) → None

Check if a tensor represents a quaternion.

Parameters:: q – A tensor representing a quaternion.

pymomentum.quaternion.check_and_normalize_weights(quaternions: Tensor, weights_in: Tensor | None = None) → Tensor

Check and normalize the weights for blending quaternions.

Parameters:

quaternions – A tensor of shape (…, k, 4) representing the quaternions to blend.
weights_in – An optional tensor of shape (…, k) representing the weights for each quaternion. If not provided, all quaternions will be weighted equally.

Returns:

A tensor of shape (…, k) representing the normalized weights.

pymomentum.quaternion.conjugate(q: Tensor) → Tensor

Conjugate a quaternion.

Parameters:: q – A quaternion ((x, y, z), w)).
Returns:: The conjugate.

pymomentum.quaternion.euler_xyz_to_quaternion(euler_xyz: Tensor) → Tensor

Convert Euler XYZ angles to a quaternion.

This function converts XYZ Euler angles to quaternions. The rotation order is X-Y-Z, meaning first rotate around X-axis, then Y-axis, then Z-axis.

Parameters:: euler_xyz – A tensor of shape (…, 3) representing the Euler XYZ angles in order [roll, pitch, yaw].
Returns:: A tensor of shape (…, 4) representing the quaternion in ((x, y, z), w) format.

pymomentum.quaternion.euler_zyx_to_quaternion(euler_zyx: Tensor) → Tensor

Convert Euler ZYX angles to a quaternion.

This function converts ZYX Euler angles (yaw-pitch-roll convention) to quaternions. The rotation order is Z-Y-X, meaning first rotate around Z-axis (yaw), then Y-axis (pitch), then X-axis (roll).

Parameters:: euler_zyx – A tensor of shape (…, 3) representing the Euler ZYX angles in order [yaw, pitch, roll].
Returns:: A tensor of shape (…, 4) representing the quaternion in ((x, y, z), w) format.

pymomentum.quaternion.from_axis_angle(axis_angle: Tensor) → Tensor

Convert an axis-angle tensor to a quaternion.

Parameters:: axis_angle – A tensor of shape (…, 3) representing the axis-angle.
Returns:: A tensor of shape (…, 4) representing the quaternion in ((x, y, z), w) format.

pymomentum.quaternion.from_rotation_matrix(matrices: Tensor, eta: float = 1e-06) → Tensor

Convert a rotation matrix to a quaternion using numerically stable method.

This implementation uses the robust algorithm that computes all four quaternion component candidates and selects the best-conditioned one, ensuring numerical stability across all rotation matrix configurations.

Parameters:

matrices – A tensor of shape (…, 3, 3) representing the rotation matrices.
eta – Numerical precision threshold (unused, kept for compatibility).

Returns:

A tensor of shape (…, 4) representing the quaternions in ((x, y, z), w) format.

pymomentum.quaternion.from_two_vectors(v1: Tensor, v2: Tensor) → Tensor

Construct a quaternion that rotates one vector into another.

Parameters:

v1 – The initial vector.
v2 – The target vector.

Returns:

A quaternion representing the rotation from v1 to v2.

pymomentum.quaternion.identity(size: Sequence[int] | None = None, device: device | None = None, dtype: dtype = torch.float32) → Tensor

Create a quaternion identity tensor.

Parameters:

sizes – A tuple of integers representing the size of the quaternion tensor.
device – The device on which to create the tensor.

Returns:

A quaternion identity tensor with the specified sizes and device.

pymomentum.quaternion.inverse(q: Tensor) → Tensor

Compute the inverse of a quaternion.

Uses numerical clamping to avoid division by very small numbers, improving numerical stability for near-zero quaternions.

Parameters:: q – A quaternion ((x, y, z), w)).
Returns:: The inverse.

pymomentum.quaternion.multiply(q1: Tensor, q2: Tensor) → Tensor

Multiply two quaternions together.

Normalizes input quaternions before multiplication for numerical stability. For performance-critical code where quaternions are guaranteed to be normalized, use multiply_assume_normalized().

Parameters:

q1 – A quaternion ((x, y, z), w)).
q2 – A quaternion ((x, y, z), w)).

Returns:

The normalized product q1*q2.

pymomentum.quaternion.multiply_assume_normalized(q1: Tensor, q2: Tensor) → Tensor

Multiply two quaternions together, assuming they are already normalized.

This is a performance-optimized version of multiply() that skips normalization of the input quaternions. Use this only when you are certain both quaternions are already normalized.

Parameters:

q1 – A normalized quaternion ((x, y, z), w)).
q2 – A normalized quaternion ((x, y, z), w)).

Returns:

The product q1*q2.

pymomentum.quaternion.multiply_backprop(q1: Tensor, q2: Tensor, grad_q: Tensor) → Tuple[Tensor, Tensor]

Custom backpropagation for quaternion multiplication.

Computes gradients for quaternion multiplication with proper handling of normalization.

This version normalizes the input quaternions. For performance-critical code where quaternions are guaranteed to be normalized, use multiply_backprop_assume_normalized().

Parameters:

q1 – The first quaternion tensor of shape (…, 4).
q2 – The second quaternion tensor of shape (…, 4).
grad_q – The gradient from the output of shape (…, 4).

Returns:

A tuple of (grad_q1, grad_q2) representing gradients with respect to the first and second quaternions respectively.

pymomentum.quaternion.multiply_backprop_assume_normalized(q1: Tensor, q2: Tensor, grad_q: Tensor) → Tuple[Tensor, Tensor]

Custom backpropagation for quaternion multiplication assuming unit quaternions.

Computes gradients for quaternion multiplication when both input quaternions are assumed to be normalized. This is more efficient than the general case but should only be used when quaternions are guaranteed to be unit quaternions.

Parameters:

q1 – The first normalized quaternion tensor of shape (…, 4).
q2 – The second normalized quaternion tensor of shape (…, 4).
grad_q – The gradient from the output of shape (…, 4).

Returns:

A tuple of (grad_q1, grad_q2) representing gradients with respect to the first and second quaternions respectively.

pymomentum.quaternion.normalize(q: Tensor) → Tensor

Normalize a quaternion.

Parameters:: q – A quaternion ((x, y, z), w)).
Returns:: The normalized quaternion.

pymomentum.quaternion.normalize_backprop(q: Tensor, grad: Tensor) → Tensor

Custom backpropagation for quaternion normalization.

This function computes gradients for quaternion normalization in a numerically stable way, avoiding potential issues with automatic differentiation when quaternions are near zero norm.

Parameters:

q – The input quaternion tensor of shape (…, 4).
grad – The gradient from the output of shape (…, 4).

Returns:

The gradient with respect to the input quaternion q.

pymomentum.quaternion.quaternion_to_xyz_euler(q: Tensor, eps: float = 1e-06) → Tensor

Parameters:: eps – a small number to avoid calling asin(1) or asin(-1). Should not be smaller than 1e-6 as this can cause NaN gradients for some models.

pymomentum.quaternion.rotate_vector(q: Tensor, v: Tensor) → Tensor

Rotate a vector by a quaternion.

Normalizes the input quaternion before rotation for numerical stability. For performance-critical code where quaternions are guaranteed to be normalized, use rotate_vector_assume_normalized().

Parameters:

q – (nBatch x k x 4) tensor with the quaternions in ((x, y, z), w) format.
v – (nBatch x k x 3) vector.

Returns:

(nBatch x k x 3) rotated vectors.

pymomentum.quaternion.rotate_vector_assume_normalized(q: Tensor, v: Tensor) → Tensor

Rotate a vector by a quaternion, assuming the quaternion is already normalized.

This is a performance-optimized version of rotate_vector() that skips normalization of the input quaternion. Use this only when you are certain the quaternion is already normalized.

Parameters:

q – (nBatch x k x 4) tensor with normalized quaternions in ((x, y, z), w) format.
v – (nBatch x k x 3) vector.

Returns:

(nBatch x k x 3) rotated vectors.

pymomentum.quaternion.rotate_vector_backprop(q: Tensor, v: Tensor, grad: Tensor) → Tuple[Tensor, Tensor]

Custom backpropagation for quaternion vector rotation.

Computes gradients for the quaternion rotation operation using the Euler-Rodrigues formula.

This version normalizes the input quaternion. For performance-critical code where quaternions are guaranteed to be normalized, use rotate_vector_backprop_assume_normalized().

Parameters:

q – The quaternion tensor of shape (…, 4).
v – The vector tensor of shape (…, 3).
grad – The gradient from the output of shape (…, 3).

Returns:

A tuple of (grad_q, grad_v) representing gradients with respect to the quaternion and vector respectively.

pymomentum.quaternion.rotate_vector_backprop_assume_normalized(q: Tensor, v: Tensor, grad: Tensor) → Tuple[Tensor, Tensor]

Custom backpropagation for quaternion vector rotation assuming unit quaternions.

This is a performance-optimized version of rotate_vector_backprop() that assumes the input quaternion is already normalized. Use this only when you are certain the quaternion is normalized to avoid numerical issues.

Parameters:

q – The normalized quaternion tensor of shape (…, 4).
v – The vector tensor of shape (…, 3).
grad – The gradient from the output of shape (…, 3).

Returns:

A tuple of (grad_q, grad_v) representing gradients with respect to the quaternion and vector respectively.

pymomentum.quaternion.slerp(q0: Tensor, q1: Tensor, t: Tensor) → Tensor

Perform spherical linear interpolation (slerp) between two quaternions.

Parameters:

q0 – The starting quaternion.
q1 – The ending quaternion.
t – The interpolation parameter, where 0 <= t <= 1. t=0 corresponds to q0, t=1 corresponds to q1.

Returns:

The interpolated quaternion.

pymomentum.quaternion.split(q: Tensor) → tuple[Tensor, Tensor]

Split a quaternion into its scalar and vector parts.

Parameters:: q – A tensor representing a quaternion.
Returns:: The scalar and vector parts of the quaternion.

pymomentum.quaternion.to_rotation_matrix(q: Tensor) → Tensor

Convert quaternions to 3x3 rotation matrices.

Parameters:: q – (nBatch x k x 4) tensor with the quaternions in ((x, y, z), w) format.
Returns:: (nBatch x k x 3 x 3) tensor with 3x3 rotation matrices.

pymomentum.quaternion.to_rotation_matrix_assume_normalized(q: Tensor) → Tensor

Convert quaternions to 3x3 rotation matrices.

Parameters:: q – (nBatch x k x 4) tensor with the quaternions in ((x, y, z), w) format.
Returns:: (nBatch x k x 3 x 3) tensor with 3x3 rotation matrices.