From 7c830d2530e78bcc96327e00bee22bbeb09b9e35 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 28 Oct 2026 14:35:24 +0100 Subject: [PATCH 1/2] Make wrapping into the domain and a patch's node grid public. --- CHANGELOG.md | 5 ++ src/patchwork/cells.py | 15 +----- src/patchwork/mesh.py | 63 +++++++++++++++++++++--- src/patchwork/nodal.py | 15 ++---- src/patchwork/tests/unit/test_markers.py | 26 ++++++++++ 5 files changed, 92 insertions(+), 32 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b599a2..66e2705 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/patchwork/cells.py b/src/patchwork/cells.py index 2290b3a..1d3fee4 100644 --- a/src/patchwork/cells.py +++ b/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 diff --git a/src/patchwork/mesh.py b/src/patchwork/mesh.py index 42995e3..2770575 100644 --- a/src/patchwork/mesh.py +++ b/src/patchwork/mesh.py @@ -1008,6 +1008,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 +1155,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( diff --git a/src/patchwork/nodal.py b/src/patchwork/nodal.py index 5b2eea7..c26936e 100644 --- a/src/patchwork/nodal.py +++ b/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 diff --git a/src/patchwork/tests/unit/test_markers.py b/src/patchwork/tests/unit/test_markers.py index d38fa96..670236c 100644 --- a/src/patchwork/tests/unit/test_markers.py +++ b/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) From 9e79609eecb55a4f3877698b6d200fa85f38cb30 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 28 Oct 2026 14:36:55 +0100 Subject: [PATCH 2/2] Profile AMRHierarchy.locate. --- src/patchwork/mesh.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/patchwork/mesh.py b/src/patchwork/mesh.py index 2770575..96df3a4 100644 --- a/src/patchwork/mesh.py +++ b/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`.