Add non-periodic grid gradient boundary mode

CHANGELOG.md

@@ -10,6 +10,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Adds an opt-in ``boundary="one_sided"`` mode to
+  ``rectilinear_grid_gradient`` and ``uniform_grid_gradient``. The default
+  periodic behavior is unchanged, while bounded domains can use one-sided
+  boundary stencils with Torch/Warp forward and autograd parity.
 - Adds dimension-generic volume mesh generation for implicit domains to
   `physicsnemo.mesh.generate`. `mesh_implicit_domain` meshes
   `{x : phi(x) < 0}`, clipped to the bounding box (box faces are honored

physicsnemo/nn/functional/derivatives/_boundary.py

@@ -0,0 +1,137 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+
+BoundaryMode = Literal["periodic", "one_sided"]
+
+
+def normalize_boundary(boundary: str, *, function_name: str) -> BoundaryMode:
+    """Validate and normalize a grid-gradient boundary mode."""
+    if not isinstance(boundary, str):
+        raise TypeError(f"{function_name} boundary must be a string")
+    if boundary not in ("periodic", "one_sided"):
+        raise ValueError(
+            f'{function_name} boundary must be "periodic" or "one_sided", '
+            f"got {boundary!r}"
+        )
+    return boundary
+
+
+def _finite_difference_weights(
+    nodes: torch.Tensor,
+    evaluation_point: torch.Tensor,
+    derivative_order: int,
+) -> torch.Tensor:
+    """Return finite-difference weights for arbitrary one-dimensional nodes."""
+    offsets = nodes - evaluation_point
+    powers = torch.arange(nodes.numel(), device=nodes.device, dtype=nodes.dtype)
+    factorials = torch.exp(torch.lgamma(powers + 1.0))
+    system = offsets.unsqueeze(0).pow(powers.unsqueeze(1)) / factorials.unsqueeze(1)
+    rhs = torch.zeros_like(nodes)
+    rhs[derivative_order] = 1.0
+    return torch.linalg.solve(system, rhs)
+
+
+def _weighted_boundary_value(
+    field: torch.Tensor,
+    *,
+    axis: int,
+    indices: torch.Tensor,
+    weights: torch.Tensor,
+) -> torch.Tensor:
+    """Evaluate a one-sided stencil and retain the differentiated axis."""
+    values = torch.index_select(field, dim=axis, index=indices)
+    view_shape = [1] * field.ndim
+    view_shape[axis] = weights.numel()
+    return (values * weights.view(view_shape)).sum(dim=axis, keepdim=True)
+
+
+def apply_one_sided_boundaries(
+    output: torch.Tensor,
+    *,
+    field: torch.Tensor,
+    coordinates: tuple[torch.Tensor, ...],
+    derivative_order: int,
+    stencil_size: int,
+    boundary_width: int = 1,
+    function_name: str,
+) -> torch.Tensor:
+    """Replace periodic boundary rows with arbitrary-grid one-sided stencils."""
+    corrected: list[torch.Tensor] = []
+    for axis, coordinates_axis in enumerate(coordinates):
+        axis_size = field.shape[axis]
+        if axis_size < stencil_size:
+            raise ValueError(
+                f"{function_name} boundary='one_sided' requires at least "
+                f"{stencil_size} points along axis {axis}, got {axis_size}"
+            )
+
+        weight_dtype = torch.float64 if field.dtype == torch.float64 else torch.float32
+        nodes = coordinates_axis.to(device=field.device, dtype=weight_dtype)
+        left_indices = torch.arange(stencil_size, device=field.device)
+        right_indices = torch.arange(
+            axis_size - stencil_size, axis_size, device=field.device
+        )
+        left_nodes = torch.index_select(nodes, 0, left_indices)
+        right_nodes = torch.index_select(nodes, 0, right_indices)
+        left_values = []
+        right_values = []
+        for offset in range(boundary_width):
+            left_weights = _finite_difference_weights(
+                left_nodes, nodes[offset], derivative_order
+            ).to(dtype=field.dtype)
+            right_weights = _finite_difference_weights(
+                right_nodes,
+                nodes[axis_size - boundary_width + offset],
+                derivative_order,
+            ).to(dtype=field.dtype)
+            left_values.append(
+                _weighted_boundary_value(
+                    field,
+                    axis=axis,
+                    indices=left_indices,
+                    weights=left_weights,
+                )
+            )
+            right_values.append(
+                _weighted_boundary_value(
+                    field,
+                    axis=axis,
+                    indices=right_indices,
+                    weights=right_weights,
+                )
+            )
+
+        left = torch.cat(left_values, dim=axis)
+        right = torch.cat(right_values, dim=axis)
+        interior_slices = [slice(None)] * field.ndim
+        interior_slices[axis] = slice(boundary_width, -boundary_width)
+        interior = output[axis][tuple(interior_slices)]
+        corrected.append(torch.cat((left, interior, right), dim=axis))
+
+    return torch.stack(corrected, dim=0)
+
+
+def uniform_coordinates(
+    field: torch.Tensor, spacing: tuple[float, ...]
+) -> tuple[torch.Tensor, ...]:
+    """Build per-axis coordinates for uniform-grid boundary stencils."""
+    dtype = torch.float64 if field.dtype == torch.float64 else torch.float32
+    return tuple(
+        torch.arange(size, device=field.device, dtype=dtype) * dx
+        for size, dx in zip(field.shape, spacing, strict=True)
+    )
+
+
+__all__ = [
+    "BoundaryMode",
+    "apply_one_sided_boundaries",
+    "normalize_boundary",
+    "uniform_coordinates",
+]

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

@@ -20,6 +20,7 @@
 
 import torch
 
+from .._boundary import apply_one_sided_boundaries, normalize_boundary
 from .._request_utils import (
     compose_derivative_outputs,
     normalize_derivative_orders,
@@ -41,14 +42,16 @@ def rectilinear_grid_gradient_torch(
     periods: float | Sequence[float] | None = None,
     derivative_order: int = 1,
     include_mixed: bool = False,
+    boundary: str = "periodic",
 ) -> torch.Tensor:
-    """Compute periodic first or pure second derivatives on rectilinear grids."""
+    """Compute first or pure second derivatives on rectilinear grids."""
     ### Validate field and coordinate inputs.
     validate_field(field)
     derivative_order = validate_derivative_request(
         derivative_order=derivative_order,
         include_mixed=include_mixed,
     )
+    boundary = normalize_boundary(boundary, function_name="rectilinear_grid_gradient")
 
     coords_tuple, period_tuple = validate_and_normalize_coordinates(
         field=field,
@@ -86,7 +89,17 @@ def rectilinear_grid_gradient_torch(
         gradients.append(grad_axis)
 
     ### Stack per-axis derivative terms into (dims, *field.shape).
-    return torch.stack(gradients, dim=0)
+    output = torch.stack(gradients, dim=0)
+    if boundary == "one_sided":
+        output = apply_one_sided_boundaries(
+            output,
+            field=field,
+            coordinates=coords_tuple,
+            derivative_order=derivative_order,
+            stencil_size=2 + derivative_order,
+            function_name="rectilinear_grid_gradient",
+        )
+    return output
 
 
 def rectilinear_grid_gradient_torch_multi(
@@ -95,6 +108,7 @@ def rectilinear_grid_gradient_torch_multi(
     periods: float | Sequence[float] | None = None,
     derivative_orders: int | Sequence[int] = 1,
     include_mixed: bool = False,
+    boundary: str = "periodic",
 ) -> torch.Tensor:
     """Compute first/second/mixed derivatives from a unified request."""
     requested_orders = normalize_derivative_orders(
@@ -123,6 +137,7 @@ def rectilinear_grid_gradient_torch_multi(
                 periods=periods,
                 derivative_order=derivative_order,
                 include_mixed=False,
+                boundary=boundary,
             )
         ),
     )

physicsnemo/nn/functional/derivatives/rectilinear_grid_gradient/_warp_impl/op.py

@@ -23,6 +23,7 @@
 
 from physicsnemo.core.function_spec import FunctionSpec
 
+from ..._boundary import apply_one_sided_boundaries, normalize_boundary
 from ..._request_utils import (
     compose_derivative_outputs,
     normalize_derivative_orders,
@@ -704,8 +705,9 @@ def rectilinear_grid_gradient_warp(
     periods: float | Sequence[float] | None = None,
     derivative_order: int = 1,
     include_mixed: bool = False,
+    boundary: str = "periodic",
 ) -> torch.Tensor:
-    """Compute periodic first or pure second derivatives on rectilinear grids.
+    """Compute first or pure second derivatives on rectilinear grids.
 
     Notes
     -----
@@ -718,6 +720,7 @@ def rectilinear_grid_gradient_warp(
         derivative_order=derivative_order,
         include_mixed=include_mixed,
     )
+    boundary = normalize_boundary(boundary, function_name="rectilinear_grid_gradient")
 
     coords_tuple, period_tuple = validate_and_normalize_coordinates(
         field=field,
@@ -728,15 +731,15 @@ def rectilinear_grid_gradient_warp(
     )
 
     if field.ndim == 1:
-        return rectilinear_grid_gradient_1d_impl(
+        output = rectilinear_grid_gradient_1d_impl(
             field,
             coords_tuple[0],
             float(period_tuple[0]),
             int(derivative_order),
             bool(include_mixed),
         )
-    if field.ndim == 2:
-        return rectilinear_grid_gradient_2d_impl(
+    elif field.ndim == 2:
+        output = rectilinear_grid_gradient_2d_impl(
             field,
             coords_tuple[0],
             coords_tuple[1],
@@ -745,17 +748,29 @@ def rectilinear_grid_gradient_warp(
             int(derivative_order),
             bool(include_mixed),
         )
-    return rectilinear_grid_gradient_3d_impl(
-        field,
-        coords_tuple[0],
-        coords_tuple[1],
-        coords_tuple[2],
-        float(period_tuple[0]),
-        float(period_tuple[1]),
-        float(period_tuple[2]),
-        int(derivative_order),
-        bool(include_mixed),
-    )
+    else:
+        output = rectilinear_grid_gradient_3d_impl(
+            field,
+            coords_tuple[0],
+            coords_tuple[1],
+            coords_tuple[2],
+            float(period_tuple[0]),
+            float(period_tuple[1]),
+            float(period_tuple[2]),
+            int(derivative_order),
+            bool(include_mixed),
+        )
+
+    if boundary == "one_sided":
+        output = apply_one_sided_boundaries(
+            output,
+            field=field,
+            coordinates=coords_tuple,
+            derivative_order=derivative_order,
+            stencil_size=2 + derivative_order,
+            function_name="rectilinear_grid_gradient",
+        )
+    return output
 
 
 def rectilinear_grid_gradient_warp_multi(
@@ -764,6 +779,7 @@ def rectilinear_grid_gradient_warp_multi(
     periods: float | Sequence[float] | None = None,
     derivative_orders: int | Sequence[int] = 1,
     include_mixed: bool = False,
+    boundary: str = "periodic",
 ) -> torch.Tensor:
     """Compute multiple derivative families, fusing first+second when possible.
 
@@ -773,6 +789,7 @@ def rectilinear_grid_gradient_warp_multi(
     ordering and autograd behavior.
     """
     validate_field(field)
+    boundary = normalize_boundary(boundary, function_name="rectilinear_grid_gradient")
     coords_tuple, period_tuple = validate_and_normalize_coordinates(
         field=field,
         coordinates=coordinates,
@@ -796,7 +813,7 @@ def rectilinear_grid_gradient_warp_multi(
     )
 
     ### Fused no-mixed path with custom-op backward for combined first+second.
-    if not mixed_terms and requested_orders == (1, 2):
+    if boundary == "periodic" and not mixed_terms and requested_orders == (1, 2):
         if field.ndim == 1:
             fused = rectilinear_derivatives_1d_fused_no_mixed_impl(
                 field,
@@ -842,6 +859,7 @@ def rectilinear_grid_gradient_warp_multi(
                 periods=period_tuple,
                 derivative_order=derivative_order,
                 include_mixed=False,
+                boundary=boundary,
             )
         ),
     )

physicsnemo/nn/functional/derivatives/rectilinear_grid_gradient/rectilinear_grid_gradient.py

@@ -29,7 +29,7 @@
 
 
 class RectilinearGridGradient(FunctionSpec):
-    r"""Compute periodic gradients on rectilinear grids with nonuniform spacing.
+    r"""Compute gradients on rectilinear grids with nonuniform spacing.
 
     This functional computes first-order and/or second-order
     derivatives of a scalar field on a 1D/2D/3D rectilinear grid where each
@@ -86,6 +86,10 @@ class RectilinearGridGradient(FunctionSpec):
         Include mixed second derivatives when requesting second derivatives.
         Mixed terms are appended in axis-pair order ``(x,y)``, ``(x,z)``,
         ``(y,z)``.
+    boundary : {"periodic", "one_sided"}, optional
+        Boundary treatment. ``"periodic"`` preserves wrap-around behavior.
+        ``"one_sided"`` uses one-sided stencils at the first and last point
+        of each axis while retaining central differences in the interior.
     implementation : {"warp", "torch"} or None
         Explicit backend selection. When ``None``, dispatch selects by rank.
 
@@ -117,6 +121,7 @@ def warp_forward(
         periods: float | Sequence[float] | None = None,
         derivative_orders: int | Sequence[int] = 1,
         include_mixed: bool = False,
+        boundary: str = "periodic",
     ) -> torch.Tensor:
         """Dispatch rectilinear gradients to the Warp backend."""
         return rectilinear_grid_gradient_warp_multi(
@@ -125,6 +130,7 @@ def warp_forward(
             periods=periods,
             derivative_orders=derivative_orders,
             include_mixed=include_mixed,
+            boundary=boundary,
         )
 
     @FunctionSpec.register(name="torch", rank=1, baseline=True)
@@ -134,6 +140,7 @@ def torch_forward(
         periods: float | Sequence[float] | None = None,
         derivative_orders: int | Sequence[int] = 1,
         include_mixed: bool = False,
+        boundary: str = "periodic",
     ) -> torch.Tensor:
         """Dispatch rectilinear gradients to eager PyTorch."""
         return rectilinear_grid_gradient_torch_multi(
@@ -142,6 +149,7 @@ def torch_forward(
             periods=periods,
             derivative_orders=derivative_orders,
             include_mixed=include_mixed,
+            boundary=boundary,
         )
 
     @classmethod