Permutation connection

meshmode/discretization/connection/__init__.py

@@ -57,6 +57,7 @@
 )
 from meshmode.discretization.connection.refinement import make_refinement_connection
 from meshmode.discretization.connection.same_mesh import make_same_mesh_connection
+from meshmode.discretization.connection.permute import make_element_permutation_connection
 
 
 logger = logging.getLogger(__name__)
@@ -81,6 +82,7 @@
     "make_partition_connection",
     "make_refinement_connection",
     "make_same_mesh_connection",
+    "make_element_permutation_connection",
 ]
 
 __doc__ = """
@@ -102,6 +104,10 @@
 ---------------------
 .. autofunction:: make_same_mesh_connection
 
+Mesh permutations
+---------------------
+.. autofunction:: make_element_permutation_connection
+
 Restriction to faces
 --------------------
 .. autodata:: FACE_RESTR_INTERIOR

meshmode/discretization/connection/permute.py

@@ -0,0 +1,90 @@
+__copyright__ = "Copyright (C) 2014 Andreas Kloeckner"
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+import numpy as np
+
+from arraycontext.metadata import NameHint
+
+from meshmode.discretization.connection.direct import (
+    DirectDiscretizationConnection,
+    DiscretizationConnectionElementGroup,
+    IdentityDiscretizationConnection,
+    InterpolationBatch,
+)
+from meshmode.transform_metadata import (
+    DiscretizationElementAxisTag,
+)
+
+
+# {{{ permute mesh elements
+
+def make_element_permutation_connection(actx, from_discr):
+    """Build discretization connection containing a permuted mesh.
+
+    :arg actx: an :class:`~arraycontext.ArrayContext`.
+    :arg from_discr: a :class:`Discretization`.
+    """
+
+    # Reorder the local mesh using Metis
+    from meshmode.distributed import get_reordering_by_pymetis
+    perm, iperm = get_reordering_by_pymetis(from_discr.mesh)
+
+    # Create target discretization
+    to_discr = from_discr.copy()
+
+    groups = []
+    for igrp, (fgrp, tgrp) in enumerate(
+        zip(from_discr.groups, to_discr.groups, strict=True)):
+        from arraycontext.metadata import NameHint
+        all_elements = actx.tag(NameHint(f"all_el_ind_grp{igrp}"),
+                                actx.tag_axis(0,
+                                              DiscretizationElementAxisTag(),
+                                              actx.from_numpy(
+                                                  np.arange(
+                                                      fgrp.nelements,
+                                                      dtype=np.intp))))
+        all_elements = actx.freeze(all_elements)
+        # Permuted elements
+        all_elements_perm = actx.tag(NameHint(f"all_el_ind_perm_grp{igrp}"),
+                                     actx.tag_axis(0,
+                                                   DiscretizationElementAxisTag(),
+                                                   actx.from_numpy(perm)))
+        all_elements_perm = actx.freeze(all_elements_perm)
+        ibatch = InterpolationBatch(
+                from_group_index=igrp,
+                from_element_indices=all_elements,
+                to_element_indices=all_elements_perm,  ## CHK??
+                result_unit_nodes=tgrp.unit_nodes,
+                to_element_face=None)
+
+        groups.append(
+                DiscretizationConnectionElementGroup([ibatch]))
+
+    return DirectDiscretizationConnection(
+        from_discr,
+        to_discr,
+        groups,
+        is_surjective=True)
+
+# }}}
+
+# vim: foldmethod=marker

meshmode/distributed.py

@@ -445,6 +445,49 @@ def get_partition_by_pymetis(mesh, num_parts, *, connectivity="facial", **kwargs
     return np.array(p)
 
 
+# FIXME: Move somewhere else, since it's not strictly limited to distributed?
+def get_reordering_by_pymetis(mesh, *, connectivity="facial", **kwargs):
+    """Return a reordered mesh created by :mod:`pymetis`.
+
+    :arg mesh: A :class:`meshmode.mesh.Mesh` instance
+    :arg connectivity: the adjacency graph to be used for partitioning. Either
+        ``"facial"`` or ``"nodal"`` (based on vertices).
+    :arg kwargs: Passed unmodified to :func:`pymetis.nested_dissection`.
+    :returns: a :class:`numpy.ndarray` with the permutation in p[0] and inverse
+        permutation in p[1]
+    """
+
+    if connectivity == "facial":
+        # shape: (2, n_el_pairs)
+        neighbor_el_pairs = np.hstack([
+                np.array([
+                    fagrp.elements + mesh.base_element_nrs[fagrp.igroup],
+                    fagrp.neighbors + mesh.base_element_nrs[fagrp.ineighbor_group]
+                    ])
+                for fagrp_list in mesh.facial_adjacency_groups
+                for fagrp in fagrp_list
+                if isinstance(fagrp, InteriorAdjacencyGroup)
+                ])
+        sorted_neighbor_el_pairs = neighbor_el_pairs[
+                :, np.argsort(neighbor_el_pairs[0])]
+        xadj = np.searchsorted(
+                sorted_neighbor_el_pairs[0],
+                np.arange(mesh.nelements+1))
+        adjncy = sorted_neighbor_el_pairs[1]
+
+    elif connectivity == "nodal":
+        xadj = mesh.nodal_adjacency.neighbors_starts.tolist()
+        adjncy = mesh.nodal_adjacency.neighbors.tolist()
+
+    else:
+        raise ValueError("invalid value of connectivity")
+
+    from pymetis import nested_dissection
+    p = nested_dissection(xadj=xadj, adjncy=adjncy, **kwargs)
+
+    return np.array(p)
+
+
 def membership_list_to_map(
             membership_list: np.ndarray[Any, Any]
         ) -> Mapping[Hashable, np.ndarray]: