Enable solving multi-domain problems involving codim-0 submeshes (#3478)

Author: ksagiyam

demos/saddle_point_pc/saddle_point_systems.py.rst

@@ -180,7 +180,7 @@ Finally, at each mesh size, we print out the number of cells in the
 mesh and the number of iterations the solver took to converge ::
 
     #
-        print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber())
+        print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber())
 
 The resulting convergence is unimpressive:
 
@@ -282,7 +282,7 @@ applying the action of blocks, so we can use a block matrix format. ::
     for n in range(8):
         solver, w = build_problem(n, parameters, block_matrix=True)
         solver.solve()
-        print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber())
+        print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber())
 
 The resulting convergence is algorithmically good, however, the larger
 problems still take a long time.
@@ -367,7 +367,7 @@ Let's see what happens. ::
     for n in range(8):
         solver, w = build_problem(n, parameters, block_matrix=True)
         solver.solve()
-        print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber())
+        print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber())
 
 This is much better, the problem takes much less time to solve and
 when observing the iteration counts for inverting :math:`S` we can see
@@ -422,7 +422,7 @@ and so we no longer need a flexible Krylov method. ::
     for n in range(8):
         solver, w = build_problem(n, parameters, block_matrix=True)
         solver.solve()
-        print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber())
+        print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber())
 
 This results in the following GMRES iteration counts
 
@@ -487,7 +487,7 @@ variable. We can provide it as an :class:`~.AuxiliaryOperatorPC` via a python pr
     for n in range(8):
         solver, w = build_problem(n, parameters, aP=None, block_matrix=False)
         solver.solve()
-        print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber())
+        print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber())
 
 This actually results in slightly worse convergence than the diagonal
 approximation we used above.
@@ -571,7 +571,7 @@ Let's see what the iteration count looks like now. ::
     for n in range(8):
         solver, w = build_problem(n, parameters, aP=riesz, block_matrix=True)
         solver.solve()
-        print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber())
+        print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber())
 
 ============== ==================
  Mesh elements  GMRES iterations

firedrake/assemble.py

@@ -162,8 +162,6 @@ def hermite_stride(bcnodes):
                 # take intersection of facet nodes, and add it to bcnodes
                 # i, j, k can also be strings.
                 bcnodes1 = []
-                if len(s) > 1 and not isinstance(self._function_space.finat_element, (finat.Lagrange, finat.GaussLobattoLegendre)):
-                    raise TypeError("Currently, edge conditions have only been tested with CG Lagrange elements")
                 for ss in s:
                     # intersection of facets
                     # Edge conditions have only been tested with Lagrange elements.

firedrake/bcs.py

@@ -566,6 +566,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 MeshSequence support.
+        mesh = mesh.unique()
         # Handle extruded mesh
         tmesh = mesh.topology
         if mesh.extruded:
@@ -835,6 +837,8 @@ def get_timestepping_history(self, mesh, name):
     @PETSc.Log.EventDecorator("SaveFunctionSpace")
     def _save_function_space(self, V):
         mesh = V.mesh()
+        # TODO: Add general MeshSequence 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)
@@ -910,10 +914,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 MeshSequence 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 MeshSequence support.
+        mesh = mesh.unique()
         mesh_key = self._generate_mesh_key_from_names(mesh.name,
                                                       mesh.topology._distribution_name,
                                                       mesh.topology._permutation_name)
@@ -1299,6 +1307,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 MeshSequence 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/checkpointing.py

@@ -43,6 +43,7 @@
 
 import firedrake
 from firedrake.petsc import PETSc
+from firedrake.mesh import MeshSequenceGeometry
 
 
 @PETSc.Log.EventDecorator()
@@ -53,8 +54,11 @@ def get_function_space(dm):
     :raises RuntimeError: if no function space was found.
     """
     info = dm.getAttr("__fs_info__")
-    meshref, element, indices, (name, names), boundary_sets = info
-    mesh = meshref()
+    meshref_tuple, element, indices, (name, names), boundary_sets = info
+    if len(meshref_tuple) == 1:
+        mesh = meshref_tuple[0]()
+    else:
+        mesh = MeshSequenceGeometry([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)
@@ -80,8 +84,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:
@@ -92,11 +94,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()
     boundary_sets = tuple(V_.boundary_set for V_ in V)
-    info = (weakref.ref(mesh), element, tuple(reversed(indices)), (V.name, names), boundary_sets)
+    info = (tuple(weakref.ref(m) for m in mesh), element, tuple(reversed(indices)), (V.name, names), boundary_sets)
     dm.setAttr("__fs_info__", info)
 
 
@@ -414,7 +417,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)

firedrake/dmhooks.py

@@ -92,7 +92,7 @@ class EnsembleFunctionSpaceBase:
     - Dual ensemble objects: :class:`EnsembleDualSpace` and :class:`~firedrake.ensemble.ensemble_function.EnsembleCofunction`.
     """
     def __init__(self, local_spaces: Collection, ensemble: Ensemble):
-        meshes = set(V.mesh() for V in local_spaces)
+        meshes = set(V.mesh().unique() for V in local_spaces)
         nlocal_meshes = len(meshes)
         max_local_meshes = ensemble.ensemble_comm.allreduce(nlocal_meshes, MPI.MAX)
         if max_local_meshes > 1:

firedrake/ensemble/ensemble_functionspace.py

@@ -593,16 +593,20 @@ def _at(self, arg, *args, **kwargs):
 
         tolerance = kwargs.get('tolerance', None)
         mesh = self.function_space().mesh()
+        if len(set(mesh)) == 1:
+            mesh_unique = mesh.unique()
+        else:
+            raise NotImplementedError("Not implemented for general mixed meshes")
         if tolerance is None:
-            tolerance = mesh.tolerance
+            tolerance = mesh_unique.tolerance
         else:
-            mesh.tolerance = tolerance
+            mesh_unique.tolerance = tolerance
 
         # Handle f._at(0.3)
         if not arg.shape:
             arg = arg.reshape(-1)
 
-        if mesh.variable_layers:
+        if mesh_unique.variable_layers:
             raise NotImplementedError("Point evaluation not implemented for variable layers")
 
         # Validate geometric dimension
@@ -778,7 +782,7 @@ def evaluate(self, function: Function) -> np.ndarray | Tuple[np.ndarray, ...]:
         if function.function_space().ufl_element().family() == "Real":
             return function.dat.data_ro
 
-        function_mesh = function.function_space().mesh()
+        function_mesh = function.function_space().mesh().unique()
         if function_mesh is not self.mesh:
             raise ValueError("Function mesh must be the same Mesh object as the PointEvaluator mesh.")
         if coord_changed := function_mesh.coordinates.dat.dat_version != self.mesh._saved_coordinate_dat_version:

firedrake/function.py

@@ -4,6 +4,7 @@
 API is functional, rather than object-based, to allow for simple
 backwards-compatibility, argument checking, and dispatch.
 """
+import itertools
 import ufl
 import finat.ufl
 
@@ -253,6 +254,8 @@ def MixedFunctionSpace(spaces, name=None, mesh=None):
         :class:`finat.ufl.mixedelement.MixedElement`, ignored otherwise.
 
     """
+    from firedrake.mesh import MeshSequenceGeometry
+
     if isinstance(spaces, finat.ufl.FiniteElementBase):
         # Build the spaces if we got a mixed element
         assert type(spaces) is finat.ufl.MixedElement and mesh is not None
@@ -267,22 +270,15 @@ def rec(eles):
                     sub_elements.append(ele)
         rec(spaces.sub_elements)
         spaces = [FunctionSpace(mesh, element) for element in sub_elements]
-
-    # Check that function spaces are on the same mesh
-    meshes = [space.mesh() for space in spaces]
-    for i in range(1, len(meshes)):
-        if meshes[i] is not meshes[0]:
-            raise ValueError("All function spaces must be defined on the same mesh!")
-
+    # Flatten MeshSequences.
+    meshes = list(itertools.chain(*[space.mesh() for space in spaces]))
     try:
         cls, = set(type(s) for s in spaces)
     except ValueError:
         # Neither primal nor dual
         # We had not implemented something in between, so let's make it primal
         cls = impl.WithGeometry
 
-    # Select mesh
-    mesh = meshes[0]
     # Get topological spaces
     spaces = tuple(s.topological for s in flatten(spaces))
     # Error checking
@@ -296,10 +292,9 @@ def rec(eles):
         else:
             raise ValueError("Can't make mixed space with %s" % type(space))
 
-    new = impl.MixedFunctionSpace(spaces, name=name)
-    if mesh is not mesh.topology:
-        new = cls.create(new, mesh)
-    return new
+    mixed_mesh_geometry = MeshSequenceGeometry(meshes)
+    new = impl.MixedFunctionSpace(spaces, mixed_mesh_geometry.topology, name=name)
+    return cls.create(new, mixed_mesh_geometry)
 
 
 @PETSc.Log.EventDecorator("CreateFunctionSpace")