Mesh: differentiable dense displacement and sparse control-point morphing#1786
Conversation
Greptile SummaryThis PR adds differentiable mesh morphing to PhysicsNeMo via two new operations:
Important Files Changed
Reviews (1): Last reviewed commit: "Mesh: differentiable dense displacement ..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
Review in-progress! Found a few edge-case bugs; added failing tests to demonstrate. (Will follow up shortly with comments) |
peterdsharpe
left a comment
peterdsharpe
left a comment
There was a problem hiding this comment.
Follow-up to the review above with the chunked-backward analysis (the earlier review went out without its summary, so it's included here).
Summary of the full review — 15 inline comments: 6 correctness items, 1 dispatch-diagnostics item, 8 maintainability items. I pushed two commits (ce5463b, b9d6d5b) adding failing regression tests that pin down the confirmed correctness items; CI is expected to be red on exactly these seven until they're addressed:
test_torch_control_chunked_gradients_match_unchunked_on_ties— cross-chunk tie gradient corruption (comment below; one-line fix suggested)test_torch_compile_rejects_invalid_python_scalars_like_eagertest_infinite_tensor_radius_agrees_across_backendstest_mesh_point_data_path_through_leaf_has_actionable_diagnostictest_mesh_displace_rejects_weights_resolving_to_tensordicttest_mesh_morph_rejects_non_tensor_control_pointstest_domain_morph_clones_domain_global_data
Overall this is a very carefully engineered piece of numerics — the overflow-safe normalized distances and the detached-reference compensation algebra are impressive. The correctness items are all in the seams (chunk boundaries, compile-vs-eager, backend parity, API-boundary validation) rather than the core math.
|
Heads up @mehdiataei : there is a change coming to how we wrap warp streams in physicsnemo, seen here: #1771 tl;dr - we will have a small custom wrapper to replace ScopedStream. If your pr merges before #1771 we may want to patch it after. Want to make sure you're aware. You can read the summary in #1771 to understand the motivation but it's stream safety issues. |
peterdsharpe
left a comment
peterdsharpe
left a comment
There was a problem hiding this comment.
Second batch: API/design review (the first batch covered correctness). The implementation quality here is genuinely high — strict validation with actionable errors, radius required with no scale-presuming default, displacements-not-destinations stated everywhere, and the raw Shepard field evaluator deliberately kept private (easy to add later, impossible to remove). The comments below are all about the public contract, since that's the part we maintain forever; five of them carry GitHub suggestions that can be applied directly.
Summary of asks, roughly by weight:
- Keep
DisplacePoints/MorphPointsout of the top-levelnn.functionalnamespace (subpackage export matches every other spec class). - Drop
amount(linearity makes it strictly redundant) and renameweights→point_weights. - Stub the
Mesh.displace/Mesh.morphmethod docstrings per thetranslateconvention — the four full copies are already drifting. - Drop the Warp backend for
DisplacePoints— auto-dispatch never selects it, and a fused axpy doesn't earn ~8 kernels of permanent contract. - Curate the new
transformationspackage exports deliberately (include or excluderotation_matrix/scale_matrix; absolute imports like the siblings). - Add
kernel="wendland_c2"to the morph APIs to name the algorithm at call sites and reserve the extension point.
One open question rather than an ask: DomainMesh.morph exists but DomainMesh.displace doesn't — the weights-as-common-key mechanism would transfer directly, so if the asymmetry is intentional (dense per-component displacement doesn't share a world-space field), a line in the DomainMesh docs saying so would preempt the feature request.
|
Thanks @peterdsharpe for the review! All addressed. Please take a look. |
peterdsharpe
left a comment
peterdsharpe
left a comment
There was a problem hiding this comment.
Looking great - thank you for addressing those issues! Just a few last things to flag, in addition to the unresolved threads (one on the f32/f64 kernel deduplication, and one on the dispatch-fallback duplication), adding just a few more cleanup things.
We should probably add a DomainMesh tutorial showing off some of these cool morph features, but let's save that for a follow-up PR - this one is already beefy and lands the numerical core. Looking forward to having this merged!
…hing (Torch + Warp) Signed-off-by: Mehdi Ataei <ataei8@gmail.com>
Each test captures a confirmed defect found while reviewing the dense-displacement / control-point morphing implementation: - torch.compile silently skips the Python-scalar radius/amount validation that eager execution enforces, producing wrong geometry instead of the eager ValueError - the Torch and Warp backends return opposite geometry for a tensor radius containing +inf (zero vs. full displacement) - point_data paths that descend into a leaf tensor surface raw tensordict internals instead of the designed KeyError diagnostic - a weights key resolving to a nested TensorDict silently produces a Mesh whose points attribute is a TensorDict rather than a tensor - Mesh.morph reports AttributeError instead of a descriptive TypeError for non-tensor control points, unlike DomainMesh.morph - DomainMesh.morph aliases the source domain's global_data instead of cloning it like every other DomainMesh operation
A query point exactly equidistant from two controls in different control chunks receives wrong point/control/radius gradients from the chunked eager path: the reference selection in _compact_shepard_query_chunk uses strict '<' while the differentiable running minimum uses torch.minimum, whose backward splits the gradient 50/50 across exact ties, breaking the detached-reference compensation algebra. Unchunked, Warp, and central finite differences all agree; the chunked path is the outlier. Forward values and displacement gradients match, so the corruption is silent.
Signed-off-by: Mehdi Ataei <ataei8@gmail.com>
Signed-off-by: Mehdi Ataei <ataei8@gmail.com>
Signed-off-by: Mehdi Ataei <ataei8@gmail.com>
Signed-off-by: Mehdi Ataei <ataei8@gmail.com>
Auto-dispatch now resolves the device-based default and delegates to FunctionSpec.dispatch instead of duplicating its explicit-implementation branch. When CUDA inputs fall back to Torch because Warp is unavailable, the standard one-time RuntimeWarning from _warn_fallback now fires instead of silently selecting the slower backend; the dispatch test asserts the warning on the fallback path.
Rewrite the float32/float64 kernel twins as generic Warp kernels (typing.Any annotations with type()-derived scalar constants) and instantiate both precisions with wp.overload, so one numerical definition serves both dtypes by construction. The exported kernel names and launch signatures are unchanged, so op.py needs no edits. The only behavioral cleanup is the no-active-control scan: instead of comparing against per-precision maximum-float sentinels, the reference minimum initializes to 2.0 (active distances always satisfy q < 1) and reference_index == -1 is the authoritative no-active flag. Forward values, auxiliaries, and gradients are unchanged; the full parity, gradcheck, opcheck, and CUDA-graph suites pass on both precisions.
peterdsharpe
left a comment
peterdsharpe
left a comment
There was a problem hiding this comment.
Looks good to me! Thanks for working through the review process!
|
/ok to test 8e29ea8 |
Signed-off-by: Mehdi Ataei <ataei8@gmail.com>
|
/ok to test 4e6cb98 |
for more information, see https://pre-commit.ci
|
/ok to test b5f56d3 |
Signed-off-by: Mehdi Ataei <ataei8@gmail.com>
|
/ok to test 24a9348 |
7 participants
PhysicsNeMo Pull Request
Description
Adds differentiable mesh morphing to physicsnemo.nn.functional and the mesh API:
Both ops support float32/float64, unbatched and batched inputs, and first-order gradients through all tensor inputs (points, controls, displacements, radii, amount, weights). Auto dispatch uses Torch on CPU and Warp on CUDA. The Warp dense-displacement backend is opt-in only (Torch is default everywhere) and exists for backend parity and benchmarking.
Tests cover analytic derivatives, gradcheck on both backends, opcheck, CUDA-graph capture, and torch.compile fullgraph. Includes API docs, a tutorial section, and benchmark registry entries.
Checklist
Checklist
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.