Mesh Adaptivity via DMPlexTransform

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/dmhooks.py

@@ -446,18 +446,81 @@ def coarsen(dm, comm):
     return cdm
 
 
+def _adaptively_refine(dm, comm):
+    from firedrake.mg.adaptive_hierarchy import AdaptiveMeshHierarchy
+    from firedrake.mg.ufl_utils import refine
+    from firedrake.mg.utils import get_level
+
+    ctx = get_appctx(dm)
+    current_solution = ctx._x
+    mesh = current_solution.function_space().mesh()
+    hierarchy, level = get_level(mesh)
+    if hierarchy is None:
+        hierarchy = AdaptiveMeshHierarchy(mesh)
+        level = 0
+    if not isinstance(hierarchy, AdaptiveMeshHierarchy):
+        raise RuntimeError("Adaptive SNES refinement requires an AdaptiveMeshHierarchy")
+
+    if level == len(hierarchy) - 1:
+        if ctx._marking_callback is None:
+            raise RuntimeError("Adaptive SNES refinement requires setting a marking_callback")
+        solver = None
+        markers = ctx._marking_callback(solver, current_solution)
+        if not isinstance(markers, (firedrake.Function, firedrake.Cofunction)):
+            raise TypeError(
+                f"marking callback must return a Function or Cofunction, not a {type(markers).__name__}"
+            )
+        M = markers.function_space()
+        if M.mesh() is not mesh:
+            raise ValueError("marking callback must return markers on the current solution mesh")
+        if M.finat_element.space_dimension() != 1:
+            raise ValueError("marking callback must return a DG0 Function or Cofunction")
+
+        hierarchy.add_mesh(mesh.refine_marked_elements(markers))
+
+    coefficient_mapping = {}
+    refined_ctx = refine(ctx, refine, coefficient_mapping=coefficient_mapping)
+    parent = get_parent(dm)
+    coarsener = get_ctx_coarsener(dm)
+    # Get all DMs from the refined problem
+    dms = [refined_ctx._problem.u_restrict.function_space().dm]
+    for value in coefficient_mapping.values():
+        if isinstance(value, (firedrake.Function, firedrake.Cofunction)):
+            dm = value.function_space().dm
+            if dm not in dms:
+                dms.append(dm)
+    # Attach refined context
+    for dm in dms:
+        add_hook(parent, setup=partial(push_parent, dm, parent),
+                 teardown=partial(pop_parent, dm, parent),
+                 call_setup=True)
+        add_hook(parent, setup=partial(push_ctx_coarsener, dm, coarsener),
+                 teardown=partial(pop_ctx_coarsener, dm, coarsener),
+                 call_setup=True)
+        add_hook(parent, setup=partial(push_appctx, dm, refined_ctx),
+                 teardown=partial(pop_appctx, dm, refined_ctx),
+                 call_setup=True)
+    return refined_ctx._problem.dm
+
+
 @PETSc.Log.EventDecorator()
 def refine(dm, comm):
     """Callback to refine a DM.
 
     :arg DM: The DM to refine.
     :arg comm: The communicator for the new DM (ignored)
     """
+    ctx = get_appctx(dm)
+    if ctx._marking_callback is not None:
+        return _adaptively_refine(dm, comm)
+
     from firedrake.mg.utils import get_level
     V = get_function_space(dm)
     if V is None:
         raise RuntimeError("No functionspace found on DM")
     hierarchy, level = get_level(V.mesh())
+    if hierarchy is None:
+        raise RuntimeError("No mesh hierarchy available")
     if level >= len(hierarchy) - 1:
         raise RuntimeError("Cannot refine finest DM")
     if hasattr(V, "_fine"):

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)