Robust multigrid transfer

firedrake/__init__.py

@@ -102,7 +102,9 @@ def init_petsc():
     NonNestedHierarchy, SemiCoarsenedExtrudedHierarchy,
     prolong, restrict, inject, TransferManager,
     OpenCascadeMeshHierarchy, AdaptiveMeshHierarchy,
-    AdaptiveTransferManager
+    AdaptiveTransferManager,
+    CoarsePatchTransferManager,
+    FinePatchTransferManager,
 )
 from firedrake.norms import errornorm, norm  # noqa: F401
 from firedrake.nullspace import VectorSpaceBasis, MixedVectorSpaceBasis  # noqa: F401

firedrake/mg/__init__.py

@@ -9,3 +9,6 @@
 from firedrake.mg.opencascade_mh import OpenCascadeMeshHierarchy  # noqa F401
 from firedrake.mg.adaptive_hierarchy import AdaptiveMeshHierarchy  # noqa F401
 from firedrake.mg.adaptive_transfer_manager import AdaptiveTransferManager  # noqa: F401
+from firedrake.mg.robust_transfer_manager import (  # noqa: F401
+    CoarsePatchTransferManager, FinePatchTransferManager,
+)

firedrake/mg/mesh.py

@@ -10,6 +10,7 @@
 from functools import cached_property
 
 from firedrake import utils
+from firedrake.petsc import PETSc
 from firedrake.cython import mgimpl as impl
 from .utils import set_level
 
@@ -28,6 +29,8 @@ class HierarchyBase(object):
     :arg refinements_per_level: number of mesh refinements each
        multigrid level should "see".
     :arg nested: Is this mesh hierarchy nested?
+    :arg coarse_facet_label: Optional subdomain ID to label the coarse facets on
+       each level of the hierarchy.
 
     .. note::
 
@@ -36,14 +39,15 @@ class HierarchyBase(object):
        :func:`ExtrudedMeshHierarchy`, or :func:`NonNestedHierarchy`.
     """
     def __init__(self, meshes, coarse_to_fine_cells, fine_to_coarse_cells,
-                 refinements_per_level=1, nested=False):
+                 refinements_per_level=1, nested=False, coarse_facet_label=None):
         petsctools.cite("Mitchell2016")
         self._meshes = tuple(meshes)
         self.meshes = tuple(meshes[::refinements_per_level])
         self.coarse_to_fine_cells = coarse_to_fine_cells
         self.fine_to_coarse_cells = fine_to_coarse_cells
         self.refinements_per_level = refinements_per_level
         self.nested = nested
+        self._coarse_facet_label = coarse_facet_label
         for level, m in enumerate(meshes):
             set_level(m, self, Fraction(level, refinements_per_level))
         for level, m in enumerate(self):
@@ -78,7 +82,8 @@ def MeshHierarchy(mesh, refinement_levels,
                   netgen_flags=False,
                   reorder=None,
                   distribution_parameters=None, callbacks=None,
-                  mesh_builder=firedrake.Mesh):
+                  mesh_builder=firedrake.Mesh,
+                  coarse_facet_label=None):
     """Build a hierarchy of meshes by uniformly refining a coarse mesh.
 
     Parameters
@@ -108,6 +113,10 @@ def MeshHierarchy(mesh, refinement_levels,
         callback receives the refined DM (and the current level).
     mesh_builder
         Function to turn a DM into a ``Mesh``. Used by pyadjoint.
+    coarse_facet_label : int | None
+        Optional subdomain ID to label the coarse facets on each
+        level of the hierarchy.
+
     Returns
     -------
     A :py:class:`HierarchyBase` object representing the
@@ -138,11 +147,29 @@ def MeshHierarchy(mesh, refinement_levels,
     else:
         before = after = lambda dm, i: None
     for i in range(refinement_levels*refinements_per_level):
+        if coarse_facet_label is not None:
+            # Create a temporary label on all the facets of the coarse dm
+            # to label every coarse facet on the fine dm
+            fstart, fend = cdm.getHeightStratum(1)
+            iset = PETSc.IS().createStride(fend-fstart, first=fstart, comm=cdm.comm)
+            cdm.createLabel("temp_label")
+            label = cdm.getLabel("temp_label")
+            label.setStratumIS(1, iset)
+
         if i % refinements_per_level == 0:
             before(cdm, i)
         rdm = cdm.refine()
         if i % refinements_per_level == 0:
             after(rdm, i)
+
+        if coarse_facet_label is not None:
+            # Move coarse_facet_label into FACE_SETS_LABEL
+            iset = rdm.getLabel("temp_label").getStratumIS(1)
+            label = rdm.getLabel(dmcommon.FACE_SETS_LABEL)
+            label.setStratumIS(coarse_facet_label, iset)
+            rdm.removeLabel("temp_label")
+            cdm.removeLabel("temp_label")
+
         dms.append(rdm)
         cdm = rdm
         # Fix up coords if refining embedded circle or sphere
@@ -191,7 +218,8 @@ def MeshHierarchy(mesh, refinement_levels,
     fine_to_coarse_cells = dict((Fraction(i, refinements_per_level), f2c)
                                 for i, f2c in enumerate(fine_to_coarse_cells))
     return HierarchyBase(meshes, coarse_to_fine_cells, fine_to_coarse_cells,
-                         refinements_per_level, nested=True)
+                         refinements_per_level, nested=True,
+                         coarse_facet_label=coarse_facet_label)
 
 
 def ExtrudedMeshHierarchy(base_hierarchy, height, base_layer=-1, refinement_ratio=2, layers=None,