TransferManager: support a redistributed AdaptiveMeshHierarchy

demos/adaptive_multigrid/adaptive_multigrid.py.rst

@@ -6,8 +6,7 @@ Contributed by Anurag Rao.
 
 The purpose of this demo is to show how to use Firedrake's multigrid solver on a hierarchy of adaptively refined Netgen meshes.
 We will first have a look at how to use the :class:`~.AdaptiveMeshHierarchy` to construct the mesh hierarchy with Netgen meshes, then we will consider a solution to the Poisson problem on an L-shaped domain.
-Finally, we will show how to use the :class:`~.AdaptiveMeshHierarchy` and :class:`~.AdaptiveTransferManager` to construct a scalable solver. The :class:`~.AdaptiveMeshHierarchy` contains information of the mesh hierarchy and the parent child relations between the meshes.
-The :class:`~.AdaptiveTransferManager` deals with the transfer operator logic across any given levels in the hierarchy.
+Finally, we will show how to use the :class:`~.AdaptiveMeshHierarchy` to construct a scalable solver. The :class:`~.AdaptiveMeshHierarchy` contains information of the mesh hierarchy and the parent child relations between the meshes.
 We begin by importing the necessary libraries ::
 
    from firedrake import *
@@ -34,10 +33,9 @@ It is important to convert the initial Netgen mesh into a Firedrake mesh before
    :align: center
    :alt: Initial mesh.
 
-We will also initialize the :class:`~.AdaptiveTransferManager` here: ::
+We initialize the :class:`~.AdaptiveMeshHierarchy` here: ::
 
    amh = AdaptiveMeshHierarchy(mesh)
-   atm = AdaptiveTransferManager()
 
 Poisson Problem
 ---------------
@@ -62,15 +60,12 @@ Our approach strongly follows the similar problem in this `lecture course <https
 
       problem = LinearVariationalProblem(a, L, uh, bcs)
       solver = LinearVariationalSolver(problem, solver_parameters=params)
-
-      solver.set_transfer_manager(atm)
       solver.solve()
 
       its = solver.snes.getLinearSolveIterations()
       return uh, its
 
-Note the code after the construction of the :class:`~.LinearVariationalProblem`. To use the :class:`~.AdaptiveMeshHierarchy` with the existing Firedrake solver, we have to set the :class:`~.AdaptiveTransferManager` as the transfer manager of the multigrid solver.
-Since we are using linear Lagrange elements, we will employ Jacobi as the multigrid relaxation, which we define with ::
+To use the :class:`~.AdaptiveMeshHierarchy` in a multigrid solver, we just set the usual multigrid solver parameters. Since we are using linear Lagrange elements, we will employ Jacobi as the multigrid relaxation, which we define with ::
 
    solver_params = {
       "mat_type": "matfree",

firedrake/cython/mgimpl.pyx

@@ -97,6 +97,9 @@ def coarse_to_fine_nodes(Vc, Vf, np.ndarray coarse_to_fine_cells):
                     k = 0
                     for l in range(fine_cell_per_coarse_cell):
                         fine = coarse_to_fine_cells[i, l]
+                        if fine < 0:
+                            k += fine_per_cell * ratio
+                            continue
                         for layer in range(ratio):
                             fine_layer = coarse_layer * ratio + layer
                             for m in range(fine_per_cell):
@@ -107,6 +110,9 @@ def coarse_to_fine_nodes(Vc, Vf, np.ndarray coarse_to_fine_cells):
                 k = 0
                 for l in range(fine_cell_per_coarse_cell):
                     fine = coarse_to_fine_cells[i, l]
+                    if fine < 0:
+                        k += fine_per_cell
+                        continue
                     for m in range(fine_per_cell):
                         coarse_to_fine_map[node, k] = fine_map[fine, m]
                         k += 1
@@ -149,6 +155,8 @@ def fine_to_coarse_nodes(Vf, Vc, np.ndarray fine_to_coarse_cells):
 
     for i in range(fine_cells):
         for l, coarse_cell in enumerate(fine_to_coarse_cells[i, :]):
+            if coarse_cell < 0:
+                continue
             for j in range(fine_per_cell):
                 node = fine_map[i, j]
                 if extruded:

firedrake/mesh.py

@@ -2900,7 +2900,7 @@ def __iter__(self):
     def unique(self):
         return self
 
-    def refine_marked_elements(self, mark, netgen_flags=None):
+    def refine_marked_elements(self, mark, netgen_flags=None, redistribute=True):
         """Refine a mesh using a DG0 marking function.
 
         This method requires that the mesh has been constructed from a
@@ -2909,6 +2909,8 @@ def refine_marked_elements(self, mark, netgen_flags=None):
         :arg mark: the marking function which is a Firedrake DG0 function
             with the number of refinements on each cell.
         :arg netgen_flags: the dictionary of flags to be passed to ngsPETSc.
+        :arg redistribute: if true, redistribute the refined mesh when the
+            parent-owned child distribution is poorly balanced.
 
         It includes the option:
             - refine_faces, which is a boolean specifying if you want to refine faces.
@@ -2918,50 +2920,9 @@ def refine_marked_elements(self, mark, netgen_flags=None):
 
         if not hasattr(self, "netgen_mesh"):
             raise ValueError("Adaptive refinement requires a netgen mesh.")
-        if netgen_flags is None:
-            netgen_flags = self.netgen_flags
-        tdim = self.topological_dimension
-        if tdim not in {2, 3}:
-            raise NotImplementedError("No implementation for dimension other than 2 and 3.")
-        with mark.dat.vec as mvec:
-            if self.sfBC_orig is None:
-                cstart, cend = self.topology_dm.getHeightStratum(0)
-                cellNum = list(map(self._cell_numbering.getOffset, range(cstart, cend)))
-                mark_np = mvec.getArray()[cellNum]
-            else:
-                sfBCInv = self.sfBC_orig.createInverse()
-                _, mvec0 = self.topology_dm.distributeField(sfBCInv,
-                                                            self._cell_numbering,
-                                                            mvec)
-                mark_np = mvec0.getArray()
-        max_refs = 0 if mark_np.size == 0 else int(mark_np.max())
-        # Create a copy of the netgen mesh
-        netgen_mesh = self.netgen_mesh.Copy()
-        refine_faces = netgen_flags.get("refine_faces", False)
-        for r in range(max_refs):
-            cells = netgen_mesh.Elements3D() if tdim == 3 else netgen_mesh.Elements2D()
-            cells.NumPy()["refine"] = (mark_np[:len(cells)] > 0)
-            if tdim == 3:
-                faces = netgen_mesh.Elements2D()
-                faces.NumPy()["refine"] = refine_faces
-            netgen_mesh.Refine(adaptive=True)
-            mark_np -= 1
-            if r < max_refs - 1:
-                parents = netgen_mesh.parentelements if tdim == 3 else netgen_mesh.parentsurfaceelements
-                parents = parents.NumPy()["i"]
-                num_fine_cells = parents.shape[0]
-                num_coarse_cells = mark_np.size
-                indices = np.arange(num_fine_cells, dtype=PETSc.IntType)
-                while (indices >= num_coarse_cells).any():
-                    fine_cells = (indices >= num_coarse_cells)
-                    indices[fine_cells] = parents[indices[fine_cells]]
-                mark_np = mark_np[indices]
-
-        return Mesh(netgen_mesh,
-                    reorder=self._did_reordering,
-                    distribution_parameters=self._distribution_parameters,
-                    comm=self.comm,
-                    netgen_flags=netgen_flags)
+
+        from firedrake.netgen import refine_marked_elements
+        return refine_marked_elements(self, mark, netgen_flags, redistribute)
 
     @PETSc.Log.EventDecorator()
     def curve_field(self, order, permutation_tol=1e-8, cg_field=None):

firedrake/mg/adaptive_hierarchy.py

@@ -1,8 +1,14 @@
+from collections import defaultdict
+from fractions import Fraction
+
 from firedrake.mesh import MeshGeometry
 from firedrake.cofunction import Cofunction
 from firedrake.function import Function
-from firedrake.mg import HierarchyBase
-from firedrake.mg.utils import set_level
+from firedrake.mg.mesh import HierarchyBase
+from firedrake.mg.utils import set_level, set_refine_level
+from firedrake.netgen import (adaptive_cell_maps, adaptive_netgen_cells,
+                              adaptive_parent_cell_numbers,
+                              make_adaptive_refined_mesh)
 
 __all__ = ["AdaptiveMeshHierarchy"]
 
@@ -17,27 +23,84 @@ class AdaptiveMeshHierarchy(HierarchyBase):
         The coarsest mesh in the hierarchy.
     nested: bool
         A flag to indicate whether the meshes are nested.
+    redistribute: bool
+        If ``True``, keep adaptively refined meshes redistributed when
+        this is needed to avoid empty ranks and use an internal
+        parent-owned mesh for transfer operators.
 
     """
-    def __init__(self, base_mesh: MeshGeometry, nested: bool = True):
+    def __init__(self, base_mesh: MeshGeometry, nested: bool = True,
+                 redistribute: bool = True):
         self.meshes = []
         self._meshes = []
+        self.coarse_to_fine_cells = {}
+        self.fine_to_coarse_cells = {Fraction(0, 1): None}
+        self.refinements_per_level = 1
         self.nested = nested
+        self.redistribute = redistribute
+        self._shared_data_cache = defaultdict(dict)
         self.add_mesh(base_mesh)
 
-    def add_mesh(self, mesh: MeshGeometry):
+    def add_mesh(self, mesh: MeshGeometry,
+                 coarse_to_fine_cells=None,
+                 fine_to_coarse_cells=None):
         """
         Adds a mesh into the hierarchy.
 
         Parameters
         ----------
         mesh
             The mesh to be added to the finest level.
+        coarse_to_fine_cells
+            Optional map from cells on the previous finest level to cells on
+            ``mesh``.
+        fine_to_coarse_cells
+            Optional map from cells on ``mesh`` to cells on the previous
+            finest level.
         """
         level = len(self.meshes)
+        if level > 0 and (coarse_to_fine_cells is None or fine_to_coarse_cells is None):
+            redist = getattr(mesh, "redist", None)
+            transfer_mesh = redist.orig if redist is not None else mesh
+            parent_cell_numbers = getattr(transfer_mesh, "_adaptive_parent_cell_numbers", None)
+            if parent_cell_numbers is None:
+                parent_cell_numbers = getattr(mesh, "_adaptive_parent_cell_numbers", None)
+            if parent_cell_numbers is None:
+                parent_cell_numbers = adaptive_parent_cell_numbers(self.meshes[-1], mesh)
+
+            if parent_cell_numbers is not None:
+                if not getattr(transfer_mesh, "_adaptive_parent_owned", False):
+                    if mesh.comm.size == 1:
+                        mesh._adaptive_parent_cell_numbers = parent_cell_numbers
+                        mesh._adaptive_parent_owned = True
+                        transfer_mesh = mesh
+                    else:
+                        mesh = make_adaptive_refined_mesh(
+                            self.meshes[-1], mesh.netgen_mesh, mesh.netgen_flags,
+                            parent_cell_numbers, self.redistribute,
+                        )
+                        redist = getattr(mesh, "redist", None)
+                        transfer_mesh = redist.orig if redist is not None else mesh
+                        parent_cell_numbers = transfer_mesh._adaptive_parent_cell_numbers
+
+                coarse_to_fine_cells, fine_to_coarse_cells = adaptive_cell_maps(
+                    self.meshes[-1], transfer_mesh, parent_cell_numbers
+                )
+
         self._meshes.append(mesh)
         self.meshes.append(mesh)
         set_level(mesh, self, level)
+        set_refine_level(mesh, level)
+        redist = getattr(mesh, "redist", None)
+
+        if hasattr(mesh, "netgen_mesh"):
+            mesh._adaptive_netgen_num_cells = len(adaptive_netgen_cells(mesh))
+            if redist is not None:
+                redist.orig._adaptive_netgen_num_cells = mesh._adaptive_netgen_num_cells
+
+        if level > 0 and coarse_to_fine_cells is not None and fine_to_coarse_cells is not None:
+            self.coarse_to_fine_cells[Fraction(level - 1, 1)] = coarse_to_fine_cells
+            self.fine_to_coarse_cells[Fraction(level, 1)] = fine_to_coarse_cells
 
     def adapt(self, eta: Function | Cofunction, theta: float):
         """
@@ -80,6 +143,6 @@ def adapt(self, eta: Function | Cofunction, theta: float):
         markers = Function(M)
         markers.dat.data_wo[should_refine] = 1
 
-        refined_mesh = mesh.refine_marked_elements(markers)
+        refined_mesh = mesh.refine_marked_elements(markers, redistribute=self.redistribute)
         self.add_mesh(refined_mesh)
-        return refined_mesh
+        return self.meshes[-1]

firedrake/mg/adaptive_transfer_manager.py

@@ -2,54 +2,18 @@
 This module contains the AdaptiveTransferManager used to perform
 transfer operations on AdaptiveMeshHierarchies
 """
+import warnings
 from firedrake.mg.embedded import TransferManager
-from firedrake.ufl_expr import action, TrialFunction
-from firedrake.interpolation import interpolate
-
 
 __all__ = ("AdaptiveTransferManager",)
 
 
-class AdaptiveTransferManager(TransferManager):
+def AdaptiveTransferManager(*args, **kwargs):
     """
     TransferManager for adaptively refined mesh hierarchies
     """
-    def __init__(self, *, native_transfers=None, use_averaging=True):
-        if native_transfers is not None:
-            raise NotImplementedError("Custom transfers not implemented.")
-        super().__init__(native_transfers=native_transfers, use_averaging=use_averaging)
-        self.cache = {}
-
-    def get_interpolator(self, Vc, Vf):
-        from firedrake.assemble import assemble
-        key = (Vc, Vf)
-        try:
-            return self.cache[key]
-        except KeyError:
-            Iexpr = interpolate(TrialFunction(Vc), Vf)
-            # TODO reusable matfree Interpolator
-            I = assemble(Iexpr, mat_type="aij")
-            return self.cache.setdefault(key, I)
-
-    def forward(self, uc, uf):
-        from firedrake.assemble import assemble
-        Vc = uc.function_space()
-        Vf = uf.function_space()
-        I = self.get_interpolator(Vc, Vf)
-        return assemble(action(I, uc), tensor=uf)
-
-    def adjoint(self, uf, uc):
-        from firedrake.assemble import assemble
-        Vc = uc.function_space().dual()
-        Vf = uf.function_space().dual()
-        I = self.get_interpolator(Vc, Vf)
-        return assemble(action(uf, I), tensor=uc)
-
-    def prolong(self, uf, uc):
-        return self.forward(uf, uc)
-
-    def inject(self, uc, uf):
-        return self.forward(uc, uf)
-
-    def restrict(self, uc, uf):
-        return self.adjoint(uc, uf)
+    warnings.warn(
+        "The ``AdaptiveTransferManager`` class is deprecated and will be removed in a future release. "
+        "Please use the ``TransferManager`` class instead.", FutureWarning
+    )
+    return TransferManager(*args, **kwargs)

firedrake/mg/interface.py

@@ -68,10 +68,13 @@ def prolong(coarse, fine):
     meshes = hierarchy._meshes
     for j in range(repeat):
         next_level += 1
-        if j == repeat - 1 and not needs_quadrature:
+        fine_mesh = meshes[next_level]
+        redist = getattr(fine_mesh, "redist", None)
+        transfer_mesh = redist.orig if redist is not None else fine_mesh
+        if j == repeat - 1 and not needs_quadrature and redist is None:
             fine = finest
         else:
-            fine = Function(Vf.reconstruct(mesh=meshes[next_level]))
+            fine = Function(Vf.reconstruct(mesh=transfer_mesh))
         Vf = fine.function_space()
         Vc = coarse.function_space()
         compose_map = lambda u: utils.fine_node_to_coarse_node_map(Vf, u.function_space())
@@ -106,8 +109,15 @@ def prolong(coarse, fine):
 
         if needs_quadrature:
             # Transfer to the actual target space
-            new_fine = finest if j == repeat-1 else Function(Vfinest.reconstruct(mesh=meshes[next_level]))
+            target_mesh = transfer_mesh if redist is not None else meshes[next_level]
+            new_fine = (finest if j == repeat-1 and redist is None
+                        else Function(Vfinest.reconstruct(mesh=target_mesh)))
             fine = new_fine.interpolate(fine)
+        if redist is not None:
+            target = (finest if j == repeat - 1
+                      else Function(Vfinest.reconstruct(mesh=fine_mesh)))
+            redist.orig2redist(fine, target)
+            fine = target
         coarse = fine
     return fine
 
@@ -145,9 +155,19 @@ def restrict(fine_dual, coarse_dual):
     coarsest = coarse_dual.zero()
     meshes = hierarchy._meshes
     for j in range(repeat):
+        fine_mesh = meshes[next_level]
+        redist = getattr(fine_mesh, "redist", None)
+        if redist is not None:
+            fine_dual_transfer = Cofunction(
+                fine_dual.function_space().reconstruct(mesh=redist.orig)
+            )
+            redist.redist2orig(fine_dual, fine_dual_transfer)
+            fine_dual = fine_dual_transfer
         if needs_quadrature:
             # Transfer to the quadrature source space
-            fine_dual = Function(Vq.reconstruct(mesh=meshes[next_level])).interpolate(fine_dual)
+            fine_dual = Function(
+                Vq.reconstruct(mesh=fine_dual.function_space().mesh())
+            ).interpolate(fine_dual)
 
         next_level -= 1
         if j == repeat - 1:
@@ -238,13 +258,28 @@ def inject(fine, coarse):
     Vcoarsest = coarsest.function_space()
     meshes = hierarchy._meshes
     for j in range(repeat):
+        fine_mesh = meshes[next_level]
+        redist = getattr(fine_mesh, "redist", None)
+        if redist is not None:
+            fine_transfer = Function(
+                fine.function_space().reconstruct(mesh=redist.orig)
+            )
+            redist.redist2orig(fine, fine_transfer)
+            fine = fine_transfer
         next_level -= 1
         if j == repeat - 1 and not needs_quadrature:
             coarse = coarsest
         else:
             coarse = Function(Vc.reconstruct(mesh=meshes[next_level]))
         Vc = coarse.function_space()
         Vf = fine.function_space()
+
+        # FIXME
+        from firedrake.mg.adaptive_hierarchy import AdaptiveMeshHierarchy
+        adaptive = isinstance(hierarchy, AdaptiveMeshHierarchy)
+        if adaptive and (dg or Vc.mesh().comm.size > 1):
+            return coarsest.interpolate(fine)
+
         if not dg:
             compose_map = lambda u: utils.coarse_node_to_fine_node_map(Vc, u.function_space())
             node_locations = utils.physical_node_locations(Vc)