From 4e6f699e09342488f8f6cac72ee1cb5f637cdd95 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Tue, 15 Mar 2022 17:52:47 -0500
Subject: [PATCH 01/33] simplify partition_mesh interface

---
 meshmode/distributed.py     |  3 +--
 meshmode/mesh/processing.py | 25 +++++++++++++++++++++++-
 test/test_partition.py      | 38 +++++++++++++++++++------------------
 3 files changed, 45 insertions(+), 21 deletions(-)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index 5d5173014..cca094635 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -106,8 +106,7 @@ def send_mesh_parts(self, mesh, part_per_element, num_parts):
         assert self.is_mananger_rank()
 
         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_per_element)[0]
 
         local_part = None
 
diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 92f85001a..ace23f46f 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -431,7 +431,7 @@ 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, part_per_element, part_num):
     """
     :arg mesh: A :class:`~meshmode.mesh.Mesh` to be partitioned.
     :arg part_per_element: A :class:`numpy.ndarray` containing one
@@ -502,6 +502,29 @@ def partition_mesh(mesh, part_per_element, part_num):
 
     return part_mesh, queried_elems
 
+
+def partition_mesh(mesh, part_per_element):
+    """
+    :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.
+
+    :returns: A pair ``(part_meshes, part_to_global_maps)``, where *part_meshes*
+        is a list of :class:`~meshmode.mesh.Mesh` where each is a partition of
+        *mesh*, and *part_to_global_maps* is a list of :class:`numpy.ndarray`
+        mapping element numbers on each mesh partition to ones in *mesh*.
+    """
+    num_parts = np.max(part_per_element) + 1
+    result_pairs = [
+        _get_mesh_part(mesh, part_per_element, i)
+        for i in range(num_parts)]
+
+    return (
+        [part_mesh for part_mesh, _ in result_pairs],
+        [part_to_global for _, part_to_global in result_pairs])
+
+
 # }}}
 
 
diff --git a/test/test_partition.py b/test/test_partition.py
index 0247f2058..9fa4eb294 100644
--- a/test/test_partition.py
+++ b/test/test_partition.py
@@ -99,8 +99,7 @@ def f(x):
                     connectivity=part_method)
 
         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_per_element)[0]
 
         connected_parts = set()
         for i_local_part, part_mesh in enumerate(part_meshes):
@@ -228,23 +227,21 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio
     has_cross_rank_adj = _check_for_cross_rank_adj(mesh, 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, part_elem_to_global_elem_maps = partition_mesh(
+        mesh, part_per_element)
 
     assert mesh.nelements == np.sum(
-        [new_meshes[i][0].nelements for i in range(num_parts)]), \
+        [part_meshes[i].nelements for i in range(num_parts)]), \
         "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_meshes[i], BTAG_ALL) for i in range(num_parts)]), \
         "part_mesh has the wrong number of BTAG_ALL boundaries"
 
     connected_parts = set()
-    for i_local_part, (part_mesh, _) in enumerate(new_meshes):
+    for i_local_part in range(num_parts):
         from meshmode.distributed import get_connected_partitions
-        neighbors = get_connected_partitions(part_mesh)
+        neighbors = get_connected_partitions(part_meshes[i_local_part])
         for i_remote_part in neighbors:
             connected_parts.add((i_local_part, i_remote_part))
 
@@ -253,10 +250,11 @@ 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]
+                fagrp for fagrp in part_mesh.facial_adjacency_groups[igrp]
                 if isinstance(fagrp, InterPartitionAdjacencyGroup)]
             for ipagrp in ipagrps:
                 for i, (elem, face) in enumerate(
@@ -266,7 +264,8 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio
     ipagrp_count = 0
 
     for part_num in range(num_parts):
-        part, part_to_global = new_meshes[part_num]
+        part = part_meshes[part_num]
+        part_to_global = part_elem_to_global_elem_maps[part_num]
         for grp_num in range(len(part.groups)):
             ipagrps = [
                 fagrp for fagrp in part.facial_adjacency_groups[grp_num]
@@ -282,7 +281,8 @@ 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]
+                    n_part_to_global = part_elem_to_global_elem_maps[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.
@@ -302,7 +302,8 @@ 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]
+                            n_part_to_global = (
+                                part_elem_to_global_elem_maps[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 "
@@ -334,9 +335,10 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio
 
     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"
 

From aef20240f44d56cb46c3311ccafc4d7c164ea7de Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Tue, 15 Mar 2022 20:51:02 -0500
Subject: [PATCH 02/33] add get_connected_partitions back to docs

---
 meshmode/distributed.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index cca094635..2b7488f91 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -4,6 +4,7 @@
 .. autoclass:: MPIBoundaryCommSetupHelper
 
 .. autofunction:: get_partition_by_pymetis
+.. autofunction:: get_connected_partitions
 .. autofunction:: get_inter_partition_tags
 
 .. autoclass:: RemoteGroupInfo

From 8e094b0d44c5319a0be9a735b478bdcdc51d08e0 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Tue, 15 Mar 2022 20:51:07 -0500
Subject: [PATCH 03/33] use opaque partition ID instead of number in
 partition_mesh

---
 meshmode/distributed.py     | 112 +++++-------
 meshmode/mesh/__init__.py   |  19 ++-
 meshmode/mesh/processing.py | 332 +++++++++++++++++++-----------------
 test/test_partition.py      |  44 +++--
 4 files changed, 258 insertions(+), 249 deletions(-)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index 2b7488f91..469e10196 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -4,8 +4,8 @@
 .. autoclass:: MPIBoundaryCommSetupHelper
 
 .. autofunction:: get_partition_by_pymetis
+.. autofunction:: membership_list_to_sets
 .. autofunction:: get_connected_partitions
-.. autofunction:: get_inter_partition_tags
 
 .. autoclass:: RemoteGroupInfo
 .. autoclass:: make_remote_group_infos
@@ -48,8 +48,7 @@
         Mesh,
         InteriorAdjacencyGroup,
         InterPartitionAdjacencyGroup,
-        BoundaryTag,
-        BTAG_PARTITION,
+        PartitionID,
 )
 
 from meshmode.discretization import ElementGroupFactory
@@ -106,13 +105,16 @@ def send_mesh_parts(self, mesh, part_per_element, num_parts):
 
         assert self.is_mananger_rank()
 
+        from meshmode.distributed import membership_list_to_sets
+        part_num_to_elements = membership_list_to_sets(part_per_element)
+
         from meshmode.mesh.processing import partition_mesh
-        parts = partition_mesh(mesh, part_per_element)[0]
+        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:
@@ -156,7 +158,8 @@ class RemoteGroupInfo:
 
 
 def make_remote_group_infos(
-        actx: ArrayContext, local_btag: BoundaryTag,
+        actx: ArrayContext,
+        remote_part_id: PartitionID,
         bdry_conn: DirectDiscretizationConnection
         ) -> Sequence[RemoteGroupInfo]:
     local_vol_mesh = bdry_conn.from_discr.mesh
@@ -168,7 +171,7 @@ def make_remote_group_infos(
                 inter_partition_adj_groups=[
                     fagrp for fagrp in local_vol_mesh.facial_adjacency_groups[igrp]
                     if isinstance(fagrp, InterPartitionAdjacencyGroup)
-                    and fagrp.boundary_tag == local_btag],
+                    and fagrp.boundary_tag.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,10 +191,6 @@ 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.
@@ -207,15 +206,14 @@ 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: PartitionID
+    remote_part_id: PartitionID
     remote_rank: int
     local_boundary_connection: DirectDiscretizationConnection
 
@@ -265,9 +263,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 +289,7 @@ def __enter__(self):
 
         # to know when we're done
         self.pending_recv_identifiers = {
-                (irbi.remote_rank, irbi.remote_part_id)
+                irbi.remote_part_id
                 for irbi in self.inter_rank_bdry_info}
 
         self.send_reqs = [
@@ -302,7 +299,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]
@@ -341,36 +338,35 @@ 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)
+        remote_part_id_to_irbi = {
+                irbi.remote_part_id: irbi for irbi in self.inter_rank_bdry_info}
+        assert len(remote_part_id_to_irbi) == len(self.inter_rank_bdry_info)
 
         for i_src_rank, recvd in zip(
                 source_ranks, data):
-            (recvd_remote_part_id, recvd_local_part_id,
+            (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 = remote_part_id_to_irbi[remote_part_id]
             assert i_src_rank == irbi.remote_rank
-            assert recvd_remote_part_id == irbi.remote_part_id
+            assert local_part_id == irbi.local_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)
+            remote_to_local_bdry_conns[remote_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))
 
-            self.pending_recv_identifiers.remove((i_src_rank, recvd_remote_part_id))
+            self.pending_recv_identifiers.remove(remote_part_id)
 
         if not self.pending_recv_identifiers:
             MPI.Request.waitall(self.send_reqs)
@@ -381,7 +377,8 @@ def complete_some(self):
 # }}}
 
 
-def get_partition_by_pymetis(mesh, num_parts, *, connectivity="facial", **kwargs):
+def get_partition_by_pymetis(
+        mesh, num_parts, *, connectivity="facial", return_sets=False, **kwargs):
     """Return a mesh partition created by :mod:`pymetis`.
 
     :arg mesh: A :class:`meshmode.mesh.Mesh` instance
@@ -426,39 +423,24 @@ def get_partition_by_pymetis(mesh, num_parts, *, connectivity="facial", **kwargs
     from pymetis import part_graph
     _, p = part_graph(num_parts, xadj=xadj, adjncy=adjncy, **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`.
-    """
-    assert mesh.facial_adjacency_groups is not None
-    # internal and deprecated, remove in July 2022
+    if return_sets:
+        return membership_list_to_sets(np.array(p))
+    else:
+        return np.array(p)
 
-    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")
 
+def membership_list_to_sets(membership_list):
     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: set(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`.
-    """
+def get_connected_partitions(mesh: Mesh) -> "Set[PartitionID]":
+    """For a local mesh part in *mesh*, determine the set of connected partitions."""
     assert mesh.facial_adjacency_groups is not None
 
     return {
-            grp.boundary_tag
+            grp.boundary_tag.part_id
             for fagrp_list in mesh.facial_adjacency_groups
             for grp in fagrp_list
             if isinstance(grp, InterPartitionAdjacencyGroup)}
diff --git a/meshmode/mesh/__init__.py b/meshmode/mesh/__init__.py
index 1d0e16aa0..11675044d 100644
--- a/meshmode/mesh/__init__.py
+++ b/meshmode/mesh/__init__.py
@@ -67,6 +67,7 @@
 # {{{ element tags
 
 BoundaryTag = Hashable
+PartitionID = Hashable
 
 
 class BTAG_NONE:  # noqa: N801
@@ -109,22 +110,22 @@ 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 partition 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: PartitionID):
+        self.part_id = 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 +133,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
diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index ace23f46f..045a40237 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -80,25 +80,30 @@ 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, part_id_to_elements, part_id_to_part_index, element_id_dtype):
     """
     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.
+    :arg nelements: The number of elements in the global mesh.
+    :arg part_id_to_elements: A :class:`dict` mapping partition identifiers to
+        sets of elements.
+    :arg part_id_to_part_index: A mapping from partition 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 partition index of the element and
+        ``global_elem_to_part_elem[ielement, 1]`` gives its partition-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 = np.sort(list(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
 
@@ -183,20 +188,23 @@ def _filter_mesh_groups(mesh, selected_elements, vertex_id_dtype):
 
 
 def _get_connected_partitions(
-        mesh, part_per_element, global_elem_to_part_elem):
+        mesh, part_id_to_part_index, global_elem_to_part_elem, self_part_id):
     """
     Find the partitions that are connected to the current partition.
 
     :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 partition 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 partition indices and partition-wide element indices. See
+        :fun:`_compute_global_elem_to_part_elem`` for details.
+    :arg self_part_id: The identifier of the partition currently being created.
+
+    :returns: A :class:`set` of identifiers of the neighboring partitions.
     """
-    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,45 +217,50 @@ def _get_connected_partitions(
             elem_base_j = mesh.base_element_nrs[jgrp]
 
             elements_are_local = global_elem_to_part_elem[facial_adj.elements
-                        + elem_base_i] >= 0
+                        + elem_base_i, 0] == self_part_index
             neighbors_are_nonlocal = global_elem_to_part_elem[facial_adj.neighbors
-                        + elem_base_j] < 0
+                        + 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])
+                    + 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):
+            self_part_index, self_mesh_groups, global_group_to_self_group,
+            self_mesh_group_elem_base):
     r"""
     Create local-to-local 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
+    :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element
+        indices to partition indices and partition-wide element indices. See
+        :fun:`_compute_global_elem_to_part_elem`` for details.
+    :arg self_part_index: The index of the partition currently being created, in
+        the range ``[0, num_parts)``.
+    :arg self_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_group_to_self_group: An array mapping groups in *mesh* to groups in
+        *self_mesh_groups* (or `None` if the group is not local).
+    :arg self_mesh_group_elem_base: An array containing the starting partition-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.
     """
-    local_to_local_adjacency_groups = [[] for _ in part_mesh_groups]
+    local_to_local_adjacency_groups = [[] 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:
+        i_self_grp = global_group_to_self_group[igrp]
+        if i_self_grp is None:
             continue
 
         int_grps = [
@@ -257,35 +270,35 @@ 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:
+            j_self_grp = global_group_to_self_group[jgrp]
+            if j_self_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
+                        + elem_base_i, 0] == self_part_index
             neighbors_are_local = global_elem_to_part_elem[facial_adj.neighbors
-                        + elem_base_j] >= 0
+                        + elem_base_j, 0] == self_part_index
 
             adj_indices, = np.where(elements_are_local & neighbors_are_local)
 
             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[i_self_grp]
+                self_elem_base_j = self_mesh_group_elem_base[j_self_grp]
 
                 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(
+                local_to_local_adjacency_groups[i_self_grp].append(
                     InteriorAdjacencyGroup(
-                        igroup=i_part_grp,
-                        ineighbor_group=j_part_grp,
+                        igroup=i_self_grp,
+                        ineighbor_group=j_self_grp,
                         elements=elements,
                         element_faces=element_faces,
                         neighbors=neighbors,
@@ -296,23 +309,25 @@ def _create_local_to_local_adjacency_groups(mesh, global_elem_to_part_elem,
 
 
 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):
+        mesh, part_id_to_part_index, global_elem_to_part_elem, self_part_id,
+        self_mesh_groups, global_group_to_self_group, self_mesh_group_elem_base,
+        connected_parts):
     """
     Create non-local 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 partition 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 partition indices and partition-wide element indices. See
+        :fun:`_compute_global_elem_to_part_elem`` for details.
+    :arg self_part_id: The identifier of the partition currently being created.
+    :arg self_mesh_groups: An array of `~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_group_to_self_group: An array mapping groups in *mesh* to groups in
+        *self_mesh_groups* (or `None` if the group is not local).
+    :arg self_mesh_group_elem_base: An array containing the starting partition-wide
+        element index for each group in *self_mesh_groups*.
     :arg connected_parts: A :class:`set` containing the partitions connected to
         the current one.
 
@@ -320,14 +335,13 @@ def _create_nonlocal_adjacency_groups(
         instances corresponding to the entries in *mesh.facial_adjacency_groups* that
         have non-local 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]
+    nonlocal_adj_groups = [[] 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:
+        i_self_grp = global_group_to_self_group[igrp]
+        if i_self_grp is None:
             continue
 
         int_grps = [
@@ -343,29 +357,31 @@ 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_local = (
+                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))
+                    & (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[i_self_grp]
 
                     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(
+                    nonlocal_adj_groups[i_self_grp].append(
                         InterPartitionAdjacencyGroup(
-                            igroup=i_part_grp,
-                            boundary_tag=BTAG_PARTITION(i_neighbor_part),
+                            igroup=i_self_grp,
+                            boundary_tag=BTAG_PARTITION(neighbor_part_id),
                             elements=elements,
                             element_faces=element_faces,
                             neighbors=neighbors,
@@ -375,31 +391,33 @@ def _create_nonlocal_adjacency_groups(
     return nonlocal_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, global_elem_to_part_elem, self_part_index,
+            self_mesh_groups, global_group_to_self_group, self_mesh_group_elem_base):
     """
     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 partition indices and partition-wide element indices. See
+        :fun:`_compute_global_elem_to_part_elem`` for details.
+    :arg self_part_index: The index of the partition currently being created, in
+        the range ``[0, num_parts)``.
+    :arg self_mesh_groups: An array of `~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_group_to_self_group: An array mapping groups in *mesh* to groups in
+        *self_mesh_groups* (or `None` if the group is not local).
+    :arg self_mesh_group_elem_base: An array containing the starting partition-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 = [[] 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:
+        i_self_grp = global_group_to_self_group[igrp]
+        if i_self_grp is None:
             continue
 
         bdry_grps = [
@@ -409,21 +427,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[i_self_grp]
                 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[i_self_grp].append(
                 BoundaryAdjacencyGroup(
-                    igroup=i_part_grp,
+                    igroup=i_self_grp,
                     boundary_tag=bdry_grp.boundary_tag,
                     elements=elements,
                     element_faces=element_faces))
@@ -431,99 +450,98 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, part_mesh_groups,
     return bdry_adj_groups
 
 
-def _get_mesh_part(mesh, part_per_element, part_num):
+def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
     """
     :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 partition identifiers to
+        sets of elements.
+    :arg self_part_id: The partition 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 partition of *mesh*.
 
     .. versionadded:: 2017.1
     """
-    assert len(part_per_element) == mesh.nelements, (
-        "part_per_element must have shape (mesh.nelements,)")
+    nelements = sum(len(elems) for elems in part_id_to_elements.values())
+    assert nelements == mesh.nelements, \
+        "counts of partitioned elements must add up to mesh.nelements"
 
-    # Contains the indices of the elements requested.
-    queried_elems, = np.where(part_per_element == part_num)
+    part_id_to_part_index = {
+        part_id: part_index
+        for part_id, part_index in zip(
+            part_id_to_elements.keys(),
+            range(len(part_id_to_elements)))}
 
-    global_elem_to_part_elem = _compute_global_elem_to_part_elem(part_per_element,
-                {part_num}, mesh.element_id_dtype)
+    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)
+    self_elements_sorted = np.sort(list(part_id_to_elements[self_part_id]))
+    self_mesh_groups, global_group_to_self_group, required_vertex_indices =\
+                _filter_mesh_groups(
+                    mesh, self_elements_sorted,
+                    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 i_self_grp, grp in enumerate(self_mesh_groups):
+        self_mesh_group_elem_base[i_self_grp] = el_nr
         el_nr += grp.nelements
 
     connected_parts = _get_connected_partitions(
-        mesh, part_per_element, global_elem_to_part_elem)
+        mesh, part_id_to_part_index, global_elem_to_part_elem,
+        self_part_index)
 
     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)
+                mesh, global_elem_to_part_elem, self_part_index, self_mesh_groups,
+                global_group_to_self_group, 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)
+                mesh, part_id_to_part_index, global_elem_to_part_elem, self_part_id,
+                self_mesh_groups, global_group_to_self_group,
+                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)
+                mesh, global_elem_to_part_elem, self_part_index, self_mesh_groups,
+                global_group_to_self_group, self_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))]
+    self_facial_adj_groups = [
+        local_to_local_adj_groups[i_self_grp]
+        + nonlocal_adj_groups[i_self_grp]
+        + boundary_adj_groups[i_self_grp]
+        for i_self_grp in range(len(self_mesh_groups))]
 
     from meshmode.mesh import Mesh
-    part_mesh = Mesh(
-            part_vertices,
-            part_mesh_groups,
-            facial_adjacency_groups=part_facial_adj_groups,
+    self_mesh = Mesh(
+            self_vertices,
+            self_mesh_groups,
+            facial_adjacency_groups=self_facial_adj_groups,
             is_conforming=mesh.is_conforming)
 
-    return part_mesh, queried_elems
+    return self_mesh
 
 
-def partition_mesh(mesh, part_per_element):
+def partition_mesh(mesh, part_id_to_elements):
     """
     :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.
-
-    :returns: A pair ``(part_meshes, part_to_global_maps)``, where *part_meshes*
-        is a list of :class:`~meshmode.mesh.Mesh` where each is a partition of
-        *mesh*, and *part_to_global_maps* is a list of :class:`numpy.ndarray`
-        mapping element numbers on each mesh partition to ones in *mesh*.
-    """
-    num_parts = np.max(part_per_element) + 1
-    result_pairs = [
-        _get_mesh_part(mesh, part_per_element, i)
-        for i in range(num_parts)]
-
-    return (
-        [part_mesh for part_mesh, _ in result_pairs],
-        [part_to_global for _, part_to_global in result_pairs])
+    :arg part_id_to_elements: A :class:`dict` mapping partition identifiers to
+        sets of elements.
 
+    :returns: A :class:`dict` mapping partition identifiers to instances of
+        :class:`~meshmode.mesh.Mesh` that represent the corresponding partition of
+        *mesh*.
+    """
+    return {
+        part_id: _get_mesh_part(mesh, part_id_to_elements, part_id)
+        for part_id in part_id_to_elements.keys()}
 
 # }}}
 
diff --git a/test/test_partition.py b/test/test_partition.py
index 9fa4eb294..52d8a0d89 100644
--- a/test/test_partition.py
+++ b/test/test_partition.py
@@ -98,19 +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_sets
+        part_num_to_elements = membership_list_to_sets(part_per_element)
+
         from meshmode.mesh.processing import partition_mesh
-        part_meshes = partition_mesh(mesh, part_per_element)[0]
+        part_meshes = partition_mesh(mesh, part_num_to_elements)
 
         connected_parts = set()
-        for i_local_part, part_mesh in enumerate(part_meshes):
+        for i_local_part, part_mesh in part_meshes.items():
             from meshmode.distributed import get_connected_partitions
             neighbors = get_connected_partitions(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,
@@ -145,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(
@@ -153,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)
@@ -226,16 +229,18 @@ 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_sets
+    part_num_to_elements = membership_list_to_sets(part_per_element)
+
     from meshmode.mesh.processing import partition_mesh
-    part_meshes, part_elem_to_global_elem_maps = partition_mesh(
-        mesh, part_per_element)
+    part_meshes = partition_mesh(mesh, part_num_to_elements)
 
     assert mesh.nelements == np.sum(
-        [part_meshes[i].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(part_meshes[i], 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()
@@ -263,16 +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_meshes[part_num]
-        part_to_global = part_elem_to_global_elem_maps[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)]
             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 +290,6 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio
                     n_meshwide_elem = ipagrp.neighbors[idx]
                     n_face = ipagrp.neighbor_faces[idx]
                     n_part = part_meshes[n_part_num]
-                    n_part_to_global = part_elem_to_global_elem_maps[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.
@@ -291,7 +298,7 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio
                     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]
+                        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,11 @@ 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 = (
-                                part_elem_to_global_elem_maps[n_part_num])
-                            p_meshwide_elem = part_to_global[elem + elem_base]
-                            p_meshwide_n_elem = n_part_to_global[n_meshwide_elem]
+                            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, ("InterPartitionAdjacencyGroup is not "
                         "consistent")
 

From 2f41937d62ad97f591c58605b329328912f8e839 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Wed, 16 Mar 2022 00:09:51 -0500
Subject: [PATCH 04/33] use self/other instead of local/nonlocal

---
 meshmode/mesh/processing.py | 61 +++++++++++++++++++------------------
 1 file changed, 32 insertions(+), 29 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 045a40237..2ec6a6753 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -216,15 +216,15 @@ 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
+            elements_are_self = global_elem_to_part_elem[facial_adj.elements
                         + elem_base_i, 0] == self_part_index
-            neighbors_are_nonlocal = global_elem_to_part_elem[facial_adj.neighbors
+            neighbors_are_other = global_elem_to_part_elem[facial_adj.neighbors
                         + elem_base_j, 0] != self_part_index
 
             connected_part_indices.update(
                 global_elem_to_part_elem[
                     facial_adj.neighbors[
-                        elements_are_local & neighbors_are_nonlocal]
+                        elements_are_self & neighbors_are_other]
                     + elem_base_j, 0])
 
     return {
@@ -233,11 +233,11 @@ def _get_connected_partitions(
         if part_index in connected_part_indices}
 
 
-def _create_local_to_local_adjacency_groups(mesh, global_elem_to_part_elem,
+def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
             self_part_index, self_mesh_groups, global_group_to_self_group,
             self_mesh_group_elem_base):
     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` that maps global element
@@ -248,15 +248,16 @@ def _create_local_to_local_adjacency_groups(mesh, global_elem_to_part_elem,
     :arg self_mesh_groups: An array of :class:`~meshmode.mesh.ElementGroup` instances
         representing the partitioned mesh groups.
     :arg global_group_to_self_group: An array mapping groups in *mesh* to groups in
-        *self_mesh_groups* (or `None` if the group is not local).
+        *self_mesh_groups* (or `None` if the group is not part of the current
+        partition).
     :arg self_mesh_group_elem_base: An array containing the starting partition-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 self_mesh_groups]
+    self_to_self_adjacency_groups = [[] for _ in self_mesh_groups]
 
     for igrp, facial_adj_list in enumerate(mesh.facial_adjacency_groups):
         i_self_grp = global_group_to_self_group[igrp]
@@ -277,12 +278,12 @@ def _create_local_to_local_adjacency_groups(mesh, global_elem_to_part_elem,
             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
+            elements_are_self = global_elem_to_part_elem[facial_adj.elements
                         + elem_base_i, 0] == self_part_index
-            neighbors_are_local = global_elem_to_part_elem[facial_adj.neighbors
+            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:
                 self_elem_base_i = self_mesh_group_elem_base[i_self_grp]
@@ -295,7 +296,7 @@ def _create_local_to_local_adjacency_groups(mesh, global_elem_to_part_elem,
                             adj_indices] + elem_base_j, 1] - self_elem_base_j
                 neighbor_faces = facial_adj.neighbor_faces[adj_indices]
 
-                local_to_local_adjacency_groups[i_self_grp].append(
+                self_to_self_adjacency_groups[i_self_grp].append(
                     InteriorAdjacencyGroup(
                         igroup=i_self_grp,
                         ineighbor_group=j_self_grp,
@@ -305,15 +306,15 @@ def _create_local_to_local_adjacency_groups(mesh, global_elem_to_part_elem,
                         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(
+def _create_self_to_other_adjacency_groups(
         mesh, part_id_to_part_index, global_elem_to_part_elem, self_part_id,
         self_mesh_groups, global_group_to_self_group, self_mesh_group_elem_base,
         connected_parts):
     """
-    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_id_to_part_index: A mapping from partition identifiers to indices in
@@ -325,7 +326,8 @@ def _create_nonlocal_adjacency_groups(
     :arg self_mesh_groups: An array of `~meshmode.mesh.ElementGroup` instances
         representing the partitioned mesh groups.
     :arg global_group_to_self_group: An array mapping groups in *mesh* to groups in
-        *self_mesh_groups* (or `None` if the group is not local).
+        *self_mesh_groups* (or `None` if the group is not part of the current
+        partition).
     :arg self_mesh_group_elem_base: An array containing the starting partition-wide
         element index for each group in *self_mesh_groups*.
     :arg connected_parts: A :class:`set` containing the partitions connected to
@@ -333,11 +335,11 @@ def _create_nonlocal_adjacency_groups(
 
     :returns: A list of lists of `~meshmode.mesh.InterPartitionAdjacencyGroup`
         instances corresponding to the entries in *mesh.facial_adjacency_groups* that
-        have non-local adjacency.
+        have self-to-other adjacency.
     """
     self_part_index = part_id_to_part_index[self_part_id]
 
-    nonlocal_adj_groups = [[] for _ in self_mesh_groups]
+    self_to_other_adj_groups = [[] for _ in self_mesh_groups]
 
     for igrp, facial_adj_list in enumerate(mesh.facial_adjacency_groups):
         i_self_grp = global_group_to_self_group[igrp]
@@ -357,7 +359,7 @@ def _create_nonlocal_adjacency_groups(
             global_elements = facial_adj.elements + elem_base_i
             global_neighbors = facial_adj.neighbors + elem_base_j
 
-            elements_are_local = (
+            elements_are_self = (
                 global_elem_to_part_elem[global_elements, 0] == self_part_index)
 
             neighbor_part_indices = global_elem_to_part_elem[global_neighbors, 0]
@@ -365,7 +367,7 @@ def _create_nonlocal_adjacency_groups(
             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
+                    elements_are_self
                     & (neighbor_part_indices == neighbor_part_index))
 
                 if len(adj_indices) > 0:
@@ -378,7 +380,7 @@ def _create_nonlocal_adjacency_groups(
                         global_neighbors[adj_indices], 1]
                     neighbor_faces = facial_adj.neighbor_faces[adj_indices]
 
-                    nonlocal_adj_groups[i_self_grp].append(
+                    self_to_other_adj_groups[i_self_grp].append(
                         InterPartitionAdjacencyGroup(
                             igroup=i_self_grp,
                             boundary_tag=BTAG_PARTITION(neighbor_part_id),
@@ -388,7 +390,7 @@ def _create_nonlocal_adjacency_groups(
                             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, self_part_index,
@@ -405,7 +407,8 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
     :arg self_mesh_groups: An array of `~meshmode.mesh.ElementGroup` instances
         representing the partitioned mesh groups.
     :arg global_group_to_self_group: An array mapping groups in *mesh* to groups in
-        *self_mesh_groups* (or `None` if the group is not local).
+        *self_mesh_groups* (or `None` if the group is not part of the current
+        partition).
     :arg self_mesh_group_elem_base: An array containing the starting partition-wide
         element index for each group in *self_mesh_groups*.
 
@@ -476,7 +479,7 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
         mesh.element_id_dtype)
 
     # Create new mesh groups that mimic the original mesh's groups but only contain
-    # the local partition's elements
+    # the current partition's elements
     self_elements_sorted = np.sort(list(part_id_to_elements[self_part_id]))
     self_mesh_groups, global_group_to_self_group, required_vertex_indices =\
                 _filter_mesh_groups(
@@ -499,11 +502,11 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
         mesh, part_id_to_part_index, global_elem_to_part_elem,
         self_part_index)
 
-    local_to_local_adj_groups = _create_local_to_local_adjacency_groups(
+    self_to_self_adj_groups = _create_self_to_self_adjacency_groups(
                 mesh, global_elem_to_part_elem, self_part_index, self_mesh_groups,
                 global_group_to_self_group, self_mesh_group_elem_base)
 
-    nonlocal_adj_groups = _create_nonlocal_adjacency_groups(
+    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, global_group_to_self_group,
                 self_mesh_group_elem_base, connected_parts)
@@ -512,10 +515,10 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
                 mesh, global_elem_to_part_elem, self_part_index, self_mesh_groups,
                 global_group_to_self_group, self_mesh_group_elem_base)
 
-    # Combine local/nonlocal/boundary adjacency groups
+    # Combine adjacency groups
     self_facial_adj_groups = [
-        local_to_local_adj_groups[i_self_grp]
-        + nonlocal_adj_groups[i_self_grp]
+        self_to_self_adj_groups[i_self_grp]
+        + self_to_other_adj_groups[i_self_grp]
         + boundary_adj_groups[i_self_grp]
         for i_self_grp in range(len(self_mesh_groups))]
 

From efdc4b5f3a0ee1b1852b601a6ed1832c8a5674fc Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Wed, 16 Mar 2022 12:58:21 -0500
Subject: [PATCH 05/33] fix compatibility

---
 meshmode/mesh/__init__.py | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/meshmode/mesh/__init__.py b/meshmode/mesh/__init__.py
index 11675044d..543899fad 100644
--- a/meshmode/mesh/__init__.py
+++ b/meshmode/mesh/__init__.py
@@ -117,8 +117,23 @@ class BTAG_PARTITION(BTAG_NO_BOUNDARY):  # noqa: N801
 
     .. versionadded:: 2017.1
     """
-    def __init__(self, part_id: PartitionID):
-        self.part_id = part_id
+    def __init__(self, part_id: PartitionID, 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_id))

From 5475dfd93d0c1c7079e357ab49251ef2925e38c9 Mon Sep 17 00:00:00 2001
From: Matt Smith <mjsmith1@gmail.com>
Date: Thu, 17 Mar 2022 12:43:26 -0500
Subject: [PATCH 06/33] remove unnecessary import

Co-authored-by: Alex Fikl <alexfikl@gmail.com>
---
 meshmode/distributed.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index 469e10196..51ed4a69c 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -105,7 +105,6 @@ def send_mesh_parts(self, mesh, part_per_element, num_parts):
 
         assert self.is_mananger_rank()
 
-        from meshmode.distributed import membership_list_to_sets
         part_num_to_elements = membership_list_to_sets(part_per_element)
 
         from meshmode.mesh.processing import partition_mesh

From f671225c9ff376117c459d2fe505c20029ffc250 Mon Sep 17 00:00:00 2001
From: Matt Smith <mjsmith1@gmail.com>
Date: Thu, 17 Mar 2022 12:48:02 -0500
Subject: [PATCH 07/33] eliminate fun

Co-authored-by: Alex Fikl <alexfikl@gmail.com>
---
 meshmode/mesh/processing.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 2ec6a6753..8e83a509f 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -197,7 +197,7 @@ def _get_connected_partitions(
         the range ``[0, num_parts)``.
     :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element
         indices to partition indices and partition-wide element indices. See
-        :fun:`_compute_global_elem_to_part_elem`` for details.
+        :func:`_compute_global_elem_to_part_elem`` for details.
     :arg self_part_id: The identifier of the partition currently being created.
 
     :returns: A :class:`set` of identifiers of the neighboring partitions.
@@ -242,7 +242,7 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
     :arg mesh: A :class:`~meshmode.mesh.Mesh` representing the unpartitioned mesh.
     :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element
         indices to partition indices and partition-wide element indices. See
-        :fun:`_compute_global_elem_to_part_elem`` for details.
+        :func:`_compute_global_elem_to_part_elem`` for details.
     :arg self_part_index: The index of the partition currently being created, in
         the range ``[0, num_parts)``.
     :arg self_mesh_groups: An array of :class:`~meshmode.mesh.ElementGroup` instances
@@ -321,7 +321,7 @@ def _create_self_to_other_adjacency_groups(
         the range ``[0, num_parts)``.
     :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element
         indices to partition indices and partition-wide element indices. See
-        :fun:`_compute_global_elem_to_part_elem`` for details.
+        :func:`_compute_global_elem_to_part_elem`` for details.
     :arg self_part_id: The identifier of the partition currently being created.
     :arg self_mesh_groups: An array of `~meshmode.mesh.ElementGroup` instances
         representing the partitioned mesh groups.
@@ -401,7 +401,7 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
     :arg mesh: A :class:`~meshmode.mesh.Mesh` representing the unpartitioned mesh.
     :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element
         indices to partition indices and partition-wide element indices. See
-        :fun:`_compute_global_elem_to_part_elem`` for details.
+        :func:`_compute_global_elem_to_part_elem`` for details.
     :arg self_part_index: The index of the partition currently being created, in
         the range ``[0, num_parts)``.
     :arg self_mesh_groups: An array of `~meshmode.mesh.ElementGroup` instances

From 87a9aa8ad15a35ca59820871e6bf119382fa5ec6 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Thu, 17 Mar 2022 12:49:17 -0500
Subject: [PATCH 08/33] stamp out remaining traces of fun

---
 meshmode/mesh/processing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 8e83a509f..7e16cbd1f 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -657,7 +657,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

From efa7f0ecb2b85f1a3375d0a57731ecfbeeebef85 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Thu, 17 Mar 2022 13:06:00 -0500
Subject: [PATCH 09/33] rename membership_list_to_sets ->
 membership_list_to_map and store index sets as numpy arrays instead of python
 sets

---
 meshmode/distributed.py     | 15 ++++++++++-----
 meshmode/mesh/processing.py |  5 ++---
 test/test_partition.py      |  8 ++++----
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index 51ed4a69c..8b186dae9 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -4,7 +4,7 @@
 .. autoclass:: MPIBoundaryCommSetupHelper
 
 .. autofunction:: get_partition_by_pymetis
-.. autofunction:: membership_list_to_sets
+.. autofunction:: membership_list_to_map
 .. autofunction:: get_connected_partitions
 
 .. autoclass:: RemoteGroupInfo
@@ -105,7 +105,7 @@ def send_mesh_parts(self, mesh, part_per_element, num_parts):
 
         assert self.is_mananger_rank()
 
-        part_num_to_elements = membership_list_to_sets(part_per_element)
+        part_num_to_elements = membership_list_to_map(part_per_element)
 
         from meshmode.mesh.processing import partition_mesh
         parts = partition_mesh(mesh, part_num_to_elements)
@@ -423,14 +423,19 @@ def get_partition_by_pymetis(
     _, p = part_graph(num_parts, xadj=xadj, adjncy=adjncy, **kwargs)
 
     if return_sets:
-        return membership_list_to_sets(np.array(p))
+        return membership_list_to_map(np.array(p))
     else:
         return np.array(p)
 
 
-def membership_list_to_sets(membership_list):
+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`).
+    """
     return {
-        entry: set(np.where(membership_list == entry)[0])
+        entry: np.where(membership_list == entry)[0]
         for entry in set(membership_list)}
 
 
diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 7e16cbd1f..1ce378df8 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -100,7 +100,7 @@ def _compute_global_elem_to_part_elem(
     """
     global_elem_to_part_elem = np.empty((nelements, 2), dtype=element_id_dtype)
     for part_id in part_id_to_elements.keys():
-        elements = np.sort(list(part_id_to_elements[part_id]))
+        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)
@@ -480,10 +480,9 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
 
     # Create new mesh groups that mimic the original mesh's groups but only contain
     # the current partition's elements
-    self_elements_sorted = np.sort(list(part_id_to_elements[self_part_id]))
     self_mesh_groups, global_group_to_self_group, required_vertex_indices =\
                 _filter_mesh_groups(
-                    mesh, self_elements_sorted,
+                    mesh, part_id_to_elements[self_part_id],
                     mesh.vertex_id_dtype)
 
     self_part_index = part_id_to_part_index[self_part_id]
diff --git a/test/test_partition.py b/test/test_partition.py
index 52d8a0d89..33403e3e6 100644
--- a/test/test_partition.py
+++ b/test/test_partition.py
@@ -98,8 +98,8 @@ def f(x):
             part_per_element = get_partition_by_pymetis(mesh, num_parts,
                     connectivity=part_method)
 
-        from meshmode.distributed import membership_list_to_sets
-        part_num_to_elements = membership_list_to_sets(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
         part_meshes = partition_mesh(mesh, part_num_to_elements)
@@ -229,8 +229,8 @@ 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_sets
-    part_num_to_elements = membership_list_to_sets(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
     part_meshes = partition_mesh(mesh, part_num_to_elements)

From 4c0e631661283e0b708aca2d7875dd69ad997420 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Thu, 17 Mar 2022 14:48:11 -0500
Subject: [PATCH 10/33] remove return_sets option from get_partition_by_pymetis

---
 meshmode/distributed.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index 8b186dae9..5cf469973 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -376,8 +376,7 @@ def complete_some(self):
 # }}}
 
 
-def get_partition_by_pymetis(
-        mesh, num_parts, *, connectivity="facial", return_sets=False, **kwargs):
+def get_partition_by_pymetis(mesh, num_parts, *, connectivity="facial", **kwargs):
     """Return a mesh partition created by :mod:`pymetis`.
 
     :arg mesh: A :class:`meshmode.mesh.Mesh` instance
@@ -422,10 +421,7 @@ def get_partition_by_pymetis(
     from pymetis import part_graph
     _, p = part_graph(num_parts, xadj=xadj, adjncy=adjncy, **kwargs)
 
-    if return_sets:
-        return membership_list_to_map(np.array(p))
-    else:
-        return np.array(p)
+    return np.array(p)
 
 
 def membership_list_to_map(membership_list):

From cdae47f00de1f055ddf4383586078b718c037e35 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Thu, 17 Mar 2022 22:08:21 -0500
Subject: [PATCH 11/33] fix bugs in MPIBoundaryCommSetupHelper

---
 meshmode/distributed.py | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index 5cf469973..397ca9c9d 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -288,7 +288,7 @@ def __enter__(self):
 
         # to know when we're done
         self.pending_recv_identifiers = {
-                irbi.remote_part_id
+                (irbi.local_part_id, irbi.remote_part_id)
                 for irbi in self.inter_rank_bdry_info}
 
         self.send_reqs = [
@@ -337,9 +337,10 @@ def complete_some(self):
 
         remote_to_local_bdry_conns = {}
 
-        remote_part_id_to_irbi = {
-                irbi.remote_part_id: irbi for irbi in self.inter_rank_bdry_info}
-        assert len(remote_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}
+        assert len(part_ids_to_irbi) == len(self.inter_rank_bdry_info)
 
         for i_src_rank, recvd in zip(
                 source_ranks, data):
@@ -351,11 +352,15 @@ def complete_some(self):
 
             # Connect local_mesh to remote_mesh
             from meshmode.discretization.connection import make_partition_connection
-            irbi = remote_part_id_to_irbi[remote_part_id]
+            irbi = part_ids_to_irbi[local_part_id, remote_part_id]
             assert i_src_rank == irbi.remote_rank
-            assert local_part_id == irbi.local_part_id
 
-            remote_to_local_bdry_conns[remote_part_id] = (
+            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,
@@ -365,7 +370,7 @@ def complete_some(self):
                         group_factory=self.bdry_grp_factory),
                     remote_group_infos=remote_group_infos))
 
-            self.pending_recv_identifiers.remove(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)

From 03df4cae2ba6022538412661b9eccaae37d6d9d2 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Fri, 18 Mar 2022 10:38:54 -0500
Subject: [PATCH 12/33] flake8

---
 meshmode/distributed.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index 397ca9c9d..b0f0c78b6 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -38,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 (

From aa0067c8f0634e34c0f4cb8bacb9da1a449c2fc1 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 21 Mar 2022 09:54:09 -0500
Subject: [PATCH 13/33] fix bug

---
 meshmode/mesh/processing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 1ce378df8..6e1591ed0 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -499,7 +499,7 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
 
     connected_parts = _get_connected_partitions(
         mesh, part_id_to_part_index, global_elem_to_part_elem,
-        self_part_index)
+        self_part_id)
 
     self_to_self_adj_groups = _create_self_to_self_adjacency_groups(
                 mesh, global_elem_to_part_elem, self_part_index, self_mesh_groups,

From 7cb188fe601e744fdec193f541bd25fda26d2e71 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Fri, 1 Apr 2022 10:38:23 -0500
Subject: [PATCH 14/33] add a couple of fixmes

---
 meshmode/discretization/connection/opposite_face.py | 1 +
 meshmode/distributed.py                             | 3 +++
 2 files changed, 4 insertions(+)

diff --git a/meshmode/discretization/connection/opposite_face.py b/meshmode/discretization/connection/opposite_face.py
index 25bb669c7..4e08a3ea5 100644
--- a/meshmode/discretization/connection/opposite_face.py
+++ b/meshmode/discretization/connection/opposite_face.py
@@ -521,6 +521,7 @@ 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):
     """
diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index b0f0c78b6..dad61158b 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -148,6 +148,9 @@ 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-partition
+# connections between two parts on the same rank.
 @dataclass
 class RemoteGroupInfo:
     inter_partition_adj_groups: List[InterPartitionAdjacencyGroup]

From e276355cb42245be1a44749fc50dfe79f69b52ed Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Thu, 21 Apr 2022 15:26:31 -0500
Subject: [PATCH 15/33] handle groupless mesh case in dim/ambient_dim

---
 meshmode/mesh/__init__.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/meshmode/mesh/__init__.py b/meshmode/mesh/__init__.py
index 5b8e39aa0..df951648d 100644
--- a/meshmode/mesh/__init__.py
+++ b/meshmode/mesh/__init__.py
@@ -1095,11 +1095,15 @@ def set_if_not_present(name, from_name=None):
 
     @property
     def ambient_dim(self):
+        if not self.groups:
+            return None
         from pytools import single_valued
         return single_valued(grp.nodes.shape[0] for grp in self.groups)
 
     @property
     def dim(self):
+        if not self.groups:
+            return None
         from pytools import single_valued
         return single_valued(grp.dim for grp in self.groups)
 

From cac1fec145a5c0ae0cf06b7150e2dee2e74b5a80 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 2 May 2022 12:42:12 -0500
Subject: [PATCH 16/33] Revert "handle groupless mesh case in dim/ambient_dim"

not a good solution
---
 meshmode/mesh/__init__.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/meshmode/mesh/__init__.py b/meshmode/mesh/__init__.py
index df951648d..5b8e39aa0 100644
--- a/meshmode/mesh/__init__.py
+++ b/meshmode/mesh/__init__.py
@@ -1095,15 +1095,11 @@ def set_if_not_present(name, from_name=None):
 
     @property
     def ambient_dim(self):
-        if not self.groups:
-            return None
         from pytools import single_valued
         return single_valued(grp.nodes.shape[0] for grp in self.groups)
 
     @property
     def dim(self):
-        if not self.groups:
-            return None
         from pytools import single_valued
         return single_valued(grp.dim for grp in self.groups)
 

From 078eab4e754d780a3e9490f6bef3d11ef344a56e Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 2 May 2022 12:45:33 -0500
Subject: [PATCH 17/33] disable removal of empty mesh groups in partitioning

---
 meshmode/mesh/processing.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 6e1591ed0..45f9a13ba 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -134,8 +134,8 @@ def _filter_mesh_groups(mesh, selected_elements, vertex_id_dtype):
     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
+        # if end_idx == start_idx:
+        #     continue
 
         new_group_to_old_group.append(igrp)
         filtered_group_elements.append(

From 5413209f0dd7cadc41bdf35ea0475e3b471947d8 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Thu, 19 May 2022 14:16:47 -0500
Subject: [PATCH 18/33] fix some issues with partitioning docs

---
 meshmode/mesh/processing.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 45f9a13ba..32a0eda8b 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -112,9 +112,10 @@ def _filter_mesh_groups(mesh, selected_elements, vertex_id_dtype):
     """
     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.
+    :arg vertex_id_dtype: The vertex index data type.
     :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),
@@ -245,8 +246,8 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
         :func:`_compute_global_elem_to_part_elem`` for details.
     :arg self_part_index: The index of the partition currently being created, in
         the range ``[0, num_parts)``.
-    :arg self_mesh_groups: An array of :class:`~meshmode.mesh.ElementGroup` instances
-        representing the partitioned mesh groups.
+    :arg self_mesh_groups: An array of :class:`~meshmode.mesh.MeshElementGroup`
+        instances representing the partitioned mesh groups.
     :arg global_group_to_self_group: An array mapping groups in *mesh* to groups in
         *self_mesh_groups* (or `None` if the group is not part of the current
         partition).
@@ -323,7 +324,7 @@ def _create_self_to_other_adjacency_groups(
         indices to partition indices and partition-wide element indices. See
         :func:`_compute_global_elem_to_part_elem`` for details.
     :arg self_part_id: The identifier of the partition currently being created.
-    :arg self_mesh_groups: An array of `~meshmode.mesh.ElementGroup` instances
+    :arg self_mesh_groups: An array of `~meshmode.mesh.MeshElementGroup` instances
         representing the partitioned mesh groups.
     :arg global_group_to_self_group: An array mapping groups in *mesh* to groups in
         *self_mesh_groups* (or `None` if the group is not part of the current
@@ -404,7 +405,7 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
         :func:`_compute_global_elem_to_part_elem`` for details.
     :arg self_part_index: The index of the partition currently being created, in
         the range ``[0, num_parts)``.
-    :arg self_mesh_groups: An array of `~meshmode.mesh.ElementGroup` instances
+    :arg self_mesh_groups: An array of `~meshmode.mesh.MeshElementGroup` instances
         representing the partitioned mesh groups.
     :arg global_group_to_self_group: An array mapping groups in *mesh* to groups in
         *self_mesh_groups* (or `None` if the group is not part of the current

From e3e32988ced8081b31ae70ee71cfae321d765600 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 20 Jun 2022 10:44:57 -0500
Subject: [PATCH 19/33] remove empty-group filtering

---
 meshmode/mesh/processing.py | 127 ++++++++++++------------------------
 1 file changed, 41 insertions(+), 86 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 32a0eda8b..6ecf1bf05 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -116,52 +116,38 @@ def _filter_mesh_groups(mesh, selected_elements, vertex_id_dtype):
     :arg selected_elements: A sorted array of indices of elements to be included in
         the filtered groups.
     :arg vertex_id_dtype: The vertex index data type.
-    :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*.
+    :returns: A tuple ``(new_groups, required_vertex_indices)``, where *new_groups*
+        is made up of groups from *groups* 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)]
+            mesh.groups[igrp].vertex_indices[
+                    filtered_group_elements[igrp], :]
+            for igrp in range(len(mesh.groups))]
 
-    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)
+    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)
@@ -178,14 +164,13 @@ 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(
@@ -235,8 +220,7 @@ def _get_connected_partitions(
 
 
 def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
-            self_part_index, self_mesh_groups, global_group_to_self_group,
-            self_mesh_group_elem_base):
+            self_part_index, self_mesh_groups, self_mesh_group_elem_base):
     r"""
     Create self-to-self facial adjacency groups for a partitioned mesh.
 
@@ -248,9 +232,6 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
         the range ``[0, num_parts)``.
     :arg self_mesh_groups: An array of :class:`~meshmode.mesh.MeshElementGroup`
         instances representing the partitioned mesh groups.
-    :arg global_group_to_self_group: An array mapping groups in *mesh* to groups in
-        *self_mesh_groups* (or `None` if the group is not part of the current
-        partition).
     :arg self_mesh_group_elem_base: An array containing the starting partition-wide
         element index for each group in *self_mesh_groups*.
 
@@ -261,10 +242,6 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
     self_to_self_adjacency_groups = [[] for _ in self_mesh_groups]
 
     for igrp, facial_adj_list in enumerate(mesh.facial_adjacency_groups):
-        i_self_grp = global_group_to_self_group[igrp]
-        if i_self_grp is None:
-            continue
-
         int_grps = [
             grp for grp in facial_adj_list
             if isinstance(grp, InteriorAdjacencyGroup)]
@@ -272,10 +249,6 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
         for facial_adj in int_grps:
             jgrp = facial_adj.ineighbor_group
 
-            j_self_grp = global_group_to_self_group[jgrp]
-            if j_self_grp is None:
-                continue
-
             elem_base_i = mesh.base_element_nrs[igrp]
             elem_base_j = mesh.base_element_nrs[jgrp]
 
@@ -287,8 +260,8 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
             adj_indices, = np.where(elements_are_self & neighbors_are_self)
 
             if len(adj_indices) > 0:
-                self_elem_base_i = self_mesh_group_elem_base[i_self_grp]
-                self_elem_base_j = self_mesh_group_elem_base[j_self_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, 1] - self_elem_base_i
@@ -297,10 +270,10 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
                             adj_indices] + elem_base_j, 1] - self_elem_base_j
                 neighbor_faces = facial_adj.neighbor_faces[adj_indices]
 
-                self_to_self_adjacency_groups[i_self_grp].append(
+                self_to_self_adjacency_groups[igrp].append(
                     InteriorAdjacencyGroup(
-                        igroup=i_self_grp,
-                        ineighbor_group=j_self_grp,
+                        igroup=igrp,
+                        ineighbor_group=jgrp,
                         elements=elements,
                         element_faces=element_faces,
                         neighbors=neighbors,
@@ -312,8 +285,7 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
 
 def _create_self_to_other_adjacency_groups(
         mesh, part_id_to_part_index, global_elem_to_part_elem, self_part_id,
-        self_mesh_groups, global_group_to_self_group, self_mesh_group_elem_base,
-        connected_parts):
+        self_mesh_groups, self_mesh_group_elem_base, connected_parts):
     """
     Create self-to-other adjacency groups for the partitioned mesh.
 
@@ -326,9 +298,6 @@ def _create_self_to_other_adjacency_groups(
     :arg self_part_id: The identifier of the partition currently being created.
     :arg self_mesh_groups: An array of `~meshmode.mesh.MeshElementGroup` instances
         representing the partitioned mesh groups.
-    :arg global_group_to_self_group: An array mapping groups in *mesh* to groups in
-        *self_mesh_groups* (or `None` if the group is not part of the current
-        partition).
     :arg self_mesh_group_elem_base: An array containing the starting partition-wide
         element index for each group in *self_mesh_groups*.
     :arg connected_parts: A :class:`set` containing the partitions connected to
@@ -343,10 +312,6 @@ def _create_self_to_other_adjacency_groups(
     self_to_other_adj_groups = [[] for _ in self_mesh_groups]
 
     for igrp, facial_adj_list in enumerate(mesh.facial_adjacency_groups):
-        i_self_grp = global_group_to_self_group[igrp]
-        if i_self_grp is None:
-            continue
-
         int_grps = [
             grp for grp in facial_adj_list
             if isinstance(grp, InteriorAdjacencyGroup)]
@@ -372,7 +337,7 @@ def _create_self_to_other_adjacency_groups(
                     & (neighbor_part_indices == neighbor_part_index))
 
                 if len(adj_indices) > 0:
-                    self_elem_base_i = self_mesh_group_elem_base[i_self_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, 1] - self_elem_base_i
@@ -381,9 +346,9 @@ def _create_self_to_other_adjacency_groups(
                         global_neighbors[adj_indices], 1]
                     neighbor_faces = facial_adj.neighbor_faces[adj_indices]
 
-                    self_to_other_adj_groups[i_self_grp].append(
+                    self_to_other_adj_groups[igrp].append(
                         InterPartitionAdjacencyGroup(
-                            igroup=i_self_grp,
+                            igroup=igrp,
                             boundary_tag=BTAG_PARTITION(neighbor_part_id),
                             elements=elements,
                             element_faces=element_faces,
@@ -395,7 +360,7 @@ def _create_self_to_other_adjacency_groups(
 
 
 def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
-            self_mesh_groups, global_group_to_self_group, self_mesh_group_elem_base):
+            self_mesh_groups, self_mesh_group_elem_base):
     """
     Create boundary groups for partitioned mesh.
 
@@ -407,9 +372,6 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
         the range ``[0, num_parts)``.
     :arg self_mesh_groups: An array of `~meshmode.mesh.MeshElementGroup` instances
         representing the partitioned mesh groups.
-    :arg global_group_to_self_group: An array mapping groups in *mesh* to groups in
-        *self_mesh_groups* (or `None` if the group is not part of the current
-        partition).
     :arg self_mesh_group_elem_base: An array containing the starting partition-wide
         element index for each group in *self_mesh_groups*.
 
@@ -420,10 +382,6 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
     bdry_adj_groups = [[] for _ in self_mesh_groups]
 
     for igrp, facial_adj_list in enumerate(mesh.facial_adjacency_groups):
-        i_self_grp = global_group_to_self_group[igrp]
-        if i_self_grp is None:
-            continue
-
         bdry_grps = [
             grp for grp in facial_adj_list
             if isinstance(grp, BoundaryAdjacencyGroup)]
@@ -436,7 +394,7 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
                 == self_part_index)
 
             if len(adj_indices) > 0:
-                self_elem_base = self_mesh_group_elem_base[i_self_grp]
+                self_elem_base = self_mesh_group_elem_base[igrp]
                 elements = global_elem_to_part_elem[bdry_grp.elements[adj_indices]
                             + elem_base, 1] - self_elem_base
                 element_faces = bdry_grp.element_faces[adj_indices]
@@ -444,9 +402,9 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
                 elements = np.empty(0, dtype=mesh.element_id_dtype)
                 element_faces = np.empty(0, dtype=mesh.face_id_dtype)
 
-            bdry_adj_groups[i_self_grp].append(
+            bdry_adj_groups[igrp].append(
                 BoundaryAdjacencyGroup(
-                    igroup=i_self_grp,
+                    igroup=igrp,
                     boundary_tag=bdry_grp.boundary_tag,
                     elements=elements,
                     element_faces=element_faces))
@@ -481,10 +439,8 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
 
     # Create new mesh groups that mimic the original mesh's groups but only contain
     # the current partition's elements
-    self_mesh_groups, global_group_to_self_group, required_vertex_indices =\
-                _filter_mesh_groups(
-                    mesh, part_id_to_elements[self_part_id],
-                    mesh.vertex_id_dtype)
+    self_mesh_groups, required_vertex_indices = _filter_mesh_groups(
+        mesh, part_id_to_elements[self_part_id], mesh.vertex_id_dtype)
 
     self_part_index = part_id_to_part_index[self_part_id]
 
@@ -494,8 +450,8 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
 
     self_mesh_group_elem_base = [0 for _ in self_mesh_groups]
     el_nr = 0
-    for i_self_grp, grp in enumerate(self_mesh_groups):
-        self_mesh_group_elem_base[i_self_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(
@@ -504,23 +460,22 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
 
     self_to_self_adj_groups = _create_self_to_self_adjacency_groups(
                 mesh, global_elem_to_part_elem, self_part_index, self_mesh_groups,
-                global_group_to_self_group, self_mesh_group_elem_base)
+                self_mesh_group_elem_base)
 
     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, global_group_to_self_group,
-                self_mesh_group_elem_base, connected_parts)
+                self_mesh_groups, self_mesh_group_elem_base, connected_parts)
 
     boundary_adj_groups = _create_boundary_groups(
                 mesh, global_elem_to_part_elem, self_part_index, self_mesh_groups,
-                global_group_to_self_group, self_mesh_group_elem_base)
+                self_mesh_group_elem_base)
 
     # Combine adjacency groups
     self_facial_adj_groups = [
-        self_to_self_adj_groups[i_self_grp]
-        + self_to_other_adj_groups[i_self_grp]
-        + boundary_adj_groups[i_self_grp]
-        for i_self_grp in range(len(self_mesh_groups))]
+        self_to_self_adj_groups[igrp]
+        + self_to_other_adj_groups[igrp]
+        + boundary_adj_groups[igrp]
+        for igrp in range(len(self_mesh_groups))]
 
     from meshmode.mesh import Mesh
     self_mesh = Mesh(

From f54e47aa535b6b6eb4a7103f86e371dd964dda6d Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 27 Jun 2022 09:49:04 -0500
Subject: [PATCH 20/33] clarify part vs. partition terminology

---
 .../connection/opposite_face.py               | 13 ++-
 meshmode/distributed.py                       | 63 ++++++++-------
 meshmode/mesh/__init__.py                     | 21 +++--
 meshmode/mesh/processing.py                   | 81 +++++++++----------
 test/test_partition.py                        | 38 ++++-----
 5 files changed, 109 insertions(+), 107 deletions(-)

diff --git a/meshmode/discretization/connection/opposite_face.py b/meshmode/discretization/connection/opposite_face.py
index 4e08a3ea5..19051ac8a 100644
--- a/meshmode/discretization/connection/opposite_face.py
+++ b/meshmode/discretization/connection/opposite_face.py
@@ -525,17 +525,16 @@ def make_opposite_face_connection(actx, volume_to_bdry_conn):
 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
 
@@ -557,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).
@@ -565,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 dad61158b..4cafdf423 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -5,7 +5,7 @@
 
 .. autofunction:: get_partition_by_pymetis
 .. autofunction:: membership_list_to_map
-.. autofunction:: get_connected_partitions
+.. autofunction:: get_connected_parts
 
 .. autoclass:: RemoteGroupInfo
 .. autoclass:: make_remote_group_infos
@@ -47,8 +47,8 @@
 from meshmode.mesh import (
         Mesh,
         InteriorAdjacencyGroup,
-        InterPartitionAdjacencyGroup,
-        PartitionID,
+        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()
@@ -119,7 +119,7 @@ def send_mesh_parts(self, mesh, part_per_element, num_parts):
             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()
 
@@ -149,11 +149,11 @@ 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-partition
-# connections between two parts on the same rank.
+# 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
@@ -161,7 +161,7 @@ class RemoteGroupInfo:
 
 def make_remote_group_infos(
         actx: ArrayContext,
-        remote_part_id: PartitionID,
+        remote_part_id: PartID,
         bdry_conn: DirectDiscretizationConnection
         ) -> Sequence[RemoteGroupInfo]:
     local_vol_mesh = bdry_conn.from_discr.mesh
@@ -170,9 +170,9 @@ 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)
+                    if isinstance(fagrp, InterPartAdjacencyGroup)
                     and fagrp.boundary_tag.part_id == remote_part_id],
                 vol_elem_indices=np.concatenate([
                     actx.to_numpy(batch.from_element_indices)
@@ -195,11 +195,11 @@ class InterRankBoundaryInfo:
     """
     .. 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
 
@@ -214,15 +214,15 @@ class InterRankBoundaryInfo:
     .. automethod:: __init__
     """
 
-    local_part_id: PartitionID
-    remote_part_id: PartitionID
+    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
         """
@@ -313,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
 
@@ -393,7 +392,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
@@ -443,15 +442,21 @@ def membership_list_to_map(membership_list):
         for entry in set(membership_list)}
 
 
-def get_connected_partitions(mesh: Mesh) -> "Set[PartitionID]":
-    """For a local mesh part in *mesh*, determine the set of connected partitions."""
+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.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 5b8e39aa0..c5c21f49b 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,7 +67,7 @@
 # {{{ element tags
 
 BoundaryTag = Hashable
-PartitionID = Hashable
+PartID = Hashable
 
 
 class BTAG_NONE:  # noqa: N801
@@ -110,14 +110,14 @@ 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 identifier of the adjacent mesh is given
+    another :class:`Mesh`. The part identifier of the adjacent mesh is given
     by ``part_id``.
 
     .. attribute:: part_id
 
     .. versionadded:: 2017.1
     """
-    def __init__(self, part_id: PartitionID, part_nr=None):
+    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. "
@@ -770,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
@@ -800,7 +800,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
@@ -809,8 +809,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
 
@@ -832,7 +831,7 @@ def __eq__(self, other):
             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(
@@ -1647,7 +1646,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 6ecf1bf05..1e8dde34b 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -35,7 +35,7 @@
     BTAG_PARTITION,
     InteriorAdjacencyGroup,
     BoundaryAdjacencyGroup,
-    InterPartitionAdjacencyGroup
+    InterPartAdjacencyGroup
 )
 
 from meshmode.mesh.tools import AffineMap
@@ -83,20 +83,19 @@ def find_group_indices(groups, meshwide_elems):
 def _compute_global_elem_to_part_elem(
         nelements, part_id_to_elements, part_id_to_part_index, element_id_dtype):
     """
-    Create a map from global element index to partition-wide element index for
-    a set of partitions.
+    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 partition identifiers to
+    :arg part_id_to_elements: A :class:`dict` mapping part identifiers to
         sets of elements.
-    :arg part_id_to_part_index: A mapping from partition identifiers to indices in
+    :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` ``global_elem_to_part_elem`` of shape
         ``(nelements, 2)``, where ``global_elem_to_part_elem[ielement, 0]`` gives
-        the partition index of the element and
-        ``global_elem_to_part_elem[ielement, 1]`` gives its partition-wide element
-        index.
+        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.empty((nelements, 2), dtype=element_id_dtype)
     for part_id in part_id_to_elements.keys():
@@ -173,20 +172,20 @@ def _filter_mesh_groups(mesh, selected_elements, vertex_id_dtype):
     return new_groups, required_vertex_indices
 
 
-def _get_connected_partitions(
+def _get_connected_parts(
         mesh, part_id_to_part_index, global_elem_to_part_elem, self_part_id):
     """
-    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_id_to_part_index: A mapping from partition identifiers to indices in
-        the range ``[0, num_parts)``.
+    :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 partition indices and partition-wide element indices. See
+        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 partition currently being created.
+    :arg self_part_id: The identifier of the part currently being created.
 
-    :returns: A :class:`set` of identifiers of the neighboring partitions.
+    :returns: A :class:`set` of identifiers of the neighboring parts.
     """
     self_part_index = part_id_to_part_index[self_part_id]
 
@@ -226,13 +225,13 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
 
     :arg mesh: A :class:`~meshmode.mesh.Mesh` representing the unpartitioned mesh.
     :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element
-        indices to partition indices and partition-wide element indices. See
+        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 partition currently being created, in
-        the range ``[0, num_parts)``.
+    :arg self_part_index: The index of the part currently being created, in the
+        range ``[0, num_parts)``.
     :arg self_mesh_groups: An array of :class:`~meshmode.mesh.MeshElementGroup`
         instances representing the partitioned mesh groups.
-    :arg self_mesh_group_elem_base: An array containing the starting partition-wide
+    :arg self_mesh_group_elem_base: An array containing the starting part-wide
         element index for each group in *self_mesh_groups*.
 
     :returns: A list of lists of `~meshmode.mesh.InteriorAdjacencyGroup` instances
@@ -290,21 +289,21 @@ def _create_self_to_other_adjacency_groups(
     Create self-to-other adjacency groups for the partitioned mesh.
 
     :arg mesh: A :class:`~meshmode.mesh.Mesh` representing the unpartitioned mesh.
-    :arg part_id_to_part_index: A mapping from partition identifiers to indices in
-        the range ``[0, num_parts)``.
+    :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 partition indices and partition-wide element indices. See
+        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 partition currently being created.
+    :arg self_part_id: The identifier of the part currently being created.
     :arg self_mesh_groups: An array of `~meshmode.mesh.MeshElementGroup` instances
         representing the partitioned mesh groups.
-    :arg self_mesh_group_elem_base: An array containing the starting partition-wide
+    :arg self_mesh_group_elem_base: An array containing the starting part-wide
         element index for each group in *self_mesh_groups*.
-    :arg connected_parts: A :class:`set` containing the partitions connected to
+    :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
+    :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.
     """
     self_part_index = part_id_to_part_index[self_part_id]
@@ -347,7 +346,7 @@ def _create_self_to_other_adjacency_groups(
                     neighbor_faces = facial_adj.neighbor_faces[adj_indices]
 
                     self_to_other_adj_groups[igrp].append(
-                        InterPartitionAdjacencyGroup(
+                        InterPartAdjacencyGroup(
                             igroup=igrp,
                             boundary_tag=BTAG_PARTITION(neighbor_part_id),
                             elements=elements,
@@ -366,13 +365,13 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
 
     :arg mesh: A :class:`~meshmode.mesh.Mesh` representing the unpartitioned mesh.
     :arg global_elem_to_part_elem: A :class:`numpy.ndarray` that maps global element
-        indices to partition indices and partition-wide element indices. See
+        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 partition currently being created, in
-        the range ``[0, num_parts)``.
+    :arg self_part_index: The index of the part currently being created, in the
+        range ``[0, num_parts)``.
     :arg self_mesh_groups: An array of `~meshmode.mesh.MeshElementGroup` instances
         representing the partitioned mesh groups.
-    :arg self_mesh_group_elem_base: An array containing the starting partition-wide
+    :arg self_mesh_group_elem_base: An array containing the starting part-wide
         element index for each group in *self_mesh_groups*.
 
     :returns: A list of lists of `~meshmode.mesh.BoundaryAdjacencyGroup` instances
@@ -415,11 +414,11 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
 def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
     """
     :arg mesh: A :class:`~meshmode.mesh.Mesh` to be partitioned.
-    :arg part_id_to_elements: A :class:`dict` mapping partition identifiers to
+    :arg part_id_to_elements: A :class:`dict` mapping part identifiers to
         sets of elements.
-    :arg self_part_id: The partition identifier of the mesh to return.
+    :arg self_part_id: The part identifier of the mesh to return.
 
-    :returns: A :class:`~meshmode.mesh.Mesh` containing a partition of *mesh*.
+    :returns: A :class:`~meshmode.mesh.Mesh` containing a part of *mesh*.
 
     .. versionadded:: 2017.1
     """
@@ -438,7 +437,7 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
         mesh.element_id_dtype)
 
     # Create new mesh groups that mimic the original mesh's groups but only contain
-    # the current partition's elements
+    # 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)
 
@@ -454,7 +453,7 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
         self_mesh_group_elem_base[igrp] = el_nr
         el_nr += grp.nelements
 
-    connected_parts = _get_connected_partitions(
+    connected_parts = _get_connected_parts(
         mesh, part_id_to_part_index, global_elem_to_part_elem,
         self_part_id)
 
@@ -490,11 +489,11 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
 def partition_mesh(mesh, part_id_to_elements):
     """
     :arg mesh: A :class:`~meshmode.mesh.Mesh` to be partitioned.
-    :arg part_id_to_elements: A :class:`dict` mapping partition identifiers to
-        sets of elements.
+    :arg part_id_to_elements: A :class:`dict` mapping part identifiers to sets of
+        elements.
 
-    :returns: A :class:`dict` mapping partition identifiers to instances of
-        :class:`~meshmode.mesh.Mesh` that represent the corresponding partition of
+    :returns: A :class:`dict` mapping part identifiers to instances of
+        :class:`~meshmode.mesh.Mesh` that represent the corresponding part of
         *mesh*.
     """
     return {
diff --git a/test/test_partition.py b/test/test_partition.py
index b0bd98ee1..3cdae413f 100644
--- a/test/test_partition.py
+++ b/test/test_partition.py
@@ -39,7 +39,7 @@
     BTAG_ALL,
     InteriorAdjacencyGroup,
     BoundaryAdjacencyGroup,
-    InterPartitionAdjacencyGroup
+    InterPartAdjacencyGroup
 )
 
 import pytest
@@ -106,8 +106,8 @@ def f(x):
 
         connected_parts = set()
         for i_local_part, part_mesh in part_meshes.items():
-            from meshmode.distributed import get_connected_partitions
-            neighbors = get_connected_partitions(part_mesh)
+            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))
 
@@ -136,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
 
@@ -195,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),
@@ -207,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
@@ -217,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")
@@ -245,8 +245,8 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio
 
     connected_parts = set()
     for i_local_part in range(num_parts):
-        from meshmode.distributed import get_connected_partitions
-        neighbors = get_connected_partitions(part_meshes[i_local_part])
+        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))
 
@@ -260,7 +260,7 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio
         for igrp in range(len(part_mesh.groups)):
             ipagrps = [
                 fagrp for fagrp in part_mesh.facial_adjacency_groups[igrp]
-                if isinstance(fagrp, InterPartitionAdjacencyGroup)]
+                if isinstance(fagrp, InterPartAdjacencyGroup)]
             for ipagrp in ipagrps:
                 for i, (elem, face) in enumerate(
                         zip(ipagrp.elements, ipagrp.element_faces)):
@@ -277,7 +277,7 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio
         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_id
@@ -297,7 +297,7 @@ 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)
+                        if isinstance(fagrp, InterPartAdjacencyGroup)
                         and fagrp.boundary_tag.part_id == part_num]
                     found_reverse_adj = False
                     for n_ipagrp in n_ipagrps:
@@ -314,7 +314,7 @@ def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitio
                             p_meshwide_n_elem = (
                                 part_elem_to_global_elem[n_part_num][
                                     n_meshwide_elem])
-                    assert found_reverse_adj, ("InterPartitionAdjacencyGroup is not "
+                    assert found_reverse_adj, ("InterPartAdjacencyGroup is not "
                         "consistent")
 
                     p_grp_num = find_group_indices(mesh.groups, p_meshwide_elem)
@@ -339,7 +339,7 @@ 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
@@ -405,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)
@@ -579,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([

From 4e84a59d0f5d8f03d156391451edd762b6187b4a Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 27 Jun 2022 11:40:25 -0500
Subject: [PATCH 21/33] change some asserts to exceptions

---
 meshmode/distributed.py     | 4 +++-
 meshmode/mesh/processing.py | 5 +++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index 4cafdf423..f6b56eeef 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -342,7 +342,9 @@ def complete_some(self):
         part_ids_to_irbi = {
                 (irbi.local_part_id, irbi.remote_part_id): irbi
                 for irbi in self.inter_rank_bdry_info}
-        assert len(part_ids_to_irbi) == len(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):
diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 1e8dde34b..3e252c8bf 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -423,8 +423,9 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
     .. versionadded:: 2017.1
     """
     nelements = sum(len(elems) for elems in part_id_to_elements.values())
-    assert nelements == mesh.nelements, \
-        "counts of partitioned elements must add up to mesh.nelements"
+    if nelements != mesh.nelements:
+        raise ValueError(
+            "counts of partitioned elements don't add up to mesh.nelements")
 
     part_id_to_part_index = {
         part_id: part_index

From f93ddfa4a98ed58c8b7f8d81d9a725c7b5d40d02 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 27 Jun 2022 11:47:48 -0500
Subject: [PATCH 22/33] add a couple of FIXMEs

---
 meshmode/distributed.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index f6b56eeef..c8ada638d 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -385,6 +385,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`.
 
@@ -444,6 +445,7 @@ def membership_list_to_map(membership_list):
         for entry in set(membership_list)}
 
 
+# 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

From e76d2cb3e16f360fd5085965573e94737c5fea2e Mon Sep 17 00:00:00 2001
From: Matt Smith <mjsmith1@gmail.com>
Date: Mon, 27 Jun 2022 12:20:38 -0500
Subject: [PATCH 23/33] cosmetic change
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Andreas Klöckner <inform@tiker.net>
---
 meshmode/distributed.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index c8ada638d..871733784 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -346,8 +346,7 @@ def complete_some(self):
             raise ValueError(
                 "duplicate local/remote part pair in inter_rank_bdry_info")
 
-        for i_src_rank, recvd in zip(
-                source_ranks, data):
+        for i_src_rank, recvd in zip(source_ranks, data):
             (remote_part_id, local_part_id,
                     remote_bdry_mesh, remote_group_infos) = recvd
 

From 74425fef4e4ae7906309135abdd8a7eb4adccb24 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 27 Jun 2022 12:35:15 -0500
Subject: [PATCH 24/33] detect elements that belong to multiple parts

---
 meshmode/mesh/processing.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 3e252c8bf..f0f6bb764 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -422,10 +422,13 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
 
     .. versionadded:: 2017.1
     """
-    nelements = sum(len(elems) for elems in part_id_to_elements.values())
-    if nelements != mesh.nelements:
-        raise ValueError(
-            "counts of partitioned elements don't add up to mesh.nelements")
+    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

From 11bd23bd60ad2fab37773bf34f778b8bea759e47 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 27 Jun 2022 13:09:57 -0500
Subject: [PATCH 25/33] add return_parts argument to partition_mesh

---
 meshmode/mesh/processing.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index f0f6bb764..17e509417 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -490,19 +490,24 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
     return self_mesh
 
 
-def partition_mesh(mesh, part_id_to_elements):
+def partition_mesh(mesh, part_id_to_elements, return_parts):
     """
     :arg mesh: A :class:`~meshmode.mesh.Mesh` to be partitioned.
     :arg part_id_to_elements: A :class:`dict` mapping part identifiers to sets 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 = part_id_to_elements.keys()
+
     return {
         part_id: _get_mesh_part(mesh, part_id_to_elements, part_id)
-        for part_id in part_id_to_elements.keys()}
+        for part_id in return_parts}
 
 # }}}
 

From 8ee9e1d9a3161d4c13c8f54cb1e40946acc92478 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 27 Jun 2022 10:41:45 -0500
Subject: [PATCH 26/33] add type hints to mesh partitioning

---
 meshmode/mesh/processing.py | 96 ++++++++++++++++++++++---------------
 1 file changed, 58 insertions(+), 38 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index b41b223c4..6b8cfdeba 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, Any, Tuple, Dict, List, Set
 
 from dataclasses import dataclass
 
@@ -32,7 +32,10 @@
 import modepy as mp
 
 from meshmode.mesh import (
+    MeshElementGroup,
+    Mesh,
     BTAG_PARTITION,
+    PartID,
     InteriorAdjacencyGroup,
     BoundaryAdjacencyGroup,
     InterPartAdjacencyGroup
@@ -81,14 +84,17 @@ def find_group_indices(groups, meshwide_elems):
 # {{{ partition_mesh
 
 def _compute_global_elem_to_part_elem(
-        nelements, part_id_to_elements, part_id_to_part_index, element_id_dtype):
+        nelements: int,
+        part_id_to_elements: Dict[PartID, np.ndarray],
+        part_id_to_part_index: Dict[PartID, int],
+        element_id_dtype: Any) -> np.ndarray:
     """
     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 part identifiers to
-        sets of elements.
+    :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.
@@ -107,7 +113,10 @@ def _compute_global_elem_to_part_elem(
     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: Any) -> Tuple[List, np.ndarray]:
     """
     Create new mesh groups containing a selected subset of elements.
 
@@ -174,7 +183,10 @@ def _filter_mesh_groups(mesh, selected_elements, vertex_id_dtype):
 
 
 def _get_connected_parts(
-        mesh, part_id_to_part_index, global_elem_to_part_elem, self_part_id):
+        mesh: Mesh,
+        part_id_to_part_index: Dict[PartID, int],
+        global_elem_to_part_elem: np.ndarray,
+        self_part_id: PartID) -> "Set[PartID]":
     """
     Find the parts that are connected to the current part.
 
@@ -219,8 +231,12 @@ def _get_connected_parts(
         if part_index in connected_part_indices}
 
 
-def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
-            self_part_index, self_mesh_groups, self_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 self-to-self facial adjacency groups for a partitioned mesh.
 
@@ -230,9 +246,9 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
         :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: An array of :class:`~meshmode.mesh.MeshElementGroup`
+    :arg self_mesh_groups: A list of :class:`~meshmode.mesh.MeshElementGroup`
         instances representing the partitioned mesh groups.
-    :arg self_mesh_group_elem_base: An array containing the starting part-wide
+    :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
@@ -284,8 +300,13 @@ def _create_self_to_self_adjacency_groups(mesh, global_elem_to_part_elem,
 
 
 def _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):
+        mesh: Mesh,
+        part_id_to_part_index: Dict[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 self-to-other adjacency groups for the partitioned mesh.
 
@@ -296,9 +317,9 @@ def _create_self_to_other_adjacency_groups(
         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: An array of `~meshmode.mesh.MeshElementGroup` instances
+    :arg self_mesh_groups: A list of `~meshmode.mesh.MeshElementGroup` instances
         representing the partitioned mesh groups.
-    :arg self_mesh_group_elem_base: An array containing the starting part-wide
+    :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.
@@ -359,8 +380,12 @@ def _create_self_to_other_adjacency_groups(
     return self_to_other_adj_groups
 
 
-def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
-            self_mesh_groups, self_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.
 
@@ -370,9 +395,9 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
         :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: An array of `~meshmode.mesh.MeshElementGroup` instances
+    :arg self_mesh_groups: A list of `~meshmode.mesh.MeshElementGroup` instances
         representing the partitioned mesh groups.
-    :arg self_mesh_group_elem_base: An array containing the starting part-wide
+    :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
@@ -412,11 +437,14 @@ def _create_boundary_groups(mesh, global_elem_to_part_elem, self_part_index,
     return bdry_adj_groups
 
 
-def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
+def _get_mesh_part(
+        mesh: Mesh,
+        part_id_to_elements: Dict[PartID, np.ndarray],
+        self_part_id: PartID) -> Mesh:
     """
     :arg mesh: A :class:`~meshmode.mesh.Mesh` to be partitioned.
-    :arg part_id_to_elements: A :class:`dict` mapping part identifiers to
-        sets of elements.
+    :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 :class:`~meshmode.mesh.Mesh` containing a part of *mesh*.
@@ -433,9 +461,7 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
 
     part_id_to_part_index = {
         part_id: part_index
-        for part_id, part_index in zip(
-            part_id_to_elements.keys(),
-            range(len(part_id_to_elements)))}
+        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,
@@ -481,21 +507,21 @@ def _get_mesh_part(mesh, part_id_to_elements, self_part_id):
         + boundary_adj_groups[igrp]
         for igrp in range(len(self_mesh_groups))]
 
-    from meshmode.mesh import Mesh
-    self_mesh = Mesh(
+    return Mesh(
             self_vertices,
             self_mesh_groups,
             facial_adjacency_groups=self_facial_adj_groups,
             is_conforming=mesh.is_conforming)
 
-    return self_mesh
-
 
-def partition_mesh(mesh, part_id_to_elements, return_parts):
+def partition_mesh(
+        mesh: Mesh,
+        part_id_to_elements: Dict[PartID, np.ndarray],
+        return_parts: Optional[Sequence[PartID]] = None) -> "Dict[PartID, Mesh]":
     """
     :arg mesh: A :class:`~meshmode.mesh.Mesh` to be partitioned.
-    :arg part_id_to_elements: A :class:`dict` mapping part identifiers to sets of
-        elements.
+    :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.
 
@@ -504,7 +530,7 @@ def partition_mesh(mesh, part_id_to_elements, return_parts):
         *mesh*.
     """
     if return_parts is None:
-        return_parts = part_id_to_elements.keys()
+        return_parts = list(part_id_to_elements.keys())
 
     return {
         part_id: _get_mesh_part(mesh, part_id_to_elements, part_id)
@@ -702,8 +728,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]
@@ -841,7 +865,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,
@@ -899,7 +922,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,
@@ -1189,8 +1211,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):

From 5fefcea03c58a085794b78c799b8783d3b681610 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Mon, 27 Jun 2022 15:59:56 -0500
Subject: [PATCH 27/33] Fix some annotations in mesh.processing

---
 meshmode/mesh/processing.py | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 6b8cfdeba..b17b0c818 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, Any, Tuple, Dict, List, Set
+from typing import Optional, Union, Any, Tuple, Dict, List, Set, Sequence
 
 from dataclasses import dataclass
 
@@ -36,6 +36,7 @@
     Mesh,
     BTAG_PARTITION,
     PartID,
+    FacialAdjacencyGroup,
     InteriorAdjacencyGroup,
     BoundaryAdjacencyGroup,
     InterPartAdjacencyGroup
@@ -255,7 +256,8 @@ def _create_self_to_self_adjacency_groups(
         corresponding to the entries in *mesh.facial_adjacency_groups* that
         have self-to-self adjacency.
     """
-    self_to_self_adjacency_groups = [[] for _ in self_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):
         int_grps = [
@@ -330,7 +332,8 @@ def _create_self_to_other_adjacency_groups(
     """
     self_part_index = part_id_to_part_index[self_part_id]
 
-    self_to_other_adj_groups = [[] for _ in self_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):
         int_grps = [
@@ -404,7 +407,8 @@ def _create_boundary_groups(
         corresponding to the entries in *mesh.facial_adjacency_groups* that have
         boundary faces.
     """
-    bdry_adj_groups = [[] for _ in self_mesh_groups]
+    bdry_adj_groups: List[List[BoundaryAdjacencyGroup]] = [
+            [] for _ in self_mesh_groups]
 
     for igrp, facial_adj_list in enumerate(mesh.facial_adjacency_groups):
         bdry_grps = [
@@ -500,12 +504,16 @@ def _get_mesh_part(
                 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 = [
-        self_to_self_adj_groups[igrp]
-        + self_to_other_adj_groups[igrp]
-        + boundary_adj_groups[igrp]
-        for igrp in range(len(self_mesh_groups))]
+            _gather_grps(igrp) for igrp in range(len(self_mesh_groups))]
 
     return Mesh(
             self_vertices,

From 58751d0c5edd368610830f93c53bb47555043606 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 27 Jun 2022 16:59:07 -0500
Subject: [PATCH 28/33] add explicit part_id attribute to
 InterPartAdjacencyGroup

---
 meshmode/distributed.py     | 4 ++--
 meshmode/mesh/__init__.py   | 7 +++++++
 meshmode/mesh/processing.py | 1 +
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/meshmode/distributed.py b/meshmode/distributed.py
index 871733784..d99e90dc8 100644
--- a/meshmode/distributed.py
+++ b/meshmode/distributed.py
@@ -173,7 +173,7 @@ def make_remote_group_infos(
                 inter_part_adj_groups=[
                     fagrp for fagrp in local_vol_mesh.facial_adjacency_groups[igrp]
                     if isinstance(fagrp, InterPartAdjacencyGroup)
-                    and fagrp.boundary_tag.part_id == remote_part_id],
+                    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]),
@@ -450,7 +450,7 @@ def get_connected_parts(mesh: Mesh) -> "Set[PartID]":
     assert mesh.facial_adjacency_groups is not None
 
     return {
-            grp.boundary_tag.part_id
+            grp.part_id
             for fagrp_list in mesh.facial_adjacency_groups
             for grp in fagrp_list
             if isinstance(grp, InterPartAdjacencyGroup)}
diff --git a/meshmode/mesh/__init__.py b/meshmode/mesh/__init__.py
index c5c21f49b..69f57e683 100644
--- a/meshmode/mesh/__init__.py
+++ b/meshmode/mesh/__init__.py
@@ -784,6 +784,10 @@ class InterPartAdjacencyGroup(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.
@@ -819,6 +823,7 @@ class InterPartAdjacencyGroup(BoundaryAdjacencyGroup):
     .. versionadded:: 2017.1
     """
 
+    part_id: PartID
     neighbors: np.ndarray
     neighbor_faces: np.ndarray
     aff_map: AffineMap
@@ -826,6 +831,7 @@ class InterPartAdjacencyGroup(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)
@@ -837,6 +843,7 @@ def as_python(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),
diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index b17b0c818..32e08fa64 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -374,6 +374,7 @@ def _create_self_to_other_adjacency_groups(
                         InterPartAdjacencyGroup(
                             igroup=igrp,
                             boundary_tag=BTAG_PARTITION(neighbor_part_id),
+                            part_id=neighbor_part_id,
                             elements=elements,
                             element_faces=element_faces,
                             neighbors=neighbors,

From 1829309c84fa86ea6c48628ea455bec39156e400 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Mon, 27 Jun 2022 17:01:34 -0500
Subject: [PATCH 29/33] cosmetic change

---
 meshmode/mesh/processing.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index 32e08fa64..ae4dbf46e 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -151,10 +151,10 @@ def _filter_mesh_groups(
     # {{{ filter vertex indices
 
     filtered_vertex_indices = [
-            mesh.groups[igrp].vertex_indices[
+            grp.vertex_indices[
                     filtered_group_elements[igrp], :]
-            for igrp in range(len(mesh.groups))
-            if mesh.groups[igrp].vertex_indices is not None]
+            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])

From d4214a2a5bef4c068065f4d565a3ca5c9b955890 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Wed, 29 Jun 2022 18:55:02 -0500
Subject: [PATCH 30/33] fix some type annotations

---
 meshmode/mesh/processing.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index ae4dbf46e..c9179753f 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -88,7 +88,7 @@ def _compute_global_elem_to_part_elem(
         nelements: int,
         part_id_to_elements: Dict[PartID, np.ndarray],
         part_id_to_part_index: Dict[PartID, int],
-        element_id_dtype: Any) -> np.ndarray:
+        element_id_dtype: np.dtype[Any]) -> np.ndarray:
     """
     Create a map from global element index to part-wide element index for a set of
     parts.
@@ -117,7 +117,7 @@ def _compute_global_elem_to_part_elem(
 def _filter_mesh_groups(
         mesh: Mesh,
         selected_elements: np.ndarray,
-        vertex_id_dtype: Any) -> Tuple[List, np.ndarray]:
+        vertex_id_dtype: np.dtype[Any]) -> Tuple[List[MeshElementGroup], np.ndarray]:
     """
     Create new mesh groups containing a selected subset of elements.
 

From 2a7e743688769916dbe3254ae3a1e8ffec42154a Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Wed, 29 Jun 2022 18:55:56 -0500
Subject: [PATCH 31/33] Dict -> Mapping

---
 meshmode/mesh/processing.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index c9179753f..fea0173e9 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, Any, Tuple, Dict, List, Set, Sequence
+from typing import Optional, Union, Any, Tuple, Mapping, List, Set, Sequence
 
 from dataclasses import dataclass
 
@@ -86,8 +86,8 @@ def find_group_indices(groups, meshwide_elems):
 
 def _compute_global_elem_to_part_elem(
         nelements: int,
-        part_id_to_elements: Dict[PartID, np.ndarray],
-        part_id_to_part_index: Dict[PartID, int],
+        part_id_to_elements: Mapping[PartID, np.ndarray],
+        part_id_to_part_index: Mapping[PartID, int],
         element_id_dtype: np.dtype[Any]) -> np.ndarray:
     """
     Create a map from global element index to part-wide element index for a set of
@@ -185,7 +185,7 @@ def _filter_mesh_groups(
 
 def _get_connected_parts(
         mesh: Mesh,
-        part_id_to_part_index: Dict[PartID, int],
+        part_id_to_part_index: Mapping[PartID, int],
         global_elem_to_part_elem: np.ndarray,
         self_part_id: PartID) -> "Set[PartID]":
     """
@@ -303,7 +303,7 @@ def _create_self_to_self_adjacency_groups(
 
 def _create_self_to_other_adjacency_groups(
         mesh: Mesh,
-        part_id_to_part_index: Dict[PartID, int],
+        part_id_to_part_index: Mapping[PartID, int],
         global_elem_to_part_elem: np.ndarray,
         self_part_id: PartID,
         self_mesh_groups: List[MeshElementGroup],
@@ -444,7 +444,7 @@ def _create_boundary_groups(
 
 def _get_mesh_part(
         mesh: Mesh,
-        part_id_to_elements: Dict[PartID, np.ndarray],
+        part_id_to_elements: Mapping[PartID, np.ndarray],
         self_part_id: PartID) -> Mesh:
     """
     :arg mesh: A :class:`~meshmode.mesh.Mesh` to be partitioned.
@@ -525,8 +525,8 @@ def _gather_grps(igrp: int) -> List[FacialAdjacencyGroup]:
 
 def partition_mesh(
         mesh: Mesh,
-        part_id_to_elements: Dict[PartID, np.ndarray],
-        return_parts: Optional[Sequence[PartID]] = None) -> "Dict[PartID, 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

From fbe156bbae8278ef9ddef09a5c82b0f6fc666d14 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Wed, 29 Jun 2022 19:02:05 -0500
Subject: [PATCH 32/33] fix outdated argument reference in docstring

---
 meshmode/mesh/processing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index fea0173e9..fa3cef00c 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -126,7 +126,7 @@ def _filter_mesh_groups(
         the filtered 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 *groups* containing only elements from
+        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*.

From f55588d0e4ea8cbc65c60536ec4e0f5d95baa438 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Thu, 30 Jun 2022 11:23:50 -0500
Subject: [PATCH 33/33] try using just dtype

dtype[Any] requires python 3.9
---
 meshmode/mesh/processing.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index fa3cef00c..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, Any, Tuple, Mapping, List, Set, Sequence
+from typing import Optional, Union, Tuple, Mapping, List, Set, Sequence
 
 from dataclasses import dataclass
 
@@ -88,7 +88,7 @@ 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[Any]) -> np.ndarray:
+        element_id_dtype: np.dtype) -> np.ndarray:
     """
     Create a map from global element index to part-wide element index for a set of
     parts.
@@ -117,7 +117,7 @@ def _compute_global_elem_to_part_elem(
 def _filter_mesh_groups(
         mesh: Mesh,
         selected_elements: np.ndarray,
-        vertex_id_dtype: np.dtype[Any]) -> Tuple[List[MeshElementGroup], np.ndarray]:
+        vertex_id_dtype: np.dtype) -> Tuple[List[MeshElementGroup], np.ndarray]:
     """
     Create new mesh groups containing a selected subset of elements.