[codex] Add uniform grid vector calculus functionals#1724
[codex] Add uniform grid vector calculus functionals#1724loliverhennigh wants to merge 2 commits into
Conversation
|
Hey @ktangsali, I had a little bit of time this week and started flushing out some of the gradient/functionals work in PhysicsNemo. I’m working my way toward the mesh gradient functionals and trying to get some real acceleration there. This PR adds uniform-grid divergence, curl, and Laplacian functionals with fused Warp kernels plus Torch references and tests: No rush, but when you get a chance, can you take a quick look at this? |
loliverhennigh
force-pushed
the
grid-vector-calculus-functionals
branch
2 times, most recently
from
07d1557 to
a1b62e6
Compare
Greptile SummaryAdds three new uniform-grid vector calculus functionals —
Important Files Changed
Reviews (1): Last reviewed commit: "Address vector calculus functional revie..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| vector_field: torch.Tensor, | ||
| spacing: float | Sequence[float] = 1.0, | ||
| order: int = 2, | ||
| ) -> torch.Tensor: | ||
| """Compute periodic uniform-grid divergence with PyTorch tensor ops.""" | ||
| grid_ndim = validate_vector_field(vector_field) |
There was a problem hiding this comment.
The positivity check for dx lives inside the main accumulation loop, so if a later spacing entry is invalid the partial derivative for earlier axes has already been computed before the error is raised. Both the curl and (separately) the warp implementations validate all spacing values up-front before any computation. Bringing divergence and laplacian in line with that pattern would make the three implementations consistent and avoids wasted computation on bad inputs.
| @wp.func | ||
| def _wrap_plus1(i: int, n: int) -> int: | ||
| return (i + 1) % n | ||
|
|
||
|
|
||
| @wp.func | ||
| def _wrap_minus1(i: int, n: int) -> int: | ||
| return (i + n - 1) % n | ||
|
|
||
|
|
||
| @wp.func | ||
| def _wrap_plus2(i: int, n: int) -> int: | ||
| return (i + 2) % n | ||
|
|
||
|
|
||
| @wp.func | ||
| def _wrap_minus2(i: int, n: int) -> int: | ||
| return (i + n - 2) % n |
There was a problem hiding this comment.
_warp_impl.py files
_wrap_plus1, _wrap_minus1, _wrap_plus2, and _wrap_minus2 are defined identically in uniform_grid_divergence/_warp_impl.py, uniform_grid_curl/_warp_impl.py, and uniform_grid_laplacian/_warp_impl.py. The implementations already share the public-facing utilities (_launch_dim, _normalize_spacing, _warp_launch_context, etc.) from ..uniform_grid_gradient._warp_impl.utils. Consolidating the wrap helpers into that shared utils module would eliminate the duplication and make future changes (e.g., a different periodic-wrap strategy) a single-site edit.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| @torch.library.custom_op( | ||
| "physicsnemo::uniform_grid_divergence_warp_impl", mutates_args=() | ||
| ) | ||
| def uniform_grid_divergence_impl( | ||
| vector_field: torch.Tensor, | ||
| spacing_meta: torch.Tensor, | ||
| order: int, | ||
| ) -> torch.Tensor: | ||
| """Evaluate uniform-grid divergence with fused Warp kernels.""" | ||
| grid_ndim = validate_vector_field(vector_field) | ||
| spacing_tuple = tuple(float(v) for v in spacing_meta.tolist()) | ||
| _validate_positive_spacing(spacing_tuple) | ||
| order = _validate_order(int(order)) | ||
| orig_dtype = vector_field.dtype | ||
| vector_field_fp32 = _to_fp32_contiguous(vector_field) | ||
| output_fp32 = torch.empty( | ||
| vector_field_fp32.shape[1:], | ||
| device=vector_field_fp32.device, | ||
| dtype=torch.float32, | ||
| ) | ||
| _launch_divergence_forward( | ||
| vector_field_fp32=vector_field_fp32, | ||
| spacing_tuple=spacing_tuple[:grid_ndim], | ||
| order=order, | ||
| output_fp32=output_fp32, | ||
| ) | ||
| return _restore_dtype(output_fp32, orig_dtype) | ||
|
|
There was a problem hiding this comment.
uniform_grid_divergence_warp (the public wrapper) calls validate_vector_field, _validate_positive_spacing, and _validate_order, then forwards to uniform_grid_divergence_impl, which repeats all three checks. This double-validation pattern is present in the curl and laplacian warp paths as well. While correct, it means every Warp call pays validation overhead twice. If the custom op is ever exercised directly (e.g., through torch.compile graph capture or export), the second validation acts as a safety net; but in the normal dispatch path it is redundant.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
loliverhennigh
force-pushed
the
grid-vector-calculus-functionals
branch
from
a1b62e6 to
f7bbfc6
Compare
|
/blossom-ci |
1 similar comment
|
/blossom-ci |
loliverhennigh
force-pushed
the
grid-vector-calculus-functionals
branch
from
f7bbfc6 to
bdee451
Compare
1 participant
Summary
Adds uniform-grid vector calculus functionals for periodic Cartesian grids:
uniform_grid_divergenceuniform_grid_curluniform_grid_laplacianEach functional follows the existing
FunctionSpecpattern with a Torch reference backend and a fused Warp backend. The Warp paths launch dedicated custom kernels instead of composing throughuniform_grid_gradient, so the accelerated backend avoids intermediate derivative tensors.This also exposes
MeshPoissonDiskSampleandMeshToVoxelFractionin the functional benchmark registry and adds derivative docs entries with figures for the three new uniform-grid vector calculus functionals.Notes
The Laplacian compare hook uses a small absolute tolerance for fourth-order CUDA parity. Fourth-order second-derivative stencils subtract nearly equal float32 values, and the fused Warp kernel and composed Torch reference take different rounding paths on CUDA. Analytic tests still cover physical finite-difference accuracy separately.
Validation
Local CPU/MPS-host validation:
uv run --frozen pytest -q test/nn/functional/derivatives/test_uniform_grid_divergence.py test/nn/functional/derivatives/test_uniform_grid_curl.py test/nn/functional/derivatives/test_uniform_grid_laplacian.py40 passeduv run --frozen pytest -q test/nn/functional359 passed, 49 skippeduv run --frozen ruff check physicsnemo/nn/functional/derivatives/uniform_grid_divergence physicsnemo/nn/functional/derivatives/uniform_grid_curl physicsnemo/nn/functional/derivatives/uniform_grid_laplacian test/nn/functional/derivatives/test_uniform_grid_divergence.py test/nn/functional/derivatives/test_uniform_grid_curl.py test/nn/functional/derivatives/test_uniform_grid_laplacian.py benchmarks/physicsnemo/nn/functional/registry.py physicsnemo/nn/functional/__init__.py physicsnemo/nn/functional/derivatives/__init__.pygit diff --checkRTX 4090 validation:
2.11.0+cu130, Warp1.13.080 passedRTX 4090 fused Warp timing sanity check, compile time excluded:
(2, 1024, 1024)(2, 1024, 1024)(1024, 1024)(3, 128, 128, 128)(3, 128, 128, 128)(128, 128, 128)