RobustTransferManager

Author: pbrubeck

firedrake/__init__.py

@@ -102,7 +102,9 @@ def init_petsc():
     NonNestedHierarchy, SemiCoarsenedExtrudedHierarchy, SubmeshHierarchy,
     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/functionspace.py

@@ -336,7 +336,17 @@ def RestrictedFunctionSpace(function_space, boundary_set=[], name=None):
         An optional name for the function space.
 
     """
-    return impl.WithGeometry(impl.RestrictedFunctionSpace(function_space,
-                                                          boundary_set=boundary_set,
-                                                          name=name),
-                             function_space.mesh())
+    if not isinstance(function_space, (impl.WithGeometry, impl.FiredrakeDualSpace)):
+        raise TypeError(f"Can't make restricted space with {type(function_space).__name__}")
+
+    make_space = type(function_space)
+    mesh = function_space.mesh()
+
+    if function_space.boundary_set:
+        boundary_set = frozenset(boundary_set) | function_space.boundary_set
+        function_space = function_space.function_space
+
+    return make_space(impl.RestrictedFunctionSpace(function_space.topological,
+                                                   boundary_set=boundary_set,
+                                                   name=name),
+                      mesh)

firedrake/mg/__init__.py

@@ -10,3 +10,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/kernels.py

@@ -345,7 +345,7 @@ def restrict_kernel(Vf, Vc):
 
 
 def inject_kernel(Vf, Vc):
-    if Vc.finat_element.is_dg():
+    if Vc.finat_element.is_dg() and Vf.ufl_element() == Vc.ufl_element():
         hierarchy, level = utils.get_level(Vc.mesh())
         if Vf.extruded:
             assert Vc.extruded

firedrake/mg/mesh.py

@@ -11,6 +11,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
 
@@ -29,6 +30,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::
 
@@ -37,14 +40,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):
@@ -79,7 +83,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
@@ -109,6 +114,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
@@ -139,11 +148,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
@@ -192,7 +219,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,