`soma.geometry`#

The soma.geometry package contains the lower-level geometry, rigging, interpolation, and Warp-accelerated kernels used internally by SOMALayer and PoseInversion.

This section is useful when you want to build custom fitting or retargeting workflows without going through the higher-level layer wrappers.

Transform and rig utilities#

Rotation and rigid-transform helpers used across SOMA geometry modules.

soma.geometry.transforms.compute_covariance(A, B, virtual_normal=True, eps=1e-8)#

Compute covariance matrix H = A^T @ B for rotation estimation.

Parameters:

A (Tensor) – Target vectors (…, N, 3)
B (Tensor) – Source vectors (…, N, 3)
virtual_normal (bool) – If True, add synthetic normal correspondence for conditioning
eps (float) – Small constant for numerical stability

Returns:

H – Covariance matrix (…, 3, 3)

Return type:

Tensor

soma.geometry.transforms.kabsch(H)#

Compute rotation matrix from covariance using Kabsch algorithm (SVD).

Parameters:: H (Tensor) – Covariance matrix (…, 3, 3)
Returns:: R – Rotation matrix (…, 3, 3) with det(R) = 1
Return type:: Tensor

soma.geometry.transforms.newton_schulz(H, num_iters=30, eps=1e-8)#

Compute rotation matrix from covariance using Newton-Schulz iteration.

This is primarily a reference implementation for testing and comparing against the Warp-accelerated Newton-Schulz kernel. Production callers should pair Newton-Schulz with validity checks or a reference-gauge policy when the covariance is rank deficient.

Parameters:

H (Tensor) – Covariance matrix (…, 3, 3)
num_iters (int) – Number of iterations (default 30)
eps (float) – Small constant for numerical stability

Returns:

R – Rotation matrix (…, 3, 3) with det(R) = 1

Return type:

Tensor

Note

Convergence depends on conditioning of H. Ill-conditioned matrices may require more iterations or may not converge to high precision.

soma.geometry.transforms.regularize_covariance_with_reference( H, reference_rotation=None, prior_strength=AUTO_ROTATION_PRIOR_STRENGTH, rank_threshold=AUTO_ROTATION_RANK_THRESHOLD, eps=1e-8, )#

Add a weak reference-gauge prior to a Procrustes covariance matrix.

soma.geometry.transforms.rotation_matrices_are_valid(R, det_tol=1e-2, orthogonality_tol=1e-2)#

Return a boolean mask for finite right-handed orthonormal rotations.

soma.geometry.transforms.rodrigues_rotation(a, b, eps=1e-8)#

Compute rotation matrix that aligns vector b to vector a.

Uses the shortest arc rotation approach similar to SciPy’s align_vectors.

Parameters:

a (Tensor) – Target vector (…, 3)
b (Tensor) – Source vector (…, 3)
eps (float) – Small constant for numerical stability

Returns:

R – Rotation matrix (…, 3, 3) such that R @ b ≈ a

Return type:

Tensor

soma.geometry.transforms.align_vectors(A, B, eps=1e-8, method='auto')#

SciPy-compatible: return rotation C such that C @ b ≈ a. Supports broadcasting across leading batch dims. Inputs: (…, N, 3).

Parameters:

A (Tensor) – Target vectors (…, N, 3)
B (Tensor) – Source vectors (…, N, 3)
eps (float) – Small constant for numerical stability
method (Literal['kabsch', 'newton-schulz', 'auto']) – ‘auto’, ‘kabsch’ (SVD-based), or ‘newton-schulz’ (iterative)

soma.geometry.transforms.SE3_from_Rt(R, t)#

autograd-safe SE(3) transform construction from rotation R and translation t. R: (…, 3, 3) t: (…, 3) Returns: T (…, 4, 4)

soma.geometry.transforms.SE3_inverse(T)#

Invert SE(3) transform(s) in homogeneous coordinates.

Parameters:: T (Tensor) – (…, 4, 4) torch.Tensor
Returns:: Tinv – (…, 4, 4)
Return type:: Tensor

soma.geometry.transforms.euler_xyz_to_matrix(euler_xyz)#

Convert XYZ Euler angles to rotation matrices.

Parameters:: euler_xyz (Tensor) – (…, 3) angles in radians, ordered as X, Y, Z.
Returns:: (…, 3, 3) rotation matrices.
Return type:: Tensor

soma.geometry.transforms.matrix_to_euler_xyz(R)#

Convert rotation matrices to XYZ Euler angles.

Parameters:: R (Tensor) – (…, 3, 3) rotation matrices.
Returns:: (…, 3) angles in radians, ordered as X, Y, Z.
Return type:: Tensor

soma.geometry.transforms.quaternion_normalize_xyzw(quaternion, eps=1e-12)#

Normalize XYZW quaternions.

Parameters:

quaternion (Tensor) – (…, 4) quaternions ordered as x, y, z, w.
eps (float) – Small constant used when normalizing quaternions.

Returns:

(…, 4) unit quaternions ordered as x, y, z, w.

Return type:

Tensor

soma.geometry.transforms.quaternion_conjugate_xyzw(quaternion)#

Return the conjugate of XYZW quaternions.

soma.geometry.transforms.quaternion_multiply_xyzw(a, b)#

Hamilton product of XYZW quaternions.

Parameters:

a (Tensor) – (…, 4) left quaternions ordered as x, y, z, w.
b (Tensor) – (…, 4) right quaternions ordered as x, y, z, w.

Returns:

(…, 4) product quaternions ordered as x, y, z, w.

Return type:

Tensor

soma.geometry.transforms.quaternion_xyzw_to_matrix(quaternion, eps=1e-12)#

Convert XYZW quaternions to rotation matrices.

Parameters:

quaternion (Tensor) – (…, 4) quaternions ordered as x, y, z, w.
eps (float) – Small constant used when normalizing quaternions.

Returns:

(…, 3, 3) rotation matrices.

Return type:

Tensor

soma.geometry.transforms.matrix_to_quaternion_xyzw(R, eps=1e-12)#

Convert rotation matrices to XYZW unit quaternions.

Parameters:

R (Tensor) – (…, 3, 3) rotation matrices.
eps (float) – Small constant used when normalizing quaternions.

Returns:

(…, 4) quaternions ordered as x, y, z, w with non-negative w.

Return type:

Tensor

soma.geometry.transforms.matrix_to_quaternion_xyzw_stable(R, eps=1e-12)#

Convert rotation matrices to XYZW quaternions with finite sqrt gradients.

This follows the same signed-branch convention as matrix_to_quaternion_xyzw but adds a tiny value inside each square root. It is intended for code paths that need gradients through matrix-to-quaternion conversion near zero-valued branches.

soma.geometry.transforms.single_axis_rotation_matrices(angles, axis, axis_signs)#

Create rotation matrices for angles around one shared local axis.

soma.geometry.transforms.quaternion_half_angle_xyzw(quaternion, eps=1e-12)#

Return the principal half-angle quaternion for XYZW rotations.

q and -q represent the same rotation. This helper first chooses the representation with non-negative w and then computes the normalized square-root quaternion sqrt(q) using [v, w + 1]. This is useful when extracting twist angles near 180 degrees without an unstable full-angle projection.

soma.geometry.transforms.quaternion_twist_angle_xyzw(quaternion, axis_ids=0, eps=1e-12)#

Extract signed twist angles around local axes from XYZW quaternions.

Parameters:

quaternion (Tensor) – (…, 4) quaternions ordered as x, y, z, w.
axis_ids (int | Tensor) – scalar axis id or tensor broadcastable to quaternion.shape[:-1]. Axis ids use 0=x, 1=y, 2=z.
eps (float) – Small constant used when normalizing quaternions.

Returns:

quaternion.shape[:-1] twist angles in radians.

Return type:

Tensor

soma.geometry.transforms.matrix_to_rotvec(R, eps=1e-6)#

(…,3,3) rotation matrices -> (…,3) rotation vectors (axis * angle). Robust for small angles and near-pi.

soma.geometry.transforms.rotvec_to_matrix(rotvec, eps=1e-8)#

(…,3) rotation vectors -> (…,3,3) rotation matrices. Robust near zero.

soma.geometry.transforms.rotation_6d_to_matrix(d6)#

Convert 6D rotation representation by Zhou et al. [1] to rotation matrix.

Uses Gram-Schmidt orthogonalization per Section B of [1].

Parameters:: d6 (Tensor) – 6D rotation representation, of size (*, 6).
Returns:: Batch of rotation matrices of size (*, 3, 3).
Return type:: Tensor

[1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. On the Continuity of Rotation Representations in Neural Networks. IEEE Conference on Computer Vision and Pattern Recognition, 2019. Retrieved from http://arxiv.org/abs/1812.07035

Skeleton hierarchy utilities for converting between local and world frames.

soma.geometry.rig_utils.joint_world_to_local( joint_world_transforms, joint_parent_ids, return_inverse=False, )#

Convert world space joint transforms to local space (relative to parent).

Parameters:

joint_world_transforms (Tensor) – (J, M, M) or (B, J, M, M) where M=4 or 3
joint_parent_ids (Sequence[int] | Tensor) – (J,) int array

Returns:

joint_local_transforms – same shape as joint_world_transforms

Return type:

Tensor | tuple[Tensor, Tensor]

soma.geometry.rig_utils.joint_local_to_world(joint_local_transforms, joint_parent_ids)#

Convert local space joint transforms (relative to parent) to world space.

Parameters:

joint_local_transforms (Tensor) – (J, M, M) or (B, J, M, M) where M=4 or 3
joint_parent_ids (Sequence[int] | Tensor) – (J,) int array or Python list[int] (list preferred for torch.export)

Returns:

joint_world_transforms – same shape as joint_local_transforms

Return type:

Tensor

soma.geometry.rig_utils.compute_skeleton_levels(joint_parent_ids, device=None)#

Group joints by tree depth for level-order forward kinematics.

Parameters:

joint_parent_ids (Sequence[int] | Tensor) – (J,) int array, tensor, or Python list of parent indices.
device (device | None) – torch device for the returned index tensors. Defaults to the device of joint_parent_ids when it is a tensor, otherwise "cpu".

Returns:

List of (joint_ids, parent_ids) tensor pairs, one per depth level. Level 0 contains the root(s); levels 1+ contain joints whose parents are at strictly lower levels.

Return type:

list[tuple[Tensor, Tensor]]

soma.geometry.rig_utils.joint_local_to_world_levelorder(joint_local_transforms, levels)#

Batched level-order forward kinematics.

Equivalent to joint_local_to_world but processes all joints at the same tree depth in a single batched bmm, reducing ~J sequential kernel launches to ~D (number of depth levels, typically 10-12 for humanoids).

Parameters:

joint_local_transforms (Tensor) – (J, M, M) or (B, J, M, M) where M = 3 or 4.
levels (list[tuple[Tensor, Tensor]]) – precomputed output of compute_skeleton_levels.

Returns:

joint_world_transforms – same shape as joint_local_transforms.

Return type:

Tensor

soma.geometry.rig_utils.precompute_joint_orient(joint_orient, joint_parent_ids)#

Precompute tensors for apply_joint_orient_local.

Parameters:

joint_orient (Tensor) – (J, 3, 3) world-space orientation per joint.
joint_parent_ids (Sequence[int] | Tensor) – (J,) int array or Python list of parent indices.

Returns:

(orient, orient_parent_T) – both (J, 3, 3) tensors ready to be passed to apply_joint_orient_local.

Return type:

tuple[Tensor, Tensor]

soma.geometry.rig_utils.apply_joint_orient_local(local_rotations, orient, orient_parent_T)#

Apply joint orient as a per-joint local operation (no FK loop).

Mathematically equivalent to:

world = joint_local_to_world(local_rotations, parent_ids)
world = world @ orient
local = joint_world_to_local(world, parent_ids)

but computed as a single parallel matmul:

R_out[j] = orient_parent_T[j] @ R_in[j] @ orient[j]

Parameters:

local_rotations (Tensor) – (B, J, 3, 3) local-space rotation matrices.
orient (Tensor) – (J, 3, 3) from precompute_joint_orient.
orient_parent_T (Tensor) – (J, 3, 3) from precompute_joint_orient.

Returns:

(B, J, 3, 3) oriented local rotations.

Return type:

Tensor

soma.geometry.rig_utils.remove_joint_orient_local(local_rotations, orient, orient_parent_T)#

Remove joint orient — inverse of apply_joint_orient_local.

Converts absolute local rotations (with orient baked in) back to T-pose-relative rotations:

R_rel[j] = orient_parent_T[j]^T @ R_abs[j] @ orient[j]^T

This is used to convert PoseInversion output (absolute) into the relative convention expected by SOMALayer.pose(absolute_pose=False).

Parameters:

local_rotations (Tensor) – (B, J, 3, 3) absolute local rotations.
orient (Tensor) – (J, 3, 3) from precompute_joint_orient.
orient_parent_T (Tensor) – (J, 3, 3) from precompute_joint_orient.

Returns:

(B, J, 3, 3) T-pose-relative rotations.

Return type:

Tensor

soma.geometry.rig_utils.get_joint_children_ids(joint_parent_ids)#

Given a list of joint parent IDs, return a list of lists of joint children IDs.

soma.geometry.rig_utils.get_joint_descendents(joint_parent_ids, joint_id)#

Given a list of joint parent IDs, return a list of all descendents of a given joint ID.

soma.geometry.rig_utils.get_body_part_vertex_ids( skinning_weights, joint_parent_ids, root_joint_id, include_root=True, weight_threshold=0.01, )#

Get the vertex IDs influenced by a body part defined by a root joint and its descendents.

Parameters:

skinning_weights (Tensor) – (V, J) array of skinning weights
joint_parent_ids (Sequence[int] | Tensor) – (J,) int array of joint parent indices
root_joint_id (int) – int, joint ID of the body part root
include_root (bool) – bool, whether to include vertices influenced by the root joint itself
weight_threshold (float) – float, minimum skinning weight to consider a vertex influenced by a joint

Returns:

vertex_ids – list of int, vertex IDs influenced by the body part

Return type:

list[int]

class soma.geometry.rig_utils.PoseMirror_SOMA(joint_names, root_name='Root')#

Bases: object

Mirrors world-space SOMA skeletal poses across the sagittal plane (YZ plane).

This class reflects joint positions and orientations to swap the “Left” and “Right” sides of the body, effectively creating a mirror image of the pose. It handles the complexities of retaining a valid Right-Handed Coordinate System (RHCS) after reflection.

Rig Assumptions:

World Up: +Y
World Forward: +Z
Bone Axis: Local +X axis points toward the child joint (along the bone).
Naming Convention: Symmetrical joints must start with “Left” or “Right”.

Mathematical Logic:

Global Reflection: Reflects world positions across the YZ plane ($$x o -x$$).
Chirality Correction: Flips the local Y-axis to restore the coordinate system determinant to +1 (fixing the “inside-out” matrix caused by reflection).
Axis Realignment: Applies 180-degree rotations to realign the bone vectors: - Limbs: Rotated 180° about Y to ensure Local +X points to the child. - Center: Rotated 180° about X to correct roll/twist on the mirror plane.
Swap: Swaps the memory of Left and Right joint indices.

Features:

Supports both NumPy arrays and PyTorch tensors (CPU/GPU).
Auto-detects backend and caches device-specific matrices.
Vectorized operations (no loops during execution).

Parameters:

joint_names (list[str]) – List of joint names in the order of the pose array.
root_name (str) – Name of the root joint. Defaults to “Root”.

class soma.geometry.rig_utils.PoseMirror_MHR(param_names, negate_params=_DEFAULT_MHR_NEGATE_PARAMS)#

Bases: object

Mirror MHR parameters across the sagittal plane (YZ plane).

Handles both pose parameters (l_ / r_ prefixed) and scale parameters (scale_l_ / scale_r_ prefixed), swapping left/right counterparts and negating parameters that reverse sign under reflection.

Features:

Supports both NumPy arrays and PyTorch tensors (CPU/GPU).
Pre-computes swap indices and sign vectors.
Vectorized execution.

Parameters:

param_names (list[str]) – List of MHR parameter names (pose and/or scale).
negate_params (set[str]) – Set of parameter names to negate (multiply by -1).

Skinning and skeleton fitting#

PyTorch linear blend skinning and Rodrigues conversion helpers.

soma.geometry.lbs.batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32)#

Calculates the rotation matrices for a batch of rotation vectors :param rot_vecs: array of N axis-angle vectors :type rot_vecs: torch.tensor Nx3

Returns:: R (torch.tensor Nx3x3) – The rotation matrices for the given axis-angle parameters
Return type:: Tensor

soma.geometry.lbs.lbs(bind_shape, skinning_weights, target_pose_world)#

Linear Blend Skinning (PyTorch).

Supported:

target batched only: bind_shape: (V, 3) skinning_weights: (V, J) # STRICT target_pose_world: (…, J, 4, 4)
all batched with same leading dims: bind_shape: (…, V, 3) skinning_weights: (V, J) # STRICT target_pose_world: (…, J, 4, 4)
fully unbatched: bind_shape: (V, 3) skinning_weights: (V, J) target_pose_world: (J, 4, 4)

Returns:: blended – (…, V, 3) # mirrors the active batch dims (from A or B)
Return type:: Tensor

Batched forward kinematics and linear blend skinning utilities.

class soma.geometry.batched_skinning.FKTopology( parent_ids, target_joint_indices=None, joint_orient=None, global_translation_joint_idx=None, bind_world_transforms=None, )#

Bases: object

Optional FK-only source topology for a target LBS rig.

BatchedSkinning always owns one target topology for LBS. A source topology describes a related FK rig whose world transforms can be expanded into the target topology by caller-owned logic.

soma.geometry.batched_skinning.topk_skinning( W, K=8, weight_eps=1e-12, sort_desc=True, pad_index=-1, dtype_idx=torch.int32, dtype_w=torch.float32, )#

Convert dense skinning weights (N, J) -> sparse top-K jointIndices/jointWeights.

Parameters:

W (Tensor) – (N, J) float tensor of skinning weights per vertex.
K (int) – number of influences per vertex to keep.
weight_eps (float) – prune tiny weights; also avoids div-by-zero in normalization.
sort_desc (bool) – sort the chosen K influences by descending weight (stable output).
pad_index (int) – index used when fewer than K nonzero weights exist for a vertex.
dtype_idx (dtype) – dtype for the joint-index output.
dtype_w (dtype) – dtype for the joint-weight output.

Returns:

idx_mat – (N, K) int32 w_mat: (N, K) float32

Return type:

tuple[Tensor, Tensor]

Skeleton fitting utilities for adapting a bind rig to a new rest shape.

class soma.geometry.skeleton_transfer.SkeletonTransfer( joint_parent_ids, bind_world_transforms, bind_shape, skinning_weights, rbf_kernel='linear', vertex_ids_to_exclude=None, freeze_rotations=None, skip_endjoints=True, use_sparse_rbf_matrix=True, use_warp_for_rotations=True, rotation_method='auto', skip_inverse_lbs=False, root_joint_idx=1, )#

Bases: Module

Fit joint positions and orientations to a new rest-shape surface.

Initialize a SkeletonTransfer instance for fitting a skeleton to new shapes. :param joint_parent_ids: (J,) int array of joint parent indices :param bind_world_transforms: (J, 4, 4) array of joint bind poses in world space :param bind_shape: (V, 3) array of vertex positions in bind pose :param skinning_weights: (V, J) array of skinning weights :param rbf_kernel: type of RBF kernel to use for joint position regression :param vertex_ids_to_exclude: (V,) int array of vertex ids to exclude from the joint position regressors :param freeze_rotations: list of joint ids to freeze to bind pose (for Warp mode) :param skip_endjoints: bool, whether to skip rotation fitting for end joints (for Warp mode) :param use_sparse_rbf_matrix: bool, whether to use a sparse RBF matrix for joint position regression :param use_warp_for_rotations: bool, whether to use Warp-based rotation fitting (requires Warp) :param rotation_method: str, rotation extraction method (‘auto’, ‘kabsch’, or ‘newton-schulz’) :param skip_inverse_lbs: bool, whether to skip Inverse LBS (skinned vertex fitting) and use identity R_init :param root_joint_idx: index of the root joint (0 for hand models, 1 for full-body SOMA)

update_bind(bind_world_transforms, bind_shape)#

Update bind-pose data without rebuilding structural caches.

This is much faster than constructing a new SkeletonTransfer and is suitable when only the identity (shape) changes but the topology, skinning weights, and skeleton structure remain the same.

fit(target_shapes)#

Fit the skeleton to new shapes by adjusting joint positions and orientations. :param target_shapes: (B, V, 3) or (V, 3) array of new vertex positions

Returns:: target_bind_world_transforms – (B, J, 4, 4) or (J, 4, 4) array of new bind poses in world space
Return type:: Tensor

fit_joint_positions(target_shapes)#

Fit the skeleton to new shapes by adjusting joint positions. :param target_shapes: (B, V, 3) or (V, 3) array of new vertex positions

Returns:: new_joint_positions – (B, J, 3) or (J, 3) array of new joint positions
Return type:: Tensor

fit_joint_rotations( new_joint_positions, target_shapes, )#

Fit the skeleton to new positions by adjusting joint orientations. :param new_joint_positions: (B, J, 3) or (J, 3) array of new joint positions :param target_shapes: (B, V, 3) or (V, 3) array of new vertex positions

Returns:: world_bind_pose – (B, J, 4, 4) or (J, 4, 4) array of new bind poses in world space
Return type:: Tensor

fit_rotations_warp( new_joint_positions, target_shapes, )#

Warp-accelerated version of fit_joint_rotations using GPU-parallel alignment.

Parameters:

new_joint_positions (Tensor) – (B, J, 3) or (J, 3) array of new joint positions
target_shapes (Tensor) – (B, V, 3) or (V, 3) array of new vertex positions

Returns:

world_bind_pose – (B, J, 4, 4) or (J, 4, 4) array of new bind poses in world space

Return type:

Tensor

Interpolation and mesh utilities#

Barycentric interpolation for mesh deformation transfer.

Example usage:

>>> import torch
>>> from soma.geometry.barycentric_interp import BarycentricInterpolator
>>>
>>> # Create interpolator
>>> interp = BarycentricInterpolator(V_src, F_src, V_dst)
>>>
>>> # Apply to deformed source mesh
>>> V_dst_deformed = interp(V_src_deformed)
>>>
>>> # Works with batches too
>>> V_dst_batch = interp(V_src_batch)  # (batch, n_verts, 3)

soma.geometry.barycentric_interp.fabricate_tet(p0, p1, p2, normal_scale='area')#

Fabricate a tetrahedron from triangle vertices by adding a point perpendicular to the triangle plane.

Parameters:

p0 (ndarray) – First triangle vertex, shape (…, 3)
p1 (ndarray) – Second triangle vertex, shape (…, 3)
p2 (ndarray) – Third triangle vertex, shape (…, 3)
normal_scale (str) – "area" uses the legacy raw cross-product normal. "edge" uses a unit normal scaled by local mean edge length.

Returns:

p3 – Fourth point of tetrahedron, shape (…, 3)

Return type:

ndarray

soma.geometry.barycentric_interp.compute_barycentric_coords_3d(p, v0, v1, v2, v3)#

Compute 3D barycentric coordinates for point p in tetrahedron (v0, v1, v2, v3).

Parameters:

p (ndarray) – Query points, shape (…, 3)
v0 (ndarray) – First tetrahedron vertex, shape (…, 3)
v1 (ndarray) – Second tetrahedron vertex, shape (…, 3)
v2 (ndarray) – Third tetrahedron vertex, shape (…, 3)
v3 (ndarray) – Fourth tetrahedron vertex, shape (…, 3)

Returns:

Barycentric coordinates, shape (…, 4)

Return type:

ndarray

soma.geometry.barycentric_interp.barycentric_interpolation(V_tet, F_tet, face_ids, bary_coords)#

Interpolate vertices using precomputed barycentric coordinates.

Parameters:

V_tet (Tensor) – Tetrahedralized vertices, shape (batch, n_verts + n_faces, 3) or (n_verts + n_faces, 3)
F_tet (Tensor) – Tetrahedron face indices, shape (n_faces, 4)
face_ids (Tensor) – Closest face indices for each target point, shape (n_target,)
bary_coords (Tensor) – Barycentric coordinates, shape (n_target, 4)

Returns:

Interpolated vertices, shape (batch, n_target, 3) or (n_target, 3)

Return type:

Tensor

class soma.geometry.barycentric_interp.BarycentricInterpolator( V_src, F_src, V_dst, correspondence_path=None, tet_normal_scale='area', )#

Bases: Module

Barycentric interpolation for transferring deformations from source to target mesh. Uses tetrahedral mesh construction and 3D barycentric coordinates.

Initialize barycentric interpolator.

Parameters:

V_src (Tensor) – Source mesh vertices, (n_src_verts, 3) - torch tensor
F_src (Tensor) – Source mesh faces, (n_faces, 3) - torch tensor
V_dst (Tensor) – Target mesh vertices, (n_dst_verts, 3) - torch tensor
correspondence_path (str | Path | None) – Optional path to load precomputed correspondence
tet_normal_scale (str) – Height scale for fabricated tetrahedra. "area" preserves legacy behavior; "edge" uses a local edge-length normal height.

compute_correspondence()#

Compute correspondence using trimesh for closest point queries.

load_correspondence(path)#

Load precomputed correspondence from file.

save_correspondence(path)#

Save computed correspondence to file.

forward(V_src_deformed)#

Apply barycentric interpolation to transfer deformation.

Parameters:: V_src_deformed (Tensor) – Deformed source vertices, shape (batch, n_src_verts, 3) or (n_src_verts, 3)
Returns:: Deformed target vertices, shape (batch, n_dst_verts, 3) or (n_dst_verts, 3)
Return type:: Tensor

class soma.geometry.interpolate.RadialBasisFunction( source_control_points, kernel='thin_plate_spline', kernel_params=None, include_polynomial=True, )#

Bases: object

RBF interpolator with precomputed system matrix (PyTorch).

get_basis_weights(query_point)#

Compute linear weights w such that interpolated_pos = sum(w_i * source_point_i).

Parameters:: query_point – (D,) vector representing the static template position of the joint (or point) being interpolated.
Returns:: (N,) vector of weights corresponding to the source_control_points.

interpolate( target_control_positions, query_points, )#

Interpolate query points from new control-point positions.

Parameters:

target_control_positions – (N, D) or (B, N, D).
query_points – (M, D).

Returns:

(M, D) if input was (N, D); (B, M, D) if input was (B, N, D).

Sparse Laplacian construction and Laplacian mesh editing helpers.

soma.geometry.laplacian.cotangent_weights(V, F)#

Compute cotangent weights for mesh Laplacian.

Parameters:

V (Tensor) – Vertices (n_verts, 3)
F (Tensor) – Faces (n_faces, 3)

Returns:

Sparse COO indices and values for cotangent Laplacian matrix

Return type:

tuple[Tensor, Tensor, Tensor]

soma.geometry.laplacian.build_cotangent_laplacian(V, F)#

Build cotangent Laplacian matrix in sparse CSR format.

Parameters:

V (Tensor) – Vertices (n_verts, 3)
F (Tensor) – Faces (n_faces, 3)

Returns:

Sparse CSR tensor of shape (n_verts, n_verts)

Return type:

Tensor

soma.geometry.laplacian.build_uniform_laplacian(F, n_verts, device=None, dtype=None)#

Build the uniform (graph) Laplacian: L = D - A, where A is the binary adjacency matrix and D is the diagonal degree matrix.

Unlike the cotangent Laplacian this is geometry-independent – it only depends on mesh connectivity.

Parameters:

F (Tensor) – Faces (n_faces, 3), int tensor
n_verts (int) – Number of vertices in the mesh
device (device | None) – Torch device (inferred from F if None)
dtype (dtype | None) – Value dtype (default float32)

Returns:

Sparse CSR tensor of shape (n_verts, n_verts)

Return type:

Tensor

soma.geometry.laplacian.power_laplacian(L, order)#

Compute L^order for higher-order Laplacian.

Parameters:

L (Tensor) – Sparse Laplacian matrix (CSR format)
order (int) – Power to raise L to

Returns:

L^order as sparse CSR tensor

Return type:

Tensor

class soma.geometry.laplacian.LaplacianMesh( V, F, mask_anchors, order=1, constraint_mode='hard', soft_weight=1e-5, jitter=0.0, solver='cholespy', )#

Bases: Module

Pure PyTorch implementation of Laplacian mesh editing. Supports both hard and soft constraints for deformation.

Initialize Laplacian mesh editor.

Parameters:

V (Tensor) – Reference vertices (n_verts, 3)
F (Tensor) – Faces (n_faces, 3)
mask_anchors (Tensor) – Boolean mask (n_verts,) - True for anchors, False for free
order (int) – Order of Laplacian (1 or 2)
constraint_mode (Literal['hard', 'soft']) – “hard” or “soft”
soft_weight (float) – Weight for soft constraints (only used if constraint_mode=”soft”)
jitter (float) – Small value added to diagonal for numerical stability
solver (Literal['cholespy', 'pytorch']) – “cholespy” (pre-factored Cholesky via torch.linalg.cholesky) or “pytorch” (torch.sparse.spsolve, requires cuDSS on CUDA)

solve(Vdef)#

Solve for deformed mesh given constraint vertices.

Parameters:: Vdef (Tensor) – Target vertices (B, n, 3) or (n, 3) For hard constraints: only values at constrained vertices are used For soft constraints: constrained vertices pull the solution
Returns:: Deformed vertices (B, n, 3) or (n, 3)
Return type:: Tensor

Warp backends#

Warp-based Linear Blend Skinning with torch.export support.

Changes from original: - Added @torch.library.custom_op decorators for torch.export compatibility - Registered fake implementations for shape inference during export - Wrapped forward/backward logic in custom operator functions - Public API (linear_blend_skinning) unchanged, fully backward compatible - Supports dynamic K (any number of bone influences)

Original implementation used torch.autograd.Function (not exportable). Current implementation uses torch.library.custom_op (exportable via torch.export).

Example usage:

from soma.geometry.lbs_warp import linear_blend_skinning

# Prepare inputs vertices = torch.rand(4, 1000, 3, device=’cuda’) # (batch, num_verts, 3) bone_weights = torch.rand(1000, 8, device=’cuda’) # (num_verts, K) - sparse bone_indices = torch.randint(0, 20, (1000, 8), device=’cuda’) # (num_verts, K) bone_transforms = torch.eye(4, device=’cuda’).unsqueeze(0).unsqueeze(0).expand(4, 20, 4, 4)

# Apply skinning posed_verts = linear_blend_skinning(vertices, bone_weights, bone_indices, bone_transforms)

# With gradients bone_transforms.requires_grad_(True) posed_verts = linear_blend_skinning(vertices, bone_weights, bone_indices, bone_transforms) posed_verts.sum().backward() # Compute gradients

Export example:

class LBSModule(torch.nn.Module):

def __init__(self, bone_weights, bone_indices):

super().__init__() self.register_buffer(‘bone_weights’, bone_weights) self.register_buffer(‘bone_indices’, bone_indices)

def forward(self, vertices, bone_transforms):

return linear_blend_skinning(vertices, self.bone_weights,: self.bone_indices, bone_transforms)

module = LBSModule(bone_weights, bone_indices) exported = torch.export.export(module, (vertices, bone_transforms)) torch.export.save(exported, “lbs.pt2”)

soma.geometry.lbs_warp.get_kernel( max_bones_count, vertices_scalar_dtype, weights_dtype, indices_dtype, vec3_dtype=wp.vec3, mat44_dtype=wp.mat44, )#

Generate the kernel for linear blend skinning.

soma.geometry.lbs_warp.linear_blend_skinning( vertices, bone_weights, bone_indices, bone_transforms, )#

Apply linear blend skinning to vertices using bone weights and transforms.

This function is compatible with both regular PyTorch usage and torch.export.

Parameters:

vertices (Tensor) – (batch_size, num_vertices, 3) vertex positions
bone_weights (Tensor) – (num_vertices, max_bones_per_vertex) bone weights
bone_indices (Tensor) – (num_vertices, max_bones_per_vertex) bone indices
bone_transforms (Tensor) – (batch_size, num_bones, 4, 4) bone transformation matrices

Returns:

Transformed vertices of shape (batch_size, num_vertices, 3)

Return type:

Tensor

Warp-based implementations of geometric transformations for improved GPU performance.

class soma.geometry.align_vectors_warp.RodriguesRotationWarp(*args, **kwargs)#

Bases: Function

Differentiable Rodrigues rotation using Warp kernel.

Computes rotation matrices from pairs of vectors using Rodrigues formula. SciPy-compatible: Returns R such that R @ B ≈ A. Fully differentiable with respect to input vectors.

static forward(ctx, A, B)#

Forward pass using Rodrigues formula.

Parameters:

A – (N, 3) - target vectors
B – (N, 3) - source vectors

Returns:

R – (N, 3, 3) - rotation matrices such that R @ B ≈ A

static backward(ctx, grad_output)#: Backward pass using Warp’s tape.

class soma.geometry.align_vectors_warp.AlignVectorsWarp(*args, **kwargs)#

Bases: Function

Differentiable Warp-based align_vectors using tape for gradients.

For N>=2 point sets using Kabsch algorithm (covariance + SVD). For N=1, use rodrigues_rotation_warp() instead.

Forward and backward passes both use Warp for full GPU acceleration.

static forward( ctx, A_flat, B_flat, offsets, counts, method='auto', )#

Forward pass using Warp kernel with tape recording.

Parameters:

A_flat – (total_vectors, 3) - flattened target vectors
B_flat – (total_vectors, 3) - flattened source vectors
offsets – (batch_size,) int32 - start index for each batch
counts – (batch_size,) int32 - number of vectors per batch (must be >= 2)
method – str - rotation extraction method (‘auto’, ‘kabsch’, or ‘newton-schulz’)

Returns:

R – (batch_size, 3, 3) - rotation matrices

static backward(ctx, grad_output)#: Backward pass using Warp’s tape.

soma.geometry.align_vectors_warp.rodrigues_rotation_warp(A, B)#

Compute rotation matrices from vector pairs using Rodrigues formula.

Differentiable Warp-based implementation. SciPy-compatible: Returns rotation R such that R @ B ≈ A.

Parameters:

A (Tensor) – (N, 3) torch tensor - target vectors
B (Tensor) – (N, 3) torch tensor - source vectors

Returns:

R – (N, 3, 3) rotation matrices such that R @ B ≈ A

Return type:

Tensor

Example

>>> A = torch.randn(100, 3, requires_grad=True)  # target
>>> B = torch.randn(100, 3, requires_grad=True)  # source
>>> R = rodrigues_rotation_warp(A, B)  # shape: (100, 3, 3)
>>> # Verify: R @ B ≈ A
>>> loss = R.pow(2).sum()
>>> loss.backward()

soma.geometry.align_vectors_warp.align_vectors_warp(A_flat, B_flat, offsets, counts, method='auto')#

Differentiable Warp-based batch alignment of vector sets with ragged support.

Uses auto, Kabsch, or Newton-Schulz rotation extraction for N>=2 point sets. For N=1 cases, use rodrigues_rotation_warp() instead.

Uses Warp kernel with tape for both forward and backward passes. Fully GPU-accelerated including gradient computation.

SciPy-compatible: returns rotation R such that R @ B ≈ A. Matches the convention in geometry.transforms.align_vectors().

Parameters:

A_flat (Tensor) – (total_vectors, 3) torch tensor - flattened target vectors
B_flat (Tensor) – (total_vectors, 3) torch tensor - flattened source vectors
offsets (Tensor) – (batch_size,) torch tensor (int32) - start index for each batch
counts (Tensor) – (batch_size,) torch tensor (int32) - number of vectors per batch (must be >= 2)
method (Literal['auto', 'kabsch', 'newton-schulz']) – str - rotation extraction method (‘auto’, ‘kabsch’, or ‘newton-schulz’, default: ‘auto’)

Returns:

R – (batch_size, 3, 3) rotation matrices such that R @ B ≈ A

Return type:

Tensor

Note

Differentiable: gradients flow through both forward and backward passes
Uses Warp’s tape mechanism for efficient gradient computation
Supports arbitrary batch sizes and variable vector counts (ragged batches)
For single vector pairs (N=1), use rodrigues_rotation_warp() instead

Example

>>> # Single batch with 10 vectors
>>> A = torch.randn(10, 3, requires_grad=True)
>>> B = torch.randn(10, 3, requires_grad=True)
>>> offsets = torch.tensor([0], dtype=torch.int32)
>>> counts = torch.tensor([10], dtype=torch.int32)
>>> R = align_vectors_warp(A, B, offsets, counts)  # shape: (1, 3, 3)
>>> loss = R.pow(2).sum()
>>> loss.backward()  # Gradients computed via Warp tape

>>> # Ragged batch: 5 vectors, 8 vectors, 3 vectors
>>> A_flat = torch.randn(16, 3, requires_grad=True)  # 5+8+3
>>> B_flat = torch.randn(16, 3, requires_grad=True)
>>> offsets = torch.tensor([0, 5, 13], dtype=torch.int32)
>>> counts = torch.tensor([5, 8, 3], dtype=torch.int32)
>>> R = align_vectors_warp(A_flat, B_flat, offsets, counts)  # shape: (3, 3, 3)

class soma.geometry.align_vectors_warp.ParallelRodriguesKabschWarp(*args, **kwargs)#

Bases: Function

Parallel execution of Rodrigues (N=1) and Kabsch (N>=2) rotations using two CUDA streams.

This function launches two independent Warp kernels in parallel: - Stream 1: Rodrigues rotation for single-child joints (N=1) - Stream 2: Covariance + Kabsch for multi-child joints (N>=2)

Both streams execute concurrently for improved performance when both types of joints exist.

static forward( ctx, rodrigues_tgt, rodrigues_src, kabsch_A_flat, kabsch_B_flat, kabsch_offsets, kabsch_counts, method='auto', )#

Forward pass with parallel kernel launches.

Parameters:

rodrigues_tgt – (N1, 3) target vectors for Rodrigues
rodrigues_src – (N1, 3) source vectors for Rodrigues
kabsch_A_flat – (N2_total, 3) flattened target vectors for Kabsch
kabsch_B_flat – (N2_total, 3) flattened source vectors for Kabsch
kabsch_offsets – (N2_batches,) offsets for Kabsch ragged batch
kabsch_counts – (N2_batches,) counts for Kabsch ragged batch
method – str - rotation extraction method (‘auto’, ‘kabsch’, or ‘newton-schulz’, default: ‘auto’)

Returns:

rodrigues_R – (N1, 3, 3) rotation matrices from Rodrigues kabsch_R: (N2_batches, 3, 3) rotation matrices from Kabsch

static backward( ctx, grad_rodrigues_R, grad_kabsch_R, )#: Backward pass using Warp’s tape.

soma.geometry.align_vectors_warp.parallel_rodrigues_kabsch_warp( rodrigues_tgt, rodrigues_src, kabsch_A_flat, kabsch_B_flat, kabsch_offsets, kabsch_counts, method='auto', )#

Execute Rodrigues and Kabsch rotations in parallel using two CUDA streams.

Parameters:

rodrigues_tgt (Tensor) – (N1, 3) target vectors for Rodrigues (can be None if no N=1 joints)
rodrigues_src (Tensor) – (N1, 3) source vectors for Rodrigues (can be None if no N=1 joints)
kabsch_A_flat (Tensor) – (N2_total, 3) flattened target vectors for Kabsch (can be None if no N>=2 joints)
kabsch_B_flat (Tensor) – (N2_total, 3) flattened source vectors for Kabsch (can be None if no N>=2 joints)
kabsch_offsets (Tensor) – (N2_batches,) offsets for Kabsch ragged batch
kabsch_counts (Tensor) – (N2_batches,) counts for Kabsch ragged batch
method (Literal['auto', 'kabsch', 'newton-schulz']) – str - rotation extraction method (‘auto’, ‘kabsch’, or ‘newton-schulz’, default: ‘auto’)

Returns:

rodrigues_R – (N1, 3, 3) rotation matrices from Rodrigues (empty tensor if no N=1 joints) kabsch_R: (N2_batches, 3, 3) rotation matrices from Kabsch (empty tensor if no N>=2 joints)

Return type:

tuple[Tensor, Tensor]

Example

>>> # Joint with single child (Rodrigues)
>>> tgt1 = torch.randn(10, 3, device='cuda', requires_grad=True)
>>> src1 = torch.randn(10, 3, device='cuda', requires_grad=True)
>>>
>>> # Joints with multiple children (Kabsch)
>>> A_flat = torch.randn(30, 3, device='cuda', requires_grad=True)  # 10+15+5
>>> B_flat = torch.randn(30, 3, device='cuda', requires_grad=True)
>>> offsets = torch.tensor([0, 10, 25], dtype=torch.int32, device='cuda')
>>> counts = torch.tensor([10, 15, 5], dtype=torch.int32, device='cuda')
>>>
>>> # Execute in parallel
>>> R1, R2 = parallel_rodrigues_kabsch_warp(tgt1, src1, A_flat, B_flat, offsets, counts)
>>> # R1.shape: (10, 3, 3), R2.shape: (3, 3, 3)

class soma.geometry.chamfer_warp.ChamferLoss#

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward( src_points, target_verts, target_faces=None, refit=False, )#

Compute Chamfer loss from source points to target mesh.

Parameters:

src_points – (B, N, 3) or (N, 3) source query points.
target_verts – (B, M, 3) or (M, 3) target vertices (supports batching).
target_faces – (F, 3) static topology (optional). If None, treats target as a point cloud by creating degenerate triangles (one per vertex).
refit – if True, rebuild the cached target meshes.

Fused Warp kernel for inverse-LBS refit: LBS + SE3_inv + src/tgt + covariance.

Replaces per-level: 2 LBS launches + ~10 PyTorch ops + 2 alignment launches With: 1 fused launch + 1 Newton-Schulz launch = 2 kernel launches per level.

soma.geometry.fused_refit_warp.fused_refit_level( bind_verts_cat, sub_bw_cat, sub_bi_cat, non_bw_cat, non_bi_cat, sub_w_sum_cat, arm_vids_cat, joint_offsets, joint_counts, joint_id_per_vert, joint_indices, D, W, target, J_level, rotation_method='auto', reference_rotations=None, auto_prior_strength=0.05, )#

Run fused LBS + covariance + rotation projection for one skeleton level.

bind_verts_cat can be unbatched (V_total, 3) (single identity) or batched (B, V_total, 3) (per-sample identity). Unbatched is expanded to (B, V_total, 3) for the kernel.

Returns: R_all (B, J_level, 3, 3) world-space rotations.

soma.geometry#

Transform and rig utilities#

Skinning and skeleton fitting#

Interpolation and mesh utilities#

Warp backends#

`soma.geometry`#