Mesh: Add differentiable shrinkwrap deformation by mehdiataei · Pull Request #1886 · NVIDIA/physicsnemo

mehdiataei · 2026-07-29T15:33:29Z

Summary

Add differentiable nearest-surface projection through shrinkwrap_points and Mesh.shrinkwrap
Provide a Torch reference search and an accelerated NVIDIA Warp search for safe float32 geometry
Replay the selected point-to-triangle projection with PyTorch autograd for source, target, weight, and offset gradients
Support masks, floating point weights, signed normal offsets, bounded search, batched source points, and Torch or Warp dispatch
Add user-facing API documentation, benchmarks, adversarial coverage, and two triangulated visual examples

Motivation

Mesh optimization workflows often need to conform selected vertices to a surface constraint without discarding the rest of the optimized geometry. This adds a connectivity-preserving projection that can participate in adjoint-based optimization with respect to both source and target coordinates.

Nearest-face selection is discrete. With that selection fixed, the projection is first-order differentiable through the continuous geometry.

API

repaired = source.shrinkwrap(
    target,
    point_weights=selection,
    offset=0.0,
    max_distance=None,
)

The tensor-level API is available as physicsnemo.nn.functional.shrinkwrap_points.

Visual examples

Validation

942 passed, 9 skipped across deformation, shrinkwrap, newly landed deformation-energy, mesh-wrapper, and functional-benchmark tests
Ruff lint and format checks
Docstring coverage
Markdown lint
License headers
Added-file size checks
Import contracts
Torch CPU and Warp CUDA geometry and adjoint agreement in both visual examples

copy-pr-bot · 2026-07-29T15:33:33Z

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 · 2026-07-29T15:33:43Z

CODEOWNERS review map

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

⏳ @loliverhennigh — 10 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/_shrinkwrap_torch_impl.py
physicsnemo/nn/functional/geometry/deform/_warp_impl/init.py
physicsnemo/nn/functional/geometry/deform/_warp_impl/shrinkwrap_kernels.py
physicsnemo/nn/functional/geometry/deform/_warp_impl/shrinkwrap_op.py
physicsnemo/nn/functional/geometry/deform/shrinkwrap.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/shrinkwrap_panel_conformance.py
examples/minimal/mesh/shrinkwrap_solid_surface.py
physicsnemo/mesh/README.md
physicsnemo/mesh/mesh.py
physicsnemo/mesh/transformations/deform/init.py
physicsnemo/mesh/transformations/deform/shrinkwrap.py
test/mesh/transformations/test_shrinkwrap.py

No CODEOWNER

CHANGELOG.md
docs/img/mesh/shrinkwrap_panel_conformance.png
docs/img/mesh/shrinkwrap_panel_conformance.py
docs/img/mesh/shrinkwrap_solid_surface.png
docs/img/mesh/shrinkwrap_solid_surface.py
test/nn/functional/geometry/test_point_shrinkwrap.py

Comment /codeowners-info to refresh.

greptile-apps · 2026-07-29T15:40:06Z

Greptile Summary

Adds a differentiable nearest-surface shrinkwrap deformation.

Introduces Torch and Warp nearest-face search implementations with PyTorch projection replay for autograd.
Exposes tensor-level and Mesh APIs supporting masks, weights, offsets, bounded searches, and batched source points.
Adds public documentation, benchmarks, tests, and visual examples.

Important Files Changed

Filename	Overview
physicsnemo/nn/functional/geometry/deform/_shrinkwrap_torch_impl.py	Implements robust Torch search and differentiable replay, but pre-normalization edge subtraction rejects valid extreme-coordinate triangles.
physicsnemo/nn/functional/geometry/deform/shrinkwrap.py	Adds validation, dispatch, and the public functional API; its validity check exposes the extreme-coordinate rejection.
physicsnemo/nn/functional/geometry/deform/_warp_impl/shrinkwrap_op.py	Adds Warp search with safety fallbacks, while repeating the overflow-prone raw edge subtraction in its safety check.
physicsnemo/mesh/transformations/deform/shrinkwrap.py	Adds a thin mesh-aware wrapper that preserves source connectivity and data.
test/nn/functional/geometry/test_point_shrinkwrap.py	Adds broad numerical, gradient, backend, and compilation coverage but does not exercise opposite-sign extremes within a target triangle.

_{Reviews (1): Last reviewed commit: "Mesh: Add differentiable shrinkwrap defo..." | Re-trigger Greptile}

Signed-off-by: Mehdi Ataei <ataei8@gmail.com>

megnvidia · 2026-07-30T19:33:19Z

+    )
+    wrapped.square().mean().backward()
+
+Source points can have shape ``(N, 3)`` or ``(B, N, 3)``. One triangle target


Suggested change

Source points can have shape ``(N, 3)`` or ``(B, N, 3)``. One triangle target

Source points use one of these shapes:

- ``(N, 3)``

- ``(B, N, 3)``

One triangle target is shared across the source batches. Keep the following

projection rules in mind:

- A positive ``offset`` follows target face winding.

- A finite ``max_distance`` leaves points at or beyond the cutoff unchanged.

then delete lines 325-327

megnvidia · 2026-07-30T19:34:41Z

+unchanged.
+
+Nearest-face selection, closest-feature changes, and distance gating are
+discrete. Between those transitions, gradients propagate through source


Suggested change

discrete. Between those transitions, gradients propagate through source

discrete. Between those transitions, gradients propagate through the following

values:

- Source points

- Selected target vertices

- Floating-point weights

- A tensor-valued ``offset``

Target connectivity is not differentiable.

megnvidia · 2026-07-30T19:34:57Z

+
+Nearest-face selection, closest-feature changes, and distance gating are
+discrete. Between those transitions, gradients propagate through source
+points, selected target vertices, floating point weights, and a tensor-valued


Suggested change

points, selected target vertices, floating point weights, and a tensor-valued

megnvidia · 2026-07-30T19:35:09Z

+Nearest-face selection, closest-feature changes, and distance gating are
+discrete. Between those transitions, gradients propagate through source
+points, selected target vertices, floating point weights, and a tensor-valued
+``offset``. Target connectivity is not differentiable.


Suggested change

``offset``. Target connectivity is not differentiable.

megnvidia · 2026-07-30T19:36:22Z

+unsafe coordinate magnitudes or face geometry.
+
+Shrinkwrap performs data-dependent validation and nearest-face search setup.
+CUDA executions with either backend are not supported inside CUDA Graph


Suggested change

CUDA executions with either backend are not supported inside CUDA Graph

Do not run CUDA executions with either backend inside CUDA Graph capture.

megnvidia · 2026-07-30T19:36:31Z

+
+Shrinkwrap performs data-dependent validation and nearest-face search setup.
+CUDA executions with either backend are not supported inside CUDA Graph
+capture.


Suggested change

capture.

megnvidia · 2026-07-30T19:44:01Z

+adjacent faces. Their closest point is the same, but a nonzero face-normal
+offset can differ.
+
+Float64 targets use the Torch search because Warp searches in float32. Safe


Suggested change

Float64 targets use the Torch search because Warp searches in float32. Safe

Search backend rules differ by dtype and geometry:

- ``float64`` targets use the Torch search because Warp searches in ``float32``.

- Warp searches safe ``float32`` coordinates unchanged.

- Warp falls back to Torch for unsafe coordinate magnitudes or face geometry.

megnvidia · 2026-07-30T19:47:17Z

+
+Torch provides the reference nearest-face search. Warp accelerates that search
+on CPU and CUDA. Both backends evaluate the selected projection with PyTorch
+in the input dtype. Exact ties at target edges or vertices can select different


on the fence about whether dtype should be formatted as dtype or not. Prevailing advice would be to stay consistent within this file, so I will leave that as is.

megnvidia · 2026-07-30T20:03:34Z

  and Warp provides first-order GPU kernels.
+- Adds differentiable nearest-surface shrinkwrap through
+  `shrinkwrap_points` and `Mesh.shrinkwrap`. Torch provides the reference
+  search, NVIDIA Warp accelerates float32 nearest-face queries on CPU and CUDA,


Suggested change

search, NVIDIA Warp accelerates float32 nearest-face queries on CPU and CUDA,

search, NVIDIA Warp accelerates ``float32`` nearest-face queries on CPU and CUDA.

megnvidia · 2026-07-30T20:04:14Z

+- Adds differentiable nearest-surface shrinkwrap through
+  `shrinkwrap_points` and `Mesh.shrinkwrap`. Torch provides the reference
+  search, NVIDIA Warp accelerates float32 nearest-face queries on CPU and CUDA,
+  and both backends replay the selected projection with PyTorch autograd.


Suggested change

and both backends replay the selected projection with PyTorch autograd.

Both backends replay the selected projection with PyTorch autograd.

peterdsharpe

A few comments

peterdsharpe · 2026-07-30T18:05:23Z

+        return False
+
+    target_f32 = target_points.detach()
+    target_magnitude_is_safe = torch.isfinite(target_f32).all() & (


[bug][P1] This magnitude check does not make Warp's closest-point search numerically safe. Warp's closest_point_to_triangle first forms O(s²) dot products and then products such as d1 * d4 - d3 * d2 (O(s⁴)), so coordinates well below _WARP_SAFE_QUERY_MAGNITUDE can still overflow; the inverse occurs at small scales.

For example, with s = 1e12,

target / s = [ [-4, -4, -4], [5, -5, 4], [1, -2, -4], [-1, 3, 2], [4, -2, -1], [-4, -2, -5], ] faces = [[0, 1, 2], [3, 4, 5]] query / s = [[0, -2, -4]]

all as float32, _float32_target_search_is_safe returns true, but Torch selects face 0 while Warp—and therefore the default CUDA path—selects face 1. Measuring the replayed outputs in float64 shows Warp's projection is 2.408× farther from the query. Explicit Warp on CPU reports no hit. In randomized targets admitted by this gate, I saw 2,011/4,608 face mismatches at scale 1e18 and 3,781/4,608 at 1e-10.

Could we either uniformly rescale into a range proven safe for all Warp intermediates, or conservatively fall back to Torch outside verified lower and upper edge-scale bounds? Please add CPU and CUDA regression cases at both ends of the accepted scale range.

peterdsharpe · 2026-07-30T18:05:23Z

+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Render selected cap repair on a broad optimized lid dome.


This file is effectively a second copy of examples/minimal/mesh/shrinkwrap_solid_surface.py, let's deduplicate.

peterdsharpe · 2026-07-31T13:58:52Z

@@ -0,0 +1,487 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.


This file breaks repo conventions; every other file in this folder is an ipynb.

peterdsharpe · 2026-07-31T13:58:58Z

@@ -0,0 +1,832 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.


This file breaks repo conventions; every other file in this folder is an ipynb.

peterdsharpe · 2026-07-31T14:00:45Z

 path. Because its backward uses atomic accumulation at shared vertices, Warp
 gradient results can have small run-to-run floating-point differences.

+Nearest-Surface Shrinkwrap


I doesn't look like shrinkwrap (the operation) is defined here; can we add a definition for users?

Mesh: Add differentiable shrinkwrap deformation

c120846

mehdiataei marked this pull request as ready for review July 29, 2026 15:37

mehdiataei requested review from loliverhennigh, megnvidia and peterdsharpe as code owners July 29, 2026 15:37

greptile-apps Bot reviewed Jul 29, 2026

View reviewed changes

Comment thread physicsnemo/nn/functional/geometry/deform/_shrinkwrap_torch_impl.py

mehdiataei assigned peterdsharpe Jul 30, 2026

Fix shrinkwrap at extreme target coordinates

f5593c4

Signed-off-by: Mehdi Ataei <ataei8@gmail.com>

megnvidia approved these changes Jul 30, 2026

View reviewed changes

peterdsharpe reviewed Jul 31, 2026

View reviewed changes

-Source points can have shape ``(N, 3)`` or ``(B, N, 3)``. One triangle target
+Source points use one of these shapes:
+- ``(N, 3)``
+- ``(B, N, 3)``
+One triangle target is shared across the source batches. Keep the following
+projection rules in mind:
+- A positive ``offset`` follows target face winding.
+- A finite ``max_distance`` leaves points at or beyond the cutoff unchanged.

-discrete. Between those transitions, gradients propagate through source
+discrete. Between those transitions, gradients propagate through the following
+values:
+- Source points
+- Selected target vertices
+- Floating-point weights
+- A tensor-valued ``offset``
+Target connectivity is not differentiable.

	CUDA executions with either backend are not supported inside CUDA Graph
	Do not run CUDA executions with either backend inside CUDA Graph capture.

-Float64 targets use the Torch search because Warp searches in float32. Safe
+Search backend rules differ by dtype and geometry:
+- ``float64`` targets use the Torch search because Warp searches in ``float32``.
+- Warp searches safe ``float32`` coordinates unchanged.
+- Warp falls back to Torch for unsafe coordinate magnitudes or face geometry.

	search, NVIDIA Warp accelerates float32 nearest-face queries on CPU and CUDA,
	search, NVIDIA Warp accelerates ``float32`` nearest-face queries on CPU and CUDA.

	and both backends replay the selected projection with PyTorch autograd.
	Both backends replay the selected projection with PyTorch autograd.

		@@ -0,0 +1,487 @@
		# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.

		@@ -0,0 +1,832 @@
		# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.

Uh oh!

Conversation

mehdiataei commented Jul 29, 2026

Summary

Motivation

API

Visual examples

Validation

Uh oh!

copy-pr-bot Bot commented Jul 29, 2026

Uh oh!

github-actions Bot commented Jul 29, 2026

CODEOWNERS review map

No CODEOWNER

Uh oh!

Uh oh!

greptile-apps Bot commented Jul 29, 2026

Greptile Summary

Important Files Changed

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

peterdsharpe left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants