introduce MixedMesh

Author: ksagiyam

firedrake/assemble.py

@@ -1020,7 +1020,7 @@ def parloop_builders(self):
                     self._bcs,
                     local_kernel,
                     subdomain_id,
-                    self.all_integer_subdomain_ids,
+                    self.all_integer_subdomain_ids[local_kernel.kinfo.domain_number],
                     diagonal=self.diagonal,
                 )
             )
@@ -1037,15 +1037,15 @@ 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 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()):
+        #    domain = extract_unique_domain(o)
+        #    if domain is not None and domain.topology != topology:
+        #        raise NotImplementedError("Assembly with multiple meshes is not supported")
 
         if isinstance(self._form, ufl.Form):
             kernels = tsfc_interface.compile_form(
@@ -1345,7 +1345,7 @@ def _make_maps_and_regions(self):
             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:
+             for local_kernel in self._all_local_kernels:
                 i, j = local_kernel.indices
                 # Make Sparsity independent of _iterset, which can be a Subset, for better reusability.
                 integral_type = local_kernel.kinfo.integral_type
@@ -1565,8 +1565,13 @@ def __init__(self, form, local_knl, subdomain_id, all_integer_subdomain_ids, dia
         self._diagonal = diagonal
         self._unroll = unroll
 
+        self._active_coordinates = _FormHandler.iter_active_coordinates(form, local_knl.kinfo)
+        self._active_cell_orientations = _FormHandler.iter_active_cell_orientations(form, local_knl.kinfo)
+        self._active_cell_sizes = _FormHandler.iter_active_cell_sizes(form, local_knl.kinfo)
         self._active_coefficients = _FormHandler.iter_active_coefficients(form, local_knl.kinfo)
         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._map_arg_cache = {}
         # Cache for holding :class:`op2.MapKernelArg` instances.
@@ -1580,8 +1585,13 @@ def build(self):
                        for arg in self._kinfo.arguments]
 
         # we should use up all of the coefficients and constants
+        assert_empty(self._active_coordinates)
+        assert_empty(self._active_cell_orientations)
+        assert_empty(self._active_cell_sizes)
         assert_empty(self._active_coefficients)
         assert_empty(self._constants)
+        assert_empty(self._active_exterior_facets)
+        assert_empty(self._active_interior_facets)
 
         iteration_regions = {"exterior_facet_top": op2.ON_TOP,
                              "exterior_facet_bottom": op2.ON_BOTTOM,
@@ -1606,7 +1616,8 @@ def _integral_type(self):
 
     @cached_property
     def _mesh(self):
-        return self._form.ufl_domains()[self._kinfo.domain_number]
+        all_meshes = extract_domains(self._form)
+        return all_meshes[self._kinfo.domain_number]
 
     @cached_property
     def _needs_subset(self):
@@ -1711,7 +1722,22 @@ def _as_global_kernel_arg_output(_, self):
 
 @_as_global_kernel_arg.register(kernel_args.CoordinatesKernelArg)
 def _as_global_kernel_arg_coordinates(_, self):
-    V = self._mesh.coordinates.function_space()
+    coord = next(self._active_coordinates)
+    V = coord.function_space()
+    return self._make_dat_global_kernel_arg(V)
+
+
+@_as_global_kernel_arg.register(kernel_args.CellOrientationsKernelArg)
+def _as_global_kernel_arg_cell_orientations(_, self):
+    c = next(self._active_cell_orientations)
+    V = c.function_space()
+    return self._make_dat_global_kernel_arg(V)
+
+
+@_as_global_kernel_arg.register(kernel_args.CellSizesKernelArg)
+def _as_global_kernel_arg_cell_sizes(_, self):
+    c = next(self._active_cell_sizes)
+    V = c.function_space()
     return self._make_dat_global_kernel_arg(V)
 
 
@@ -1739,19 +1765,15 @@ def _as_global_kernel_arg_constant(_, self):
     return op2.GlobalKernelArg((value_size,))
 
 
-@_as_global_kernel_arg.register(kernel_args.CellSizesKernelArg)
-def _as_global_kernel_arg_cell_sizes(_, self):
-    V = self._mesh.cell_sizes.function_space()
-    return self._make_dat_global_kernel_arg(V)
-
-
 @_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,))
 
 
 @_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,))
 
 
@@ -1764,12 +1786,6 @@ def _as_global_kernel_arg_cell_facet(_, self):
     return op2.DatKernelArg((num_facets, 2))
 
 
-@_as_global_kernel_arg.register(kernel_args.CellOrientationsKernelArg)
-def _as_global_kernel_arg_cell_orientations(_, self):
-    V = self._mesh.cell_orientations().function_space()
-    return self._make_dat_global_kernel_arg(V)
-
-
 @_as_global_kernel_arg.register(LayerCountKernelArg)
 def _as_global_kernel_arg_layer_count(_, self):
     return op2.GlobalKernelArg((1,))
@@ -1803,8 +1819,13 @@ def __init__(self, form, bcs, local_knl, subdomain_id,
         self._diagonal = diagonal
         self._bcs = bcs
 
+        self._active_coordinates = _FormHandler.iter_active_coordinates(form, local_knl.kinfo)
+        self._active_cell_orientations = _FormHandler.iter_active_cell_orientations(form, local_knl.kinfo)
+        self._active_cell_sizes = _FormHandler.iter_active_cell_sizes(form, local_knl.kinfo)
         self._active_coefficients = _FormHandler.iter_active_coefficients(form, local_knl.kinfo)
         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)
 
     def build(self, tensor):
         """Construct the parloop.
@@ -1931,7 +1952,8 @@ def _indexed_tensor(self):
 
     @cached_property
     def _mesh(self):
-        return self._form.ufl_domains()[self._kinfo.domain_number]
+        all_meshes = extract_domains(self._form)
+        return all_meshes[self._kinfo.domain_number]
 
     @cached_property
     def _iterset(self):
@@ -2004,8 +2026,22 @@ def _as_parloop_arg_output(_, self):
 
 @_as_parloop_arg.register(kernel_args.CoordinatesKernelArg)
 def _as_parloop_arg_coordinates(_, self):
-    func = self._mesh.coordinates
-    map_ = self._get_map(func.function_space())
+    func = next(self._active_coordinates)
+    map_ = self._get_map(func.function_space())  #Compose!
+    return op2.DatParloopArg(func.dat, map_)
+
+
+@_as_parloop_arg.register(kernel_args.CellOrientationsKernelArg)
+def _as_parloop_arg_cell_orientations(_, self):
+    func = next(self._active_cell_orientations)
+    map_ = self._get_map(func.function_space())  #Compose!
+    return op2.DatParloopArg(func.dat, map_)
+
+
+@_as_parloop_arg.register(kernel_args.CellSizesKernelArg)
+def _as_parloop_arg_cell_sizes(_, self):
+    func = next(self._active_cell_sizes)
+    map_ = self._get_map(func.function_space())  #Compose!
     return op2.DatParloopArg(func.dat, map_)
 
 
@@ -2025,28 +2061,26 @@ def _as_parloop_arg_constant(arg, self):
     return op2.GlobalParloopArg(const.dat)
 
 
-@_as_parloop_arg.register(kernel_args.CellOrientationsKernelArg)
-def _as_parloop_arg_cell_orientations(_, self):
-    func = self._mesh.cell_orientations()
-    m = self._get_map(func.function_space())
-    return op2.DatParloopArg(func.dat, m)
-
-
-@_as_parloop_arg.register(kernel_args.CellSizesKernelArg)
-def _as_parloop_arg_cell_sizes(_, self):
-    func = self._mesh.cell_sizes
-    m = self._get_map(func.function_space())
-    return op2.DatParloopArg(func.dat, m)
-
-
 @_as_parloop_arg.register(kernel_args.ExteriorFacetKernelArg)
 def _as_parloop_arg_exterior_facet(_, self):
-    return op2.DatParloopArg(self._mesh.exterior_facets.local_facet_dat)
+    mesh, local_facet_dat = next(self._active_exterior_facets)
+    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 == "exterior_facets"
+    return op2.DatParloopArg(local_facet_dat, m)
 
 
 @_as_parloop_arg.register(kernel_args.InteriorFacetKernelArg)
 def _as_parloop_arg_interior_facet(_, self):
-    return op2.DatParloopArg(self._mesh.interior_facets.local_facet_dat)
+    mesh, local_facet_dat = next(self._active_interior_facets)
+    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 == "interior_facets"
+    return op2.DatParloopArg(local_facet_dat, m)
 
 
 @_as_parloop_arg.register(CellFacetKernelArg)
@@ -2068,6 +2102,27 @@ def _as_parloop_arg_layer_count(_, self):
 class _FormHandler:
     """Utility class for inspecting forms and local kernels."""
 
+    @staticmethod
+    def iter_active_coordinates(form, kinfo):
+        """Yield the form coordinates referenced in ``kinfo``."""
+        all_meshes = extract_domains(form)
+        for i in kinfo.active_domain_numbers.coordinates:
+            yield all_meshes[i].coordinates
+
+    @staticmethod
+    def iter_active_cell_orientations(form, kinfo):
+        """Yield the form cell orientations referenced in ``kinfo``."""
+        all_meshes = extract_domains(form)
+        for i in kinfo.active_domain_numbers.cell_orientations:
+            yield all_meshes[i].cell_orientations()
+
+    @staticmethod
+    def iter_active_cell_sizes(form, kinfo):
+        """Yield the form cell sizes referenced in ``kinfo``."""
+        all_meshes = extract_domains(form)
+        for i in kinfo.active_domain_numbers.cell_sizes:
+            yield all_meshes[i].cell_sizes
+
     @staticmethod
     def iter_active_coefficients(form, kinfo):
         """Yield the form coefficients referenced in ``kinfo``."""
@@ -2087,6 +2142,22 @@ def iter_constants(form, kinfo):
             for constant_index in kinfo.constant_numbers:
                 yield all_constants[constant_index]
 
+    @staticmethod
+    def iter_active_exterior_facets(form, kinfo):
+        """Yield the form exterior facets referenced in ``kinfo``."""
+        all_meshes = extract_domains(form)
+        for i in kinfo.active_domain_numbers.exterior_facets:
+            mesh = all_meshes[i]
+            yield mesh, mesh.exterior_facets.local_facet_dat
+
+    @staticmethod
+    def iter_active_interior_facets(form, kinfo):
+        """Yield the form interior facets referenced in ``kinfo``."""
+        all_meshes = extract_domains(form)
+        for i in kinfo.active_domain_numbers.interior_facets:
+            mesh = all_meshes[i]
+            yield mesh, mesh.interior_facets.local_facet_dat
+
     @staticmethod
     def index_function_spaces(form, indices):
         """Return the function spaces of the form's arguments, indexed

firedrake/checkpointing.py

@@ -565,6 +565,8 @@ def save_mesh(self, mesh, distribution_name=None, permutation_name=None):
         :kwarg distribution_name: the name under which distribution is saved; if `None`, auto-generated name will be used.
         :kwarg permutation_name: the name under which permutation is saved; if `None`, auto-generated name will be used.
         """
+        # TODO: Add general MixedMesh support.
+        mesh = mesh.unique()
         mesh.init()
         # Handle extruded mesh
         tmesh = mesh.topology
@@ -827,6 +829,8 @@ def get_timestepping_history(self, mesh, name):
     @PETSc.Log.EventDecorator("SaveFunctionSpace")
     def _save_function_space(self, V):
         mesh = V.mesh()
+        # TODO: Add general MixedMesh support.
+        mesh = mesh.unique()
         if isinstance(V.topological, impl.MixedFunctionSpace):
             V_name = self._generate_function_space_name(V)
             base_path = self._path_to_mixed_function_space(mesh.name, V_name)
@@ -902,10 +906,12 @@ def save_function(self, f, idx=None, name=None, timestepping_info={}):
             each index.
         """
         V = f.function_space()
-        mesh = V.mesh()
         if name:
             g = Function(V, val=f.dat, name=name)
             return self.save_function(g, idx=idx, timestepping_info=timestepping_info)
+        mesh = V.mesh()
+        # TODO: Add general MixedMesh support.
+        mesh = mesh.unique()
         # -- Save function space --
         self._save_function_space(V)
         # -- Save function --
@@ -1224,6 +1230,8 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters):
 
     @PETSc.Log.EventDecorator("LoadFunctionSpace")
     def _load_function_space(self, mesh, name):
+        # TODO: Add general MixedMesh support.
+        mesh = mesh.unique()
         mesh.init()
         mesh_key = self._generate_mesh_key_from_names(mesh.name,
                                                       mesh.topology._distribution_name,
@@ -1301,6 +1309,8 @@ def load_function(self, mesh, name, idx=None):
             be loaded with idx only when it was saved with idx.
         :returns: the loaded :class:`~.Function`.
         """
+        # TODO: Add general MixedMesh support.
+        mesh = mesh.unique()
         tmesh = mesh.topology
         if name in self._get_mixed_function_name_mixed_function_space_name_map(mesh.name):
             V_name = self._get_mixed_function_name_mixed_function_space_name_map(mesh.name)[name]

firedrake/dmhooks.py

@@ -43,7 +43,7 @@
 
 import firedrake
 from firedrake.petsc import PETSc
-
+from firedrake.mesh import MixedMeshGeometry
 
 @PETSc.Log.EventDecorator()
 def get_function_space(dm):
@@ -53,8 +53,11 @@ def get_function_space(dm):
     :raises RuntimeError: if no function space was found.
     """
     info = dm.getAttr("__fs_info__")
-    meshref, element, indices, (name, names) = info
-    mesh = meshref()
+    meshref_tuple, element, indices, (name, names) = info
+    if len(meshref_tuple) == 1:
+        mesh = meshref_tuple[0]()
+    else:
+        mesh = MixedMeshGeometry(*(meshref() for meshref in meshref_tuple))
     if mesh is None:
         raise RuntimeError("Somehow your mesh was collected, this should never happen")
     V = firedrake.FunctionSpace(mesh, element, name=name)
@@ -78,8 +81,6 @@ def set_function_space(dm, V):
        This stores the information necessary to make a function space given a DM.
 
     """
-    mesh = V.mesh()
-
     indices = []
     names = []
     while V.parent is not None:
@@ -90,11 +91,12 @@ def set_function_space(dm, V):
             assert V.index is None
             indices.append(V.component)
         V = V.parent
+    mesh = V.mesh()
     if len(V) > 1:
         names = tuple(V_.name for V_ in V)
     element = V.ufl_element()
 
-    info = (weakref.ref(mesh), element, tuple(reversed(indices)), (V.name, names))
+    info = (tuple(weakref.ref(m) for m in mesh), element, tuple(reversed(indices)), (V.name, names))
     dm.setAttr("__fs_info__", info)
 
 
@@ -412,7 +414,9 @@ def coarsen(dm, comm):
     """
     from firedrake.mg.utils import get_level
     V = get_function_space(dm)
-    hierarchy, level = get_level(V.mesh())
+    # TODO: Think harder.
+    m, = set(m_ for m_ in V.mesh())
+    hierarchy, level = get_level(m)
     if level < 1:
         raise RuntimeError("Cannot coarsen coarsest DM")
     coarsen = get_ctx_coarsener(dm)