Tidy up Submesh (#4999)

* Tidy up Submesh

* Make everything a kwarg

Author: pbrubeck

firedrake/mesh.py

@@ -4797,22 +4797,27 @@ def SubDomainData(geometric_expr):
     return op2.Subset(m.cell_set, indices)
 
 
-def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo=False, reorder=None, comm=None):
+def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ignore_halo=False, reorder=None, comm=None):
     """Construct a submesh from a given mesh.
 
     Parameters
     ----------
     mesh : MeshGeometry
         Parent mesh (`MeshGeometry`).
-    subdim : int
+    subdim : int | None
         Topological dimension of the submesh.
+        Defaults to ``mesh.topological_dimension``.
     subdomain_id : int | None
         Subdomain ID representing the submesh.
-        `None` defines the submesh owned by the sub-communicator.
+        If `None` the submesh will cover the entire domain.
+        This is useful to obtain a codim-1 submesh over all facets or
+        a submesh over a different communicator.
     label_name : str | None
         Name of the label to search ``subdomain_id`` in.
+        Defaults to ``'Cell Sets'`` or ``'Face Sets'`` depending on ``subdim``.
     name : str |  None
         Name of the submesh.
+        Defaults to ``mesh.name + "_submesh"``·
     ignore_halo : bool
         Whether to exclude the halo from the submesh.
     reorder : bool | None
@@ -4855,24 +4860,24 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo=
         raise NotImplementedError("Can not create a submesh of an ``ExtrudedMesh``")
     elif isinstance(mesh.topology, VertexOnlyMeshTopology):
         raise NotImplementedError("Can not create a submesh of a ``VertexOnlyMesh``")
+    if subdim is None:
+        subdim = mesh.topological_dimension
     plex = mesh.topology_dm
     dim = plex.getDimension()
-    if subdim not in [dim, dim - 1]:
+    if subdim not in {dim, dim - 1}:
         raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})")
-    if label_name is None:
+    if subdomain_id is None:
+        if label_name is not None:
+            raise ValueError("subdomain_id=None requires label_name=None.")
+        # Select all entities
+        label_name = "depth"
+        subdomain_id = subdim
+    elif label_name is None:
         if subdim == dim:
             label_name = dmcommon.CELL_SETS_LABEL
         elif subdim == dim - 1:
             label_name = dmcommon.FACE_SETS_LABEL
-    if subdomain_id is None:
-        # Filter the plex with PETSc's default label (cells owned by comm)
-        if label_name != dmcommon.CELL_SETS_LABEL:
-            raise ValueError("subdomain_id == None requires label_name == CELL_SETS_LABEL.")
-        subplex, sf = plex.filter(sanitizeSubMesh=True, ignoreHalo=ignore_halo, comm=comm)
-        dmcommon.submesh_update_facet_labels(plex, subplex)
-        dmcommon.submesh_correct_entity_classes(plex, subplex, sf)
-    else:
-        subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, ignore_halo, comm=comm)
+    subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, ignore_halo, comm=comm)
 
     comm = comm or mesh.comm
     name = name or _generate_default_submesh_name(mesh.name)

firedrake/mg/mesh.py

@@ -126,9 +126,8 @@ def MeshHierarchy(mesh, refinement_levels,
     # Effectively "invert" addOverlap().
     # -- The resulting plex is to have the identical data structure as the one before addOverlap().
     #    This is algorithmically guaranteed.
-    dm_cell_type, = mesh.dm_cell_types
     tdim = mesh.topology_dm.getDimension()
-    cdm = dmcommon.submesh_create(mesh.topology_dm, tdim, "celltype", dm_cell_type, True)
+    cdm = dmcommon.submesh_create(mesh.topology_dm, tdim, "depth", tdim, True)
     cdm.removeLabel("pyop2_core")
     cdm.removeLabel("pyop2_owned")
     cdm.removeLabel("pyop2_ghost")

firedrake/mg/netgen.py

@@ -265,8 +265,7 @@ def NetgenHierarchy(mesh, levs, flags, distribution_parameters=None):
         )
         mesh = reconstruct_mesh(mesh, coordinates)
     # Make a plex (cdm) without overlap.
-    dm_cell_type, = mesh.dm_cell_types
-    cdm = dmcommon.submesh_create(mesh.topology_dm, tdim, "celltype", dm_cell_type, True)
+    cdm = dmcommon.submesh_create(mesh.topology_dm, tdim, "depth", tdim, True)
     cdm.removeLabel("pyop2_core")
     cdm.removeLabel("pyop2_owned")
     cdm.removeLabel("pyop2_ghost")

firedrake/preconditioners/bddc.py

@@ -199,7 +199,7 @@ def local_mesh(mesh):
             return cache[key]
         except KeyError:
             if mesh.comm.size > 1:
-                submesh = Submesh(mesh, mesh.topological_dimension, None, ignore_halo=True, reorder=False, comm=COMM_SELF)
+                submesh = Submesh(mesh, ignore_halo=True, comm=COMM_SELF)
             else:
                 submesh = None
             return cache.setdefault(key, submesh)

tests/firedrake/submesh/test_submesh_comm.py

@@ -28,10 +28,9 @@ def assert_local_equality(A, Asub, V, Vsub):
 @pytest.mark.parametrize("reorder", [False, True])
 @pytest.mark.parametrize("ignore_halo", [False, True])
 def test_create_submesh_comm_self(reorder, ignore_halo):
-    subdomain_id = None
     nx = 4
     mesh = UnitSquareMesh(nx, nx, quadrilateral=True, reorder=reorder, distribution_parameters={"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)})
-    submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=ignore_halo, reorder=reorder, comm=COMM_SELF)
+    submesh = Submesh(mesh, ignore_halo=ignore_halo, reorder=reorder, comm=COMM_SELF)
     assert submesh.submesh_parent is mesh
     assert submesh.comm.size == 1
     # Submesh on COMM_SELF should not have halo
@@ -49,13 +48,12 @@ def test_create_submesh_comm_self(reorder, ignore_halo):
 @pytest.mark.parametrize("family,degree", [("DG", 0), ("CG", 1)])
 @pytest.mark.parametrize("reorder", [False, True])
 def test_assemble_submesh_comm_self(family, degree, reorder):
-    subdomain_id = None
     nx = 6
     ny = 5
     px = -np.cos(np.linspace(0, np.pi, nx))
     py = -np.cos(np.linspace(0, np.pi, ny))
     mesh = TensorRectangleMesh(px, py, reorder=reorder)
-    submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=True, reorder=reorder, comm=COMM_SELF)
+    submesh = Submesh(mesh, ignore_halo=True, reorder=reorder, comm=COMM_SELF)
 
     Vsub = FunctionSpace(submesh, family, degree)
     Asub = assemble(inner(TrialFunction(Vsub), TestFunction(Vsub))*dx)

tests/firedrake/submesh/test_submesh_facet.py

@@ -122,3 +122,15 @@ def test_submesh_facet_corner_case_2():
     facet_value = 999
     mesh = RelabeledMesh(mesh, [facet_function], [facet_value])
     _ = Submesh(mesh, mesh.topological_dimension - 1, facet_value)
+
+
+def test_submesh_facet_all_facets():
+    mesh = UnitCubeMesh(2, 2, 2)
+    submesh1 = Submesh(mesh, mesh.topological_dimension - 1)
+
+    V = FunctionSpace(mesh, "HDiv Trace", 0)
+    facet_function = Function(V).assign(1)
+    facet_value = 999
+    rmesh = RelabeledMesh(mesh, [facet_function], [facet_value])
+    submesh2 = Submesh(rmesh, mesh.topological_dimension - 1, facet_value)
+    assert submesh2.cell_set.size == submesh1.cell_set.size