Add differentiable Sobolev mesh deformation

[mehdiataei]

Summary

• Add sobolev_deform_points, SobolevDeformPoints, and Mesh.sobolev_deform as public APIs.
• Implement a matrix-free uniform-mass P1 Helmholtz filter with Jacobi-preconditioned conjugate-gradient solves.
• Add a CUDA Warp backend for segments, triangles, and tetrahedra. CUDA selects Warp by default when available, while CPU and higher-dimensional simplices use Torch.
• Provide first-order implicit differentiation through the displacement and source vertices. The Warp backward performs an adjoint CG solve and an analytic P1 geometry vector-Jacobian product.
• Use deterministic point-incidence gathers, periodic true-residual replacement, and per-batch normalization for reliable float32 and small-magnitude solves.
• Support optional fixed vertices, batched point tensors with shared topology, isolated points, and exact zero-length or empty-topology dense displacement.
• Add API documentation, 2D and 3D optimization examples, reproducible figures, benchmark registration, and focused tests.

Motivation

Direct optimization of dense mesh vertices can produce vertex-scale oscillations in shape updates and adjoints. This adds a mesh-aware deformation parameterization that filters those high-frequency modes while preserving broad shape changes and fixed-point constraints.

The Warp backend moves P1 assembly, matrix-free operator application, CG, and the analytic geometry pullback onto CUDA. Warp's optim.linear module disables generated backward kernels, so the implementation differentiates the converged linear system with an explicit adjoint solve instead of differentiating the CG iteration history.

User impact

Users can optimize candidate vertex coordinates through mesh.sobolev_deform(...) and receive smoother reverse-mode sensitivities. The smoothing length uses the same physical units as the mesh coordinates.

Torch is the default on CPU. Warp is the default on CUDA for segment, triangle, and tetrahedron meshes when it is available. Both backends support first-order gradients and torch.compile. CUDA Graph capture and higher-order derivatives at positive smoothing lengths are not supported.

Validation

• CUDA_VISIBLE_DEVICES=0 uv run --no-sync pytest -q test/nn/functional/geometry/test_point_sobolev.py test/mesh/transformations/test_sobolev.py
  • 50 passed
  • 18 expected skips for CUDA-only cases on CPU
• Torch and Warp forward and gradient parity passed for segments, triangles, and tetrahedra in float32 and float64.
• Tiny-field relative-tolerance, true-residual restart, empty-topology identity, fixed and isolated vertices, FakeTensor, CUDA Graph rejection, and full-graph compile tests passed.
• CUDA Inductor forward and backward completed successfully.
• Both standalone Sobolev optimization examples completed on CUDA. The 3D example also completed on CPU.
• Ruff check and format, Interrogate, Markdownlint, license checks, large-file checks, and import-linter passed through pre-commit.

[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.

[github-actions]

CODEOWNERS review map

Current for commit bf8979c87aab. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @loliverhennigh — 7 file(s)

• benchmarks/physicsnemo/nn/functional/registry.py
• docs/api/nn/functionals/geometry.rst
• physicsnemo/nn/functional/init.py
• physicsnemo/nn/functional/geometry/init.py
• physicsnemo/nn/functional/geometry/deform/init.py
• physicsnemo/nn/functional/geometry/deform/_sobolev_torch_impl.py
• physicsnemo/nn/functional/geometry/deform/sobolev.py

⏳ @megnvidia — 2 file(s)

• docs/api/mesh/transformations.rst
• docs/api/nn/functionals/geometry.rst

⏳ @peterdsharpe — 9 file(s)

• docs/api/mesh/transformations.rst
• examples/minimal/mesh/README.md
• examples/minimal/mesh/sobolev_shape_optimization.py
• examples/minimal/mesh/sobolev_surface_shape_optimization.py
• physicsnemo/mesh/README.md
• physicsnemo/mesh/mesh.py
• physicsnemo/mesh/transformations/deform/init.py
• physicsnemo/mesh/transformations/deform/sobolev.py
• test/mesh/transformations/test_sobolev.py

No CODEOWNER

• CHANGELOG.md
• docs/img/mesh/sobolev_adjoint_field.png
• docs/img/mesh/sobolev_adjoint_field.py
• docs/img/mesh/sobolev_adjoint_field_3d.png
• docs/img/mesh/sobolev_adjoint_field_3d.py
• test/nn/functional/geometry/test_point_sobolev.py

Comment /codeowners-info to refresh.

[mehdiataei]

@loliverhennigh Can you please review this? Thanks!

[greptile-apps]

Greptile Summary

Adds differentiable Sobolev mesh deformation across the functional and Mesh APIs.

• Implements matrix-free Torch and CUDA Warp backends with implicit first-order differentiation.
• Supports fixed and isolated vertices, batched points, empty topology, and exact zero-length displacement.
• Adds public exports, documentation, examples, benchmarks, and focused backend and compilation tests.

Important Files Changed

Filename	Overview
physicsnemo/nn/functional/geometry/deform/_sobolev_torch_impl.py	Implements P1 assembly, matrix-free PCG, implicit differentiation, and a direct zero-length branch that resolves the previously reported compiled-call behavior.
physicsnemo/nn/functional/geometry/deform/_warp_impl/sobolev_op.py	Implements CUDA Warp dispatch, solves, and analytic geometry pullback for supported simplex meshes.
physicsnemo/nn/functional/geometry/deform/sobolev.py	Defines the public multi-backend API, input normalization, validation, dispatch policy, and benchmark inputs.
physicsnemo/mesh/transformations/deform/sobolev.py	Integrates Sobolev deformation with Mesh while preserving the transformation API conventions.
test/nn/functional/geometry/test_point_sobolev.py	Covers numerical references, gradients, backend parity, compilation, dynamic runtime zero, fixed vertices, and edge cases.

_{Reviews (2): Last reviewed commit: "Handle dynamic zero Sobolev scales" | Re-trigger Greptile}

[loliverhennigh]

LGTM, no big comment but just wanted to confirm there is now warp implementation for this? I only see PyTorch.

[mehdiataei]

@loliverhennigh Yeah I didn't attempted the Warp backend due to the matrix-free Jacobi-PCG solves to keep it differentiable. Let me give a try I have an idea.

[mehdiataei]

@loliverhennigh I implemented the Warp backend. I'll try to upstream some of these to Warp later. Please take a look.