pymomentum.trs

TRS (Translation-Rotation-Scale) Utilities

This module provides utilities for working with TRS (Translation-Rotation-Scale) transforms in PyMomentum.

A TRS transform is represented as a tuple of three separate tensors:

Translation (t): tensor of shape […, 3] - 3D position offset
Rotation (r): tensor of shape […, 3, 3] - 3x3 rotation matrix
Scale (s): tensor of shape […, 1] - uniform scale factor

This representation provides efficient operations since rotation matrices can be directly used for transformations and inverted via transpose, avoiding quaternion to matrix conversions. It is also preferable to working with 4x4 matrices, because extracting the scale from a fully general 4x4 matrix is expensive.

Note that internally, momentum mostly uses skel_states which use quaternions rather than rotation matrices. However, many use cases (particularly for ML) prefer rotation matrices (such as the widely-used 6D rotation representation), so this library provides useful functionality for converting between the two.

Key features:

Creating TRS transforms from individual components (from_translation(), from_rotation_matrix(), from_scale())
Converting between TRS transforms and 4x4 transformation matrices (to_matrix(), from_matrix())
Performing transformations and operations (multiply(), inverse(), transform_points())
Interoperability with skeleton states (from_skeleton_state(), to_skeleton_state())
Fast inverse computation using rotation matrix transpose

Example:

Creating and using a TRS transform:

import torch
from pymomentum import trs

# Create identity transform
t, r, s = trs.identity()

# Create from translation
translation = torch.tensor([1.0, 2.0, 3.0])
t, r, s = trs.from_translation(translation)

# Transform points
points = torch.tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
transformed = trs.transform_points((t, r, s), points)

Note:

All functions in this module work with TRS transforms represented as tuples of (translation, rotation_matrix, scale) tensors.

pymomentum.trs.blend(trs_transforms: Sequence[tuple[Tensor, Tensor, Tensor]], weights: Tensor | None = None) → tuple[Tensor, Tensor, Tensor]

Blend multiple TRS transforms with the given weights.

Parameters:

trs_transforms (Sequence[TRSTransform]) – A sequence of TRS transform tuples to blend.
weights (torch.Tensor, optional) – The weights to use for blending. If not provided, equal weights are used. Should have shape [num_transforms] or […, num_transforms].

Returns:

The blended TRS transform.

Return type:

TRSTransform

pymomentum.trs.from_matrix(matrices: Tensor) → tuple[Tensor, Tensor, Tensor]

Convert 4x4 transformation matrices to TRS transforms. Assumes uniform scaling.

Parameters:: matrices (torch.Tensor) – A tensor of 4x4 matrices of shape […, 4, 4].
Returns:: TRS transform tuple (translation, rotation_matrix, scale).
Return type:: TRSTransform

pymomentum.trs.from_rotation_matrix(rotation_matrix: Tensor) → tuple[Tensor, Tensor, Tensor]

Create a TRS transform from rotation matrix.

Parameters:: rotation_matrix (torch.Tensor) – The rotation matrix component of shape […, 3, 3].
Returns:: TRS transform tuple (zero_translation, rotation_matrix, unit_scale).
Return type:: TRSTransform

pymomentum.trs.from_scale(scale: Tensor) → tuple[Tensor, Tensor, Tensor]

Create a TRS transform from scale.

Parameters:: scale (torch.Tensor) – The scale component of shape […, 1].
Returns:: TRS transform tuple (zero_translation, identity_rotation, scale).
Return type:: TRSTransform

pymomentum.trs.from_skeleton_state(skeleton_state: Tensor) → tuple[Tensor, Tensor, Tensor]

Convert skeleton state to TRS transform.

Parameters:: skeleton_state (torch.Tensor) – The skeleton state tensor of shape […, 8] containing (tx, ty, tz, rx, ry, rz, rw, s).
Returns:: TRS transform tuple (translation, rotation_matrix, scale).
Return type:: TRSTransform

pymomentum.trs.from_translation(translation: Tensor) → tuple[Tensor, Tensor, Tensor]

Create a TRS transform from translation.

Parameters:: translation (torch.Tensor) – The translation component of shape […, 3].
Returns:: TRS transform tuple (translation, identity_rotation, unit_scale).
Return type:: TRSTransform

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

Returns a TRS transform representing the identity transform.

Parameters:

size (Sequence[int], optional) – The size of each batch dimension in the output tensors. Defaults to None, which means the output will have no batch dimensions.
device (torch.device, optional) – The device on which to create the tensors. Defaults to None.
dtype (torch.dtype, optional) – The data type of the tensors. Defaults to None (float32).

Returns:

Identity TRS transform tuple.

Return type:

TRSTransform

pymomentum.trs.inverse(trs: tuple[Tensor, Tensor, Tensor]) → tuple[Tensor, Tensor, Tensor]

Compute the inverse of a TRS transform. This is efficient since rotation matrix inverse is just transpose.

Parameters:: trs (TRSTransform) – The TRS transform tuple to invert.
Returns:: The inverted TRS transform tuple.
Return type:: TRSTransform

pymomentum.trs.multiply(trs1: tuple[Tensor, Tensor, Tensor], trs2: tuple[Tensor, Tensor, Tensor]) → tuple[Tensor, Tensor, Tensor]

Multiply (compose) two TRS transforms.

Parameters:

trs1 (TRSTransform) – The first TRS transform tuple (t1, r1, s1).
trs2 (TRSTransform) – The second TRS transform tuple (t2, r2, s2).

Returns:

The composed TRS transform tuple representing trs1 * trs2.

Return type:

TRSTransform

pymomentum.trs.rotmat_from_euler_xyz(euler: Tensor) → Tensor

Create rotation matrices from XYZ Euler angles.

This function converts XYZ Euler angles to rotation matrices using the ZYX rotation order (also known as Tait-Bryan angles). The rotation is applied as follows: first around the X-axis, then Y-axis, then Z-axis.

Parameters:: euler (torch.Tensor) – Euler angles of shape […, 3] where the last dimension contains [rx, ry, rz] rotations in radians.
Returns:: Rotation matrices of shape […, 3, 3].
Return type:: torch.Tensor

pymomentum.trs.rotmat_inverse(r: Tensor) → Tensor

Compute the inverse of rotation matrices.

For orthogonal rotation matrices, the inverse is simply the transpose.

Parameters:: r (torch.Tensor) – Rotation matrices of shape […, 3, 3].
Returns:: Inverse rotation matrices of shape […, 3, 3].
Return type:: torch.Tensor

pymomentum.trs.rotmat_multiply(r1: Tensor, r2: Tensor) → Tensor

Multiply two rotation matrices.

Parameters:

r1 (torch.Tensor) – First rotation matrices of shape […, 3, 3].
r2 (torch.Tensor) – Second rotation matrices of shape […, 3, 3].

Returns:

Product r1 @ r2 of shape […, 3, 3].

Return type:

torch.Tensor

pymomentum.trs.rotmat_rotate_vector(r: Tensor, v: Tensor) → Tensor

Apply rotation matrices to vectors.

Parameters:

r (torch.Tensor) – Rotation matrices of shape […, 3, 3].
v (torch.Tensor) – Vectors of shape […, 3].

Returns:

Rotated vectors r @ v of shape […, 3].

Return type:

torch.Tensor

pymomentum.trs.slerp(trs0: tuple[Tensor, Tensor, Tensor], trs1: tuple[Tensor, Tensor, Tensor], t: Tensor) → tuple[Tensor, Tensor, Tensor]

Spherical linear interpolation between two TRS transforms.

Parameters:

trs0 (TRSTransform) – The first TRS transform tuple.
trs1 (TRSTransform) – The second TRS transform tuple.
t (torch.Tensor) – The interpolation factor where 0 <= t <= 1. t=0 corresponds to trs0, t=1 corresponds to trs1.

Returns:

The interpolated TRS transform.

Return type:

TRSTransform

pymomentum.trs.to_matrix(trs: tuple[Tensor, Tensor, Tensor]) → Tensor

Convert TRS transform to 4x4 transformation matrix.

Parameters:: trs (TRSTransform) – The TRS transform tuple to convert.
Returns:: A tensor containing 4x4 transformation matrices of shape […, 4, 4].
Return type:: torch.Tensor

pymomentum.trs.to_skeleton_state(trs: tuple[Tensor, Tensor, Tensor]) → Tensor

Convert TRS transform to skeleton state.

Parameters:: trs (TRSTransform) – The TRS transform tuple to convert.
Returns:: Skeleton state tensor of shape […, 8] containing (tx, ty, tz, rx, ry, rz, rw, s).
Return type:: torch.Tensor

pymomentum.trs.transform_points(trs: tuple[Tensor, Tensor, Tensor], points: Tensor) → Tensor

Transform 3D points by the TRS transform.

Parameters:

trs (TRSTransform) – The TRS transform tuple to use for transformation.
points (torch.Tensor) – The points to transform of shape […, 3].

Returns:

The transformed points of shape […, 3].

Return type:

torch.Tensor