support labels on COMM_SELF

Author: pbrubeck

firedrake/mesh.py

@@ -4818,24 +4818,21 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo=
         raise NotImplementedError("Can not create a submesh of a ``VertexOnlyMesh``")
     plex = mesh.topology_dm
     dim = plex.getDimension()
+    if subdim not in [dim, dim - 1]:
+        raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})")
+    if 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 is not None:
-            raise ValueError("subdomain_id == None requires label_name == None.")
-        if subdim != dim:
-            raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})")
-        subplex, _ = plex.filter(sanitizeSubMesh=True, ignoreHalo=ignore_halo, comm=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:
-        if comm is not None and comm != mesh.comm:
-            raise NotImplementedError("Submesh on sub-communicator not implemented on cell subsets.")
-        if subdim not in [dim, dim - 1]:
-            raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})")
-        if label_name is None:
-            if subdim == dim:
-                label_name = dmcommon.CELL_SETS_LABEL
-            elif subdim == dim - 1:
-                label_name = dmcommon.FACE_SETS_LABEL
         subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, ignore_halo, comm=comm)
 
     comm = comm or mesh.comm

tests/firedrake/submesh/test_submesh_comm.py

@@ -1,15 +1,25 @@
 import pytest
 import numpy as np
 from firedrake import *
+from firedrake.petsc import PETSc
+
+
+def assert_local_equality(A, B):
+    i0, j0, v0 = A.getValuesCSR()
+    i1, j1, v1 = B.getValuesCSR()
+    j0 -= A.createVecs()[0].getOwnershipRange()[0]
+    j1 -= B.createVecs()[0].getOwnershipRange()[0]
+    assert np.array_equal(i0, i1)
+    assert np.array_equal(j0, j1)
+    assert np.allclose(v0, v1)
 
 
 @pytest.mark.parallel([1, 3])
 def test_create_submesh_comm_self():
-    comm = COMM_SELF
     subdomain_id = None
     nx = 4
     mesh = UnitSquareMesh(nx, nx, quadrilateral=True, reorder=False)
-    submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=True, comm=comm, reorder=False)
+    submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=True, reorder=False, comm=COMM_SELF)
     assert submesh.submesh_parent is mesh
     assert submesh.comm.size == 1
     assert submesh.cell_set.size == mesh.cell_set.size
@@ -18,24 +28,51 @@ def test_create_submesh_comm_self():
 
 @pytest.mark.parallel([1, 3])
 def test_assemble_submesh_comm_self():
-    comm = COMM_SELF
     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=False)
-    submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=True, comm=comm, reorder=False)
+    submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=True, reorder=False, comm=COMM_SELF)
+
+    Vsub = FunctionSpace(submesh, "DG", 0)
+    Asub = assemble(inner(TrialFunction(Vsub), TestFunction(Vsub))*dx)
 
     V = FunctionSpace(mesh, "DG", 0)
     A = assemble(inner(TrialFunction(V), TestFunction(V))*dx)
+    assert_local_equality(A.petscmat, Asub.petscmat)
+
+
+@pytest.mark.parallel([1, 3])
+@pytest.mark.parametrize("label", ["some", "all"])
+def test_label_submesh_comm_self(label):
+    subdomain_id = 999
+    nx = 8
+    mesh = UnitSquareMesh(nx, nx, reorder=False)
+
+    M = FunctionSpace(mesh, "DG", 0)
+    marker = Function(M)
+    if label == "some":
+        x, y = SpatialCoordinate(mesh)
+        marker.interpolate(conditional(Or(x > 0.5, y > 0.5), 1, 0))
+    elif label == "all":
+        marker.assign(1)
+    else:
+        raise ValueError(f"Unrecognized label {label}")
+
+    mesh = RelabeledMesh(mesh, [marker], [subdomain_id])
+    submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=True, reorder=False, comm=COMM_SELF)
 
     Vsub = FunctionSpace(submesh, "DG", 0)
-    Asub = assemble(inner(TrialFunction(Vsub), TestFunction(Vsub))*dx)
+    Asub = assemble(inner(TrialFunction(Vsub), TestFunction(Vsub)) * dx)
 
-    i0, j0, v0 = A.petscmat.getValuesCSR()
-    j0 -= V.dof_dset.layout_vec.getOwnershipRange()[0]
-    i1, j1, v1 = Asub.petscmat.getValuesCSR()
-    assert np.array_equal(i0, i1)
-    assert np.array_equal(j0, j1)
-    assert np.allclose(v0, v1)
+    V = FunctionSpace(mesh, "DG", 0)
+    A = assemble(inner(TrialFunction(V), TestFunction(V)) * dx)
+    if label == "all":
+        assert_local_equality(A.petscmat, Asub.petscmat)
+    else:
+        lgmap = V.dof_dset.lgmap
+        indices = PETSc.IS().createGeneral(lgmap.apply(np.flatnonzero(marker.dat.data).astype(PETSc.IntType)))
+        Amat = A.petscmat.createSubMatrix(indices, indices)
+        assert_local_equality(Amat, Asub.petscmat)