Make wrapping into the domain and a patch's node grid public

CHANGELOG.md

@@ -7,6 +7,11 @@ format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Added
 
+- `AMRHierarchy.wrap(points)` and `wrapped(points)`: points beyond a periodic face moved back
+  into the domain, in place or as a copy (the array itself when nothing needs wrapping), and
+  `AMRHierarchy.node_grid(patch)`: a patch's nodes as an `mpifield.BlockGrid`, for compiled
+  kernels that run on a patch's node array. The private copies in `cells` and `nodal` are gone.
+
 - Solver backends (`patchwork.backends`): `NodalPoisson(backend=...)` takes a `Backend` or
   its name, and every backend solves the same assembled operator (`LinearSystem`).
   `"lu"` (`LUBackend`, the default) factorizes it with SuperLU, or with LAPACK's tridiagonal

src/patchwork/cells.py

@@ -117,17 +117,6 @@ class EmbeddedFloating:
 Body = EmbeddedDirichlet | EmbeddedFloating
 
 
-def _wrapped(hierarchy: AMRHierarchy, points: np.ndarray) -> np.ndarray:
-    """``points`` ``(ndim, k)`` moved into the domain along its periodic axes."""
-    if not any(hierarchy.periodic):
-        return points
-    points = points.copy()
-    for axis, (lower, upper) in enumerate(hierarchy.bounds):
-        if hierarchy.periodic[axis]:
-            points[axis] = lower + (points[axis] - lower) % (upper - lower)
-    return points
-
-
 class CompositeCells:
     """The leaf cells of a hierarchy, their faces and the maps to its patches.
 
@@ -216,7 +205,7 @@ def locate(self, points: np.ndarray) -> np.ndarray:
         Periodic axes wrap; a point on a face belongs to the cell above it, the
         upper faces of the domain to the last cell.
         """
-        points = _wrapped(self.hierarchy, np.asarray(points, dtype=float))
+        points = self.hierarchy.wrapped(np.asarray(points, dtype=float))
         lower = np.asarray(self.hierarchy.lower_bound, dtype=float)
         shape = np.asarray(self.owner.shape)
         lattice = np.floor((points.T - lower) / self._spacing).astype(np.int64)
@@ -769,7 +758,7 @@ def _fractions(self, outside: np.ndarray, inside: np.ndarray) -> np.ndarray:
                 fraction = np.asarray(body.fraction(start, start + step), dtype=float)
             else:
                 fraction = _surface_crossing(
-                    body.inside, start, step, lambda p: _wrapped(hierarchy, p)
+                    body.inside, start, step, hierarchy.wrapped
                 )
             result[mine] = np.maximum(fraction, self._min_fraction)
         return result

src/patchwork/mesh.py

@@ -993,6 +993,7 @@ def _points(self, coordinates: Sequence[Any]) -> tuple[tuple, tuple[int, ...]]:
             )
         return tuple(array.reshape(-1) for array in arrays), shape
 
+    @sp.profile("patchwork.locate", tags=PROFILE_TAGS)
     def locate(self, *coordinates: Any) -> Any:
         """The finest patch that holds each point: its index in :attr:`patches`.
 
@@ -1008,6 +1009,40 @@ def locate(self, *coordinates: Any) -> Any:
         points, shape = self._points(coordinates)
         return self._cells(points)[0].reshape(shape)
 
+    def wrap(self, points: Any) -> None:
+        """Move ``points`` ``(ndim, k)`` beyond a periodic face back into the domain, in place.
+
+        Only the points outside the domain along a periodic axis are touched.
+        """
+        import numpy as np
+
+        for axis, ((lower, upper), periodic) in enumerate(
+            zip(self.bounds, self.periodic, strict=True)
+        ):
+            if periodic:
+                row = points[axis]
+                out = np.flatnonzero((row < lower) | (row >= upper))
+                if out.size:
+                    row[out] = lower + np.mod(row[out] - lower, upper - lower)
+
+    def wrapped(self, points: Any) -> Any:
+        """``points`` ``(ndim, k)`` moved into the domain along its periodic axes.
+
+        The array itself if no point is beyond a periodic face, else a wrapped
+        copy (see :meth:`wrap`).
+        """
+        import numpy as np
+
+        for axis, ((lower, upper), periodic) in enumerate(
+            zip(self.bounds, self.periodic, strict=True)
+        ):
+            row = points[axis]
+            if periodic and ((row < lower) | (row >= upper)).any():
+                points = np.array(points, dtype=float)
+                self.wrap(points)
+                break
+        return points
+
     def _owners(self, found: Any, cell: tuple[Any, ...], name: str) -> Any:
         """The rank that holds the block of field ``name`` where each marker's cell starts."""
         import numpy as np
@@ -1121,25 +1156,40 @@ def migrate(
         )
         return tuple(np.asarray(array) for array in received)
 
+    def _grid(self, patch: Patch, first: Sequence[int], shape: Sequence[int]) -> Any:
+        """The :class:`mpifield.BlockGrid` of ``shape`` nodes of ``patch`` from node ``first`` on."""
+        from mpifield import BlockGrid
+
+        spacing = self.spacing(patch.level)
+        lower = tuple(
+            self.lower_bound[axis] + spacing[axis] * (patch.box[axis][0] + first[axis])
+            for axis in range(self.ndim)
+        )
+        return BlockGrid.from_geometry(lower, spacing, tuple(shape))
+
+    def node_grid(self, patch: Patch) -> Any:
+        """The nodes of ``patch`` as an :class:`mpifield.BlockGrid`: all of them, no halo.
+
+        For compiled kernels that run on a patch's node array, such as a push
+        that gathers each marker's field in its own patch. No axis is periodic:
+        a patch holds every node, both images of a periodic seam too.
+        """
+        return self._grid(patch, (0,) * self.ndim, patch.node_shape)
+
     def _stored_block(self, patch: Patch, name: str) -> tuple[Any, Any]:
         """This rank's block of field ``name`` on ``patch``, halo included, and its grid.
 
         The grid is an :class:`mpifield.BlockGrid` of the stored nodes (the first one is
         the first halo node), so mpifield's weightings run on the block as on a field's.
         """
-        from mpifield import BlockGrid
-
         array = patch.field(name)
         layout = array.layout
-        spacing = self.spacing(patch.level)
         storage = array.local_with_halos
         first = [
-            self.lower_bound[axis]
-            + spacing[axis]
-            * (patch.box[axis][0] + layout.index_bounds[axis][0] - layout.halo[axis])
+            layout.index_bounds[axis][0] - layout.halo[axis]
             for axis in range(self.ndim)
         ]
-        return storage, BlockGrid.from_geometry(tuple(first), spacing, storage.shape)
+        return storage, self._grid(patch, first, storage.shape)
 
     @sp.profile("patchwork.deposit", tags=PROFILE_TAGS)
     def deposit(

src/patchwork/nodal.py

@@ -1009,17 +1009,6 @@ class NodalPoisson:
     """
 
     # Assembly ------------------------------------------------------------------#
-    def _wrapped(self, points: np.ndarray) -> np.ndarray:
-        """``points`` ``(ndim, k)`` moved into the domain along its periodic axes."""
-        hierarchy = self.nodes.hierarchy
-        if not any(hierarchy.periodic):
-            return points
-        points = points.copy()
-        for axis, (lower, upper) in enumerate(hierarchy.bounds):
-            if hierarchy.periodic[axis]:
-                points[axis] = lower + (points[axis] - lower) % (upper - lower)
-        return points
-
     def _classify_bodies(self) -> None:
         """Which nodes are inside a body, and which one (the last listed)."""
         nodes = self.nodes
@@ -1047,7 +1036,9 @@ def _fractions(self, outside: np.ndarray, inside: np.ndarray) -> np.ndarray:
             if body.fraction is not None:
                 fraction = np.asarray(body.fraction(start, start + step), dtype=float)
             else:
-                fraction = _surface_crossing(body.inside, start, step, self._wrapped)
+                fraction = _surface_crossing(
+                    body.inside, start, step, self.nodes.hierarchy.wrapped
+                )
             result[mine] = np.maximum(fraction, self._min_fraction)
         return result
 

src/patchwork/tests/unit/test_markers.py

@@ -371,3 +371,29 @@ def test_interpolation_outside_the_domain_is_zero_and_checks_its_field() -> None
     if SIZE > 1:
         with pytest.raises(ValueError, match="rank that owns them"):
             mesh.interpolate(x, y, "q")  # every rank has all of them
+
+
+def test_points_beyond_a_periodic_face_are_wrapped_into_the_domain() -> None:
+    mesh = hierarchy()
+    inside = np.array([[0.2, 0.7], [0.0, 0.5]])
+    assert mesh.wrapped(inside) is inside  # nothing to wrap: no copy
+    points = np.array([[0.2, 0.7, 1.5], [1.25, -0.25, 1.0]])
+    wrapped = mesh.wrapped(points)
+    np.testing.assert_allclose(wrapped, [[0.2, 0.7, 1.5], [0.25, 0.75, 0.0]])
+    assert points[1, 0] == 1.25  # a copy
+    mesh.wrap(points)  # in place; x is not periodic and stays
+    np.testing.assert_array_equal(points, wrapped)
+    np.testing.assert_array_equal(mesh.locate(*wrapped), mesh.locate(*points))
+
+
+def test_the_node_grid_of_a_patch_holds_its_nodes() -> None:
+    mesh = hierarchy()
+    for patch in mesh.patches:
+        grid = mesh.node_grid(patch)
+        dx, dy = mesh.spacing(patch.level)
+        (i0, _), (j0, _) = patch.box
+        np.testing.assert_allclose(grid.bounds.lower_bound, (i0 * dx, j0 * dy))
+        np.testing.assert_allclose(grid.grid_spacing, (dx, dy))
+        assert tuple(grid.num_gridpoints) == patch.node_shape
+        assert grid.periodic == (False, False)
+        assert grid.geometry == (i0 * dx, j0 * dy, dx, dy, *patch.node_shape, 0, 0)