[codex] Add uniform grid vector calculus functionals

benchmarks/physicsnemo/nn/functional/registry.py

@@ -23,7 +23,10 @@
     MeshLSQGradient,
     RectilinearGridGradient,
     SpectralGridGradient,
+    UniformGridCurl,
+    UniformGridDivergence,
     UniformGridGradient,
+    UniformGridLaplacian,
 )
 from physicsnemo.nn.functional.fourier_spectral import (
     IRFFT,
@@ -36,6 +39,8 @@
 )
 from physicsnemo.nn.functional.geometry import (
     FarthestPointSampling,
+    MeshPoissonDiskSample,
+    MeshToVoxelFraction,
     SignedDistanceField,
 )
 from physicsnemo.nn.functional.interpolation import (
@@ -64,8 +69,13 @@
     MeshGreenGaussGradient,
     SpectralGridGradient,
     MeshlessFDDerivatives,
+    UniformGridDivergence,
+    UniformGridCurl,
+    UniformGridLaplacian,
     # Geometry.
     FarthestPointSampling,
+    MeshPoissonDiskSample,
+    MeshToVoxelFraction,
     SignedDistanceField,
     # Interpolation.
     GridToPointInterpolation,

docs/api/nn/functionals/derivatives.rst

@@ -32,3 +32,24 @@ Derivative Functionals
    :width: 100%
 
 .. autofunction:: physicsnemo.nn.functional.meshless_fd_derivatives
+
+Uniform Grid Vector Calculus
+----------------------------
+
+.. autofunction:: physicsnemo.nn.functional.uniform_grid_divergence
+
+.. figure:: /img/nn/functional/derivatives/uniform_grid_divergence.png
+   :alt: Uniform grid divergence example
+   :width: 100%
+
+.. autofunction:: physicsnemo.nn.functional.uniform_grid_curl
+
+.. figure:: /img/nn/functional/derivatives/uniform_grid_curl.png
+   :alt: Uniform grid curl example
+   :width: 100%
+
+.. autofunction:: physicsnemo.nn.functional.uniform_grid_laplacian
+
+.. figure:: /img/nn/functional/derivatives/uniform_grid_laplacian.png
+   :alt: Uniform grid Laplacian example
+   :width: 100%

physicsnemo/nn/functional/__init__.py

@@ -20,7 +20,10 @@
     meshless_fd_derivatives,
     rectilinear_grid_gradient,
     spectral_grid_gradient,
+    uniform_grid_curl,
+    uniform_grid_divergence,
     uniform_grid_gradient,
+    uniform_grid_laplacian,
 )
 from .equivariant_ops import (
     legendre_polynomials,
@@ -60,6 +63,9 @@
     "irfft2",
     "drop_path",
     "farthest_point_sampling",
+    "uniform_grid_curl",
+    "uniform_grid_divergence",
+    "uniform_grid_laplacian",
     "grid_to_point_interpolation",
     "imag",
     "interpolation",

physicsnemo/nn/functional/derivatives/__init__.py

@@ -25,15 +25,24 @@
     rectilinear_grid_gradient,
 )
 from .spectral_grid_gradient import SpectralGridGradient, spectral_grid_gradient
+from .uniform_grid_curl import UniformGridCurl, uniform_grid_curl
+from .uniform_grid_divergence import UniformGridDivergence, uniform_grid_divergence
 from .uniform_grid_gradient import UniformGridGradient, uniform_grid_gradient
+from .uniform_grid_laplacian import UniformGridLaplacian, uniform_grid_laplacian
 
 __all__ = [
+    "UniformGridCurl",
+    "UniformGridDivergence",
+    "UniformGridLaplacian",
     "MeshGreenGaussGradient",
     "MeshlessFDDerivatives",
     "MeshLSQGradient",
     "RectilinearGridGradient",
     "SpectralGridGradient",
     "UniformGridGradient",
+    "uniform_grid_curl",
+    "uniform_grid_divergence",
+    "uniform_grid_laplacian",
     "mesh_green_gauss_gradient",
     "meshless_fd_derivatives",
     "mesh_lsq_gradient",

physicsnemo/nn/functional/derivatives/uniform_grid_curl/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .uniform_grid_curl import UniformGridCurl, uniform_grid_curl
+
+__all__ = ["UniformGridCurl", "uniform_grid_curl"]

physicsnemo/nn/functional/derivatives/uniform_grid_curl/_torch_impl.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+from .utils import validate_vector_field
+
+_SUPPORTED_ORDERS = (2, 4)
+
+
+def _normalize_spacing(
+    spacing: float | Sequence[float], ndim: int
+) -> tuple[float, ...]:
+    if isinstance(spacing, (float, int)):
+        return tuple(float(spacing) for _ in range(ndim))
+    spacing_tuple = tuple(float(x) for x in spacing)
+    if len(spacing_tuple) != ndim:
+        raise ValueError(
+            f"spacing must have {ndim} entries for a {ndim}D field, got {len(spacing_tuple)}"
+        )
+    return spacing_tuple
+
+
+def _validate_order(order: int) -> int:
+    if not isinstance(order, int):
+        raise TypeError(f"order must be an integer, got {type(order)}")
+    if order not in _SUPPORTED_ORDERS:
+        raise ValueError(
+            "uniform_grid_curl supports central orders "
+            f"{list(_SUPPORTED_ORDERS)}, got order={order}"
+        )
+    return order
+
+
+def _central_derivative_order2(
+    field: torch.Tensor, axis: int, dx: float
+) -> torch.Tensor:
+    return (
+        torch.roll(field, shifts=-1, dims=axis) - torch.roll(field, shifts=1, dims=axis)
+    ) / (2.0 * dx)
+
+
+def _central_derivative_order4(
+    field: torch.Tensor, axis: int, dx: float
+) -> torch.Tensor:
+    return (
+        -torch.roll(field, shifts=-2, dims=axis)
+        + 8.0 * torch.roll(field, shifts=-1, dims=axis)
+        - 8.0 * torch.roll(field, shifts=1, dims=axis)
+        + torch.roll(field, shifts=2, dims=axis)
+    ) / (12.0 * dx)
+
+
+_DERIVATIVE_DISPATCH = {
+    2: _central_derivative_order2,
+    4: _central_derivative_order4,
+}
+
+
+def uniform_grid_curl_torch(
+    vector_field: torch.Tensor,
+    spacing: float | Sequence[float] = 1.0,
+    order: int = 2,
+) -> torch.Tensor:
+    """Compute periodic uniform-grid curl with PyTorch tensor ops."""
+    grid_ndim = validate_vector_field(vector_field)
+    spacing_tuple = _normalize_spacing(spacing, grid_ndim)
+    for dx in spacing_tuple:
+        if dx <= 0.0:
+            raise ValueError("all spacing entries must be strictly positive")
+    derivative_fn = _DERIVATIVE_DISPATCH[_validate_order(order)]
+
+    def derivative(component: int, axis: int) -> torch.Tensor:
+        return derivative_fn(vector_field[component], axis, spacing_tuple[axis])
+
+    if grid_ndim == 2:
+        return derivative(1, 0) - derivative(0, 1)
+
+    curl_x = derivative(2, 1) - derivative(1, 2)
+    curl_y = derivative(0, 2) - derivative(2, 0)
+    curl_z = derivative(1, 0) - derivative(0, 1)
+    return torch.stack((curl_x, curl_y, curl_z), dim=0)