Redistributed Submesh and MeshHierarchy

firedrake/adapt.py

@@ -7,7 +7,7 @@
 from firedrake.utils import IntType
 from firedrake.function import Function
 from firedrake.functionspace import FunctionSpace
-from firedrake.mesh import Mesh, DISTRIBUTION_PARAMETERS_NOOP
+from firedrake.mesh import Mesh, Submesh, DISTRIBUTION_PARAMETERS_NOOP
 from firedrake.netgen import _transfer_high_order_coordinates
 from firedrake.petsc import PETSc
 
@@ -72,7 +72,33 @@ def _copy_adaptive_refinement_metadata(source_mesh, target_mesh):
         target_mesh.netgen_flags = source_mesh.netgen_flags
 
 
-def refine_marked_elements(mesh, cell_marker):
+def _redistribute_adaptive_refined_mesh(coarse_mesh, refined_mesh, redistribute=True):
+    """Redistribute an adaptively refined mesh if it has empty ranks.
+
+    Parameters
+    ----------
+    coarse_mesh : firedrake.mesh.MeshGeometry
+        The mesh that was refined.
+    refined_mesh : firedrake.mesh.MeshGeometry
+        The result of refining ``coarse_mesh``.
+    redistribute : bool
+        If ``True``, redistribute ``refined_mesh`` when it has empty ranks.
+
+    Returns
+    -------
+    firedrake.mesh.MeshGeometry
+        ``refined_mesh``, or a redistributed `~firedrake.mesh.Submesh` of it.
+
+    """
+    _copy_adaptive_refinement_metadata(coarse_mesh, refined_mesh)
+    if not (redistribute and refined_mesh.has_empty_rank):
+        return refined_mesh
+    redist_mesh = Submesh(refined_mesh, redistribute=True, name=refined_mesh.name)
+    _copy_adaptive_refinement_metadata(refined_mesh, redist_mesh)
+    return redist_mesh
+
+
+def refine_marked_elements(mesh, cell_marker, redistribute=True):
     """Adaptively refine a mesh using a DG0 marking function.
 
     Positive integer marker values request repeated refinement of the
@@ -86,6 +112,9 @@ def refine_marked_elements(mesh, cell_marker):
     cell_marker
         A DG0 `~firedrake.function.Function` on ``mesh``: cells with a
         positive value ``n`` are refined ``n`` times.
+    redistribute
+        If ``True``, redistribute the refined mesh when the coarse mesh
+        has empty ranks.
 
     Returns
     -------
@@ -145,7 +174,12 @@ def refine_marked_elements(mesh, cell_marker):
                 final_mesh = _transfer_high_order_coordinates(mesh, final_mesh, order)
 
     final_mesh.topology_dm.removeLabel(PARENT_LABEL)
+    # _redistribute_adaptive_refined_mesh copies the construction metadata
+    # across, and may hand back a different mesh, so record the provenance on
+    # whichever mesh comes out of it.
+    final_mesh = _redistribute_adaptive_refined_mesh(
+        mesh, final_mesh, redistribute=redistribute
+    )
     final_mesh.adaptive_parent = mesh
     final_mesh.adaptive_cell_maps = (coarse_to_fine, fine_to_coarse)
-    _copy_adaptive_refinement_metadata(mesh, final_mesh)
     return final_mesh

firedrake/assign.py

@@ -17,12 +17,79 @@
 from firedrake.cofunction import Cofunction
 from firedrake.constant import Constant
 from firedrake.function import Function
+from firedrake.halo import _get_mtype
 from firedrake.petsc import PETSc
 from firedrake.utils import ScalarType, split_by
 
 from mpi4py import MPI
 
 
+def _submesh_point_sf(target_mesh, source_mesh):
+    """Find the point SF relating two meshes with different distributions.
+
+    Parameters
+    ----------
+    target_mesh : AbstractMeshTopology
+        The mesh being assigned to.
+    source_mesh : AbstractMeshTopology
+        The mesh being assigned from.
+
+    Returns
+    -------
+    tuple
+        The `PETSc.SF` mapping the points of the parent mesh (roots) to the
+        points of the submesh (leaves), and whether ``target_mesh`` is the
+        submesh. Both are `None` if the two meshes share their distribution,
+        in which case they are related by entity maps instead.
+
+    """
+    if target_mesh.submesh_parent is source_mesh:
+        return target_mesh.submesh_point_sf, True
+    elif source_mesh.submesh_parent is target_mesh:
+        return source_mesh.submesh_point_sf, False
+    else:
+        return None, None
+
+
+def _make_section_sf(point_sf, root_V, leaf_V):
+    """Expand a point SF into an SF relating the nodes of two function spaces.
+
+    Parameters
+    ----------
+    point_sf : PETSc.SF
+        SF mapping the points of the mesh of ``root_V`` (roots) to the points
+        of the mesh of ``leaf_V`` (leaves).
+    root_V : firedrake.functionspaceimpl.WithGeometry
+        Function space holding the root data.
+    leaf_V : firedrake.functionspaceimpl.WithGeometry
+        Function space holding the leaf data.
+
+    Returns
+    -------
+    tuple
+        The `PETSc.SF` mapping the nodes of ``root_V`` to the nodes of
+        ``leaf_V``, and the boolean array telling which nodes of ``root_V``
+        have a counterpart in ``leaf_V``.
+
+    """
+    cache = leaf_V.mesh().topology._shared_data_cache["submesh_section_sf"]
+    key = (root_V, leaf_V)
+    try:
+        return cache[key]
+    except KeyError:
+        root_section = root_V.dm.getSection()
+        leaf_section = leaf_V.dm.getSection()
+        # `distributeSection` overwrites the section it is handed, so let it
+        # build its own and only keep the root offsets it broadcasts.
+        remote_offsets, distributed_section = point_sf.distributeSection(root_section)
+        if distributed_section.getChart() != leaf_section.getChart():
+            raise RuntimeError("Point SF does not cover the nodes of the leaf function space")
+        section_sf = point_sf.createSectionSF(root_section, remote_offsets, leaf_section)
+        # A submesh only covers part of its parent, so not every root node
+        # is reduced into.
+        return cache.setdefault(key, (section_sf, section_sf.computeDegree() > 0))
+
+
 def _isconstant(expr):
     return isinstance(expr, Constant) or \
         (isinstance(expr, (Function, Cofunction)) and expr.ufl_element().family() == "Real")
@@ -298,6 +365,69 @@ def source_indices(f):
             lhs_func.dat.halo_valid = True
 
     def _assign_multi_mesh(self, lhs_func, subset, funcs, operator, allow_missing_dofs):
+        target_mesh = extract_unique_domain(lhs_func).topology
+        source_V, = set(f.function_space() for f in funcs)
+        source_mesh = source_V.mesh().topology
+        if target_mesh.submesh_shares_distribution(source_mesh):
+            self._assign_submesh(lhs_func, subset, funcs, operator, allow_missing_dofs)
+            return
+        point_sf, target_is_submesh = _submesh_point_sf(target_mesh, source_mesh)
+        if point_sf is None:
+            raise NotImplementedError(
+                "Can only assign between a redistributed mesh and its parent"
+            )
+        self._assign_redistributed(lhs_func, subset, funcs, point_sf,
+                                   target_is_submesh, allow_missing_dofs)
+
+    def _assign_redistributed(self, lhs_func, subset, funcs, point_sf,
+                              target_is_submesh, allow_missing_dofs):
+        """Assign between (co)functions on a redistributed submesh and its parent.
+
+        The nodes of the two spaces correspond one to one. The expression is
+        evaluated in the source layout. One communication then moves the
+        result into the target layout.
+        """
+        target_V = lhs_func.function_space()
+        source_V, = set(f.function_space() for f in funcs)
+        if target_is_submesh:
+            root_V, leaf_V = source_V, target_V
+        else:
+            root_V, leaf_V = target_V, source_V
+        section_sf, covered_roots = _make_section_sf(point_sf, root_V, leaf_V)
+
+        source_buffer = Function(source_V)
+        target_buffer = Function(target_V)
+        func_data = np.array([f.dat.data_ro_with_halos for f in funcs])
+        source_buffer.dat.data_wo_with_halos[...] = self._compute_rvalue(func_data)
+        mtype, _ = _get_mtype(source_buffer.dat)
+        source_data = source_buffer.dat.data_ro_with_halos
+        target_data = target_buffer.dat.data_wo_with_halos
+        if target_is_submesh:
+            section_sf.bcastBegin(mtype, source_data, target_data, MPI.REPLACE)
+            section_sf.bcastEnd(mtype, source_data, target_data, MPI.REPLACE)
+            # Every node of a submesh, including its halo, has a counterpart
+            # in the parent.
+            indices = Ellipsis if subset is None else subset.indices
+            assign_to_halos = True
+        else:
+            section_sf.reduceBegin(mtype, source_data, target_data, MPI.REPLACE)
+            section_sf.reduceEnd(mtype, source_data, target_data, MPI.REPLACE)
+            # Only the owned parent nodes that the submesh covers have been
+            # reduced into; the parent halo never is.
+            owned = covered_roots[:target_V.dof_dset.size]
+            comm = target_V.mesh().comm
+            if not comm.allreduce(owned.all(), op=MPI.LAND) and not allow_missing_dofs:
+                raise ValueError("Found assignee nodes with no matching assigner "
+                                 "nodes: run with `allow_missing_dofs=True`")
+            indices, = np.nonzero(owned)
+            if subset is not None:
+                indices = np.intersect1d(indices, subset.owned_indices)
+            target_data = target_buffer.dat.data_ro
+            assign_to_halos = False
+        self._assign_single_dat(lhs_func.dat, indices, target_data[indices], assign_to_halos)
+        lhs_func.dat.halo_valid = assign_to_halos
+
+    def _assign_submesh(self, lhs_func, subset, funcs, operator, allow_missing_dofs):
         target_mesh = extract_unique_domain(lhs_func)
         target_V = lhs_func.function_space()
         source_V, = set(f.function_space() for f in funcs)

firedrake/cython/dmcommon.pyx

@@ -4064,6 +4064,143 @@ def submesh_create(PETSc.DM dm,
     return subdm
 
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def submesh_vertex_numbering(PETSc.SF point_sf,
+                             PETSc.Section parent_numbering,
+                             PETSc.Section numbering):
+    """Inherit the universal vertex numbering of the submesh parent.
+
+    Parameters
+    ----------
+    point_sf : PETSc.SF
+        SF whose roots are the points of the parent plex and whose leaves
+        are the points of the submesh plex.
+    parent_numbering : PETSc.Section
+        Section describing the universal vertex numbering of the parent.
+    numbering : PETSc.Section
+        Section describing the universal vertex numbering of the submesh.
+
+    Returns
+    -------
+    PETSc.Section
+        Copy of ``numbering`` in which each vertex carries the universal
+        number of the corresponding vertex of the parent.
+
+    Notes
+    -----
+    Cell closures are ordered by universal vertex number, so a submesh that
+    inherits the numbering of its parent orients its entities exactly as the
+    parent does. The nodes of a function space are then in one-to-one
+    correspondence on the two meshes, even though the meshes are distributed
+    differently.
+
+    """
+    cdef:
+        PETSc.Section inherited
+        PetscInt nroots, pStart, pEnd, ppStart, ppEnd, p, dof, offset
+        np.ndarray[PetscInt, ndim=1, mode="c"] roots, leaves
+        MPI.Datatype typ
+        MPI.Op replace = MPI.REPLACE
+
+    CHKERR(PetscSFGetGraph(point_sf.sf, &nroots, NULL, NULL, NULL))
+    ppStart, ppEnd = parent_numbering.getChart()
+    if ppEnd - ppStart != nroots:
+        raise ValueError("Point SF must have one root per point of the parent plex")
+    pStart, pEnd = numbering.getChart()
+    roots = np.full(nroots, -1, dtype=IntType)
+    for p in range(ppStart, ppEnd):
+        CHKERR(PetscSectionGetDof(parent_numbering.sec, p, &dof))
+        # Points not owned by this rank carry complementary inverses.
+        if cabs(dof) > 0:
+            CHKERR(PetscSectionGetOffset(parent_numbering.sec, p, &offset))
+            roots[p - ppStart] = cabs(offset)
+    leaves = np.full(pEnd - pStart, -1, dtype=IntType)
+    try:
+        tdict = MPI.__TypeDict__
+    except AttributeError:
+        tdict = MPI._typedict
+    typ = tdict[roots.dtype.char]
+    CHKERR(PetscSFBcastBegin(point_sf.sf, typ.ob_mpi,
+                             <const void *>roots.data,
+                             <void *>leaves.data,
+                             replace.ob_mpi))
+    CHKERR(PetscSFBcastEnd(point_sf.sf, typ.ob_mpi,
+                           <const void *>roots.data,
+                           <void *>leaves.data,
+                           replace.ob_mpi))
+    inherited = numbering.clone()
+    for p in range(pStart, pEnd):
+        CHKERR(PetscSectionGetDof(inherited.sec, p, &dof))
+        if cabs(dof) > 0:
+            offset = leaves[p - pStart]
+            if offset < 0:
+                raise RuntimeError("Found a vertex with no counterpart in the submesh parent")
+            CHKERR(PetscSectionSetOffset(inherited.sec, p,
+                                         offset if dof > 0 else cneg(offset)))
+    return inherited
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def submesh_cell_orientations(PETSc.DM parent_plex,
+                              PETSc.Section parent_cell_numbering,
+                              np.ndarray parent_orientations,
+                              PETSc.SF point_sf,
+                              PETSc.DM plex,
+                              PETSc.Section cell_numbering):
+    """Inherit the cell orientations of the submesh parent.
+
+    Parameters
+    ----------
+    parent_plex : PETSc.DM
+        The parent plex.
+    parent_cell_numbering : PETSc.Section
+        Section describing the cell numbering of the parent.
+    parent_orientations : numpy.ndarray
+        Cell orientations of the parent.
+    point_sf : PETSc.SF
+        SF whose roots are the points of ``parent_plex`` and whose leaves
+        are the points of ``plex``.
+    plex : PETSc.DM
+        The submesh plex.
+    cell_numbering : PETSc.Section
+        Section describing the cell numbering of the submesh.
+
+    Returns
+    -------
+    numpy.ndarray
+        Cell orientations of the submesh.
+
+    """
+    cdef:
+        MPI.Datatype dtype
+        PETSc.Section new_section
+        PetscInt *new_values = NULL
+        PetscInt c, cStart, cEnd, l, r
+        np.ndarray orientations
+
+    try:
+        tdict = MPI.__TypeDict__
+    except AttributeError:
+        tdict = MPI._typedict
+    dtype = tdict[np.dtype(IntType).char]
+    new_section = PETSc.Section().create(comm=plex.comm)
+    CHKERR(DMPlexDistributeData(parent_plex.dm, point_sf.sf,
+                                parent_cell_numbering.sec, dtype.ob_mpi,
+                                <void *>parent_orientations.data,
+                                new_section.sec, <void **>&new_values))
+    get_height_stratum(plex.dm, 0, &cStart, &cEnd)
+    orientations = np.empty(cEnd - cStart, dtype=IntType)
+    for c in range(cStart, cEnd):
+        CHKERR(PetscSectionGetOffset(cell_numbering.sec, c, &l))
+        CHKERR(PetscSectionGetOffset(new_section.sec, c, &r))
+        orientations[l] = new_values[r]
+    if new_values != NULL:
+        CHKERR(PetscFree(new_values))
+    return orientations
+
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def submesh_correct_entity_classes(PETSc.DM dm,

firedrake/cython/petschdr.pxi

@@ -123,6 +123,7 @@ cdef extern from "petscvec.h" nogil:
 
 cdef extern from "petscis.h" nogil:
     PetscErrorCode PetscSectionGetOffset(PETSc.PetscSection,PetscInt,PetscInt*)
+    PetscErrorCode PetscSectionSetOffset(PETSc.PetscSection,PetscInt,PetscInt)
     PetscErrorCode PetscSectionGetDof(PETSc.PetscSection,PetscInt,PetscInt*)
     PetscErrorCode PetscSectionSetDof(PETSc.PetscSection,PetscInt,PetscInt)
     PetscErrorCode PetscSectionSetFieldDof(PETSc.PetscSection,PetscInt,PetscInt,PetscInt)
@@ -151,8 +152,8 @@ cdef extern from "petscsf.h" nogil:
 
     PetscErrorCode PetscSFGetGraph(PETSc.PetscSF,PetscInt*,PetscInt*,PetscInt**,PetscSFNode**)
     PetscErrorCode PetscSFSetGraph(PETSc.PetscSF,PetscInt,PetscInt,PetscInt*,PetscCopyMode,PetscSFNode*,PetscCopyMode)
-    PetscErrorCode PetscSFBcastBegin(PETSc.PetscSF,MPI.MPI_Datatype,const void*, void*,)
-    PetscErrorCode PetscSFBcastEnd(PETSc.PetscSF,MPI.MPI_Datatype,const void*, void*)
+    PetscErrorCode PetscSFBcastBegin(PETSc.PetscSF,MPI.MPI_Datatype,const void*, void*,MPI.MPI_Op)
+    PetscErrorCode PetscSFBcastEnd(PETSc.PetscSF,MPI.MPI_Datatype,const void*, void*,MPI.MPI_Op)
     PetscErrorCode PetscSFReduceBegin(PETSc.PetscSF,MPI.MPI_Datatype,const void*, void*,MPI.MPI_Op)
     PetscErrorCode PetscSFReduceEnd(PETSc.PetscSF,MPI.MPI_Datatype,const void*, void*,MPI.MPI_Op)