Enable solving multi-domain problems involving codim-1 submeshes (#4430)

Co-authored-by: Connor Ward <[email protected]>

Author: ksagiyam

firedrake/assemble.py

@@ -1661,6 +1661,7 @@ def __init__(self, form, local_knl, subdomain_id, all_integer_subdomain_ids, dia
         self._constants = _FormHandler.iter_constants(form, local_knl.kinfo)
         self._active_exterior_facets = _FormHandler.iter_active_exterior_facets(form, local_knl.kinfo)
         self._active_interior_facets = _FormHandler.iter_active_interior_facets(form, local_knl.kinfo)
+        self._active_orientations_cell = _FormHandler.iter_active_orientations_cell(form, local_knl.kinfo)
         self._active_orientations_exterior_facet = _FormHandler.iter_active_orientations_exterior_facet(form, local_knl.kinfo)
         self._active_orientations_interior_facet = _FormHandler.iter_active_orientations_interior_facet(form, local_knl.kinfo)
 
@@ -1683,6 +1684,7 @@ def build(self):
         assert_empty(self._constants)
         assert_empty(self._active_exterior_facets)
         assert_empty(self._active_interior_facets)
+        assert_empty(self._active_orientations_cell)
         assert_empty(self._active_orientations_exterior_facet)
         assert_empty(self._active_orientations_interior_facet)
 
@@ -1880,6 +1882,17 @@ def _as_global_kernel_arg_interior_facet(_, self):
         return op2.DatKernelArg((2,), m._global_kernel_arg)
 
 
+@_as_global_kernel_arg.register(kernel_args.OrientationsCellKernelArg)
+def _(_, self):
+    mesh = next(self._active_orientations_cell)
+    if mesh is self._mesh:
+        return op2.DatKernelArg((1,))
+    else:
+        m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)
+        assert integral_type == "cell"
+        return op2.DatKernelArg((1,), m._global_kernel_arg)
+
+
 @_as_global_kernel_arg.register(kernel_args.OrientationsExteriorFacetKernelArg)
 def _(_, self):
     mesh = next(self._active_orientations_exterior_facet)
@@ -1951,6 +1964,7 @@ def __init__(self, form, bcs, local_knl, subdomain_id,
         self._constants = _FormHandler.iter_constants(form, local_knl.kinfo)
         self._active_exterior_facets = _FormHandler.iter_active_exterior_facets(form, local_knl.kinfo)
         self._active_interior_facets = _FormHandler.iter_active_interior_facets(form, local_knl.kinfo)
+        self._active_orientations_cell = _FormHandler.iter_active_orientations_cell(form, local_knl.kinfo)
         self._active_orientations_exterior_facet = _FormHandler.iter_active_orientations_exterior_facet(form, local_knl.kinfo)
         self._active_orientations_interior_facet = _FormHandler.iter_active_orientations_interior_facet(form, local_knl.kinfo)
 
@@ -2216,6 +2230,17 @@ def _as_parloop_arg_interior_facet(_, self):
     return op2.DatParloopArg(mesh.interior_facets.local_facet_dat, m)
 
 
+@_as_parloop_arg.register(kernel_args.OrientationsCellKernelArg)
+def _(_, self):
+    mesh = next(self._active_orientations_cell)
+    if mesh is self._mesh:
+        m = None
+    else:
+        m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)
+        assert integral_type == "cell"
+    return op2.DatParloopArg(mesh.local_cell_orientation_dat, m)
+
+
 @_as_parloop_arg.register(kernel_args.OrientationsExteriorFacetKernelArg)
 def _(_, self):
     mesh = next(self._active_orientations_exterior_facet)
@@ -2312,6 +2337,14 @@ def iter_active_interior_facets(form, kinfo):
             mesh = all_meshes[i]
             yield mesh
 
+    @staticmethod
+    def iter_active_orientations_cell(form, kinfo):
+        """Yield the form cell orientations referenced in ``kinfo``."""
+        all_meshes = extract_domains(form)
+        for i in kinfo.active_domain_numbers.orientations_cell:
+            mesh = all_meshes[i]
+            yield mesh
+
     @staticmethod
     def iter_active_orientations_exterior_facet(form, kinfo):
         """Yield the form exterior facet orientations referenced in ``kinfo``."""

firedrake/embedding.py

@@ -5,7 +5,7 @@
 
 
 def get_embedding_dg_element(element, value_shape, broken_cg=False):
-    cell = element.cell
+    cell, = set(element.cell.cells)
     family = lambda c: "DG" if c.is_simplex else "DQ"
     if isinstance(cell, ufl.TensorProductCell):
         degree = element.degree()

firedrake/functionspaceimpl.py

@@ -13,6 +13,7 @@
 import ufl
 import finat.ufl
 
+from ufl.cell import CellSequence
 from ufl.duals import is_dual, is_primal
 from pyop2 import op2, mpi
 from pyop2.utils import as_tuple
@@ -52,6 +53,9 @@ def check_element(element, top=True):
     ValueError
         If the element is illegal.
     """
+    if isinstance(element.cell, CellSequence) and \
+       type(element) is not finat.ufl.MixedElement:
+        raise ValueError("MixedElement modifier must be outermost")
     if element.cell.cellname == "hexahedron" and \
        element.family() not in ["Q", "DQ", "Real"]:
         raise NotImplementedError("Currently can only use 'Q', 'DQ', and/or 'Real' elements on hexahedral meshes, not", element.family())

firedrake/interpolation.py

@@ -1198,7 +1198,7 @@ def get_interp_node_map(source_mesh, target_mesh, fs):
         else:
             raise ValueError("Have coefficient with unexpected mesh")
     else:
-        m_ = fs.entity_node_map(target_mesh.topology, "cell", None, None)
+        m_ = fs.entity_node_map(target_mesh.topology, "cell", "everywhere", None)
     return m_