From 4f7bc427f7370a2f2ec6b9742e8eec405d65fce6 Mon Sep 17 00:00:00 2001 From: Matt Smith Date: Thu, 30 Jun 2022 15:00:41 -0500 Subject: [PATCH] Recast `partition_mesh` in terms of generic part identifiers (#308) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * simplify partition_mesh interface * add get_connected_partitions back to docs * use opaque partition ID instead of number in partition_mesh * use self/other instead of local/nonlocal * fix compatibility * remove unnecessary import Co-authored-by: Alex Fikl * eliminate fun Co-authored-by: Alex Fikl * stamp out remaining traces of fun * rename membership_list_to_sets -> membership_list_to_map and store index sets as numpy arrays instead of python sets * remove return_sets option from get_partition_by_pymetis * fix bugs in MPIBoundaryCommSetupHelper * flake8 * fix bug * add a couple of fixmes * handle groupless mesh case in dim/ambient_dim * Revert "handle groupless mesh case in dim/ambient_dim" not a good solution * disable removal of empty mesh groups in partitioning * fix some issues with partitioning docs * remove empty-group filtering * clarify part vs. partition terminology * change some asserts to exceptions * add a couple of FIXMEs * cosmetic change Co-authored-by: Andreas Klöckner * detect elements that belong to multiple parts * add return_parts argument to partition_mesh * add type hints to mesh partitioning * Fix some annotations in mesh.processing * add explicit part_id attribute to InterPartAdjacencyGroup * cosmetic change * fix some type annotations * Dict -> Mapping * fix outdated argument reference in docstring * try using just dtype dtype[Any] requires python 3.9 Co-authored-by: Alex Fikl Co-authored-by: Andreas Klöckner --- .../connection/opposite_face.py | 14 +- meshmode/distributed.py | 170 +++--- meshmode/mesh/__init__.py | 56 +- meshmode/mesh/processing.py | 488 ++++++++++-------- test/test_partition.py | 100 ++-- 5 files changed, 447 insertions(+), 381 deletions(-) diff --git a/meshmode/discretization/connection/opposite_face.py b/meshmode/discretization/connection/opposite_face.py index 25bb669c7..19051ac8a 100644 --- a/meshmode/discretization/connection/opposite_face.py +++ b/meshmode/discretization/connection/opposite_face.py @@ -521,20 +521,20 @@ def make_opposite_face_connection(actx, volume_to_bdry_conn): # {{{ make_partition_connection +# FIXME: Consider adjusting terminology from local/remote to self/other. def make_partition_connection(actx, *, local_bdry_conn, remote_bdry_discr, remote_group_infos): """ - Connects ``local_bdry_conn`` to a neighboring partition. + Connects ``local_bdry_conn`` to a neighboring part. - :arg local_bdry_conn: A :class:`DiscretizationConnection` of the local - partition. + :arg local_bdry_conn: A :class:`DiscretizationConnection` of the local part. :arg remote_bdry_discr: A :class:`~meshmode.discretization.Discretization` - of the boundary of the remote partition. + of the boundary of the remote part. :arg remote_group_infos: An array of :class:`meshmode.distributed.RemoteGroupInfo` instances, one per remote volume element group. :returns: A :class:`DirectDiscretizationConnection` that performs data - exchange across faces from the remote partition to the local partition. + exchange across faces from the remote part to the local part. .. versionadded:: 2017.1 @@ -556,7 +556,7 @@ def make_partition_connection(actx, *, local_bdry_conn, # The code assumes that there is the same number of volume and surface groups. # # A weak reason to choose remote as the outer loop is because - # InterPartitionAdjacency refers to neighbors by global volume element + # InterPartAdjacency refers to neighbors by global volume element # numbers, and we only have enough information to resolve those to (group, # group_local_el_nr) for local elements (whereas we have no information # about remote volume elements). @@ -564,7 +564,7 @@ def make_partition_connection(actx, *, local_bdry_conn, # (See the find_group_indices below.) for rgi in remote_group_infos: - rem_ipags = rgi.inter_partition_adj_groups + rem_ipags = rgi.inter_part_adj_groups for rem_ipag in rem_ipags: i_local_grps = find_group_indices(local_vol_groups, rem_ipag.neighbors) diff --git a/meshmode/distributed.py b/meshmode/distributed.py index 5d5173014..d99e90dc8 100644 --- a/meshmode/distributed.py +++ b/meshmode/distributed.py @@ -4,7 +4,8 @@ .. autoclass:: MPIBoundaryCommSetupHelper .. autofunction:: get_partition_by_pymetis -.. autofunction:: get_inter_partition_tags +.. autofunction:: membership_list_to_map +.. autofunction:: get_connected_parts .. autoclass:: RemoteGroupInfo .. autoclass:: make_remote_group_infos @@ -37,7 +38,7 @@ from dataclasses import dataclass import numpy as np -from typing import List, Set, Union, Mapping, cast, Sequence, TYPE_CHECKING, Hashable +from typing import List, Set, Union, Mapping, cast, Sequence, TYPE_CHECKING from arraycontext import ArrayContext from meshmode.discretization.connection import ( @@ -46,9 +47,8 @@ from meshmode.mesh import ( Mesh, InteriorAdjacencyGroup, - InterPartitionAdjacencyGroup, - BoundaryTag, - BTAG_PARTITION, + InterPartAdjacencyGroup, + PartID, ) from meshmode.discretization import ElementGroupFactory @@ -94,10 +94,10 @@ def send_mesh_parts(self, mesh, part_per_element, num_parts): :arg part_per_element: A :class:`numpy.ndarray` containing one integer per element of *mesh* indicating which part of the partitioned mesh the element is to become a part of. - :arg num_parts: The number of partitions to divide the mesh into. + :arg num_parts: The number of parts to divide the mesh into. - Sends each partition to a different rank. - Returns one partition that was not sent to any other rank. + Sends each part to a different rank. + Returns one part that was not sent to any other rank. """ mpi_comm = self.mpi_comm rank = mpi_comm.Get_rank() @@ -105,20 +105,21 @@ def send_mesh_parts(self, mesh, part_per_element, num_parts): assert self.is_mananger_rank() + part_num_to_elements = membership_list_to_map(part_per_element) + from meshmode.mesh.processing import partition_mesh - parts = [partition_mesh(mesh, part_per_element, i)[0] - for i in range(num_parts)] + parts = partition_mesh(mesh, part_num_to_elements) local_part = None reqs = [] - for r, part in enumerate(parts): + for r, part in parts.items(): if r == self.manager_rank: local_part = part else: reqs.append(mpi_comm.isend(part, dest=r, tag=TAG_DISTRIBUTE_MESHES)) - logger.info("rank %d: sent all mesh partitions", rank) + logger.info("rank %d: sent all mesh parts", rank) for req in reqs: req.wait() @@ -147,16 +148,20 @@ def receive_mesh_part(self): # {{{ remote group info +# FIXME: "Remote" is perhaps not the best naming convention for this. For example, +# in a multi-volume context it may be used when constructing inter-part connections +# between two parts on the same rank. @dataclass class RemoteGroupInfo: - inter_partition_adj_groups: List[InterPartitionAdjacencyGroup] + inter_part_adj_groups: List[InterPartAdjacencyGroup] vol_elem_indices: np.ndarray bdry_elem_indices: np.ndarray bdry_faces: np.ndarray def make_remote_group_infos( - actx: ArrayContext, local_btag: BoundaryTag, + actx: ArrayContext, + remote_part_id: PartID, bdry_conn: DirectDiscretizationConnection ) -> Sequence[RemoteGroupInfo]: local_vol_mesh = bdry_conn.from_discr.mesh @@ -165,10 +170,10 @@ def make_remote_group_infos( return [ RemoteGroupInfo( - inter_partition_adj_groups=[ + inter_part_adj_groups=[ fagrp for fagrp in local_vol_mesh.facial_adjacency_groups[igrp] - if isinstance(fagrp, InterPartitionAdjacencyGroup) - and fagrp.boundary_tag == local_btag], + if isinstance(fagrp, InterPartAdjacencyGroup) + and fagrp.part_id == remote_part_id], vol_elem_indices=np.concatenate([ actx.to_numpy(batch.from_element_indices) for batch in bdry_conn.groups[igrp].batches]), @@ -188,17 +193,13 @@ def make_remote_group_infos( @dataclass(init=True, frozen=True) class InterRankBoundaryInfo: """ - .. attribute:: local_btag - - A boundary tag for the local boundary towards the remote partition. - .. attribute:: local_part_id - An opaque, hashable, picklable identifier for the local partition. + An opaque, hashable, picklable identifier for the local part. .. attribute:: remote_part_id - An opaque, hashable, picklable identifier for the remote partition. + An opaque, hashable, picklable identifier for the remote part. .. attribute:: remote_rank @@ -207,22 +208,21 @@ class InterRankBoundaryInfo: .. attribute:: local_boundary_connection A :class:`~meshmode.discretization.connection.DirectDiscretizationConnection` - from the volume onto the boundary described by :attr:`local_btag`. + from the volume onto the boundary described by + ``BTAG_PARTITION(remote_part_id)``. .. automethod:: __init__ """ - # FIXME better names? - local_btag: BoundaryTag - local_part_id: Hashable - remote_part_id: Hashable + local_part_id: PartID + remote_part_id: PartID remote_rank: int local_boundary_connection: DirectDiscretizationConnection class MPIBoundaryCommSetupHelper: """ - Helper for setting up inter-partition facial data exchange. + Helper for setting up inter-part facial data exchange. .. automethod:: __init__ .. automethod:: __enter__ @@ -240,11 +240,11 @@ def __init__(self, ], bdry_grp_factory: ElementGroupFactory): """ - :arg local_bdry_conns: A :class:`dict` mapping remote partition to + :arg local_bdry_conns: A :class:`dict` mapping remote part to `local_bdry_conn`, where `local_bdry_conn` is a :class:`~meshmode.discretization.connection.DirectDiscretizationConnection` - that performs data exchange from - the volume to the faces adjacent to partition `i_remote_part`. + that performs data exchange from the volume to the faces adjacent to + part `i_remote_part`. :arg bdry_grp_factory: Group factory to use when creating the remote-to-local boundary connections """ @@ -265,9 +265,8 @@ def __init__(self, inter_rank_bdry_info = [ InterRankBoundaryInfo( - local_btag=BTAG_PARTITION(remote_rank), - local_part_id=remote_rank, - remote_part_id=self.i_local_rank, + local_part_id=self.i_local_rank, + remote_part_id=remote_rank, remote_rank=remote_rank, local_boundary_connection=conn ) @@ -292,7 +291,7 @@ def __enter__(self): # to know when we're done self.pending_recv_identifiers = { - (irbi.remote_rank, irbi.remote_part_id) + (irbi.local_part_id, irbi.remote_part_id) for irbi in self.inter_rank_bdry_info} self.send_reqs = [ @@ -302,7 +301,7 @@ def __enter__(self): irbi.remote_part_id, irbi.local_boundary_connection.to_discr.mesh, make_remote_group_infos( - self.array_context, irbi.local_btag, + self.array_context, irbi.remote_part_id, irbi.local_boundary_connection)), dest=irbi.remote_rank) for irbi in self.inter_rank_bdry_info] @@ -314,13 +313,12 @@ def __exit__(self, type, value, traceback): def complete_some(self): """ - Returns a :class:`dict` mapping a subset of remote partitions to + Returns a :class:`dict` mapping a subset of remote parts to remote-to-local boundary connections, where a remote-to-local boundary connection is a :class:`~meshmode.discretization.connection.DirectDiscretizationConnection` - that performs data exchange across faces from partition `i_remote_part` - to the local mesh. When an empty dictionary is returned, setup is - complete. + that performs data exchange across faces from part `i_remote_part` to the + local mesh. When an empty dictionary is returned, setup is complete. """ from mpi4py import MPI @@ -341,36 +339,41 @@ def complete_some(self): remote_to_local_bdry_conns = {} - local_part_id_to_irbi = { - irbi.local_part_id: irbi for irbi in self.inter_rank_bdry_info} - assert len(local_part_id_to_irbi) == len(self.inter_rank_bdry_info) + part_ids_to_irbi = { + (irbi.local_part_id, irbi.remote_part_id): irbi + for irbi in self.inter_rank_bdry_info} + if len(part_ids_to_irbi) < len(self.inter_rank_bdry_info): + raise ValueError( + "duplicate local/remote part pair in inter_rank_bdry_info") - for i_src_rank, recvd in zip( - source_ranks, data): - (recvd_remote_part_id, recvd_local_part_id, + for i_src_rank, recvd in zip(source_ranks, data): + (remote_part_id, local_part_id, remote_bdry_mesh, remote_group_infos) = recvd logger.debug("rank %d: Received part id '%s' data from rank %d", - self.i_local_rank, recvd_local_part_id, i_src_rank) + self.i_local_rank, remote_part_id, i_src_rank) # Connect local_mesh to remote_mesh from meshmode.discretization.connection import make_partition_connection - irbi = local_part_id_to_irbi[recvd_local_part_id] + irbi = part_ids_to_irbi[local_part_id, remote_part_id] assert i_src_rank == irbi.remote_rank - assert recvd_remote_part_id == irbi.remote_part_id - remote_to_local_bdry_conns[recvd_local_part_id] \ - = make_partition_connection( - self.array_context, - local_bdry_conn=irbi.local_boundary_connection, - remote_bdry_discr=( - irbi.local_boundary_connection.to_discr.copy( - actx=self.array_context, - mesh=remote_bdry_mesh, - group_factory=self.bdry_grp_factory)), - remote_group_infos=remote_group_infos) + if self._using_old_timey_interface: + key = remote_part_id + else: + key = (remote_part_id, local_part_id) + + remote_to_local_bdry_conns[key] = ( + make_partition_connection( + self.array_context, + local_bdry_conn=irbi.local_boundary_connection, + remote_bdry_discr=irbi.local_boundary_connection.to_discr.copy( + actx=self.array_context, + mesh=remote_bdry_mesh, + group_factory=self.bdry_grp_factory), + remote_group_infos=remote_group_infos)) - self.pending_recv_identifiers.remove((i_src_rank, recvd_remote_part_id)) + self.pending_recv_identifiers.remove((local_part_id, remote_part_id)) if not self.pending_recv_identifiers: MPI.Request.waitall(self.send_reqs) @@ -381,6 +384,7 @@ def complete_some(self): # }}} +# FIXME: Move somewhere else, since it's not strictly limited to distributed? def get_partition_by_pymetis(mesh, num_parts, *, connectivity="facial", **kwargs): """Return a mesh partition created by :mod:`pymetis`. @@ -390,7 +394,7 @@ def get_partition_by_pymetis(mesh, num_parts, *, connectivity="facial", **kwargs ``"facial"`` or ``"nodal"`` (based on vertices). :arg kwargs: Passed unmodified to :func:`pymetis.part_graph`. :returns: a :class:`numpy.ndarray` with one entry per element indicating - to which partition each element belongs, with entries between ``0`` and + to which part each element belongs, with entries between ``0`` and ``num_parts-1``. .. versionchanged:: 2020.2 @@ -429,39 +433,33 @@ def get_partition_by_pymetis(mesh, num_parts, *, connectivity="facial", **kwargs return np.array(p) -def get_connected_partitions(mesh: Mesh) -> "Set[int]": - """For a local mesh part in *mesh*, determine the set of boundary - tags for connections to other parts, cf. - :class:`meshmode.mesh.InterPartitionAdjacencyGroup`. +def membership_list_to_map(membership_list): + """ + Convert a :class:`numpy.ndarray` that maps an index to a key into a + :class:`dict` that maps a key to a set of indices (with each set of indices + stored as a sorted :class:`numpy.ndarray`). """ - assert mesh.facial_adjacency_groups is not None - # internal and deprecated, remove in July 2022 - - def _get_neighbor_part_nr(btag): - if isinstance(btag, BTAG_PARTITION): - return btag.part_nr - else: - raise ValueError("unexpected inter-partition boundary tag type found") - return { - _get_neighbor_part_nr(grp.boundary_tag) - for fagrp_list in mesh.facial_adjacency_groups - for grp in fagrp_list - if isinstance(grp, InterPartitionAdjacencyGroup)} + entry: np.where(membership_list == entry)[0] + for entry in set(membership_list)} -def get_inter_partition_tags(mesh: Mesh) -> "Set[BoundaryTag]": - """For a local mesh part in *mesh*, determine the set of boundary - tags for connections to other parts, cf. - :class:`meshmode.mesh.InterPartitionAdjacencyGroup`. - """ +# FIXME: Move somewhere else, since it's not strictly limited to distributed? +def get_connected_parts(mesh: Mesh) -> "Set[PartID]": + """For a local mesh part in *mesh*, determine the set of connected parts.""" assert mesh.facial_adjacency_groups is not None return { - grp.boundary_tag + grp.part_id for fagrp_list in mesh.facial_adjacency_groups for grp in fagrp_list - if isinstance(grp, InterPartitionAdjacencyGroup)} + if isinstance(grp, InterPartAdjacencyGroup)} + +def get_connected_partitions(mesh: Mesh) -> "Set[PartID]": + warn( + "get_connected_partitions is deprecated and will stop working in June 2023. " + "Use get_connected_parts instead.", DeprecationWarning, stacklevel=2) + return get_connected_parts(mesh) # vim: foldmethod=marker diff --git a/meshmode/mesh/__init__.py b/meshmode/mesh/__init__.py index 3fcbd8cc3..69f57e683 100644 --- a/meshmode/mesh/__init__.py +++ b/meshmode/mesh/__init__.py @@ -44,7 +44,7 @@ .. autoclass:: FacialAdjacencyGroup .. autoclass:: InteriorAdjacencyGroup .. autoclass:: BoundaryAdjacencyGroup -.. autoclass:: InterPartitionAdjacencyGroup +.. autoclass:: InterPartAdjacencyGroup .. autofunction:: as_python .. autofunction:: is_true_boundary @@ -67,6 +67,7 @@ # {{{ element tags BoundaryTag = Hashable +PartID = Hashable class BTAG_NONE: # noqa: N801 @@ -109,22 +110,37 @@ class BTAG_NO_BOUNDARY: # noqa: N801 class BTAG_PARTITION(BTAG_NO_BOUNDARY): # noqa: N801 """ A boundary tag indicating that this edge is adjacent to an element of - another :class:`Mesh`. The partition number of the adjacent mesh - is given by ``part_nr``. + another :class:`Mesh`. The part identifier of the adjacent mesh is given + by ``part_id``. - .. attribute:: part_nr + .. attribute:: part_id .. versionadded:: 2017.1 """ - def __init__(self, part_nr): - self.part_nr = int(part_nr) + def __init__(self, part_id: PartID, part_nr=None): + if part_nr is not None: + from warnings import warn + warn("part_nr is deprecated and will stop working in March 2023. " + "Use part_id instead.", + DeprecationWarning, stacklevel=2) + self.part_id = int(part_nr) + else: + self.part_id = part_id + + @property + def part_nr(self): + from warnings import warn + warn("part_nr is deprecated and will stop working in March 2023. " + "Use part_id instead.", + DeprecationWarning, stacklevel=2) + return self.part_id def __hash__(self): - return hash((type(self), self.part_nr)) + return hash((type(self), self.part_id)) def __eq__(self, other): if isinstance(other, BTAG_PARTITION): - return self.part_nr == other.part_nr + return self.part_id == other.part_id else: return False @@ -132,10 +148,10 @@ def __ne__(self, other): return not self.__eq__(other) def __repr__(self): - return "<{}({})>".format(type(self).__name__, repr(self.part_nr)) + return "<{}({})>".format(type(self).__name__, repr(self.part_id)) def as_python(self): - return f"{self.__class__.__name__}({self.part_nr})" + return f"{self.__class__.__name__}({self.part_id})" class BTAG_INDUCED_BOUNDARY(BTAG_NO_BOUNDARY): # noqa: N801 @@ -754,9 +770,9 @@ def as_python(self): # {{{ partition adjacency @dataclass(frozen=True, eq=False) -class InterPartitionAdjacencyGroup(BoundaryAdjacencyGroup): +class InterPartAdjacencyGroup(BoundaryAdjacencyGroup): """ - Describes inter-partition adjacency information for one + Describes inter-part adjacency information for one :class:`MeshElementGroup`. .. attribute:: igroup @@ -768,6 +784,10 @@ class InterPartitionAdjacencyGroup(BoundaryAdjacencyGroup): The boundary tag identifier of this group. Will be an instance of :class:`~meshmode.mesh.BTAG_PARTITION`. + .. attribute:: part_id + + The identifier of the neighboring part. + .. attribute:: elements Group-local element numbers. @@ -784,7 +804,7 @@ class InterPartitionAdjacencyGroup(BoundaryAdjacencyGroup): .. attribute:: neighbors ``element_id_dtype neighbors[i]`` gives the volume element number - within the neighboring partition of the element connected to + within the neighboring part of the element connected to ``element_id_dtype elements[i]`` (which is a boundary element index). Use `~meshmode.mesh.processing.find_group_indices` to find the group that the element belongs to, then subtract ``element_nr_base`` to find the @@ -793,8 +813,7 @@ class InterPartitionAdjacencyGroup(BoundaryAdjacencyGroup): .. attribute:: neighbor_faces ``face_id_dtype global_neighbor_faces[i]`` gives face index within the - neighboring partition of the face connected to - ``element_id_dtype elements[i]`` + neighboring part of the face connected to ``element_id_dtype elements[i]`` .. attribute:: aff_map @@ -804,6 +823,7 @@ class InterPartitionAdjacencyGroup(BoundaryAdjacencyGroup): .. versionadded:: 2017.1 """ + part_id: PartID neighbors: np.ndarray neighbor_faces: np.ndarray aff_map: AffineMap @@ -811,17 +831,19 @@ class InterPartitionAdjacencyGroup(BoundaryAdjacencyGroup): def __eq__(self, other): return ( super().__eq__(other) + and np.array_equal(self.part_id, other.part_id) and np.array_equal(self.neighbors, other.neighbors) and np.array_equal(self.neighbor_faces, other.neighbor_faces) and self.aff_map == other.aff_map) def as_python(self): - if type(self) != InterPartitionAdjacencyGroup: + if type(self) != InterPartAdjacencyGroup: raise NotImplementedError(f"Not implemented for {type(self)}.") return self._as_python( igroup=self.igroup, boundary_tag=_boundary_tag_as_python(self.boundary_tag), + part_id=self.part_id, elements=_numpy_array_as_python(self.elements), element_faces=_numpy_array_as_python(self.element_faces), neighbors=_numpy_array_as_python(self.neighbors), @@ -1631,7 +1653,7 @@ def as_python(mesh, function_name="make_mesh"): FacialAdjacencyGroup, InteriorAdjacencyGroup, BoundaryAdjacencyGroup, - InterPartitionAdjacencyGroup, + InterPartAdjacencyGroup, BTAG_NONE, BTAG_ALL, BTAG_REALLY_ALL) diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py index b8ad6554f..cc0503741 100644 --- a/meshmode/mesh/processing.py +++ b/meshmode/mesh/processing.py @@ -22,7 +22,7 @@ from functools import reduce from numbers import Real -from typing import Optional, Union +from typing import Optional, Union, Tuple, Mapping, List, Set, Sequence from dataclasses import dataclass @@ -32,10 +32,14 @@ import modepy as mp from meshmode.mesh import ( + MeshElementGroup, + Mesh, BTAG_PARTITION, + PartID, + FacialAdjacencyGroup, InteriorAdjacencyGroup, BoundaryAdjacencyGroup, - InterPartitionAdjacencyGroup + InterPartAdjacencyGroup ) from meshmode.mesh.tools import AffineMap @@ -80,83 +84,80 @@ def find_group_indices(groups, meshwide_elems): # {{{ partition_mesh -def _compute_global_elem_to_part_elem(part_per_element, parts, element_id_dtype): +def _compute_global_elem_to_part_elem( + nelements: int, + part_id_to_elements: Mapping[PartID, np.ndarray], + part_id_to_part_index: Mapping[PartID, int], + element_id_dtype: np.dtype) -> np.ndarray: """ - Create a map from global element index to partition-wide element index for - a set of partitions. - - :arg part_per_element: A :class:`numpy.ndarray` mapping element indices to - partition numbers. - :arg parts: A :class:`set` of partition numbers. + Create a map from global element index to part-wide element index for a set of + parts. + + :arg nelements: The number of elements in the global mesh. + :arg part_id_to_elements: A :class:`dict` mapping a part identifier to + a sorted :class:`numpy.ndarray` of elements. + :arg part_id_to_part_index: A mapping from part identifiers to indices in + the range ``[0, num_parts)``. :arg element_id_dtype: The element index data type. - :returns: A :class:`numpy.ndarray` that maps an element's global index - to its corresponding partition-wide index if that partition belongs to - *parts* (and if not, to -1). + :returns: A :class:`numpy.ndarray` ``global_elem_to_part_elem`` of shape + ``(nelements, 2)``, where ``global_elem_to_part_elem[ielement, 0]`` gives + the part index of the element and + ``global_elem_to_part_elem[ielement, 1]`` gives its part-wide element index. """ - global_elem_to_part_elem = np.full(len(part_per_element), -1, - dtype=element_id_dtype) - for ipart in parts: - belongs_to_part = part_per_element == ipart - global_elem_to_part_elem[belongs_to_part] = ( - np.cumsum(belongs_to_part)[belongs_to_part]-1) + global_elem_to_part_elem = np.empty((nelements, 2), dtype=element_id_dtype) + for part_id in part_id_to_elements.keys(): + elements = part_id_to_elements[part_id] + global_elem_to_part_elem[elements, 0] = part_id_to_part_index[part_id] + global_elem_to_part_elem[elements, 1] = np.indices( + (len(elements),), dtype=element_id_dtype) return global_elem_to_part_elem -def _filter_mesh_groups(mesh, selected_elements, vertex_id_dtype): +def _filter_mesh_groups( + mesh: Mesh, + selected_elements: np.ndarray, + vertex_id_dtype: np.dtype) -> Tuple[List[MeshElementGroup], np.ndarray]: """ Create new mesh groups containing a selected subset of elements. - :arg groups: An array of `~meshmode.mesh.ElementGroup` instances. + :arg mesh: A `~meshmode.mesh.Mesh` instance. :arg selected_elements: A sorted array of indices of elements to be included in the filtered groups. - :returns: A tuple ``(new_groups, group_to_new_group, required_vertex_indices)``, - where *new_groups* is made up of groups from *groups* with elements not in - *selected_elements* removed (Note: empty groups are omitted), - *group_to_new_group* maps groups in *groups* to their corresponding location - in *new_groups*, and *required_vertex_indices* contains indices of all - vertices required for elements belonging to *new_groups*. + :arg vertex_id_dtype: The vertex index data type. + :returns: A tuple ``(new_groups, required_vertex_indices)``, where *new_groups* + is made up of groups from *mesh* containing only elements from + *selected_elements* (Note: resulting groups may be empty) and + *required_vertex_indices* contains indices of all vertices required for + elements belonging to *new_groups*. """ - # {{{ find n_new_groups, group_to_new_group, filtered_group_elements + # {{{ find filtered_group_elements group_elem_starts = [ np.searchsorted(selected_elements, base_element_nr) for base_element_nr in mesh.base_element_nrs ] + [len(selected_elements)] - new_group_to_old_group = [] filtered_group_elements = [] for igrp in range(len(mesh.groups)): start_idx, end_idx = group_elem_starts[igrp:igrp+2] - if end_idx == start_idx: - continue - new_group_to_old_group.append(igrp) filtered_group_elements.append( selected_elements[start_idx:end_idx] - mesh.base_element_nrs[igrp]) - n_new_groups = len(new_group_to_old_group) - - group_to_new_group = [None] * len(mesh.groups) - for i_new_grp, i_old_grp in enumerate(new_group_to_old_group): - group_to_new_group[i_old_grp] = i_new_grp - # }}} # {{{ filter vertex indices filtered_vertex_indices = [ - mesh.groups[i_old_grp].vertex_indices[ - filtered_group_elements[i_new_grp], :] - for i_new_grp, i_old_grp in enumerate(new_group_to_old_group) - if mesh.groups[i_old_grp].vertex_indices is not None] - - if n_new_groups > 0: - filtered_vertex_indices_flat = np.concatenate([indices.ravel() for indices - in filtered_vertex_indices]) - else: - filtered_vertex_indices_flat = np.empty(0, dtype=vertex_id_dtype) + grp.vertex_indices[ + filtered_group_elements[igrp], :] + for igrp, grp in enumerate(mesh.groups) + if grp.vertex_indices is not None] + + filtered_vertex_indices_flat = np.concatenate([indices.ravel() for indices + in filtered_vertex_indices]) required_vertex_indices, new_vertex_indices_flat = np.unique( filtered_vertex_indices_flat, return_inverse=True) @@ -173,31 +174,36 @@ def _filter_mesh_groups(mesh, selected_elements, vertex_id_dtype): from dataclasses import replace new_groups = [ - replace(mesh.groups[i_old_grp], - vertex_indices=new_vertex_indices[i_new_grp], - nodes=mesh.groups[i_old_grp].nodes[ - :, filtered_group_elements[i_new_grp], :].copy(), + replace(grp, + vertex_indices=new_vertex_indices[igrp], + nodes=grp.nodes[:, filtered_group_elements[igrp], :].copy(), element_nr_base=None, node_nr_base=None) - for i_new_grp, i_old_grp in enumerate(new_group_to_old_group)] + for igrp, grp in enumerate(mesh.groups)] - return new_groups, group_to_new_group, required_vertex_indices + return new_groups, required_vertex_indices -def _get_connected_partitions( - mesh, part_per_element, global_elem_to_part_elem): +def _get_connected_parts( + mesh: Mesh, + part_id_to_part_index: Mapping[PartID, int], + global_elem_to_part_elem: np.ndarray, + self_part_id: PartID) -> "Set[PartID]": """ - Find the partitions that are connected to the current partition. + Find the parts that are connected to the current part. :arg mesh: A :class:`~meshmode.mesh.Mesh` representing the unpartitioned mesh. - :arg part_per_element: A :class:`numpy.ndarray` mapping element indices to - partition numbers. - :arg global_elem_to_part_elem: A :class:`numpy.ndarray` mapping from global - element index to local partition-wide element index for local elements (and - -1 otherwise). - - :returns: A :class:`set` of indices of the neighboring partitions. + :arg part_id_to_part_index: A mapping from part identifiers to indices in the + range ``[0, num_parts)``. + :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element + indices to part indices and part-wide element indices. See + :func:`_compute_global_elem_to_part_elem`` for details. + :arg self_part_id: The identifier of the part currently being created. + + :returns: A :class:`set` of identifiers of the neighboring parts. """ - connected_parts = set() + self_part_index = part_id_to_part_index[self_part_id] + + connected_part_indices = set() for igrp, facial_adj_list in enumerate(mesh.facial_adjacency_groups): int_grps = [ @@ -209,48 +215,51 @@ def _get_connected_partitions( elem_base_i = mesh.base_element_nrs[igrp] elem_base_j = mesh.base_element_nrs[jgrp] - elements_are_local = global_elem_to_part_elem[facial_adj.elements - + elem_base_i] >= 0 - neighbors_are_nonlocal = global_elem_to_part_elem[facial_adj.neighbors - + elem_base_j] < 0 + elements_are_self = global_elem_to_part_elem[facial_adj.elements + + elem_base_i, 0] == self_part_index + neighbors_are_other = global_elem_to_part_elem[facial_adj.neighbors + + elem_base_j, 0] != self_part_index - connected_parts.update( - part_per_element[ + connected_part_indices.update( + global_elem_to_part_elem[ facial_adj.neighbors[ - elements_are_local & neighbors_are_nonlocal] - + elem_base_j]) + elements_are_self & neighbors_are_other] + + elem_base_j, 0]) - return connected_parts + return { + part_id + for part_id, part_index in part_id_to_part_index.items() + if part_index in connected_part_indices} -def _create_local_to_local_adjacency_groups(mesh, global_elem_to_part_elem, - part_mesh_groups, global_group_to_part_group, - part_mesh_group_elem_base): +def _create_self_to_self_adjacency_groups( + mesh: Mesh, + global_elem_to_part_elem: np.ndarray, + self_part_index: int, + self_mesh_groups: List[MeshElementGroup], + self_mesh_group_elem_base: List[int]) -> List[List[InteriorAdjacencyGroup]]: r""" - Create local-to-local facial adjacency groups for a partitioned mesh. + Create self-to-self facial adjacency groups for a partitioned mesh. :arg mesh: A :class:`~meshmode.mesh.Mesh` representing the unpartitioned mesh. - :arg global_elem_to_part_elem: A :class:`numpy.ndarray` mapping from global - element index to local partition-wide element index for local elements (and - -1 otherwise). - :arg part_mesh_groups: An array of :class:`~meshmode.mesh.ElementGroup` instances - representing the partitioned mesh groups. - :arg global_group_to_part_group: An array mapping groups in *mesh* to groups in - *part_mesh_groups* (or `None` if the group is not local). - :arg part_mesh_group_elem_base: An array containing the starting partition-wide - element index for each group in *part_mesh_groups*. + :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element + indices to part indices and part-wide element indices. See + :func:`_compute_global_elem_to_part_elem`` for details. + :arg self_part_index: The index of the part currently being created, in the + range ``[0, num_parts)``. + :arg self_mesh_groups: A list of :class:`~meshmode.mesh.MeshElementGroup` + instances representing the partitioned mesh groups. + :arg self_mesh_group_elem_base: A list containing the starting part-wide + element index for each group in *self_mesh_groups*. :returns: A list of lists of `~meshmode.mesh.InteriorAdjacencyGroup` instances corresponding to the entries in *mesh.facial_adjacency_groups* that - have local-to-local adjacency. + have self-to-self adjacency. """ - local_to_local_adjacency_groups = [[] for _ in part_mesh_groups] + self_to_self_adjacency_groups: List[List[InteriorAdjacencyGroup]] = [ + [] for _ in self_mesh_groups] for igrp, facial_adj_list in enumerate(mesh.facial_adjacency_groups): - i_part_grp = global_group_to_part_group[igrp] - if i_part_grp is None: - continue - int_grps = [ grp for grp in facial_adj_list if isinstance(grp, InteriorAdjacencyGroup)] @@ -258,79 +267,75 @@ def _create_local_to_local_adjacency_groups(mesh, global_elem_to_part_elem, for facial_adj in int_grps: jgrp = facial_adj.ineighbor_group - j_part_grp = global_group_to_part_group[jgrp] - if j_part_grp is None: - continue - elem_base_i = mesh.base_element_nrs[igrp] elem_base_j = mesh.base_element_nrs[jgrp] - elements_are_local = global_elem_to_part_elem[facial_adj.elements - + elem_base_i] >= 0 - neighbors_are_local = global_elem_to_part_elem[facial_adj.neighbors - + elem_base_j] >= 0 + elements_are_self = global_elem_to_part_elem[facial_adj.elements + + elem_base_i, 0] == self_part_index + neighbors_are_self = global_elem_to_part_elem[facial_adj.neighbors + + elem_base_j, 0] == self_part_index - adj_indices, = np.where(elements_are_local & neighbors_are_local) + adj_indices, = np.where(elements_are_self & neighbors_are_self) if len(adj_indices) > 0: - part_elem_base_i = part_mesh_group_elem_base[i_part_grp] - part_elem_base_j = part_mesh_group_elem_base[j_part_grp] + self_elem_base_i = self_mesh_group_elem_base[igrp] + self_elem_base_j = self_mesh_group_elem_base[jgrp] elements = global_elem_to_part_elem[facial_adj.elements[ - adj_indices] + elem_base_i] - part_elem_base_i + adj_indices] + elem_base_i, 1] - self_elem_base_i element_faces = facial_adj.element_faces[adj_indices] neighbors = global_elem_to_part_elem[facial_adj.neighbors[ - adj_indices] + elem_base_j] - part_elem_base_j + adj_indices] + elem_base_j, 1] - self_elem_base_j neighbor_faces = facial_adj.neighbor_faces[adj_indices] - local_to_local_adjacency_groups[i_part_grp].append( + self_to_self_adjacency_groups[igrp].append( InteriorAdjacencyGroup( - igroup=i_part_grp, - ineighbor_group=j_part_grp, + igroup=igrp, + ineighbor_group=jgrp, elements=elements, element_faces=element_faces, neighbors=neighbors, neighbor_faces=neighbor_faces, aff_map=facial_adj.aff_map)) - return local_to_local_adjacency_groups + return self_to_self_adjacency_groups -def _create_nonlocal_adjacency_groups( - mesh, part_per_element, global_elem_to_part_elem, part_mesh_groups, - global_group_to_part_group, part_mesh_group_elem_base, connected_parts): +def _create_self_to_other_adjacency_groups( + mesh: Mesh, + part_id_to_part_index: Mapping[PartID, int], + global_elem_to_part_elem: np.ndarray, + self_part_id: PartID, + self_mesh_groups: List[MeshElementGroup], + self_mesh_group_elem_base: List[int], + connected_parts: Set[PartID]) -> List[List[InterPartAdjacencyGroup]]: """ - Create non-local adjacency groups for the partitioned mesh. + Create self-to-other adjacency groups for the partitioned mesh. :arg mesh: A :class:`~meshmode.mesh.Mesh` representing the unpartitioned mesh. - :arg part_per_element: A :class:`numpy.ndarray` mapping element indices to - partition numbers. - :arg global_elem_to_part_elem: A :class:`numpy.ndarray` mapping from global - element index to local partition-wide element index for local elements (and - -1 otherwise). - :arg part_mesh_groups: An array of `~meshmode.mesh.ElementGroup` instances + :arg part_id_to_part_index: A mapping from part identifiers to indices in the + range ``[0, num_parts)``. + :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element + indices to part indices and part-wide element indices. See + :func:`_compute_global_elem_to_part_elem`` for details. + :arg self_part_id: The identifier of the part currently being created. + :arg self_mesh_groups: A list of `~meshmode.mesh.MeshElementGroup` instances representing the partitioned mesh groups. - :arg global_group_to_part_group: An array mapping groups in *mesh* to groups in - *part_mesh_groups* (or `None` if the group is not local). - :arg part_mesh_group_elem_base: An array containing the starting partition-wide - element index for each group in *part_mesh_groups*. - :arg connected_parts: A :class:`set` containing the partitions connected to + :arg self_mesh_group_elem_base: A list containing the starting part-wide + element index for each group in *self_mesh_groups*. + :arg connected_parts: A :class:`set` containing the parts connected to the current one. - :returns: A list of lists of `~meshmode.mesh.InterPartitionAdjacencyGroup` - instances corresponding to the entries in *mesh.facial_adjacency_groups* that - have non-local adjacency. + :returns: A list of lists of `~meshmode.mesh.InterPartAdjacencyGroup` instances + corresponding to the entries in *mesh.facial_adjacency_groups* that + have self-to-other adjacency. """ - global_elem_to_neighbor_elem = _compute_global_elem_to_part_elem( - part_per_element, connected_parts, mesh.element_id_dtype) + self_part_index = part_id_to_part_index[self_part_id] - nonlocal_adj_groups = [[] for _ in part_mesh_groups] + self_to_other_adj_groups: List[List[InterPartAdjacencyGroup]] = [ + [] for _ in self_mesh_groups] for igrp, facial_adj_list in enumerate(mesh.facial_adjacency_groups): - i_part_grp = global_group_to_part_group[igrp] - if i_part_grp is None: - continue - int_grps = [ grp for grp in facial_adj_list if isinstance(grp, InteriorAdjacencyGroup)] @@ -344,65 +349,69 @@ def _create_nonlocal_adjacency_groups( global_elements = facial_adj.elements + elem_base_i global_neighbors = facial_adj.neighbors + elem_base_j - elements_are_local = global_elem_to_part_elem[global_elements] >= 0 + elements_are_self = ( + global_elem_to_part_elem[global_elements, 0] == self_part_index) - neighbor_parts = part_per_element[global_neighbors] + neighbor_part_indices = global_elem_to_part_elem[global_neighbors, 0] - for i_neighbor_part in connected_parts: + for neighbor_part_id in connected_parts: + neighbor_part_index = part_id_to_part_index[neighbor_part_id] adj_indices, = np.where( - elements_are_local - & (neighbor_parts == i_neighbor_part)) + elements_are_self + & (neighbor_part_indices == neighbor_part_index)) if len(adj_indices) > 0: - part_elem_base_i = part_mesh_group_elem_base[i_part_grp] + self_elem_base_i = self_mesh_group_elem_base[igrp] elements = global_elem_to_part_elem[facial_adj.elements[ - adj_indices] + elem_base_i] - part_elem_base_i + adj_indices] + elem_base_i, 1] - self_elem_base_i element_faces = facial_adj.element_faces[adj_indices] - neighbors = global_elem_to_neighbor_elem[ - global_neighbors[adj_indices]] + neighbors = global_elem_to_part_elem[ + global_neighbors[adj_indices], 1] neighbor_faces = facial_adj.neighbor_faces[adj_indices] - nonlocal_adj_groups[i_part_grp].append( - InterPartitionAdjacencyGroup( - igroup=i_part_grp, - boundary_tag=BTAG_PARTITION(i_neighbor_part), + self_to_other_adj_groups[igrp].append( + InterPartAdjacencyGroup( + igroup=igrp, + boundary_tag=BTAG_PARTITION(neighbor_part_id), + part_id=neighbor_part_id, elements=elements, element_faces=element_faces, neighbors=neighbors, neighbor_faces=neighbor_faces, aff_map=facial_adj.aff_map)) - return nonlocal_adj_groups + return self_to_other_adj_groups -def _create_boundary_groups(mesh, global_elem_to_part_elem, part_mesh_groups, - global_group_to_part_group, part_mesh_group_elem_base): +def _create_boundary_groups( + mesh: Mesh, + global_elem_to_part_elem: np.ndarray, + self_part_index: PartID, + self_mesh_groups: List[MeshElementGroup], + self_mesh_group_elem_base: List[int]) -> List[List[BoundaryAdjacencyGroup]]: """ Create boundary groups for partitioned mesh. :arg mesh: A :class:`~meshmode.mesh.Mesh` representing the unpartitioned mesh. - :arg global_elem_to_part_elem: A :class:`numpy.ndarray` mapping from global - element index to local partition-wide element index for local elements (and - -1 otherwise). - :arg part_mesh_groups: An array of `~meshmode.mesh.ElementGroup` instances + :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element + indices to part indices and part-wide element indices. See + :func:`_compute_global_elem_to_part_elem`` for details. + :arg self_part_index: The index of the part currently being created, in the + range ``[0, num_parts)``. + :arg self_mesh_groups: A list of `~meshmode.mesh.MeshElementGroup` instances representing the partitioned mesh groups. - :arg global_group_to_part_group: An array mapping groups in *mesh* to groups in - *part_mesh_groups* (or `None` if the group is not local). - :arg part_mesh_group_elem_base: An array containing the starting partition-wide - element index for each group in *part_mesh_groups*. + :arg self_mesh_group_elem_base: A list containing the starting part-wide + element index for each group in *self_mesh_groups*. :returns: A list of lists of `~meshmode.mesh.BoundaryAdjacencyGroup` instances corresponding to the entries in *mesh.facial_adjacency_groups* that have boundary faces. """ - bdry_adj_groups = [[] for _ in part_mesh_groups] + bdry_adj_groups: List[List[BoundaryAdjacencyGroup]] = [ + [] for _ in self_mesh_groups] for igrp, facial_adj_list in enumerate(mesh.facial_adjacency_groups): - i_part_grp = global_group_to_part_group[igrp] - if i_part_grp is None: - continue - bdry_grps = [ grp for grp in facial_adj_list if isinstance(grp, BoundaryAdjacencyGroup)] @@ -410,21 +419,22 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, part_mesh_groups, for bdry_grp in bdry_grps: elem_base = mesh.base_element_nrs[igrp] - adj_indices, = np.where(global_elem_to_part_elem[bdry_grp.elements - + elem_base] >= 0) + adj_indices, = np.where( + global_elem_to_part_elem[bdry_grp.elements + elem_base, 0] + == self_part_index) if len(adj_indices) > 0: - part_elem_base = part_mesh_group_elem_base[i_part_grp] + self_elem_base = self_mesh_group_elem_base[igrp] elements = global_elem_to_part_elem[bdry_grp.elements[adj_indices] - + elem_base] - part_elem_base + + elem_base, 1] - self_elem_base element_faces = bdry_grp.element_faces[adj_indices] else: elements = np.empty(0, dtype=mesh.element_id_dtype) element_faces = np.empty(0, dtype=mesh.face_id_dtype) - bdry_adj_groups[i_part_grp].append( + bdry_adj_groups[igrp].append( BoundaryAdjacencyGroup( - igroup=i_part_grp, + igroup=igrp, boundary_tag=bdry_grp.boundary_tag, elements=elements, element_faces=element_faces)) @@ -432,76 +442,108 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, part_mesh_groups, return bdry_adj_groups -def partition_mesh(mesh, part_per_element, part_num): +def _get_mesh_part( + mesh: Mesh, + part_id_to_elements: Mapping[PartID, np.ndarray], + self_part_id: PartID) -> Mesh: """ :arg mesh: A :class:`~meshmode.mesh.Mesh` to be partitioned. - :arg part_per_element: A :class:`numpy.ndarray` containing one - integer per element of *mesh* indicating which part of the - partitioned mesh the element is to become a part of. - :arg part_num: The part number of the mesh to return. + :arg part_id_to_elements: A :class:`dict` mapping a part identifier to + a sorted :class:`numpy.ndarray` of elements. + :arg self_part_id: The part identifier of the mesh to return. - :returns: A tuple ``(part_mesh, part_to_global)``, where *part_mesh* - is a :class:`~meshmode.mesh.Mesh` that is a partition of mesh, and - *part_to_global* is a :class:`numpy.ndarray` mapping element - numbers on *part_mesh* to ones in *mesh*. + :returns: A :class:`~meshmode.mesh.Mesh` containing a part of *mesh*. .. versionadded:: 2017.1 """ - assert len(part_per_element) == mesh.nelements, ( - "part_per_element must have shape (mesh.nelements,)") - - # Contains the indices of the elements requested. - queried_elems, = np.where(part_per_element == part_num) - - global_elem_to_part_elem = _compute_global_elem_to_part_elem(part_per_element, - {part_num}, mesh.element_id_dtype) + element_counts = np.zeros(mesh.nelements) + for elements in part_id_to_elements.values(): + element_counts[elements] += 1 + if np.any(element_counts > 1): + raise ValueError("elements cannot belong to multiple parts") + if np.any(element_counts < 1): + raise ValueError("partition must contain all elements") + + part_id_to_part_index = { + part_id: part_index + for part_index, part_id in enumerate(part_id_to_elements.keys())} + + global_elem_to_part_elem = _compute_global_elem_to_part_elem( + mesh.nelements, part_id_to_elements, part_id_to_part_index, + mesh.element_id_dtype) # Create new mesh groups that mimic the original mesh's groups but only contain - # the local partition's elements - part_mesh_groups, global_group_to_part_group, required_vertex_indices =\ - _filter_mesh_groups(mesh, queried_elems, mesh.vertex_id_dtype) + # the current part's elements + self_mesh_groups, required_vertex_indices = _filter_mesh_groups( + mesh, part_id_to_elements[self_part_id], mesh.vertex_id_dtype) - part_vertices = np.zeros((mesh.ambient_dim, len(required_vertex_indices))) + self_part_index = part_id_to_part_index[self_part_id] + + self_vertices = np.zeros((mesh.ambient_dim, len(required_vertex_indices))) for dim in range(mesh.ambient_dim): - part_vertices[dim] = mesh.vertices[dim][required_vertex_indices] + self_vertices[dim] = mesh.vertices[dim][required_vertex_indices] - part_mesh_group_elem_base = [0 for _ in part_mesh_groups] + self_mesh_group_elem_base = [0 for _ in self_mesh_groups] el_nr = 0 - for i_part_grp, grp in enumerate(part_mesh_groups): - part_mesh_group_elem_base[i_part_grp] = el_nr + for igrp, grp in enumerate(self_mesh_groups): + self_mesh_group_elem_base[igrp] = el_nr el_nr += grp.nelements - connected_parts = _get_connected_partitions( - mesh, part_per_element, global_elem_to_part_elem) + connected_parts = _get_connected_parts( + mesh, part_id_to_part_index, global_elem_to_part_elem, + self_part_id) - local_to_local_adj_groups = _create_local_to_local_adjacency_groups( - mesh, global_elem_to_part_elem, part_mesh_groups, - global_group_to_part_group, part_mesh_group_elem_base) + self_to_self_adj_groups = _create_self_to_self_adjacency_groups( + mesh, global_elem_to_part_elem, self_part_index, self_mesh_groups, + self_mesh_group_elem_base) - nonlocal_adj_groups = _create_nonlocal_adjacency_groups( - mesh, part_per_element, global_elem_to_part_elem, - part_mesh_groups, global_group_to_part_group, - part_mesh_group_elem_base, connected_parts) + self_to_other_adj_groups = _create_self_to_other_adjacency_groups( + mesh, part_id_to_part_index, global_elem_to_part_elem, self_part_id, + self_mesh_groups, self_mesh_group_elem_base, connected_parts) boundary_adj_groups = _create_boundary_groups( - mesh, global_elem_to_part_elem, part_mesh_groups, - global_group_to_part_group, part_mesh_group_elem_base) - - # Combine local/nonlocal/boundary adjacency groups - part_facial_adj_groups = [ - local_to_local_adj_groups[i_part_grp] - + nonlocal_adj_groups[i_part_grp] - + boundary_adj_groups[i_part_grp] - for i_part_grp in range(len(part_mesh_groups))] - - from meshmode.mesh import Mesh - part_mesh = Mesh( - part_vertices, - part_mesh_groups, - facial_adjacency_groups=part_facial_adj_groups, + mesh, global_elem_to_part_elem, self_part_index, self_mesh_groups, + self_mesh_group_elem_base) + + def _gather_grps(igrp: int) -> List[FacialAdjacencyGroup]: + self_grps: Sequence[FacialAdjacencyGroup] = self_to_self_adj_groups[igrp] + other_grps: Sequence[FacialAdjacencyGroup] = self_to_other_adj_groups[igrp] + bdry_grps: Sequence[FacialAdjacencyGroup] = boundary_adj_groups[igrp] + + return list(self_grps) + list(other_grps) + list(bdry_grps) + + # Combine adjacency groups + self_facial_adj_groups = [ + _gather_grps(igrp) for igrp in range(len(self_mesh_groups))] + + return Mesh( + self_vertices, + self_mesh_groups, + facial_adjacency_groups=self_facial_adj_groups, is_conforming=mesh.is_conforming) - return part_mesh, queried_elems + +def partition_mesh( + mesh: Mesh, + part_id_to_elements: Mapping[PartID, np.ndarray], + return_parts: Optional[Sequence[PartID]] = None) -> "Mapping[PartID, Mesh]": + """ + :arg mesh: A :class:`~meshmode.mesh.Mesh` to be partitioned. + :arg part_id_to_elements: A :class:`dict` mapping a part identifier to + a sorted :class:`numpy.ndarray` of elements. + :arg return_parts: An optional list of parts to return. By default, returns all + parts. + + :returns: A :class:`dict` mapping part identifiers to instances of + :class:`~meshmode.mesh.Mesh` that represent the corresponding part of + *mesh*. + """ + if return_parts is None: + return_parts = list(part_id_to_elements.keys()) + + return { + part_id: _get_mesh_part(mesh, part_id_to_elements, part_id) + for part_id in return_parts} # }}} @@ -614,7 +656,7 @@ def get_simplex_element_flip_matrix(order, unit_nodes, permutation=None): first two barycentric coordinates. :arg order: The order of the function space on the simplex, - (see second argument in :fun:`modepy.simplex_best_available_basis`). + (see second argument in :func:`modepy.simplex_best_available_basis`). :arg unit_nodes: A np array of unit nodes with shape *(dim, nunit_nodes)*. :arg permutation: Either *None*, or a tuple of shape storing a permutation: the *i*th barycentric coordinate gets mapped to the *permutation[i]*th @@ -695,8 +737,6 @@ def perform_flips(mesh, flip_flags, skip_tests=False): flip_flags = flip_flags.astype(bool) - from meshmode.mesh import Mesh - new_groups = [] for base_element_nr, grp in zip(mesh.base_element_nrs, mesh.groups): grp_flip_flags = flip_flags[base_element_nr:base_element_nr + grp.nelements] @@ -834,7 +874,6 @@ def merge_disjoint_meshes(meshes, skip_tests=False, single_group=False): # }}} - from meshmode.mesh import Mesh return Mesh(vertices, new_groups, skip_tests=skip_tests, nodal_adjacency=nodal_adjacency, facial_adjacency_groups=facial_adjacency_groups, @@ -892,7 +931,6 @@ def split_mesh_groups(mesh, element_flags, return_subgroup_mapping=False): element_nr_base=None, node_nr_base=None, )) - from meshmode.mesh import Mesh mesh = Mesh( vertices=mesh.vertices, groups=new_groups, @@ -1182,8 +1220,6 @@ def glue_mesh_boundaries(mesh, bdry_pair_mappings_and_tols, *, use_tree=None): _match_boundary_faces(mesh, mapping, tol, use_tree=use_tree) for mapping, tol in bdry_pair_mappings_and_tols] - from meshmode.mesh import InteriorAdjacencyGroup, BoundaryAdjacencyGroup - facial_adjacency_groups = [] for igrp, old_fagrp_list in enumerate(mesh.facial_adjacency_groups): diff --git a/test/test_partition.py b/test/test_partition.py index b2b1c95aa..3cdae413f 100644 --- a/test/test_partition.py +++ b/test/test_partition.py @@ -39,7 +39,7 @@ BTAG_ALL, InteriorAdjacencyGroup, BoundaryAdjacencyGroup, - InterPartitionAdjacencyGroup + InterPartAdjacencyGroup ) import pytest @@ -98,20 +98,22 @@ def f(x): part_per_element = get_partition_by_pymetis(mesh, num_parts, connectivity=part_method) + from meshmode.distributed import membership_list_to_map + part_num_to_elements = membership_list_to_map(part_per_element) + from meshmode.mesh.processing import partition_mesh - part_meshes = [ - partition_mesh(mesh, part_per_element, i)[0] for i in range(num_parts)] + part_meshes = partition_mesh(mesh, part_num_to_elements) connected_parts = set() - for i_local_part, part_mesh in enumerate(part_meshes): - from meshmode.distributed import get_connected_partitions - neighbors = get_connected_partitions(part_mesh) + for i_local_part, part_mesh in part_meshes.items(): + from meshmode.distributed import get_connected_parts + neighbors = get_connected_parts(part_mesh) for i_remote_part in neighbors: connected_parts.add((i_local_part, i_remote_part)) from meshmode.discretization import Discretization - vol_discrs = [Discretization(actx, part_meshes[i], group_factory) - for i in range(num_parts)] + vol_discrs = [Discretization(actx, part_mesh, group_factory) + for part_mesh in part_meshes.values()] from meshmode.mesh import BTAG_PARTITION from meshmode.discretization.connection import (make_face_restriction, @@ -134,7 +136,7 @@ def f(x): remote_bdry_nelements = sum( grp.nelements for grp in remote_bdry_conn.to_discr.groups) assert bdry_nelements == remote_bdry_nelements, \ - "partitions do not have the same number of connected elements" + "parts do not have the same number of connected elements" local_bdry = local_bdry_conn.to_discr @@ -146,7 +148,7 @@ def f(x): local_bdry_conn=local_bdry_conn, remote_bdry_discr=remote_bdry, remote_group_infos=make_remote_group_infos( - actx, BTAG_PARTITION(i_local_part), remote_bdry_conn)) + actx, i_local_part, remote_bdry_conn)) # Connect from local mesh to remote mesh local_to_remote_conn = make_partition_connection( @@ -154,7 +156,7 @@ def f(x): local_bdry_conn=remote_bdry_conn, remote_bdry_discr=local_bdry, remote_group_infos=make_remote_group_infos( - actx, BTAG_PARTITION(i_remote_part), local_bdry_conn)) + actx, i_remote_part, local_bdry_conn)) check_connection(actx, remote_to_local_conn) check_connection(actx, local_to_remote_conn) @@ -193,7 +195,7 @@ def _check_for_cross_rank_adj(mesh, part_per_element): return False -@pytest.mark.parametrize(("dim", "mesh_size", "num_parts", "scramble_partitions"), +@pytest.mark.parametrize(("dim", "mesh_size", "num_parts", "scramble_parts"), [ (2, 4, 4, False), (2, 4, 4, True), @@ -205,7 +207,7 @@ def _check_for_cross_rank_adj(mesh, part_per_element): (3, 7, 32, False), ]) @pytest.mark.parametrize("num_groups", [1, 2, 7]) -def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitions): +def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_parts): np.random.seed(42) nelements_per_axis = (mesh_size,) * dim from meshmode.mesh.generation import generate_regular_rect_mesh @@ -215,7 +217,7 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio from meshmode.mesh.processing import merge_disjoint_meshes mesh = merge_disjoint_meshes(meshes) - if scramble_partitions: + if scramble_parts: part_per_element = np.random.randint(num_parts, size=mesh.nelements) else: pytest.importorskip("pymetis") @@ -227,24 +229,24 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio # adjacency (e.g., when #groups == #parts) has_cross_rank_adj = _check_for_cross_rank_adj(mesh, part_per_element) + from meshmode.distributed import membership_list_to_map + part_num_to_elements = membership_list_to_map(part_per_element) + from meshmode.mesh.processing import partition_mesh - # TODO: The same part_per_element array must be used to partition each mesh. - # Maybe the interface should be changed to guarantee this. - new_meshes = [ - partition_mesh(mesh, part_per_element, i) for i in range(num_parts)] + part_meshes = partition_mesh(mesh, part_num_to_elements) assert mesh.nelements == np.sum( - [new_meshes[i][0].nelements for i in range(num_parts)]), \ + [part_mesh.nelements for part_mesh in part_meshes.values()]), \ "part_mesh has the wrong number of elements" assert count_tags(mesh, BTAG_ALL) == np.sum( - [count_tags(new_meshes[i][0], BTAG_ALL) for i in range(num_parts)]), \ + [count_tags(part_mesh, BTAG_ALL) for part_mesh in part_meshes.values()]), \ "part_mesh has the wrong number of BTAG_ALL boundaries" connected_parts = set() - for i_local_part, (part_mesh, _) in enumerate(new_meshes): - from meshmode.distributed import get_connected_partitions - neighbors = get_connected_partitions(part_mesh) + for i_local_part in range(num_parts): + from meshmode.distributed import get_connected_parts + neighbors = get_connected_parts(part_meshes[i_local_part]) for i_remote_part in neighbors: connected_parts.add((i_local_part, i_remote_part)) @@ -253,11 +255,12 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio num_tags = np.zeros((num_parts,)) index_lookup_table = dict() - for ipart, (m, _) in enumerate(new_meshes): - for igrp in range(len(m.groups)): + for ipart in range(num_parts): + part_mesh = part_meshes[ipart] + for igrp in range(len(part_mesh.groups)): ipagrps = [ - fagrp for fagrp in m.facial_adjacency_groups[igrp] - if isinstance(fagrp, InterPartitionAdjacencyGroup)] + fagrp for fagrp in part_mesh.facial_adjacency_groups[igrp] + if isinstance(fagrp, InterPartAdjacencyGroup)] for ipagrp in ipagrps: for i, (elem, face) in enumerate( zip(ipagrp.elements, ipagrp.element_faces)): @@ -265,15 +268,19 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio ipagrp_count = 0 + part_elem_to_global_elem = { + part_num: np.sort(list(elements)) + for part_num, elements in part_num_to_elements.items()} + for part_num in range(num_parts): - part, part_to_global = new_meshes[part_num] + part = part_meshes[part_num] for grp_num in range(len(part.groups)): ipagrps = [ fagrp for fagrp in part.facial_adjacency_groups[grp_num] - if isinstance(fagrp, InterPartitionAdjacencyGroup)] + if isinstance(fagrp, InterPartAdjacencyGroup)] ipagrp_count += len(ipagrps) for ipagrp in ipagrps: - n_part_num = ipagrp.boundary_tag.part_nr + n_part_num = ipagrp.boundary_tag.part_id num_tags[n_part_num] += len(ipagrp.elements) elem_base = part.base_element_nrs[grp_num] for idx in range(len(ipagrp.elements)): @@ -282,7 +289,7 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio face = ipagrp.element_faces[idx] n_meshwide_elem = ipagrp.neighbors[idx] n_face = ipagrp.neighbor_faces[idx] - n_part, n_part_to_global = new_meshes[n_part_num] + n_part = part_meshes[n_part_num] # Hack: find_igrps expects a numpy.ndarray and returns # a numpy.ndarray. But if a single integer is fed # into find_igrps, an integer is returned. @@ -290,8 +297,8 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio n_part.groups, n_meshwide_elem)) n_ipagrps = [ fagrp for fagrp in n_part.facial_adjacency_groups[n_grp_num] - if isinstance(fagrp, InterPartitionAdjacencyGroup) - and fagrp.boundary_tag.part_nr == part_num] + if isinstance(fagrp, InterPartAdjacencyGroup) + and fagrp.boundary_tag.part_id == part_num] found_reverse_adj = False for n_ipagrp in n_ipagrps: n_elem_base = n_part.base_element_nrs[n_grp_num] @@ -302,10 +309,12 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio meshwide_elem == n_ipagrp.neighbors[n_idx] and face == n_ipagrp.neighbor_faces[n_idx]) if found_reverse_adj: - _, n_part_to_global = new_meshes[n_part_num] - p_meshwide_elem = part_to_global[elem + elem_base] - p_meshwide_n_elem = n_part_to_global[n_meshwide_elem] - assert found_reverse_adj, ("InterPartitionAdjacencyGroup is not " + p_meshwide_elem = ( + part_elem_to_global_elem[part_num][elem + elem_base]) + p_meshwide_n_elem = ( + part_elem_to_global_elem[n_part_num][ + n_meshwide_elem]) + assert found_reverse_adj, ("InterPartAdjacencyGroup is not " "consistent") p_grp_num = find_group_indices(mesh.groups, p_meshwide_elem) @@ -330,13 +339,14 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio "Tag does not give correct neighbor" assert ipagrp_count > 0 or not has_cross_rank_adj,\ - "expected at least one InterPartitionAdjacencyGroup" + "expected at least one InterPartAdjacencyGroup" for i_remote_part in range(num_parts): tag_sum = 0 - for i_local_part, (mesh, _) in enumerate(new_meshes): + for i_local_part in range(num_parts): if (i_local_part, i_remote_part) in connected_parts: - tag_sum += count_tags(mesh, BTAG_PARTITION(i_remote_part)) + tag_sum += count_tags( + part_meshes[i_local_part], BTAG_PARTITION(i_remote_part)) assert num_tags[i_remote_part] == tag_sum,\ "part_mesh has the wrong number of BTAG_PARTITION boundaries" @@ -395,8 +405,8 @@ def _test_mpi_boundary_swap(dim, order, num_groups): from meshmode.discretization import Discretization vol_discr = Discretization(actx, local_mesh, group_factory) - from meshmode.distributed import get_connected_partitions - connected_parts = get_connected_partitions(local_mesh) + from meshmode.distributed import get_connected_parts + connected_parts = get_connected_parts(local_mesh) # Check that the connectivity makes sense before doing any communication _test_connected_parts(mpi_comm, connected_parts) @@ -569,12 +579,12 @@ def f(x): # {{{ MPI pytest entrypoint @pytest.mark.mpi -@pytest.mark.parametrize("num_partitions", [3, 4]) +@pytest.mark.parametrize("num_parts", [3, 4]) @pytest.mark.parametrize("order", [2, 3]) -def test_mpi_communication(num_partitions, order): +def test_mpi_communication(num_parts, order): pytest.importorskip("mpi4py") - num_ranks = num_partitions + num_ranks = num_parts from subprocess import check_call import sys check_call([