Submesh on subcommunicator (#4761)

---------

Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>

Author: pbrubeck

firedrake/__init__.py

@@ -3,6 +3,7 @@
 # the specific version, here we are more permissive. This is to catch the
 # case where users don't update their PETSc for a really long time or
 # accidentally install a too-new release that isn't yet supported.
+# TODO RELEASE set to ">=3.25"
 PETSC_SUPPORTED_VERSIONS = ">=3.23.0"
 
 

firedrake/cython/dmcommon.pyx

@@ -3819,7 +3819,8 @@ def submesh_create(PETSc.DM dm,
                    PetscInt subdim,
                    label_name,
                    PetscInt label_value,
-                   PetscBool ignore_label_halo):
+                   PetscBool ignore_label_halo,
+                   comm=None):
     """Create submesh.
 
     Parameters
@@ -3834,12 +3835,12 @@ def submesh_create(PETSc.DM dm,
         Value in the label
     ignore_label_halo : bool
         If labeled points in the halo are ignored.
+    comm : PETSc.Comm | None
+        An optional sub-communicator to define the submesh.
 
     """
     cdef:
-        PETSc.DM subdm = PETSc.DMPlex()
         PETSc.DMLabel label, temp_label
-        PETSc.SF ownership_transfer_sf = PETSc.SF()
         char *temp_label_name = <char *>"firedrake_submesh_temp_label"
         PetscInt pStart, pEnd, p, i, stratum_size
         PETSc.PetscIS stratum_is = NULL
@@ -3863,7 +3864,11 @@ def submesh_create(PETSc.DM dm,
         CHKERR(ISRestoreIndices(stratum_is, &stratum_indices))
         CHKERR(ISDestroy(&stratum_is))
     # Make submesh using temp_label.
-    CHKERR(DMPlexFilter(dm.dm, temp_label.dmlabel, label_value, ignore_label_halo, PETSC_TRUE, &ownership_transfer_sf.sf, &subdm.dm))
+    subdm, ownership_transfer_sf = dm.filter(label=temp_label,
+                                             value=label_value,
+                                             ignoreHalo=ignore_label_halo,
+                                             sanitizeSubMesh=PETSC_TRUE,
+                                             comm=comm)
     # Destroy temp_label.
     dm.removeLabel(temp_label_name)
     subdm.removeLabel(temp_label_name)
@@ -3901,7 +3906,7 @@ def submesh_correct_entity_classes(PETSc.DM dm,
         DMLabel lbl_core, lbl_owned, lbl_ghost
         PetscBool has
 
-    if dm.comm.size == 1:
+    if subdm.comm.size == 1:
         return
     CHKERR(DMPlexGetChart(dm.dm, &pStart, &pEnd))
     CHKERR(DMPlexGetChart(subdm.dm, &subpStart, &subpEnd))

firedrake/cython/petschdr.pxi

@@ -76,7 +76,6 @@ cdef extern from "petscdmplex.h" nogil:
     PetscErrorCode DMPlexLabelComplete(PETSc.PetscDM, PETSc.PetscDMLabel)
     PetscErrorCode DMPlexDistributeOverlap(PETSc.PetscDM,PetscInt,PETSc.PetscSF*,PETSc.PetscDM*)
 
-    PetscErrorCode DMPlexFilter(PETSc.PetscDM,PETSc.PetscDMLabel,PetscInt,PetscBool,PetscBool,PETSc.PetscSF*,PETSc.PetscDM*)
     PetscErrorCode DMPlexGetSubpointIS(PETSc.PetscDM,PETSc.PetscIS*)
     PetscErrorCode DMPlexGetSubpointMap(PETSc.PetscDM,PETSc.PetscDMLabel*)
     PetscErrorCode DMPlexSetSubpointMap(PETSc.PetscDM,PETSc.PetscDMLabel)

firedrake/mesh.py

@@ -3012,7 +3012,7 @@ def make_mesh_from_coordinates(coordinates, name, tolerance=0.5):
 
 
 def make_mesh_from_mesh_topology(topology, name, tolerance=0.5):
-    """Make mesh from tpology.
+    """Make mesh from topology.
 
     Parameters
     ----------
@@ -4734,7 +4734,7 @@ def SubDomainData(geometric_expr):
     return op2.Subset(m.cell_set, indices)
 
 
-def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None):
+def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo=False, reorder=True, comm=None):
     """Construct a submesh from a given mesh.
 
     Parameters
@@ -4743,12 +4743,20 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None):
         Parent mesh (`MeshGeometry`).
     subdim : int
         Topological dimension of the submesh.
-    subdomain_id : int
+    subdomain_id : int | None
         Subdomain ID representing the submesh.
-    label_name : str
+        `None` defines the submesh owned by the sub-communicator.
+    label_name : str | None
         Name of the label to search ``subdomain_id`` in.
-    name : str
+    name : str |  None
         Name of the submesh.
+    ignore_halo : bool
+        Whether to exclude the halo from the submesh.
+    reorder : bool
+        Whether to reorder the mesh entities.
+    comm : PETSc.Comm | None
+        An optional sub-communicator to define the submesh.
+        By default, the submesh is defined on `mesh.comm`.
 
     Returns
     -------
@@ -4817,15 +4825,27 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None):
             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)
+
+    comm = comm or mesh.comm
     name = name or _generate_default_submesh_name(mesh.name)
-    subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, False)
     subplex.setName(_generate_default_mesh_topology_name(name))
     if subplex.getDimension() != subdim:
         raise RuntimeError(f"Found subplex dim ({subplex.getDimension()}) != expected ({subdim})")
     submesh = Mesh(
         subplex,
         submesh_parent=mesh,
         name=name,
+        comm=comm,
+        reorder=reorder,
         distribution_parameters={
             "partition": False,
             "overlap_type": (DistributedMeshOverlapType.NONE, 0),

setup.py

@@ -20,6 +20,7 @@
 
 
 # Ensure that the PETSc getting linked against is compatible
+# TODO RELEASE set to ">=3.25"
 petsctools.init(version_spec=">=3.23.0")
 import petsc4py
 

tests/firedrake/submesh/test_submesh_comm.py

@@ -0,0 +1,78 @@
+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():
+    subdomain_id = None
+    nx = 4
+    mesh = UnitSquareMesh(nx, nx, quadrilateral=True, 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
+    assert np.allclose(mesh.coordinates.dat.data_ro, submesh.coordinates.dat.data_ro)
+
+
+@pytest.mark.parallel([1, 3])
+def test_assemble_submesh_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, 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)
+
+    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)