Mesh: Add differentiable shape-to-shape fitting

[loliverhennigh]

PhysicsNeMo Pull Request

Description

Adds differentiable shape-to-shape fitting for prealigned 3D triangle meshes.

The new tensor APIs include:

• arap_energy, a batched as-rigid-as-possible edge-graph energy with stable first-order envelope gradients.
• point_to_mesh_distance, with a memory-bounded Torch baseline and Warp BVH nearest-face search.
• fit_template_points, a fixed-step local/global fitter with closest-surface correspondences, ARAP regularization, matrix-free conjugate gradient solves, and an implicit first-order adjoint.

Mesh.fit_template provides a connectivity-preserving object wrapper. The tensor APIs are exported from physicsnemo.nn.functional; the Mesh wrapper preserves template connectivity, ordering, attached data, and topology caches while invalidating geometry-dependent caches.

The fitting scope is intentionally prealigned, unbatched 3D triangle surfaces with one-way template-vertex-to-target-surface fitting. DomainMesh is not exposed because its heterogeneous interior and boundary components do not define one unambiguous target surface. Nearest-face, triangle-region, and local-rotation choices are discrete; gradients are branchwise and valid almost everywhere. Higher-order differentiation is outside the fitting contract.

Warp accelerates nearest-face selection for CUDA float32. Local rotations, ARAP assembly, and the implicit conjugate gradient solves remain in Torch. The target BVH is rebuilt for each call and fitting step so the public API stays stateless; CUDA Graph capture is not part of the initial contract.

Validation completed:

• 444 focused and integration tests passed with 6 expected CUDA-only skips.
• Default-Inductor compiled forward and backward match eager execution.
• Torch/Warp CPU forward and gradient parity passed.
• torch.library.opcheck, finite-difference/gradcheck, dense/conjugate-gradient parity, and Mesh preservation tests passed.
• All pre-commit hooks passed.
• A local-source Sphinx dummy build completed successfully.

CUDA hardware was not available locally. CUDA-specific Warp parity, compiled execution, and device behavior require NVIDIA CI or reviewer hardware.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards. (Not applicable; no model code is changed.)

Dependencies

No new dependencies.

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.