Mesh Adaptivity via DMPlexTransform

AGENTS.md

@@ -76,6 +76,11 @@ toolchain:
   mesh-bound data.
 * **Docstrings:** All public-facing APIs must include properly formatted `numpydoc`-style docstrings.
 * **Type Hints:** New code should include type hints on function/method signatures.
+* **Demos Are Literate Programs:** `pylit` converts each `demos/<name>/<name>.py.rst` into a `.py` that
+  `tests/firedrake/demos/test_demos_run.py` executes, so prose and code must stay in step. A paragraph
+  ending in `::` makes the indented block after it *executable*; a `.. code-block:: python` directive is
+  excluded from that rule, so its snippet renders in the docs but never runs. Prefer `::` — reach for
+  the directive only for an illustrative fragment naming things the demo never defines.
 
 ## Testing Requirements
 

demos/adaptive_multigrid/adaptive_multigrid.py.rst

@@ -1,13 +1,12 @@
-Adaptive Multigrid Methods using AdaptiveMeshHierarchy
-======================================================
+Adaptive Multigrid Methods
+==========================
 
 
 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.
+A :func:`~.MeshHierarchy` is not restricted to uniform refinement: the same object records the parent child relations between adaptively refined meshes, and grows a level at a time as the solution is resolved.
+We will first have a look at how to construct such a hierarchy from Netgen meshes, then we will consider a solution to the Poisson problem on an L-shaped domain, and finally we will use the hierarchy to construct a scalable solver.
 We begin by importing the necessary libraries ::
 
    from firedrake import *
@@ -28,16 +27,15 @@ We begin with the L-shaped domain, which we build as the union of two rectangles
    ngmsh = geo.GenerateMesh(maxh=0.5)
    mesh = Mesh(ngmsh)
 
-It is important to convert the initial Netgen mesh into a Firedrake mesh before constructing the :class:`~.AdaptiveMeshHierarchy`. To call the constructor to the hierarchy, we must pass the initial mesh. Our initial mesh looks like this:
+It is important to convert the initial Netgen mesh into a Firedrake mesh before constructing the :func:`~.MeshHierarchy`. To call the constructor to the hierarchy, we must pass the initial mesh. Our initial mesh looks like this:
 
 .. figure:: initial_mesh.png
    :align: center
    :alt: Initial mesh.
 
-We will also initialize the :class:`~.AdaptiveTransferManager` here: ::
+We initialize the :func:`~.MeshHierarchy` here. The default of zero uniform refinement levels gives a hierarchy holding just the initial mesh, which we will grow adaptively below; passing a positive number instead would start us off with that many uniformly refined levels, and the adaptive levels would stack on top of them just the same: ::
 
-   amh = AdaptiveMeshHierarchy(mesh)
-   atm = AdaptiveTransferManager()
+   mh = MeshHierarchy(mesh)
 
 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 hierarchy 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",
@@ -100,7 +95,7 @@ The initial solution is shown below.
 
 Adaptive Mesh Refinement
 ------------------------
-In this section we will discuss how to adaptively refine select elements and add the newly refined mesh into the :class:`~.AdaptiveMeshHierarchy`.
+In this section we will discuss how to adaptively refine select elements and add the newly refined mesh into the hierarchy.
 For this problem, we will be using the Babuška-Rheinbolt a-posteriori estimate for an element:
 
 .. math::
@@ -159,7 +154,7 @@ With these helper functions complete, we can solve the system iteratively. In th
    for level in range(refinements):
       print(f"level {level}")
 
-      mesh = amh[-1]
+      mesh = mh[-1]
       uh, its = solve_poisson(mesh, solver_params)
       VTKFile(f"output/adaptive_loop_{level}.pvd").write(uh)
 
@@ -175,10 +170,17 @@ With these helper functions complete, we can solve the system iteratively. In th
          rates = -numpy.diff(numpy.log(est_errors)) / numpy.diff(numpy.log(sqrt_dofs))
          print(f"  rate = {rates[-1]}")
 
-      if i != refinements - 1:
-         amh.adapt(eta, theta)
+      if level != refinements - 1:
+         mh.adapt(eta, theta)
+
+To perform Dörfler marking, refine the current mesh, and add the mesh to the hierarchy, we use the :meth:`~.HierarchyBase.adapt` method. In this method the input is the recently computed error estimator ``eta`` and the Dörfler marking parameter ``theta``. The method always performs this on the current fine mesh in the hierarchy.
+To mark cells by some other criterion, refine the finest mesh yourself and add the result, which is all that :meth:`~.HierarchyBase.adapt` does once it has marked:
+
+.. code-block:: python
+
+   mh.add_mesh(mh[-1].refine_marked_elements(markers))
 
-To perform Dörfler marking, refine the current mesh, and add the mesh to the :class:`~.AdaptiveMeshHierarchy`, we use the ``amh.adapt(eta, theta)`` method. In this method the input is the recently computed error estimator ``eta`` and the Dörfler marking parameter ``theta``. The method always performs this on the current fine mesh in the hierarchy. There is another method for adding a mesh to the hierarchy: ``amh.add_mesh(mesh)``. In this method, refinement on the mesh is performed externally by some custom procedure and the resulting mesh directly gets added to the hierarchy.
+Here ``markers`` is a DG0 function whose value on each cell is the number of times to refine it. If the mesh was instead produced by some procedure Firedrake cannot trace the parent child relations through, pass those cell maps to :meth:`~.HierarchyBase.add_mesh` explicitly.
 The meshes now refine according to the error estimator. The error estimators at levels 3,5, and 15 are shown below. Zooming into the vertex of the L-shape at level 15 shows the error indicator remains strongest there. Further refinements will focus on that area.
 
 +-------------------------------+-------------------------------+-------------------------------+

firedrake/adapt.py

@@ -0,0 +1,146 @@
+"""Adaptive mesh refinement helpers."""
+import numpy as np
+import petsctools
+
+from firedrake.cython import dmcommon
+from firedrake.cython import mgimpl as impl
+from firedrake.utils import IntType
+from firedrake.function import Function
+from firedrake.functionspace import FunctionSpace
+from firedrake.mesh import Mesh, DISTRIBUTION_PARAMETERS_NOOP
+from firedrake.netgen import _transfer_high_order_coordinates
+
+
+# PETSc's DMAdaptFlag value requesting refinement, for the adapt label.
+DM_ADAPT_REFINE = 1
+
+# Label holding, on every cell, the number of the cell of the original mesh it
+# descends from. The refinement transform propagates labels from a cell to its
+# children, so this stays relative to the original mesh however many times we
+# refine.
+PARENT_LABEL = "_adaptive_dmplex_parent"
+
+ADAPT_LABEL = "_adaptive_dmplex_adapt"
+
+
+def _adapt_marked_cells(mesh, cell_marker):
+    """Refine the cells of ``mesh`` marked by ``cell_marker`` and return the refined DMPlex."""
+    dm = mesh.topology_dm
+    ncoarse = mesh.cell_set.size
+
+    dm.createLabel(ADAPT_LABEL)
+    adapt_label = dm.getLabel(ADAPT_LABEL)
+    adapt_indicator = np.zeros(cell_marker.dat.data_ro_with_halos.shape, dtype=IntType)
+    adapt_indicator[:ncoarse] = cell_marker.dat.data_ro.real > 0
+    dmcommon.mark_points_with_function_array(
+        dm, cell_marker.function_space().dm.getSection(), 0,
+        adapt_indicator, adapt_label, DM_ADAPT_REFINE,
+    )
+
+    parameters = {"dm_plex_transform_type": "refine_sbr"}
+    try:
+        # options_prefix="" is essential
+        with petsctools.inserted_options(parameters=parameters, options_prefix=""):
+            new_dm = dm.adaptLabel(ADAPT_LABEL)
+    finally:
+        dm.removeLabel(ADAPT_LABEL)
+
+    # The transform propagates every label, including the temporary adapt
+    # label and the coarse mesh's stale pyop2_core/owned/ghost point
+    # classification. Mesh() skips recomputing that classification if it's
+    # already present, so it must be dropped here to force a fresh one for
+    # the new mesh's own point count and distribution.
+    for label in ("pyop2_core", "pyop2_owned", "pyop2_ghost", ADAPT_LABEL):
+        if new_dm.hasLabel(label):
+            new_dm.removeLabel(label)
+
+    return new_dm
+
+
+def _copy_adaptive_refinement_metadata(source_mesh, target_mesh):
+    """Copy mesh-construction metadata from a mesh onto its adaptively-derived successor."""
+    target_mesh._distribution_parameters = dict(source_mesh._distribution_parameters)
+    target_mesh._did_reordering = source_mesh._did_reordering
+    target_mesh._tolerance = source_mesh.tolerance
+    if hasattr(source_mesh, "netgen_mesh") and not hasattr(target_mesh, "netgen_mesh"):
+        target_mesh.netgen_mesh = source_mesh.netgen_mesh
+    if hasattr(source_mesh, "netgen_flags") and not hasattr(target_mesh, "netgen_flags"):
+        target_mesh.netgen_flags = source_mesh.netgen_flags
+
+
+def refine_marked_elements(mesh, cell_marker):
+    """Adaptively refine a mesh using a DG0 marking function.
+
+    Positive integer marker values request repeated refinement of the
+    corresponding cells. Curved Netgen meshes are re-curved to the
+    original coordinate degree after refinement.
+
+    Parameters
+    ----------
+    mesh
+        The mesh to refine.
+    cell_marker
+        A DG0 `~firedrake.function.Function` on ``mesh``: cells with a
+        positive value ``n`` are refined ``n`` times.
+
+    Returns
+    -------
+    MeshGeometry
+        The adaptively refined mesh, with ``adaptive_parent`` set to
+        ``mesh`` and ``adaptive_cell_maps`` set to the
+        ``(coarse_to_fine, fine_to_coarse)`` cell maps relative to it.
+
+    """
+    with cell_marker.dat.vec_ro as v:
+        _, num_refinements = v.max()
+    # Always run at least one adaptation pass, even when no cell is marked,
+    # so that a fresh mesh (with its own cell maps) is produced uniformly.
+    num_refinements = max(int(np.rint(num_refinements)), 1)
+
+    ncoarse = mesh.cell_set.size
+    coarse_dm = mesh.topology_dm
+    if coarse_dm.hasLabel(PARENT_LABEL):
+        coarse_dm.removeLabel(PARENT_LABEL)
+    coarse_dm.createLabel(PARENT_LABEL)
+    impl.set_adaptive_parent_label(coarse_dm, mesh._cell_numbering, ncoarse, PARENT_LABEL)
+
+    current_mesh = mesh
+    current_mark = cell_marker
+    try:
+        for ref in range(num_refinements):
+            new_dm = _adapt_marked_cells(current_mesh, current_mark)
+            current_mesh = Mesh(
+                new_dm,
+                dim=mesh.geometric_dimension,
+                reorder=False,
+                distribution_parameters=DISTRIBUTION_PARAMETERS_NOOP,
+                comm=mesh.comm,
+                tolerance=mesh.tolerance,
+            )
+            coarse_to_fine, fine_to_coarse = impl.adaptive_parent_child_cell_maps(
+                new_dm, current_mesh._cell_numbering, ncoarse,
+                current_mesh.cell_set.size, PARENT_LABEL,
+            )
+            if ref < num_refinements - 1:
+                # A cell asking for n refinements stays marked until n rounds
+                # have happened, so its descendants inherit n minus the number
+                # of rounds so far.
+                ancestor = fine_to_coarse[:, 0]
+                refined = ancestor >= 0
+                current_mark = Function(FunctionSpace(current_mesh, "DG", 0))
+                current_mark.dat.data_wo[refined] = \
+                    cell_marker.dat.data_ro[ancestor[refined]] - (ref + 1)
+    finally:
+        coarse_dm.removeLabel(PARENT_LABEL)
+
+    final_mesh = current_mesh
+    if hasattr(mesh, "netgen_mesh"):
+        order = mesh.coordinates.function_space().ufl_element().degree()
+        if order > 1:
+            final_mesh = _transfer_high_order_coordinates(mesh, final_mesh, order)
+
+    final_mesh.topology_dm.removeLabel(PARENT_LABEL)
+    final_mesh.adaptive_parent = mesh
+    final_mesh.adaptive_cell_maps = (coarse_to_fine, fine_to_coarse)
+    _copy_adaptive_refinement_metadata(mesh, final_mesh)
+    return final_mesh

firedrake/cython/mgimpl.pyx

@@ -6,6 +6,7 @@ import numpy as np
 from firedrake.cython import dmcommon
 from firedrake.petsc import PETSc
 from firedrake.utils import IntType
+from pyop2.mpi import MPI
 
 cimport numpy as np
 cimport petsc4py.PETSc as PETSc
@@ -97,6 +98,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 +111,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 +156,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:
@@ -192,6 +201,104 @@ def create_lgmap(PETSc.DM dm):
     return lgmap
 
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def set_adaptive_parent_label(PETSc.DM coarse_dm,
+                              PETSc.Section coarse_cell_numbering,
+                              PetscInt ncoarse,
+                              label_name):
+    """Seed each coarse cell's own Firedrake cell number onto a label.
+
+    Must be called *before* refining ``coarse_dm``. Since the refinement
+    transform propagates labels from a cell to its children, every cell of
+    every subsequent refinement then carries the number of the coarse cell it
+    descends from, which :func:`adaptive_parent_child_cell_maps` reads back.
+
+    :arg coarse_dm: the (pre-refinement) coarse mesh DMPlex.
+    :arg coarse_cell_numbering: the coarse mesh's cell numbering section.
+    :arg ncoarse: the number of coarse cells.
+    :arg label_name: name of a label already created on ``coarse_dm``,
+        to be populated with each owned cell's Firedrake cell number.
+    """
+    cdef:
+        PetscInt cStart, cEnd, c, off
+        DMLabel parent_label = NULL
+
+    label_name = label_name.encode()
+    CHKERR(DMGetLabel(coarse_dm.dm, <const char*>label_name, &parent_label))
+    cStart, cEnd = coarse_dm.getHeightStratum(0)
+    for c in range(cStart, cEnd):
+        CHKERR(PetscSectionGetOffset(coarse_cell_numbering.sec, c, &off))
+        if 0 <= off < ncoarse:
+            CHKERR(DMLabelSetValue(parent_label, c, off))
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def adaptive_parent_child_cell_maps(PETSc.DM fine_dm,
+                                    PETSc.Section fine_cell_numbering,
+                                    PetscInt ncoarse,
+                                    PetscInt nfine,
+                                    label_name):
+    """Build Firedrake-numbered parent/child cell maps from a DMPlex label.
+
+    Must be called on ``fine_dm``, a DMPlex obtained by refining the coarse
+    mesh that :func:`set_adaptive_parent_label` was seeded on, however many
+    times.
+
+    :arg fine_dm: the refined (child) mesh DMPlex.
+    :arg fine_cell_numbering: the fine mesh's cell numbering section.
+    :arg ncoarse: the number of coarse (parent) cells.
+    :arg nfine: the number of fine (child) cells.
+    :arg label_name: name of the label, on ``fine_dm``, propagated from
+        :func:`set_adaptive_parent_label`, mapping each fine cell to
+        its coarse parent's Firedrake cell number.
+    :returns: a ``(coarse_to_fine, fine_to_coarse)`` pair of Firedrake-numbered cell maps.
+    """
+    cdef:
+        PetscInt cStart, cEnd, c, off, parent, max_children
+        DMLabel parent_label = NULL
+        PetscInt[::1] child_counts
+        PetscInt[:, ::1] coarse_to_fine
+        PetscInt[:, ::1] fine_to_coarse
+
+    label_name = label_name.encode()
+    CHKERR(DMGetLabel(fine_dm.dm, <const char*>label_name, &parent_label))
+    fine_to_coarse = np.full((nfine, 1), -1, dtype=IntType)
+    child_counts = np.zeros(ncoarse, dtype=IntType)
+    cStart, cEnd = fine_dm.getHeightStratum(0)
+    for c in range(cStart, cEnd):
+        CHKERR(PetscSectionGetOffset(fine_cell_numbering.sec, c, &off))
+        if not (0 <= off < nfine):
+            continue
+        CHKERR(DMLabelGetValue(parent_label, c, &parent))
+        if 0 <= parent < ncoarse:
+            fine_to_coarse[off, 0] = parent
+            child_counts[parent] += 1
+
+    # coarse_to_fine is rectangular, so every coarse cell's row must be wide
+    # enough for its most prolific sibling. Different coarse cells can be
+    # refined a different number of times, so this varies by process; take
+    # the max across all ranks so the array shape agrees everywhere.
+    max_children = 0
+    for c in range(ncoarse):
+        if child_counts[c] > max_children:
+            max_children = child_counts[c]
+    max_children = fine_dm.comm.tompi4py().allreduce(max_children, op=MPI.MAX)
+    coarse_to_fine = np.full((ncoarse, max_children), -1, dtype=IntType)
+    # Re-walk the fine cells (in Firedrake order this time, via
+    # fine_to_coarse) appending each one to its parent's row. child_counts is
+    # reused as a per-parent write cursor, reset to zero first.
+    child_counts[:] = 0
+    for c in range(nfine):
+        parent = fine_to_coarse[c, 0]
+        if parent >= 0:
+            coarse_to_fine[parent, child_counts[parent]] = c
+            child_counts[parent] += 1
+
+    return np.asarray(coarse_to_fine), np.asarray(fine_to_coarse)
+
+
 # Exposition:
 #
 # These next functions compute maps from coarse mesh cells to fine