[codex] Add mesh vector calculus functionals#1735
Conversation
loliverhennigh
force-pushed
the
mesh-vector-calculus-functionals
branch
from
e301081 to
f820150
Compare
loliverhennigh
force-pushed
the
mesh-vector-calculus-functionals
branch
from
208bcce to
28b889f
Compare
Greptile SummaryThis PR adds five new FunctionSpec-backed mesh vector calculus functionals — LSQ divergence, curl, and double-LSQ Laplacian, plus cotangent/DEC Laplacian and divergence — each with eager PyTorch and Warp backends. It then wires
Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Remove cross-backend operator utilities, fix cotangent edge cases, expose the complete Mesh API, and merge current upstream main.
|
ORD A100-SXM4-80GB update: the original LSQ results still hold (3.36–3.90x faster through Mesh at 65k; 27–55x faster for isolated 262k training steps). Cotangent follow-up job 30030777 passed all 18 finite/output/gradient parity records (max relative L2 8.9e-7). Geometry caching makes warm public calls 8–10x faster at 65k and 19–24x faster at 262k versus the focused ablation. Warp divergence is now 1.12–1.13x faster through Mesh and up to 2.76x faster for 262k all-gradient functional training. Warp Laplacian functional forward is about 2x faster than the prior Warp path, but remains about 9% behind Torch; field-only training crosses over at 262k while all-input training remains about 25% slower. Torch remains the default solely for backwards compatibility. |
| def partition_cells( | ||
| mesh: Mesh, | ||
| seeds: Float[torch.Tensor, "n_seeds n_spatial_dims"], | ||
| implementation: Literal["cuml", "torch", "scipy"] | None = "torch", |
There was a problem hiding this comment.
[blocking] Defaulting to implementation="torch" here is a scalability regression. Previously this passed None to knn, which auto-dispatches to cuML (GPU) or scipy KDTree (CPU). The torch backend materializes the full (M, N, D) broadcast tensor (physicsnemo/nn/functional/neighbors/knn/_torch_impl.py:45), so memory goes from O(M+N) to O(M·N·D).
Concretely: partition_cells is called by the GLOBE DrivAer pipeline (examples/cfd/external_aerodynamics/globe/drivaer/dataset.py:284) with seeds = boundary.cell_centroids. At DrivAerML scale (~1e5 boundary cells × ~1e6 volume cells × 3 dims × 4 B) that intermediate is O(TB) — an instant OOM where the code works today, and a silent behavior change to a public API.
The determinism rationale is also weaker here than for the calculus functionals: a k=1 NN query returns the same result on every backend modulo distance ties, unlike floating-point reduction order in the derivative kernels. Suggest keeping None (auto-dispatch) as the default for the KNN-backed entry points and letting callers opt into "torch" explicitly.
| def find_nearest_cells( | ||
| mesh: "Mesh", | ||
| query_points: Float[torch.Tensor, "n_queries n_spatial_dims"], | ||
| implementation: Literal["cuml", "torch", "scipy"] | None = "torch", |
There was a problem hiding this comment.
[blocking] Same concern as partition_cells: this default change (also at match_points, sample_data_at_points, and Mesh.sample_data_at_points) switches existing callers from KNN auto-dispatch to the brute-force torch backend, which allocates an (n_cells, n_queries, dims) intermediate. For find_nearest_cells on a production-size mesh this OOMs where the scipy/cuML paths were fine. Suggest defaulting these five entry points back to None.
| ) | ||
|
|
||
|
|
||
| def _curl_from_jacobian( |
There was a problem hiding this comment.
_curl_from_jacobian is now dead code — this PR removed both call sites (points and cells now go through mesh_lsq_curl), and nothing in physicsnemo/ or test/ references it anymore. Safe to delete. (Its sibling _apply_cotan_laplacian_operator in laplacian.py is still used by mesh/curvature/_laplacian.py:99, so that one should stay.)
| import torch | ||
|
|
||
|
|
||
| def _safe_eps(dtype: torch.dtype) -> float: |
There was a problem hiding this comment.
_safe_eps and the ~70-line _validate_inputs are duplicated verbatim between _torch_impl.py and _warp_impl.py in both cotan packages (4 copies of the safe-eps formula, plus a 5th inline in mesh_cotan_divergence.py's _make_benchmark_case). The layering contract legitimately forbids importing mesh.utilities._tolerances.safe_eps from the functional layer, but the existing convention in mesh_lsq_gradient/ is a package-local shared utils.py — following that here would keep these from drifting. This formula encodes a subtle float16 cap that would be easy to fix in one copy and miss in the others.
| values: torch.Tensor, | ||
| ) -> torch.Tensor: | ||
| """Apply the normalized cotangent Laplacian with Warp kernels.""" | ||
| _validate_inputs( |
There was a problem hiding this comment.
Validation runs twice on this path: once here in the wrapper and again inside mesh_cotan_laplacian_impl (line 187). The divergence op validates only inside the custom op — inconsistent between the two. Each _validate_inputs call does edges.min().item() / edges.max().item(), i.e. two host-device syncs per validation on CUDA (and a graph break under torch.compile). Note the DEC Mesh.laplacian/Mesh.divergence paths previously ran sync-free inline code, so this is a new per-call cost on the default torch path too. Deduplicating so validation runs once per public call would recover most of it.
| case _: | ||
| raise ValueError( | ||
| f"Invalid {data_source=!r}. Must be 'points' or 'cells'." | ||
| f"Invalid {method=!r} (must be 'cotan', 'dec', or 'lsq') " |
There was a problem hiding this comment.
nit: this catch-all reports "Invalid method=... or data_source=..." without distinguishing which one is wrong. Splitting the two checks would give crisper errors.
| shape ``(n_entities, 3)`` in 3D. | ||
| """ | ||
|
|
||
| _COMPARE_ATOL = 8e-3 |
There was a problem hiding this comment.
nit: 8e-3 abs+rel is loose enough to hide real backend divergence (the cotan specs use 5e-6). Presumably this absorbs fp32 lstsq-vs-Warp-QR noise from the underlying LSQ solve — a one-line comment justifying the magnitude would help future readers, and applies to the divergence/laplacian specs too.
| to ``"torch"`` for backwards-compatible numerical precision. Pass | ||
| ``"warp"`` to use the accelerated ambient LSQ gradient kernels. | ||
| Intrinsic tangent-space LSQ gradients use their existing Torch | ||
| implementation. |
There was a problem hiding this comment.
This last sentence reads as a silent fallback, but the actual behavior when passing implementation="warp" with an intrinsic gradient on a codim > 0 mesh is a NotImplementedError (raised in compute_gradient_points_lsq, and covered by test_intrinsic_point_gradient_rejects_warp_implementation). The raise is the right behavior — silent fallback would hide intent — so the docstring should describe it, e.g. "requesting \"warp\" for intrinsic tangent-space gradients raises NotImplementedError; use gradient_type=\"extrinsic\" or implementation=\"torch\"". The same misleading sentence appears in compute_point_derivatives (calculus/derivatives.py:88-89).
| if intrinsic and mesh.codimension > 0: | ||
| if implementation not in (None, "torch"): | ||
| raise NotImplementedError( | ||
| "Warp implementation is not available for intrinsic tangent-space " |
There was a problem hiding this comment.
This message suggests gradient_type='extrinsic', but gradient_type is a parameter of the callers (Mesh.gradient / compute_point_derivatives) — this function's own parameter is intrinsic: bool. A direct caller of compute_gradient_points_lsq gets pointed at a kwarg that doesn't exist here. Suggest phrasing in the local vocabulary (intrinsic=False or implementation='torch'), or mentioning both spellings.
| laplacian = laplacian / dual_volumes_0.view( | ||
| -1, *([1] * (point_values.ndim - 1)) | ||
| ).clamp(min=safe_eps(dual_volumes_0.dtype)) | ||
| def compute_laplacian_points_lsq( |
There was a problem hiding this comment.
Two documentation gaps around the new LSQ Laplacians:
- This module's top docstring still describes only the intrinsic cotangent/DEC operator ("This is the cotangent Laplacian, intrinsic to the manifold"), but the module now also hosts two extrinsic double-LSQ operators.
- These two new public functions have ~3-line docstrings in a module whose house style is full NumPy-doc with the math written out (cf.
compute_laplacian_points_decjust above). The double-LSQ operator in particular deserves a Notes section on its accuracy properties (boundary bias / noise amplification from composing two first-order reconstructions) — the oldMesh.laplaciandocstring carried exactly that caveat ("a double-LSQ discretization with different accuracy properties") and it was dropped in the rewrite.
|
Top-level summary for my review above (GitHub dropped the review body on submit): Thanks — the core contribution here is solid. I hand-verified the backward formulas for both cotan custom ops (including the clamped-volume gradient mask) and ran The functional decomposition is clean: the mesh layer computes geometry (cotan weights, dual volumes, CSR adjacency) and the functional layer does math on raw tensors, which respects the import-linter layering and makes the operators benchmarkable. One blocking issue: the KNN-backed sampling APIs ( The rest is cleanup-level: dead |
…-perf # Conflicts: # physicsnemo/mesh/calculus/curl.py # physicsnemo/mesh/calculus/divergence.py
peterdsharpe
left a comment
peterdsharpe
left a comment
There was a problem hiding this comment.
All review feedback addressed — approving. Verified locally against the new head (CPU + GPU, 896 tests passing across test/mesh/calculus, test/mesh/sampling, test/mesh/remeshing, test/nn/functional/derivatives, and test/core/test_function_spec.py):
- KNN defaults reverted to auto-dispatch on all five entry points (
partition_cells,find_nearest_cells,match_points,sample_data_at_points,Mesh.sample_data_at_points) — this was my one blocking concern, and the forwarding tests were updated to match. - Dead
_curl_from_jacobianremoved; error messages now validatemethod/data_source/gradient_typeup front and speak each function's own vocabulary; DEC gradients now raiseNotImplementedErrorforimplementation="warp"instead of silently ignoring it; docstrings (module-level Laplacian, intrinsic+warp behavior, double-LSQ accuracy caveats, parity tolerances) all brought up to standard;implementation=Noneauto-dispatch now has test coverage. - The new work since last review holds up under scrutiny. I re-derived the fused Warp backward math (the per-endpoint
g[v0]/V0 − g[v1]/V1form is equivalent to the previous normalize-then-scatter formulation), rantorch.autograd.gradcheckon all differentiable inputs of the new fused forward+backward custom ops, and confirmed double backward works through the differentiable torch fallback. The geometry/topology caching split (topology keyed oncells, geometry onpoints, cleared on geometric transforms, bypassed underrequires_grad, inference-mode tensors handled) is correct, and I verified warm-call parity through the cache. The algebraicC = H·G⁻¹·Hᵀrewrite in the cotan weights checks out block-by-block. Single-syncaminmax().tolist()validation is a nice touch.
Non-blocking follow-up (leaving the inline thread open): _safe_eps still has four copies across the cotan backends and _validate_inputs is duplicated per backend file. Now that the torch and Warp implementations have diverged substantially, per-package self-containment is more defensible than it was — but a package-local utils.py (as in mesh_lsq_gradient/) would still be the tidier home. Fine to defer to a later cleanup.
Nice work on this one — the Warp backward kernels and the caching layer in particular are careful, well-tested additions.
2 participants
Summary
implementationfrom mesh-level functional callers, defaulting mesh APIs to"torch"(Noneretains KNN auto-dispatch)Validation
pytest -q test/nn/functional/derivatives test/mesh/calculus— 317 passed, 11 skippedNotes