DomainDecomposition.coarsen, refine fix, ghost-sync fixes in axpy and essential BCs

feectools/api/essential_bc.py

@@ -72,6 +72,9 @@ def apply_essential_bc_stencil(a, *, axis, ext, order, identity=False):
     if isinstance(a, StencilVector):
         V = a.space
         n = V.ndim
+        # Boundary entries may be ghost entries of neighbouring processes, which
+        # all call this function: their ghost regions are no longer up to date.
+        a.ghost_regions_in_sync = False
     elif isinstance(a, StencilMatrix):
         V = a.codomain
         n = V.ndim * 2

feectools/api/tests/test_essential_bc_ghosts.py

@@ -0,0 +1,33 @@
+import numpy as np
+import pytest
+
+from feectools.api.essential_bc import apply_essential_bc_stencil
+from feectools.ddm.cart import DomainDecomposition
+from feectools.ddm.mpi import mpi as MPI
+from feectools.fem.splines import SplineSpace
+from feectools.fem.tensor import TensorFemSpace
+
+
+@pytest.mark.parametrize("parallel", [False, pytest.param(True, marks=pytest.mark.mpi)])
+def test_essential_bc_marks_ghosts_stale(parallel):
+    """After setting boundary coefficients to zero, ghost regions (on all processes) must be refreshed."""
+    comm = MPI.COMM_WORLD if parallel else None
+    p, nc = 2, 8
+    spaces = [SplineSpace(p, grid=np.linspace(0.0, 1.0, nc + 1), periodic=False) for _ in range(2)]
+    dd = DomainDecomposition([nc, nc], [False, False], comm=comm)
+    V = TensorFemSpace(dd, *spaces).coeff_space
+
+    v = V.zeros()
+    v._data[...] = 1.0
+    v.update_ghost_regions()
+    assert v.ghost_regions_in_sync
+
+    apply_essential_bc_stencil(v, axis=0, ext=-1, order=0)
+    assert not v.ghost_regions_in_sync
+
+    v.update_ghost_regions()
+    # every copy (owned or ghost) of the boundary coefficients i0 = 0 is zero now
+    s0, pad = V.starts[0], V.pads[0]
+    for i_loc in range(v._data.shape[0]):
+        if s0 - pad + i_loc == 0:
+            assert np.all(v._data[i_loc, pad:-pad] == 0.0)

feectools/ddm/cart.py

@@ -1,5 +1,6 @@
 # coding: utf-8
 
+import copy
 import os
 import numpy as np
 import cunumpy as xp
@@ -409,42 +410,100 @@ def coords_exist( self, coords ):
     def refine(self, ncells, global_element_starts, global_element_ends):
         """ Create the new Cartesian decomposition of the refined domain.
 
+        The process topology (and its communicators) is shared with ``self``.
+
         Parameters
         ----------
         ncells : list or tuple of int
             Number of cells of refined space.
 
-        global_starts: list of list of int
-            The starts of the coefficients for every process along each direction.
+        global_element_starts : list of list of int
+            The element starts for every process along each direction.
 
-        global_ends: list of list of int
-            The ends of the coefficients for every process along each direction.
+        global_element_ends : list of list of int
+            The element ends for every process along each direction.
 
         Returns
         -------
-        domain : CartDecomposition
-            Cartesian decomposition of the refined domain.
+        domain : DomainDecomposition
+            Domain decomposition of the refined domain.
         """
 
         # Check input arguments
         assert len( ncells ) == len( self.ncells )
         assert all(nc>=snc for nc, snc in zip(ncells, self.ncells))
 
-        domain         = DomainDecomposition(self.ncells, self.periods, comm=self.comm,
-                                            global_comm=self.global_comm, num_threads=self.num_threads,
-                                            size=self.size)
-        domain._ncells = tuple ( ncells )
+        return self._with_element_partition(ncells, global_element_starts, global_element_ends)
+
+    def coarsen(self, factors):
+        """ Create the Cartesian decomposition of a coarsened domain, aligned with ``self``.
+
+        Along axis ``i`` every ``factors[i]`` consecutive cells are merged into one coarse cell.
+        Each process owns the coarse cells covering exactly its fine cells, so the process
+        topology (and its communicators) is shared with ``self``. This requires that the
+        element starts and ends+1 of every process are divisible by ``factors[i]``.
+
+        Parameters
+        ----------
+        factors : list or tuple of int
+            Coarsening factor (>= 1) along each direction.
+
+        Returns
+        -------
+        domain : DomainDecomposition
+            Domain decomposition of the coarse domain.
+        """
+
+        assert len( factors ) == self.ndim
+        assert all( isinstance(f, (int, np.integer)) and f >= 1 for f in factors )
+
+        ncells = []
+        global_element_starts = []
+        global_element_ends   = []
+        for axis, f in enumerate(factors):
+            gs = xp.asarray(self._global_element_starts[axis])
+            ge = xp.asarray(self._global_element_ends  [axis])
+            if self._ncells[axis] % f != 0 or xp.any(gs % f != 0) or xp.any((ge + 1) % f != 0):
+                raise ValueError(
+                    f"Cannot coarsen axis {axis} by a factor {f}: ncells={self._ncells[axis]}, "
+                    f"element starts={gs.tolist()}, ends={ge.tolist()} are not all aligned."
+                )
+            ncells.append(self._ncells[axis] // f)
+            global_element_starts.append(xp.array(gs // f))
+            global_element_ends  .append(xp.array((ge + 1) // f - 1))
+
+        return self._with_element_partition(ncells, global_element_starts, global_element_ends)
+
+    def _with_element_partition(self, ncells, global_element_starts, global_element_ends):
+        """ Return a copy of ``self`` with the same process topology but a new element partition.
+
+        Communicators are shared (not duplicated), hence this method is not collective.
+        """
+
+        assert len( ncells ) == self.ndim
+        for axis in range(self.ndim):
+            gs = xp.asarray(global_element_starts[axis])
+            ge = xp.asarray(global_element_ends  [axis])
+            assert len(gs) == len(ge) == self._nprocs[axis], \
+                f"Axis {axis}: need one block per process ({self._nprocs[axis]}), got {len(gs)}."
+            assert gs[0] == 0 and ge[-1] == ncells[axis] - 1, \
+                f"Axis {axis}: blocks must cover [0, {ncells[axis] - 1}]."
+            assert xp.all(ge >= gs), f"Axis {axis}: empty blocks are not allowed."
+            assert xp.all(gs[1:] == ge[:-1] + 1), f"Axis {axis}: blocks must be contiguous."
+
+        domain = copy.copy(self)
+        domain._ncells = tuple( int(n) for n in ncells )
 
         # Store arrays with all the starts and ends along each direction for every process
-        domain._global_element_starts = tuple(global_element_starts)
-        domain._global_element_ends   = tuple(global_element_ends)
+        domain._global_element_starts = list(global_element_starts)
+        domain._global_element_ends   = list(global_element_ends)
         if self.is_comm_null:return domain
 
         # Start/end values of global indices (without ghost regions)
         domain._starts = tuple( domain._global_element_starts[axis][c] for axis,c in zip(range(self._ndims), self._coords) )
         domain._ends   = tuple( domain._global_element_ends  [axis][c] for axis,c in zip(range(self._ndims), self._coords) )
 
-        domain._local_ncells = tuple(e-s+1 for s,e in zip(self._starts, self._ends))
+        domain._local_ncells = tuple(e-s+1 for s,e in zip(domain._starts, domain._ends))
         return domain
 
 #==================================================================================

feectools/ddm/tests/test_coarsen.py

@@ -0,0 +1,67 @@
+import numpy as np
+import pytest
+
+from feectools.ddm.cart import DomainDecomposition
+from feectools.ddm.mpi import mpi as MPI
+
+
+def _comm(parallel):
+    return MPI.COMM_WORLD if parallel else None
+
+
+def _as_list(arrs):
+    return [np.asarray(a).tolist() for a in arrs]
+
+
+@pytest.mark.parametrize("parallel", [False, pytest.param(True, marks=pytest.mark.mpi)])
+@pytest.mark.parametrize("periods", [(True, True, True), (False, True, False)])
+def test_coarsen(parallel, periods):
+    comm = _comm(parallel)
+    fine = DomainDecomposition([32, 16, 8], periods, comm=comm)
+    coarse = fine.coarsen([2, 4, 1])
+
+    assert coarse.ncells == (16, 4, 8)
+    assert coarse.periods == fine.periods
+    assert tuple(coarse.nprocs) == tuple(fine.nprocs)
+    assert coarse.comm_cart is fine.comm_cart
+    assert tuple(coarse.coords) == tuple(fine.coords)
+
+    # every process owns exactly the coarse cells covering its fine cells
+    for axis, f in enumerate([2, 4, 1]):
+        assert coarse.starts[axis] * f == fine.starts[axis]
+        assert (coarse.ends[axis] + 1) * f == fine.ends[axis] + 1
+        assert coarse.local_ncells[axis] * f == fine.local_ncells[axis]
+
+    # the original object is untouched
+    assert fine.ncells == (32, 16, 8)
+
+    # refining back gives the original partition
+    back = coarse.refine(
+        fine.ncells,
+        [np.asarray(s) * f for s, f in zip(coarse.global_element_starts, [2, 4, 1])],
+        [(np.asarray(e) + 1) * f - 1 for e, f in zip(coarse.global_element_ends, [2, 4, 1])],
+    )
+    assert back.ncells == fine.ncells
+    assert _as_list(back.global_element_starts) == _as_list(fine.global_element_starts)
+    assert _as_list(back.global_element_ends) == _as_list(fine.global_element_ends)
+    assert back.starts == fine.starts
+    assert back.ends == fine.ends
+    assert back.local_ncells == fine.local_ncells
+
+
+@pytest.mark.parametrize("parallel", [False, pytest.param(True, marks=pytest.mark.mpi)])
+def test_coarsen_repeated(parallel):
+    comm = _comm(parallel)
+    dd = DomainDecomposition([64, 32, 1], (True, True, True), comm=comm, mpi_dims_mask=[True, True, False])
+    for _ in range(3):
+        dd = dd.coarsen([2, 2, 1])
+    assert dd.ncells == (8, 4, 1)
+    assert dd.local_ncells[2] == 1
+
+
+def test_coarsen_misaligned():
+    dd = DomainDecomposition([6, 4, 4], (True, True, True))
+    with pytest.raises(ValueError):
+        dd.coarsen([4, 1, 1])
+    with pytest.raises(AssertionError):
+        dd.coarsen([0, 1, 1])

feectools/linalg/stencil.py

@@ -332,7 +332,7 @@ def axpy(self, a, x, y):
                 import cupy as cp
                 y._interface_data[axis, ext][:] = cp.asarray(y_int_np)
 
-        x._sync = x._sync and y._sync
+        y._sync = x._sync and y._sync
 
     #--------------------------------------
     # Other properties/methods

feectools/linalg/tests/test_axpy_ghost_sync.py

@@ -0,0 +1,31 @@
+from feectools.ddm.cart import DomainDecomposition, CartDecomposition
+from feectools.linalg.stencil import StencilVectorSpace
+
+
+def _space():
+    dd = DomainDecomposition([8, 8], [True, True])
+    cart = CartDecomposition(dd, [8, 8], [[0], [0]], [[7], [7]], [2, 2], [1, 1])
+    return StencilVectorSpace(cart)
+
+
+def test_axpy_ghost_sync():
+    """y += a*x must mark y out of sync if x is, and must not modify the flag of x."""
+    V = _space()
+    x, y = V.zeros(), V.zeros()
+
+    x.ghost_regions_in_sync = False
+    y.ghost_regions_in_sync = True
+    y.mul_iadd(2.0, x)
+    assert not y.ghost_regions_in_sync
+    assert not x.ghost_regions_in_sync
+
+    x.ghost_regions_in_sync = True
+    y.ghost_regions_in_sync = False
+    y.mul_iadd(2.0, x)
+    assert not y.ghost_regions_in_sync
+    assert x.ghost_regions_in_sync
+
+    x.ghost_regions_in_sync = True
+    y.ghost_regions_in_sync = True
+    y.mul_iadd(2.0, x)
+    assert y.ghost_regions_in_sync

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name            = "feectools"
-version         = "0.2.0"
+version         = "0.3.0"
 description     = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies."
 readme          = "README.md"
 requires-python = ">= 3.10"