diff --git a/firedrake/adapt.py b/firedrake/adapt.py index f1b6a4851b..209a912c7b 100644 --- a/firedrake/adapt.py +++ b/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.any_rank_is_empty): + 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 diff --git a/firedrake/assign.py b/firedrake/assign.py index 065fd90ff8..a0d0eb096a 100644 --- a/firedrake/assign.py +++ b/firedrake/assign.py @@ -17,12 +17,80 @@ 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. `point_sf` addresses each root on its owner, so a + # covered root is one this rank owns. + 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 +366,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: a broadcast onto the submesh, which + covers every one of its nodes, or a reduction onto the parent, which + reaches only the nodes the submesh covers. + """ + 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) + 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) + # A reduction reaches owned nodes alone, so the halo of the + # assignee is left stale for a later exchange to fill in. + 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) diff --git a/firedrake/cython/dmcommon.pyx b/firedrake/cython/dmcommon.pyx index ca12036c51..e1d82aa417 100644 --- a/firedrake/cython/dmcommon.pyx +++ b/firedrake/cython/dmcommon.pyx @@ -4076,6 +4076,144 @@ 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)) + # A global section negates the dof and the offset of a point that + # this rank does not own, so compare and store their magnitudes. + 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, + roots.data, + leaves.data, + replace.ob_mpi)) + CHKERR(PetscSFBcastEnd(point_sf.sf, typ.ob_mpi, + roots.data, + 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, + parent_orientations.data, + new_section.sec, &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, diff --git a/firedrake/cython/petschdr.pxi b/firedrake/cython/petschdr.pxi index f172e551a0..179e4f9e05 100644 --- a/firedrake/cython/petschdr.pxi +++ b/firedrake/cython/petschdr.pxi @@ -125,6 +125,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) @@ -153,8 +154,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) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index be425bf581..3097e63442 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -503,7 +503,7 @@ class AbstractMeshTopology(object, metaclass=abc.ABCMeta): """A representation of an abstract mesh topology without a concrete PETSc DM implementation""" - def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm, submesh_parent=None): + def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm, submesh_parent=None, submesh_point_sf=None): """Initialise a mesh topology. Parameters @@ -532,6 +532,11 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, Communicator. submesh_parent: AbstractMeshTopology Submesh parent. + submesh_point_sf: PETSc.PetscSF + `PETSc.SF` that pushes the points of ``submesh_parent`` to the + points of ``topology_dm``; only given if this mesh is to be + redistributed, i.e. does not inherit the parallel distribution + of ``submesh_parent``. """ utils._init() @@ -544,6 +549,14 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, self.sfXB = sfXB r"The PETSc SF that pushes the global point number slab [0, NX) to input (naive) plex." self.submesh_parent = submesh_parent + self.submesh_point_sf = submesh_point_sf + r"""The PETSc SF that pushes the points of ``submesh_parent`` to the points of this mesh. + + This is `None` whenever this mesh shares the parallel distribution of + ``submesh_parent``, in which case the points of the two meshes are + related locally by ``topology_dm.getSubpointIS()``. Ask + `is_redistributed` rather than testing it. + """ self.sfBC_orig = None # User comm self.user_comm = comm @@ -554,6 +567,9 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, self._add_overlap() if self.sfXB is not None: self.sfXC = sfXB.compose(self.sfBC) if self.sfBC else self.sfXB + if self.is_redistributed and self.sfBC: + # Push the parent points onto the redistributed plex. + self.submesh_point_sf = self.submesh_point_sf.compose(self.sfBC) dmcommon.label_facets(self.topology_dm) dmcommon.complete_facet_labels(self.topology_dm) # TODO: Allow users to set distribution name if they want to save @@ -659,6 +675,12 @@ def _topology_dm(self): warn("_topology_dm is deprecated (use topology_dm instead)", DeprecationWarning, stacklevel=2) return self.topology_dm + @cached_property + def any_rank_is_empty(self) -> bool: + """Whether any rank of this mesh owns no cells. Collective.""" + with temp_internal_comm(self.comm) as icomm: + return icomm.allreduce(self.cell_set.size == 0, op=MPI.LOR) + def ufl_cell(self): """The UFL :class:`~ufl.classes.Cell` associated with the mesh. @@ -990,6 +1012,39 @@ def submesh_youngest_common_ancestor(self, other): break return c + @property + def is_redistributed(self) -> bool: + """Whether this mesh was repartitioned instead of taking its parent's distribution.""" + return self.submesh_point_sf is not None + + def submesh_shares_distribution(self, other): + """Return whether `self` and ``other`` are related and share their distribution. + + Two meshes of the same submesh family have entity maps only under one + condition. Every mesh between them and their youngest common ancestor + must take its parent's parallel distribution. + + Parameters + ---------- + other : AbstractMeshTopology + The other mesh. + + Returns + ------- + bool + Whether the two meshes are related by entity maps. + + """ + common = self.submesh_youngest_common_ancestor(other) + if common is None: + return False + for mesh in (self, other): + while mesh is not common: + if mesh.is_redistributed: + return False + mesh = mesh.submesh_parent + return True + def submesh_map_child_parent(self, source_integral_type, source_subset_points, reverse=False): """Return the map from submesh child entities to submesh parent entities or its reverse. @@ -1084,6 +1139,7 @@ def __init__( distribution_name=None, permutation_name=None, submesh_parent=None, + submesh_point_sf=None, comm=COMM_WORLD, ): """Initialise a mesh topology. @@ -1114,6 +1170,9 @@ def __init__( Name of the entity permutation (reordering); if `None`, automatically generated. submesh_parent: MeshTopology Submesh parent. + submesh_point_sf: PETSc.PetscSF + `PETSc.SF` that pushes the points of ``submesh_parent`` to the + points of ``plex``; only given if ``plex`` is to be redistributed. comm : mpi4py.MPI.Comm Communicator. @@ -1133,7 +1192,7 @@ def __init__( # Disable auto distribution and reordering before setFromOptions is called. plex.distributeSetDefault(False) plex.reorderSetDefault(PETSc.DMPlex.ReorderDefaultFlag.FALSE) - super().__init__(plex, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm, submesh_parent=submesh_parent) + super().__init__(plex, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm, submesh_parent=submesh_parent, submesh_point_sf=submesh_point_sf) def _distribute(self): # Distribute/redistribute the dm to all ranks @@ -1231,6 +1290,71 @@ def dm_cell_types(self): """All DM.PolytopeTypes of cells in the mesh.""" return dmcommon.get_dm_cell_types(self.topology_dm) + @cached_property + def _universal_vertex_numbering(self): + """Section describing the universal (globally unique) vertex numbering. + + Cell closures are ordered by universal vertex number. A redistributed + submesh therefore takes the numbering of its parent, and orients its + entities exactly as the parent does. + """ + numbering = self._vertex_numbering.createGlobalSection(self.topology_dm.getPointSF()) + if not self.is_redistributed: + return numbering + return dmcommon.submesh_vertex_numbering( + self.submesh_point_sf, + self.submesh_parent._universal_vertex_numbering, + numbering, + ) + + @cached_property + def _quadrilateral_cell_orientations(self): + """Global orientation of each quadrilateral cell. + + Neighbouring cells must agree on the direction of the edge they + share, which is decided by a distributed algorithm whose outcome + depends on the partition. A redistributed submesh therefore + inherits the orientations of its parent rather than choosing + its own. + """ + plex = self.topology_dm + if self.is_redistributed: + return dmcommon.submesh_cell_orientations( + self.submesh_parent.topology_dm, + self.submesh_parent._cell_numbering, + self.submesh_parent._quadrilateral_cell_orientations, + self.submesh_point_sf, + plex, + self._cell_numbering, + ) + vertex_numbering = self._universal_vertex_numbering + cell_ranks = dmcommon.get_cell_remote_ranks(plex) + facet_orientations = dmcommon.quadrilateral_facet_orientations( + plex, vertex_numbering, cell_ranks) + cell_orientations = dmcommon.orientations_facet2cell( + plex, vertex_numbering, cell_ranks, + facet_orientations, self._cell_numbering) + dmcommon.exchange_cell_orientations(plex, + self._cell_numbering, + cell_orientations) + return cell_orientations + + @cached_property + def _inherits_parent_cell_closure(self) -> bool: + """Whether this mesh takes its cell closures from its submesh parent. + + A quadrilateral submesh of a hexahedral mesh is the exception. Its + own closures must follow the orientation restriction that a + quadrilateral cell carries, which the hexahedral closures do not, so + working with the parent permutes the quadrature points instead. + """ + if self.submesh_parent is None or self.is_redistributed: + return False + if len(self.submesh_parent.dm_cell_types) != 1: + return False + return not (self.submesh_parent.ufl_cell().cellname == "hexahedron" + and self.ufl_cell().cellname == "quadrilateral") + @cached_property def cell_closure(self): """2D array of ordered cell closures @@ -1242,19 +1366,11 @@ def cell_closure(self): # Cell numbering and global vertex numbering cell_numbering = self._cell_numbering - vertex_numbering = self._vertex_numbering.createGlobalSection(plex.getPointSF()) + vertex_numbering = self._universal_vertex_numbering cell = self.ufl_cell() assert tdim == cell.topological_dimension - if self.submesh_parent is not None and \ - not (self.submesh_parent.ufl_cell().cellname == "hexahedron" and cell.cellname == "quadrilateral") and \ - len(self.submesh_parent.dm_cell_types) == 1: - # Codim-1 submesh of a hex mesh (i.e. a quad submesh) can not - # inherit cell_closure from the hex mesh as the cell_closure - # must follow the special orientation restriction. This means - # that, when the quad submesh works with the parent hex mesh, - # quadrature points must be permuted (i.e. use the canonical - # quadrature point ordering based on the cone ordering). + if self._inherits_parent_cell_closure: topology = FIAT.ufc_cell(cell).get_topology() entity_per_cell = np.zeros(len(topology), dtype=IntType) for d, ents in topology.items(): @@ -1279,22 +1395,9 @@ def cell_closure(self): elif cell.cellname == "quadrilateral": petsctools.cite("Homolya2016") petsctools.cite("McRae2016") - # Quadrilateral mesh - cell_ranks = dmcommon.get_cell_remote_ranks(plex) - - facet_orientations = dmcommon.quadrilateral_facet_orientations( - plex, vertex_numbering, cell_ranks) - - cell_orientations = dmcommon.orientations_facet2cell( - plex, vertex_numbering, cell_ranks, - facet_orientations, cell_numbering) - - dmcommon.exchange_cell_orientations(plex, - cell_numbering, - cell_orientations) - return dmcommon.quadrilateral_closure_ordering( - plex, vertex_numbering, cell_numbering, cell_orientations) + plex, vertex_numbering, cell_numbering, + self._quadrilateral_cell_orientations) elif cell.cellname == "hexahedron": # TODO: Should change and use create_cell_closure() for all cell types. topology = FIAT.ufc_cell(cell).get_topology() @@ -1635,6 +1738,12 @@ def submesh_map_child_parent(self, source_integral_type, source_subset_points, r """ if self.submesh_parent is None: raise RuntimeError("Must only be called on submesh") + if self.is_redistributed: + raise NotImplementedError( + "Assembling or interpolating across a submesh and its parent " + "requires the two to share the same parallel distribution; use " + "`Function.assign` to transfer data between redistributed meshes" + ) if reverse: source = self.submesh_parent target = self @@ -1841,6 +1950,7 @@ def __init__(self, mesh, layers, periodic=False, name=None): self.cell_set = op2.ExtrudedSet(mesh.cell_set, layers=layers, extruded_periodic=periodic) # submesh self.submesh_parent = None + self.submesh_point_sf = None @cached_property def _ufl_cell(self): @@ -2960,7 +3070,7 @@ def unique(self): return self @PETSc.Log.EventDecorator() - def refine_marked_elements(self, mark): + def refine_marked_elements(self, mark, redistribute=True): """Adaptively refine a mesh using a DG0 marking function. Parameters @@ -2969,6 +3079,10 @@ def refine_marked_elements(self, mark): A DG0 `~firedrake.function.Function` on this 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 ------- MeshGeometry @@ -2978,7 +3092,7 @@ def refine_marked_elements(self, mark): :meth:`~firedrake.mg.mesh.HierarchyBase.add_mesh`. """ from firedrake.adapt import refine_marked_elements - return refine_marked_elements(self, mark) + return refine_marked_elements(self, mark, redistribute) @PETSc.Log.EventDecorator() def curve_field(self, order, permutation_tol=None, cg_field=None): @@ -3082,7 +3196,7 @@ def curve_field(self, order, permutation_tol=None, cg_field=None): @PETSc.Log.EventDecorator() -def make_mesh_from_coordinates(coordinates, name, tolerance=0.5): +def make_mesh_from_coordinates(coordinates, name, tolerance=0.5, submesh_parent=None): """Given a coordinate field build a new mesh, using said coordinate field. Parameters @@ -3093,6 +3207,8 @@ def make_mesh_from_coordinates(coordinates, name, tolerance=0.5): The name of the mesh. tolerance : numbers.Number The tolerance; see `Mesh`. + submesh_parent : MeshGeometry + The mesh this one is a submesh of, if any. comm: mpi4py.Intracomm Communicator. @@ -3121,6 +3237,7 @@ def make_mesh_from_coordinates(coordinates, name, tolerance=0.5): mesh._tolerance = tolerance mesh._did_reordering = orig_mesh._did_reordering mesh._distribution_parameters = orig_mesh._distribution_parameters + mesh.submesh_parent = submesh_parent return mesh @@ -3372,10 +3489,10 @@ def Mesh(meshfile, **kwargs): coordinates = meshfile else: coordinates = None - if coordinates is not None: - return make_mesh_from_coordinates(coordinates, name) - tolerance = kwargs.get("tolerance", 0.5) + if coordinates is not None: + return make_mesh_from_coordinates(coordinates, name, tolerance=tolerance, + submesh_parent=kwargs.get("submesh_parent")) utils._init() @@ -3425,6 +3542,7 @@ def Mesh(meshfile, **kwargs): distribution_name=kwargs.get("distribution_name"), permutation_name=kwargs.get("permutation_name"), submesh_parent=submesh_parent.topology if submesh_parent else None, + submesh_point_sf=kwargs.get("submesh_point_sf"), comm=user_comm) mesh = make_mesh_from_mesh_topology(topology, name) @@ -4897,7 +5015,41 @@ def SubDomainData(geometric_expr): return op2.Subset(m.cell_set, indices) -def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ignore_halo=False, reorder=None, comm=None): +def _make_submesh_point_sf(plex, subplex): + """Create the `PETSc.SF` relating the points of a plex and of its submesh. + + Parameters + ---------- + plex : PETSc.DMPlex + The parent plex. + subplex : PETSc.DMPlex + The submesh plex, before it is distributed. + + Returns + ------- + PETSc.SF + SF whose roots are the points of ``plex`` and whose leaves are the + points of ``subplex``. + + """ + pStart, pEnd = plex.getChart() + # Address every parent point on the rank that owns it, which the parent's + # own point SF records for its ghosts. Data reduced onto a ghost point + # would never reach the owner. + owners = np.empty((pEnd - pStart, 2), dtype=IntType) + owners[:, 0] = plex.comm.rank + owners[:, 1] = np.arange(pStart, pEnd, dtype=IntType) + if plex.isDistributed(): + _, ghosts, ghost_owners = plex.getPointSF().getGraph() + owners[ghosts] = ghost_owners + with subplex.getSubpointIS() as subpoints: + remote = owners[subpoints] + point_sf = PETSc.SF().create(comm=subplex.comm) + point_sf.setGraph(pEnd - pStart, None, remote) + return point_sf + + +def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ignore_halo=False, reorder=None, comm=None, redistribute=False): """Construct a submesh from a given mesh. Parameters @@ -4927,6 +5079,13 @@ def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ig comm : PETSc.Comm | None An optional sub-communicator to define the submesh. By default, the submesh is defined on `mesh.comm`. + redistribute : bool + Whether to repartition the submesh, instead of inheriting the + parallel distribution of ``mesh``. This implies ``ignore_halo=True``, + and is currently only supported for submeshes of co-dimension 0. + A redistributed submesh can not be assembled or interpolated + alongside its parent; use `~.Function.assign` to transfer data + between the two. Returns ------- @@ -4995,6 +5154,10 @@ def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ig >>> submesh = Submesh(mesh, ignore_halo=True, comm=COMM_SELF) + Construct a repartitioned copy of the entire mesh + + >>> submesh = Submesh(mesh, redistribute=True) + """ if not isinstance(mesh, MeshGeometry): raise TypeError("Parent mesh must be a `MeshGeometry`") @@ -5002,8 +5165,33 @@ def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ig 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``") - - subplex = dmcommon.submesh_create(mesh.topology_dm, subdim, label_name, subdomain_id, ignore_halo, comm=comm) + if redistribute: + if comm is not None: + raise NotImplementedError("Can only redistribute a submesh over the parent communicator") + # Drop the parent halo so that every point of the submesh is owned + # by exactly one rank before it is repartitioned. + ignore_halo = True + distribution_parameters = dict(mesh._distribution_parameters, partition=True) + else: + distribution_parameters = DISTRIBUTION_PARAMETERS_NOOP + + plex = mesh.topology_dm + subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, ignore_halo, comm=comm) + if redistribute and subplex.getDimension() != plex.getDimension(): + # The two meshes must be made of the same cells. Only then are their + # entities oriented consistently, and only then do their nodes + # correspond. + raise NotImplementedError("Can only redistribute a submesh of co-dimension 0") + if redistribute: + # The point correspondence must be recorded before the submesh is + # distributed, as distributing it discards the subpoint IS. + point_sf = _make_submesh_point_sf(plex, subplex) + # Repartitioning invalidates the entity classification the submesh + # takes from its parent. Drop the labels, so that they are recomputed. + for label in ("pyop2_core", "pyop2_owned", "pyop2_ghost"): + subplex.removeLabel(label) + else: + point_sf = None comm = comm or mesh.comm name = name or _generate_default_submesh_name(mesh.name) @@ -5015,14 +5203,80 @@ def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ig submesh = Mesh( subplex, submesh_parent=mesh, + submesh_point_sf=point_sf, name=name, comm=comm, reorder=reorder, - distribution_parameters=DISTRIBUTION_PARAMETERS_NOOP, + distribution_parameters=distribution_parameters, + tolerance=mesh.tolerance, ) - # Tag the relabeled mesh with the original distribution parameters - submesh._distribution_parameters = mesh._distribution_parameters - return submesh + if not redistribute: + # DISTRIBUTION_PARAMETERS_NOOP keeps Mesh() from distributing the + # submesh again. The submesh has the distribution of its parent, so it + # reports the parameters of the parent. + submesh._distribution_parameters = mesh._distribution_parameters + + if _plex_carries_parent_coordinates(mesh, submesh): + return submesh + return _submesh_with_transferred_coordinates(mesh, submesh, name) + + +def _plex_carries_parent_coordinates(mesh, submesh): + """Whether the plex of ``submesh`` already holds the coordinates of ``mesh``. + + Parameters + ---------- + mesh : MeshGeometry + The parent mesh. + submesh : MeshGeometry + The submesh, on the coordinates its plex carries. + + Returns + ------- + bool + `False` when the parent is curved or periodic, and so keeps its + coordinates in a `~firedrake.function.Function` of its own. + + """ + if len(mesh.topology.dm_cell_types) > 1: + # Such a mesh carries no coordinate Function at all. + return True + # A submesh of lower dimension has a different cell than its parent, so + # the two coordinate elements are compared on the parent's cell. + plex_element = submesh.coordinates.ufl_element().reconstruct(cell=mesh.ufl_cell()) + return mesh.coordinates.ufl_element() == plex_element + + +def _submesh_with_transferred_coordinates(mesh, submesh, name): + """Rebuild a submesh on the coordinates of its parent. + + Parameters + ---------- + mesh : MeshGeometry + The parent mesh, whose coordinates its plex does not carry. + submesh : MeshGeometry + The submesh to rebuild. + name : str + Name of the new mesh. + + Returns + ------- + MeshGeometry + A submesh of ``mesh`` on the topology of ``submesh``, carrying the + coordinates of ``mesh`` restricted to it. + + """ + if submesh.ufl_cell() != mesh.ufl_cell(): + raise NotImplementedError( + "Can only transfer the coordinates of a curved or periodic mesh " + "onto a submesh of the same dimension" + ) + import firedrake.function as function + + V = mesh.coordinates.function_space().reconstruct(mesh=submesh) + coordinates = function.Function(V).assign(mesh.coordinates) + return Mesh(coordinates, name=name, submesh_parent=mesh, + tolerance=mesh.tolerance) def coordinates_from_topology(topology: AbstractMeshTopology, element: finat.ufl.FiniteElement) -> "CoordinatelessFunction": @@ -5174,6 +5428,14 @@ def __init__(self, meshes): raise ValueError(f"Got {type(m)}") self._meshes = tuple(meshes) self.comm = meshes[0].comm + # A mesh sequence is never a submesh. + self.submesh_parent = None + self.submesh_point_sf = None + + @property + def is_redistributed(self) -> bool: + """A mesh sequence is never a submesh, so it never has a distribution of its own.""" + return False @property def topology(self): diff --git a/firedrake/mg/interface.py b/firedrake/mg/interface.py index 6c42d3c434..8e11560dac 100644 --- a/firedrake/mg/interface.py +++ b/firedrake/mg/interface.py @@ -68,10 +68,13 @@ def prolong(coarse, fine): meshes = hierarchy._meshes for j in range(repeat): next_level += 1 - if j == repeat - 1 and not needs_quadrature: + fine_mesh = meshes[next_level] + transfer_mesh = utils.transfer_mesh(fine_mesh) + last = j == repeat - 1 + if last and not needs_quadrature and transfer_mesh is fine_mesh: fine = finest else: - fine = Function(Vf.reconstruct(mesh=meshes[next_level])) + fine = Function(Vf.reconstruct(mesh=transfer_mesh)) Vf = fine.function_space() Vc = coarse.function_space() compose_map = lambda u: utils.fine_node_to_coarse_node_map(Vf, u.function_space()) @@ -110,8 +113,13 @@ def prolong(coarse, fine): if needs_quadrature: # Transfer to the actual target space - new_fine = finest if j == repeat-1 else Function(Vfinest.reconstruct(mesh=meshes[next_level])) + new_fine = (finest if last and transfer_mesh is fine_mesh + else Function(Vfinest.reconstruct(mesh=transfer_mesh))) fine = new_fine.interpolate(fine) + if transfer_mesh is not fine_mesh: + # Move onto the redistributed mesh + target = finest if last else Function(Vfinest.reconstruct(mesh=fine_mesh)) + fine = target.assign(fine) coarse = fine return fine @@ -149,9 +157,17 @@ def restrict(fine_dual, coarse_dual): coarsest = coarse_dual.zero() meshes = hierarchy._meshes for j in range(repeat): + fine_mesh = meshes[next_level] + transfer_mesh = utils.transfer_mesh(fine_mesh) + if transfer_mesh is not fine_mesh: + # Move off the redistributed mesh, so that we can restrict + Vf_transfer = fine_dual.function_space().reconstruct(mesh=transfer_mesh) + fine_dual = Function(Vf_transfer).assign(fine_dual) if needs_quadrature: # Transfer to the quadrature source space - fine_dual = Function(Vq.reconstruct(mesh=meshes[next_level])).interpolate(fine_dual) + fine_dual = Function( + Vq.reconstruct(mesh=fine_dual.function_space().mesh()) + ).interpolate(fine_dual) next_level -= 1 if j == repeat - 1: @@ -242,6 +258,12 @@ def inject(fine, coarse): Vcoarsest = coarsest.function_space() meshes = hierarchy._meshes for j in range(repeat): + fine_mesh = meshes[next_level] + transfer_mesh = utils.transfer_mesh(fine_mesh) + if transfer_mesh is not fine_mesh: + # Move off the redistributed mesh, so that we can inject + Vf_transfer = fine.function_space().reconstruct(mesh=transfer_mesh) + fine = Function(Vf_transfer).assign(fine) next_level -= 1 if j == repeat - 1 and not needs_quadrature: coarse = coarsest diff --git a/firedrake/mg/mesh.py b/firedrake/mg/mesh.py index 6451dbf239..41228264e5 100644 --- a/firedrake/mg/mesh.py +++ b/firedrake/mg/mesh.py @@ -11,9 +11,9 @@ from functools import cached_property from firedrake import utils +from firedrake.cython import dmcommon from firedrake.cython import mgimpl as impl -import firedrake.cython.dmcommon as dmcommon -from .utils import set_level +from .utils import set_level, set_dm_refine_level __all__ = ("HierarchyBase", "MeshHierarchy", "ExtrudedMeshHierarchy", "NonNestedHierarchy", "SemiCoarsenedExtrudedHierarchy", "SubmeshHierarchy") @@ -52,6 +52,8 @@ class HierarchyBase(object): Number of mesh refinements each multigrid level should "see". nested : Is this mesh hierarchy nested? + redistribute : + Redistribute adaptively refined meshes that have empty ranks? Notes ----- @@ -60,7 +62,7 @@ class HierarchyBase(object): """ def __init__(self, meshes, coarse_to_fine_cells, fine_to_coarse_cells, - refinements_per_level=1, nested=False): + refinements_per_level=1, nested=False, redistribute=True): petsctools.cite("Mitchell2016") self._meshes = list(meshes) self.meshes = self._meshes[::refinements_per_level] @@ -68,6 +70,7 @@ def __init__(self, meshes, coarse_to_fine_cells, fine_to_coarse_cells, self.fine_to_coarse_cells = fine_to_coarse_cells self.refinements_per_level = refinements_per_level self.nested = nested + self.redistribute = redistribute for level, m in enumerate(meshes): set_level(m, self, Fraction(level, refinements_per_level)) for level, m in enumerate(self): @@ -180,7 +183,8 @@ def adapt(self, eta, theta: float): markers = firedrake.Function(M) markers.dat.data_wo[eta.dat.data_ro > theta * eta_max] = 1 - return self.add_mesh(mesh.refine_marked_elements(markers)) + return self.add_mesh( + mesh.refine_marked_elements(markers, redistribute=self.redistribute)) def MeshHierarchy(mesh, refinement_levels=0, @@ -188,7 +192,8 @@ def MeshHierarchy(mesh, refinement_levels=0, netgen_flags=False, reorder=None, distribution_parameters=None, callbacks=None, - mesh_builder=firedrake.Mesh, nested=True): + mesh_builder=firedrake.Mesh, nested=True, + redistribute=True): """Build a hierarchy of meshes by uniformly refining a coarse mesh. Parameters @@ -211,6 +216,11 @@ def MeshHierarchy(mesh, refinement_levels=0, for details. If ``None``, use the same distribution parameters as were used to distribute the coarse mesh, otherwise, these options override the default. + redistribute : bool + If ``True``, redistribute refined meshes when this is needed to + avoid empty ranks. Transfer operators use an internal + parent-owned mesh before moving data to or from the redistributed + mesh. reorder : bool optional flag indicating whether to reorder the refined meshes. @@ -247,15 +257,26 @@ def MeshHierarchy(mesh, refinement_levels=0, else: before = after = lambda dm, i: None - # Refine an unoverlapped plex at each level. Keeping every dm here - # unoverlapped means overlap only ever needs to be added once, by - # mesh_builder below. + parameters = {} + if distribution_parameters is not None: + parameters.update(distribution_parameters) + else: + parameters.update(mesh._distribution_parameters) + parameters["partition"] = False + + # Refine an unoverlapped plex at each level, and redistribute the + # refined mesh whenever refining alone would leave empty ranks. Keeping + # the refined plex unoverlapped means that overlap only ever needs to be + # added once, by mesh_builder below. cdm = mesh.topology_dm if refinement_levels > 0: cdm = make_unoverlapped_dm(cdm) - cdm.setRefinementUniform(True) - dms = [cdm] + lgmaps = [(impl.create_lgmap(cdm), impl.create_lgmap(mesh.topology_dm))] + meshes = [mesh] + coarse_to_fine_cells = [] + fine_to_coarse_cells = [None] for i in range(refinement_levels*refinements_per_level): + cdm.setRefinementUniform(True) if i % refinements_per_level == 0: before(cdm, i) rdm = cdm.refine() @@ -271,19 +292,9 @@ def MeshHierarchy(mesh, refinement_levels=0, scale = mesh._radius / np.linalg.norm(coords, axis=1).reshape(-1, 1) coords *= scale - dms.append(rdm) - cdm = rdm - - # Build a mesh for each level, adding overlap here. - parameters = {} - if distribution_parameters is not None: - parameters.update(distribution_parameters) - else: - parameters.update(mesh._distribution_parameters) - parameters["partition"] = False - - meshes = [mesh] - for rdm in dms[1:]: + # The cell maps relate the refined mesh to the mesh it was refined + # from, so they must be built before it is redistributed. + rlgmap = impl.create_lgmap(rdm) fmesh = mesh_builder( rdm, dim=mesh.geometric_dimension, @@ -291,32 +302,28 @@ def MeshHierarchy(mesh, refinement_levels=0, reorder=reorder, comm=mesh.comm, ) - meshes.append(fmesh) - - # Build local-to-global maps and coarse/fine cell maps between - # consecutive levels. - lgmaps = [ - (impl.create_lgmap(dm), impl.create_lgmap(m.topology_dm)) - for dm, m in zip(dms, meshes) - ] - coarse_to_fine_cells = [] - fine_to_coarse_cells = [None] - for (coarse, fine), (clgmaps, flgmaps) in zip(zip(meshes[:-1], meshes[1:]), - zip(lgmaps[:-1], lgmaps[1:])): - c2f, f2c = impl.coarse_to_fine_cells(coarse, fine, clgmaps, flgmaps) + flgmaps = (rlgmap, impl.create_lgmap(fmesh.topology_dm)) + c2f, f2c = impl.coarse_to_fine_cells(meshes[-1], fmesh, lgmaps[-1], flgmaps) coarse_to_fine_cells.append(c2f) fine_to_coarse_cells.append(f2c) + if redistribute and fmesh.any_rank_is_empty: + fmesh = firedrake.Submesh(fmesh, redistribute=True) + cdm = make_unoverlapped_dm(fmesh.topology_dm) + lgmaps.append((impl.create_lgmap(cdm), impl.create_lgmap(fmesh.topology_dm))) + meshes.append(fmesh) + for i, m in enumerate(meshes): # Firedrake counts multigrid levels, PETSc counts refinements - m.topology_dm.setRefineLevel(i) + set_dm_refine_level(m, i) coarse_to_fine_cells = dict((Fraction(i, refinements_per_level), c2f) for i, c2f in enumerate(coarse_to_fine_cells)) fine_to_coarse_cells = dict((Fraction(i, refinements_per_level), f2c) for i, f2c in enumerate(fine_to_coarse_cells)) return HierarchyBase(meshes, coarse_to_fine_cells, fine_to_coarse_cells, - refinements_per_level, nested=nested) + refinements_per_level, nested=nested, + redistribute=redistribute) def ExtrudedMeshHierarchy(base_hierarchy, height, base_layer=-1, refinement_ratio=2, layers=None, diff --git a/firedrake/mg/utils.py b/firedrake/mg/utils.py index 042bf5edaf..53b19bd465 100644 --- a/firedrake/mg/utils.py +++ b/firedrake/mg/utils.py @@ -425,7 +425,7 @@ def restrict_preserved_nodes(fine_dual, coarse_dual): section_sf = preserved_node_sf(coarse_V, fine_dual.function_space()) if section_sf is None: return - buffer = firedrake.Function(coarse_V) + buffer = type(coarse_dual)(coarse_V) mtype, _ = _get_mtype(buffer.dat) source = fine_dual.dat.data_ro target = buffer.dat.data_wo_with_halos @@ -456,9 +456,88 @@ def physical_node_locations(V): return cache.setdefault(key, locations) +def transfer_mesh(mesh): + """Return the mesh that grid transfer operates on. + + A redistributed mesh has no cell maps relating it to the coarse mesh. + Transfers therefore go through the mesh it was redistributed from. The + values are then assigned across the two. + + Parameters + ---------- + mesh : firedrake.mesh.MeshGeometry + A mesh in a `HierarchyBase`. + + Returns + ------- + firedrake.mesh.MeshGeometry + ``mesh`` itself, or the mesh it was redistributed from. + + """ + return mesh.submesh_parent if mesh.is_redistributed else mesh + + +def _redistribution_ancestors(topology): + """Yield a mesh topology together with the topologies it was redistributed from. + + The transfer operators work on the mesh a redistributed mesh came from, + so both must carry the same multigrid level. + + Parameters + ---------- + topology : firedrake.mesh.AbstractMeshTopology + The topology to start from. + + Yields + ------ + firedrake.mesh.AbstractMeshTopology + ``topology``, then each mesh topology it was redistributed from, in + order. + + """ + yield topology + while topology.is_redistributed: + topology = topology.submesh_parent + yield topology + + +def set_dm_refine_level(mesh, level): + """Set the refinement level of a mesh and of the meshes it was redistributed from. + + Parameters + ---------- + mesh : firedrake.mesh.MeshGeometry + The mesh to set the refinement level of. + level : int + The refinement level to set. + + """ + for topology in _redistribution_ancestors(mesh.topology): + topology.topology_dm.setRefineLevel(level) + + def set_level(obj, hierarchy, level): - """Attach hierarchy and level info to an object.""" - setattr(obj.topological, "__level_info__", (hierarchy, level)) + """Attach hierarchy and level info to an object. + + Parameters + ---------- + obj : firedrake.mesh.MeshGeometry + The mesh to attach the hierarchy and level info to. The meshes it was + redistributed from take the same level, because the transfer + operators work on those. + hierarchy : HierarchyBase + The hierarchy ``obj`` belongs to. + level : Fraction + The level of ``obj`` in ``hierarchy``. + + Returns + ------- + firedrake.mesh.MeshGeometry + ``obj``, unchanged. + + """ + for topology in _redistribution_ancestors(obj.topological): + setattr(topology, "__level_info__", (hierarchy, level)) return obj diff --git a/tests/firedrake/multigrid/test_adaptive_multigrid.py b/tests/firedrake/multigrid/test_adaptive_multigrid.py index bfdc63abe1..e8f80987f0 100644 --- a/tests/firedrake/multigrid/test_adaptive_multigrid.py +++ b/tests/firedrake/multigrid/test_adaptive_multigrid.py @@ -2,7 +2,7 @@ import numpy as np from mpi4py import MPI from firedrake import * -from firedrake.mg.utils import preserved_node_sf +from firedrake.mg.utils import preserved_node_sf, transfer_mesh from firedrake.utils import complex_mode @@ -80,7 +80,7 @@ def test_refine_marked_elements_populates_cell_maps(coarse_mesh): fine_to_coarse = mh.fine_to_coarse_cells[1] assert coarse_to_fine.shape[0] == mesh.cell_set.size - assert fine_to_coarse.shape == (refined_mesh.cell_set.size, 1) + assert fine_to_coarse.shape == (transfer_mesh(refined_mesh).cell_set.size, 1) assert (fine_to_coarse >= -1).all() assert (fine_to_coarse >= 0).any() assert (coarse_to_fine >= 0).any() @@ -245,7 +245,7 @@ def _assert_adapt_after_uniform_refinement(mh): fine_to_coarse = mh.fine_to_coarse_cells[level] assert coarse_to_fine.shape[0] == mesh.cell_set.size - assert fine_to_coarse.shape == (refined_mesh.cell_set.size, 1) + assert fine_to_coarse.shape == (transfer_mesh(refined_mesh).cell_set.size, 1) # A rank may legitimately own zero local cells (e.g. more ranks than # coarse cells), leaving these arrays empty on that rank alone, so the # "some entry is valid" check must be collective, not per-rank. @@ -355,7 +355,7 @@ def _copied_nodes(mh, V): copied = 0 for level in range(len(mh) - 1): V_coarse = V.reconstruct(mesh=mh[level]) - V_fine = V.reconstruct(mesh=mh[level + 1]) + V_fine = V.reconstruct(mesh=transfer_mesh(mh[level + 1])) section_sf = preserved_node_sf(V_coarse, V_fine) if section_sf is not None: _, preserved, _ = section_sf.getGraph() diff --git a/tests/firedrake/multigrid/test_redist_mesh.py b/tests/firedrake/multigrid/test_redist_mesh.py new file mode 100644 index 0000000000..a5f2f9fa2a --- /dev/null +++ b/tests/firedrake/multigrid/test_redist_mesh.py @@ -0,0 +1,126 @@ +import numpy as np +import pytest + +from firedrake import * + + +@pytest.mark.parallel(2) +def test_redistributed_hierarchy(): + m = UnitIntervalMesh(1) + mh = MeshHierarchy(m, 1) + + assert mh[1].cell_set.size > 0 + + +@pytest.mark.parallel(4) +def test_uniform_hierarchy_no_empty_ranks(): + dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)} + base = UnitSquareMesh(1, 1, distribution_parameters=dparams) + mh = MeshHierarchy(base, 2) + + for l, m in enumerate(mh[1:]): + assert m.cell_set.size > 0 + for k, v in dparams.items(): + assert m._distribution_parameters.get(k, None) == v + + Vc = FunctionSpace(mh[l], "CG", 1) + Vf = FunctionSpace(mh[l+1], "CG", 1) + + xc, yc = SpatialCoordinate(mh[l]) + xf, yf = SpatialCoordinate(mh[l+1]) + coarse_expr = xc + 2*yc + fine_expr = xf + 2*yf + + # test prolong CG1 + coarse = Function(Vc).interpolate(coarse_expr) + fine = Function(Vf) + prolong(coarse, fine) + assert errornorm(fine_expr, fine) < 1e-12 + + # test restrict CG1 + one_coarse = Function(Vc).assign(1) + one_fine = Function(Vf) + prolong(one_coarse, one_fine) + + fine_dual = assemble(conj(TestFunction(Vf))*dx) + coarse_dual = Cofunction(Vc.dual()) + restrict(fine_dual, coarse_dual) + assert np.allclose( + assemble(action(coarse_dual, one_coarse)), + assemble(action(fine_dual, one_fine)), + rtol=1e-12, + atol=1e-12, + ) + + # test inject CG1 + coarse_injected = Function(Vc) + inject(fine, coarse_injected) + assert errornorm(coarse_expr, coarse_injected) < 1e-12 + + # test inject DG0 + Qc = FunctionSpace(mh[l], "DG", 0) + Qf = FunctionSpace(mh[l+1], "DG", 0) + fine_expr = conditional(xf > 1, 1, 0) + coarse_expr = conditional(xc > 1, 1, 0) + fine = Function(Qf).interpolate(fine_expr) + coarse_injected = Function(Qc) + inject(fine, coarse_injected) + assert np.allclose( + assemble(coarse_expr * dx), + assemble(coarse_injected * dx), + rtol=1e-12, + atol=1e-12, + ) + assert l == 0 or errornorm(coarse_expr, coarse_injected) < 1e-12 + + +@pytest.mark.parallel(3) +def test_adaptive_hierarchy_redistributes_empty_ranks(): + # Two cells spread over three ranks leave a rank with no cells. Adaptive + # refinement alone does not fix that. The refined mesh has to be + # redistributed. + dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1), + "partitioner_type": "simple"} + mesh = UnitSquareMesh(1, 1, distribution_parameters=dparams) + assert mesh.any_rank_is_empty + mh = MeshHierarchy(mesh) + + # Mark only one of the two cells, so the refined mesh keeps both a + # refined region and a region the refinement leaves alone. + M = FunctionSpace(mesh, "DG", 0) + markers = Function(M) + markers.dat.data_wo[:1] = 1 + + refined_mesh = mesh.refine_marked_elements(markers) + assert refined_mesh.topology.is_redistributed + assert not refined_mesh.any_rank_is_empty + mh.add_mesh(refined_mesh) + + V_coarse = FunctionSpace(mesh, "CG", 1) + V_fine = FunctionSpace(refined_mesh, "CG", 1) + xc, yc = SpatialCoordinate(mesh) + xf, yf = SpatialCoordinate(refined_mesh) + expr_coarse = xc + 2 * yc + expr_fine = xf + 2 * yf + + # test prolong CG1 + u_coarse = Function(V_coarse).interpolate(expr_coarse) + u_fine = Function(V_fine) + prolong(u_coarse, u_fine) + assert errornorm(expr_fine, u_fine) <= 1e-12 + + # test restrict CG1 + r_fine = assemble(conj(TestFunction(V_fine)) * dx) + r_coarse = Cofunction(V_coarse.dual()) + restrict(r_fine, r_coarse) + assert np.allclose( + assemble(action(r_coarse, u_coarse)), + assemble(action(r_fine, u_fine)), + rtol=1e-12, + atol=1e-12, + ) + + # test inject CG1 + u_coarse_injected = Function(V_coarse) + inject(u_fine, u_coarse_injected) + assert errornorm(expr_coarse, u_coarse_injected) <= 1e-12 diff --git a/tests/firedrake/submesh/test_submesh_assemble.py b/tests/firedrake/submesh/test_submesh_assemble.py index 8692857cad..7b15aff565 100644 --- a/tests/firedrake/submesh/test_submesh_assemble.py +++ b/tests/firedrake/submesh/test_submesh_assemble.py @@ -597,3 +597,34 @@ def test_submesh_assemble_facet_macroelement(): vsub = TestFunction(Vsub) aref = assemble(inner(1, vsub)*dx(submesh)) assert np.allclose(a.dat[1].data_ro, aref.dat.data_ro) + + +@pytest.mark.parallel(nprocs=[1, 3]) +def test_submesh_assemble_redistributed(): + # A redistributed submesh does not share the point numbering of its + # parent, so the entity maps that multidomain assembly needs do not exist. + dim = 2 + mesh = RectangleMesh(2, 1, 2., 1., quadrilateral=True) + x, _ = SpatialCoordinate(mesh) + DQ0 = FunctionSpace(mesh, "DQ", 0) + mesh.mark_entities(Function(DQ0).interpolate(conditional(x > 1., 1, 0)), 999) + subm = Submesh(mesh, dim, 999, redistribute=True) + V = FunctionSpace(mesh, "CG", 1) * FunctionSpace(subm, "CG", 1) + u0, u1 = split(TrialFunction(V)) + v0, v1 = split(TestFunction(V)) + dx0 = Measure("dx", domain=mesh, intersect_measures=(Measure("dx", subm),)) + with pytest.raises(NotImplementedError): + assemble(inner(u1, v0) * dx0(999), mat_type="nest") + + +@pytest.mark.parallel(nprocs=[1, 3]) +def test_submesh_interpolate_redistributed(): + dim = 2 + mesh = RectangleMesh(2, 1, 2., 1., quadrilateral=True) + x, _ = SpatialCoordinate(mesh) + DQ0 = FunctionSpace(mesh, "DQ", 0) + mesh.mark_entities(Function(DQ0).interpolate(conditional(x > 1., 1, 0)), 999) + subm = Submesh(mesh, dim, 999, redistribute=True) + f = Function(FunctionSpace(mesh, "CG", 1)).interpolate(x) + with pytest.raises(NotImplementedError): + assemble(interpolate(f, FunctionSpace(subm, "CG", 1))) diff --git a/tests/firedrake/submesh/test_submesh_assign.py b/tests/firedrake/submesh/test_submesh_assign.py index 12adf956ce..495bced9b4 100644 --- a/tests/firedrake/submesh/test_submesh_assign.py +++ b/tests/firedrake/submesh/test_submesh_assign.py @@ -307,3 +307,70 @@ def test_submesh_assign_cofunction_3_quads_2_processes(): cof_ = Cofunction(V_r.dual()).assign(cof) subset_indices = np.where(coords_r.dat.data_ro_with_halos[:, 0] > 1.999) assert np.allclose(cof_.dat.data_ro_with_halos[subset_indices], cof_r.dat.data_ro_with_halos[subset_indices]) + + +@pytest.mark.parametrize("mesh_type,family,degree", [ + ("simplex", "CG", 3), + ("simplex", "RT", 2), + ("simplex", "N1curl", 1), + ("quadrilateral", "CG", 2), + ("quadrilateral", "RTCF", 1), +]) +@pytest.mark.parallel(nprocs=[1, 3]) +def test_submesh_assign_function_redistributed(mesh_type, family, degree): + # A redistributed submesh holds the same mesh as its parent. Assigning + # between the two must therefore be exact for every element. That includes + # the elements whose DoFs depend on the orientation of their entity. + mesh = UnitSquareMesh(4, 4, quadrilateral=(mesh_type == "quadrilateral")) + submesh = Submesh(mesh, redistribute=True) + assert submesh.topology.submesh_parent is mesh.topology + assert submesh.topology.is_redistributed + + V = FunctionSpace(mesh, family, degree) + V_sub = FunctionSpace(submesh, family, degree) + x, y = SpatialCoordinate(mesh) + x_sub, y_sub = SpatialCoordinate(submesh) + if V.value_size == 1: + expr, expr_sub = sin(x) * cos(3 * y), sin(x_sub) * cos(3 * y_sub) + else: + expr = as_vector([sin(x), cos(3 * y)]) + expr_sub = as_vector([sin(x_sub), cos(3 * y_sub)]) + f = Function(V).interpolate(expr) + # -- mesh -> submesh + f_sub = Function(V_sub).assign(f) + assert np.allclose(norm(f_sub - Function(V_sub).interpolate(expr_sub)), 0) + # -- submesh -> mesh + assert np.allclose(norm(Function(V).assign(f_sub) - f), 0) + # -- the dual assignment preserves the action + cof_sub = assemble(inner(expr_sub, TestFunction(V_sub)) * dx) + cof = Cofunction(V.dual()).assign(cof_sub) + assert np.isclose(assemble(action(cof, f)), assemble(action(cof_sub, f_sub))) + + +@pytest.mark.parametrize("family,degree", [("CG", 2), ("RT", 1)]) +@pytest.mark.parallel(nprocs=[1, 3]) +def test_submesh_assign_function_redistributed_subdomain(family, degree): + mesh = UnitSquareMesh(4, 4) + x, y = SpatialCoordinate(mesh) + DG0 = FunctionSpace(mesh, "DG", 0) + mesh.mark_entities(Function(DG0).interpolate(conditional(x < 0.5, 1, 0)), 111) + submesh = Submesh(mesh, subdomain_id=111, redistribute=True) + + V = FunctionSpace(mesh, family, degree) + V_sub = FunctionSpace(submesh, family, degree) + x_sub, y_sub = SpatialCoordinate(submesh) + if V.value_size == 1: + expr, expr_sub = sin(x) * cos(3 * y), sin(x_sub) * cos(3 * y_sub) + else: + expr = as_vector([sin(x), cos(3 * y)]) + expr_sub = as_vector([sin(x_sub), cos(3 * y_sub)]) + f = Function(V).interpolate(expr) + # -- mesh -> submesh + f_sub = Function(V_sub).assign(f) + assert np.allclose(norm(f_sub - Function(V_sub).interpolate(expr_sub)), 0) + # -- submesh -> mesh: the parent has nodes outside the subdomain + with pytest.raises(ValueError): + Function(V).assign(f_sub) + g = Function(V).interpolate(expr) + g.assign(f_sub, allow_missing_dofs=True) + assert np.allclose(norm(g - f), 0) diff --git a/tests/firedrake/submesh/test_submesh_basics.py b/tests/firedrake/submesh/test_submesh_basics.py index ccd2500e88..cea5acaef4 100644 --- a/tests/firedrake/submesh/test_submesh_basics.py +++ b/tests/firedrake/submesh/test_submesh_basics.py @@ -1,4 +1,11 @@ +import os +import pytest +import numpy as np from firedrake import * +from petsc4py import PETSc + + +cwd = os.path.abspath(os.path.dirname(__file__)) def test_submesh_parent(): @@ -14,3 +21,67 @@ def test_submesh_parent(): submesh = Submesh(parent, parent.topological_dimension, cell_marker) assert submesh.topology.submesh_parent is parent.topology assert submesh.submesh_parent is parent + + +def test_submesh_redistribute_codim(): + # The entities of a submesh of non-zero co-dimension are not the entities + # of its parent, so they can not inherit its orientations. + mesh = UnitSquareMesh(2, 2) + with pytest.raises(NotImplementedError): + Submesh(mesh, subdomain_id="on_boundary", redistribute=True) + + +def _curved_mesh(nx=4, degree=2): + """Build a unit square whose curved coordinates the plex can not carry.""" + mesh = UnitSquareMesh(nx, nx) + V = VectorFunctionSpace(mesh, "CG", degree) + x, y = SpatialCoordinate(mesh) + coordinates = Function(V).interpolate(as_vector([x + 0.15 * sin(pi * y) * x * (1 - x), y])) + return Mesh(coordinates) + + +@pytest.mark.parallel([1, 2, 3]) +def test_submesh_curved_coordinates(): + # The plex carries no curved coordinates, so a submesh of the same + # dimension must take them from its parent. + mesh = _curved_mesh() + x, _ = SpatialCoordinate(mesh) + DG0 = FunctionSpace(mesh, "DG", 0) + mesh.mark_entities(Function(DG0).interpolate(conditional(x > 0.5, 1, 0)), 77) + + submesh = Submesh(mesh, mesh.topological_dimension, 77) + assert submesh.coordinates.ufl_element() == mesh.coordinates.ufl_element() + area = assemble(Constant(1.0) * dx(submesh)) + assert np.isclose(area, assemble(conditional(x > 0.5, 1.0, 0.0) * dx(mesh))) + # The same region of an affine parent would measure exactly one half. + assert not np.isclose(area, 0.5) + + +@pytest.mark.parallel([1, 2, 3]) +def test_submesh_affine_coordinates(): + # An affine parent and a submesh of lower dimension carry the same element + # up to their cell. The plex therefore already holds the right coordinates. + mesh = UnitSquareMesh(3, 3) + submesh = Submesh(mesh, 1, "on_boundary", label_name="exterior_facets") + assert submesh.coordinates.ufl_element() == \ + mesh.coordinates.ufl_element().reconstruct(cell=submesh.ufl_cell()) + assert np.isclose(assemble(Constant(1.0) * dx(submesh)), 4.0) + + +@pytest.mark.parallel([1, 2, 3]) +def test_submesh_curved_codim(): + # A parent and a submesh of lower dimension share only some of the nodes + # on a parent cell. The cell node maps can not select those nodes. + mesh = _curved_mesh() + with pytest.raises(NotImplementedError): + Submesh(mesh, 1, "on_boundary", label_name="exterior_facets") + + +@pytest.mark.parallel([1, 2]) +def test_submesh_multiple_cell_types_coordinates(): + # A mesh with several cell types carries no coordinate Function, so a + # submesh of it takes the coordinates the plex holds. + mesh = Mesh(os.path.join(cwd, "..", "meshes", "mixed_cell_unit_square.msh")) + submesh = Submesh(mesh, mesh.topological_dimension, + PETSc.DM.PolytopeType.TRIANGLE, label_name="celltype") + assert submesh.ufl_cell().cellname == "triangle"