Skip to content

Add differentiable Procrustes registration#1818

Draft
loliverhennigh wants to merge 1 commit into
mainfrom
codex/rigid-registration
Draft

Add differentiable Procrustes registration#1818
loliverhennigh wants to merge 1 commit into
mainfrom
codex/rigid-registration

Conversation

@loliverhennigh

@loliverhennigh loliverhennigh commented Jul 9, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add physicsnemo.nn.functional.procrustes under geometry.deform for orientation-preserving rigid and similarity registration of point sets with known correspondences.
  • Support unbatched (N, D) and aligned batched (B, N, D) inputs for D in {1, 2, 3}, with optional nonnegative isotropic scale estimation.
  • Provide a Torch baseline and NVIDIA Warp backend. Automatic dispatch selects Torch on CPU and Warp on CUDA; the Warp path uses torch.library.custom_op, fake registration, and registered autograd.
  • Add backend parity, gradient, opcheck, compile, CUDA Graph, dispatch, benchmark-contract, validation, and numerical edge-case tests, plus API docs, changelog coverage, and a functional visualization.

Why

This is the closed-form Procrustes/Kabsch portion of rigid-registration roadmap item NVIDIA/physicsnemo-roadmap#2452. It supplies a differentiable alignment primitive for corresponding point sets while following the deformation package and backend conventions established by #1786.

PR #1786 has now merged, so this branch is rebased directly onto its merge commit on main and the diff contains only the Procrustes work.

Scope

This PR intentionally implements Procrustes only. ICP remains a separate follow-up because soft=True does not yet define the correspondence relaxation, temperature, iteration and stopping policy, outlier handling, scale policy, gradient semantics, or transform return type.

The returned transform follows the row-vector convention:

aligned = scale * (source @ rotation.transpose(-1, -2)) + translation

Implementation notes

  • The Torch baseline uses a determinant-corrected SVD with an analytic first-order polar-factor VJP.
  • The Warp projector uses the closed-form SO(2) solution in 2D and a double-precision Davenport/Jacobi quaternion solve in 3D. Float32 inputs promote only the tiny batched 2x2/3x3 covariance projection and return the public dtype afterward.
  • The Warp VJP solves the polar-factor Sylvester equation from the saved symmetric factor. Both backends explicitly reject unsupported higher-order differentiation instead of returning a partial Hessian.

Validation

  • pytest -q test/nn/functional/geometry/test_point_procrustes.py: 93 passed, 5 skipped (CUDA-only)
  • pre-commit run --from-ref origin/main --to-ref HEAD: all hooks passed
  • Torch/Warp forward and source/target gradient parity: float32/float64, 1D/2D/3D, rigid/similarity, batched/unbatched
  • Raw Warp custom-op opcheck, fake-tensor/AOT tracing, fullgraph AOT-eager forward/backward, finite-difference gradcheck, and unsupported second-derivative behavior passed
  • Independent stress checks covered tens of thousands of random covariance matrices, reflections, repeated spectra, rank-2 elongated clouds, and extreme float32 coordinate magnitudes
  • Existing point-deformation suite: 75 passed, 3 skipped; its 3 default-Inductor compile cases cannot run from this local checkout because PyTorch's generated clang command does not quote the workspace path containing spaces

CUDA hardware was unavailable locally, so no GPU performance claim is made; CUDA dispatch and graph-capture coverage remains in the test suite for CI.

@copy-pr-bot

copy-pr-bot Bot commented Jul 9, 2026

Copy link
Copy Markdown

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.

@loliverhennigh
loliverhennigh force-pushed the codex/rigid-registration branch 2 times, most recently from 1466a0a to 53240a9 Compare July 9, 2026 21:22
Signed-off-by: Oliver Hennigh <loliverhennigh101@gmail.com>
@loliverhennigh
loliverhennigh force-pushed the codex/rigid-registration branch from 53240a9 to 0759ed7 Compare July 9, 2026 21:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant