submesh: enable cell submesh

Author: ksagiyam

firedrake/__future__.py

@@ -1,6 +1,6 @@
 from ufl.domain import as_domain, extract_unique_domain
 from ufl.algorithms import extract_arguments
-from firedrake.mesh import VertexOnlyMeshTopology
+from firedrake.mesh import MeshTopology, VertexOnlyMeshTopology
 from firedrake.interpolation import (interpolate as interpolate_old,
                                      Interpolator as InterpolatorOld,
                                      SameMeshInterpolator as SameMeshInterpolatorOld,
@@ -16,13 +16,13 @@ class Interpolator(InterpolatorOld):
     def __new__(cls, expr, V, **kwargs):
         target_mesh = as_domain(V)
         source_mesh = extract_unique_domain(expr) or target_mesh
-        if target_mesh is not source_mesh:
+        if target_mesh is source_mesh or all(isinstance(m.topology, MeshTopology) for m in [target_mesh, source_mesh]) and target_mesh.submesh_ancesters[-1] is source_mesh.submesh_ancesters[-1]:
+            return object.__new__(SameMeshInterpolator)
+        else:
             if isinstance(target_mesh.topology, VertexOnlyMeshTopology):
                 return object.__new__(SameMeshInterpolator)
             else:
                 return object.__new__(CrossMeshInterpolator)
-        else:
-            return object.__new__(SameMeshInterpolator)
 
     interpolate = InterpolatorOld._interpolate_future
 

firedrake/assemble.py

@@ -18,7 +18,7 @@
 from firedrake import (extrusion_utils as eutils, matrix, parameters, solving,
                        tsfc_interface, utils)
 from firedrake.adjoint_utils import annotate_assemble
-from firedrake.ufl_expr import extract_unique_domain
+from firedrake.ufl_expr import extract_domains
 from firedrake.bcs import DirichletBC, EquationBC, EquationBCSplit
 from firedrake.functionspaceimpl import WithGeometry, FunctionSpace, FiredrakeDualSpace
 from firedrake.functionspacedata import entity_dofs_key, entity_permutations_key
@@ -1037,15 +1037,16 @@ def local_kernels(self):
             each possible combination.
 
         """
-        #try:
-        #    topology, = set(d.topology for d in self._form.ufl_domains())
-        #except ValueError:
-        #    raise NotImplementedError("All integration domains must share a mesh topology")
+        try:
+            topology, = set(d.topology.submesh_ancesters[-1] for d in self._form.ufl_domains())
+        except ValueError:
+            raise NotImplementedError("All integration domains must share a mesh topology")
 
-        #for o in itertools.chain(self._form.arguments(), self._form.coefficients()):
-        #    domain = extract_unique_domain(o)
-        #    if domain is not None and domain.topology != topology:
-        #        raise NotImplementedError("Assembly with multiple meshes is not supported")
+        for o in itertools.chain(self._form.arguments(), self._form.coefficients()):
+            domains = extract_domains(o)
+            for domain in domains:
+                if domain is not None and domain.topology.submesh_ancesters[-1] != topology:
+                    raise NotImplementedError("Assembly with multiple meshes is not supported")
 
         if isinstance(self._form, ufl.Form):
             kernels = tsfc_interface.compile_form(
@@ -1337,20 +1338,23 @@ def _make_maps_and_regions(self):
         test, trial = self._form.arguments()
         if self._allocation_integral_types is not None:
             return ExplicitMatrixAssembler._make_maps_and_regions_default(test, trial, self._allocation_integral_types)
-        elif any(local_kernel.indices == (None, None) for local_kernel in self._all_local_kernels):
+        elif any(local_kernel.indices == (None, None) for local_kernel, _ in self._all_local_kernels):
             # Handle special cases: slate or split=False
-            assert all(local_kernel.indices == (None, None) for local_kernel in self._all_local_kernels)
+            assert all(local_kernel.indices == (None, None) for local_kernel, _ in self._all_local_kernels)
             allocation_integral_types = set(local_kernel.kinfo.integral_type
-                                            for local_kernel in self._all_local_kernels)
+                                            for local_kernel, _ in self._all_local_kernels)
             return ExplicitMatrixAssembler._make_maps_and_regions_default(test, trial, allocation_integral_types)
         else:
             maps_and_regions = defaultdict(lambda: defaultdict(set))
-             for local_kernel in self._all_local_kernels:
+            all_meshes = extract_domains(self._form)
+            for local_kernel, subdomain_id in self._all_local_kernels:
                 i, j = local_kernel.indices
+                mesh = all_meshes[local_kernel.kinfo.domain_number]  # integration domain
                 # Make Sparsity independent of _iterset, which can be a Subset, for better reusability.
                 integral_type = local_kernel.kinfo.integral_type
-                rmap_ = test.function_space().topological[i].entity_node_map(integral_type)
-                cmap_ = trial.function_space().topological[j].entity_node_map(integral_type)
+                all_subdomain_ids = self.all_integer_subdomain_ids[local_kernel.kinfo.domain_number]
+                rmap_ = test.function_space().topological[i].entity_node_map(mesh.topology, integral_type, subdomain_id, all_subdomain_ids)
+                cmap_ = trial.function_space().topological[j].entity_node_map(mesh.topology, integral_type, subdomain_id, all_subdomain_ids)
                 region = ExplicitMatrixAssembler._integral_type_region_map[integral_type]
                 maps_and_regions[(i, j)][(rmap_, cmap_)].add(region)
             return {block_indices: [map_pair + (tuple(region_set), ) for map_pair, region_set in map_pair_to_region_set.items()]
@@ -1366,8 +1370,14 @@ def _make_maps_and_regions_default(test, trial, allocation_integral_types):
         # Use outer product of component maps.
         for integral_type in allocation_integral_types:
             region = ExplicitMatrixAssembler._integral_type_region_map[integral_type]
-            for i, rmap_ in enumerate(test.function_space().topological.entity_node_map(integral_type)):
-                for j, cmap_ in enumerate(trial.function_space().topological.entity_node_map(integral_type)):
+            #for i, rmap_ in enumerate(test.function_space().topological.entity_node_map(mesh.topology, integral_type, None, None)):
+            #    for j, cmap_ in enumerate(trial.function_space().topological.entity_node_map(mesh.topology, integral_type, None, None)):
+            #        maps_and_regions[(i, j)][(rmap_, cmap_)].add(region)
+            for i, Vrow in enumerate(test.function_space()):
+                for j, Vcol in enumerate(trial.function_space()):
+                    mesh = Vrow.mesh()
+                    rmap_ = Vrow.topological.entity_node_map(mesh.topology, integral_type, None, None)
+                    cmap_ = Vcol.topological.entity_node_map(mesh.topology, integral_type, None, None)
                     maps_and_regions[(i, j)][(rmap_, cmap_)].add(region)
         return {block_indices: [map_pair + (tuple(region_set), ) for map_pair, region_set in map_pair_to_region_set.items()]
                 for block_indices, map_pair_to_region_set in maps_and_regions.items()}
@@ -1389,7 +1399,7 @@ def _all_local_kernels(self):
         When constructing sparsity, we use all parloop_builders
         that are to be used in the actual assembly.
         """
-        all_local_kernels = tuple(local_kernel for local_kernel, _ in self.local_kernels)
+        all_local_kernels = self.local_kernels
         for bc in self._bcs:
             if isinstance(bc, EquationBCSplit):
                 _assembler = type(self)(bc.f, bcs=bc.bcs, form_compiler_parameters=self._form_compiler_params, needs_zeroing=False)
@@ -1559,7 +1569,7 @@ def __init__(self, form, local_knl, subdomain_id, all_integer_subdomain_ids, dia
         self._form = form
         self._indices, self._kinfo = local_knl
         self._subdomain_id = subdomain_id
-        self._all_integer_subdomain_ids = all_integer_subdomain_ids.get(self._kinfo.integral_type, None)
+        self._all_integer_subdomain_ids = all_integer_subdomain_ids
         self._diagonal = diagonal
         self._unroll = unroll
 
@@ -1626,7 +1636,7 @@ def _needs_subset(self):
         if self._subdomain_id == "everywhere":
             return False
         elif self._subdomain_id == "otherwise":
-            return self._all_integer_subdomain_ids is not None
+            return self._all_integer_subdomain_ids.get(self._kinfo.integral_type, None) is not None
         else:
             return True
 
@@ -1646,7 +1656,7 @@ def _get_dim(self, finat_element):
 
     def _make_dat_global_kernel_arg(self, V, index=None):
         finat_element = create_element(V.ufl_element())
-        map_arg = V.topological.entity_node_map(self._integral_type)._global_kernel_arg
+        map_arg = V.topological.entity_node_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)._global_kernel_arg
         if isinstance(finat_element, finat.EnrichedElement) and finat_element.is_mixed:
             assert index is None
             subargs = tuple(self._make_dat_global_kernel_arg(Vsub, index=index)
@@ -1664,7 +1674,7 @@ def _make_mat_global_kernel_arg(self, Vrow, Vcol):
             shape = len(relem.elements), len(celem.elements)
             return op2.MixedMatKernelArg(subargs, shape)
         else:
-            rmap_arg, cmap_arg = (V.topological.entity_node_map(self._integral_type)._global_kernel_arg for V in [Vrow, Vcol])
+            rmap_arg, cmap_arg = (V.topological.entity_node_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)._global_kernel_arg for V in [Vrow, Vcol])
             # PyOP2 matrix objects have scalar dims so we flatten them here
             rdim = numpy.prod(self._get_dim(relem), dtype=int)
             cdim = numpy.prod(self._get_dim(celem), dtype=int)
@@ -1765,14 +1775,24 @@ def _as_global_kernel_arg_constant(_, self):
 
 @_as_global_kernel_arg.register(kernel_args.ExteriorFacetKernelArg)
 def _as_global_kernel_arg_exterior_facet(_, self):
-    _ = next(self._active_exterior_facets)
-    return op2.DatKernelArg((1,))
+    mesh, _ = next(self._active_exterior_facets)
+    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 == "exterior_facet"
+        return op2.DatKernelArg((1,), m._global_kernel_arg)
 
 
 @_as_global_kernel_arg.register(kernel_args.InteriorFacetKernelArg)
 def _as_global_kernel_arg_interior_facet(_, self):
-    _ = next(self._active_interior_facets)
-    return op2.DatKernelArg((2,))
+    mesh, _ = next(self._active_interior_facets)
+    if mesh is self._mesh:
+        return op2.DatKernelArg((2,))
+    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 == "interior_facet"
+        return op2.DatKernelArg((2,), m._global_kernel_arg)
 
 
 @_as_global_kernel_arg.register(CellFacetKernelArg)
@@ -1978,7 +1998,7 @@ def _iterset(self):
     def _get_map(self, V):
         """Return the appropriate PyOP2 map for a given function space."""
         assert isinstance(V, (WithGeometry, FiredrakeDualSpace, FunctionSpace))
-        return V.entity_node_map(self._integral_type)
+        return V.topological.entity_node_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)
 
     def _as_parloop_arg(self, tsfc_arg):
         """Return a :class:`op2.ParloopArg` corresponding to the provided
@@ -2066,7 +2086,7 @@ def _as_parloop_arg_exterior_facet(_, self):
         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 == "exterior_facets"
+        assert integral_type == "exterior_facet"
     return op2.DatParloopArg(local_facet_dat, m)
 
 
@@ -2077,7 +2097,7 @@ def _as_parloop_arg_interior_facet(_, self):
         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 == "interior_facets"
+        assert integral_type == "interior_facet"
     return op2.DatParloopArg(local_facet_dat, m)