combine/optimize _match_vertices and find_point_permutation implementations (#458, part 2)

* optimize vertex matching in periodicity processing

* use np.median instead of approximating with np.histogram

    turns out to be just as fast

* combine _match_vertices and mesh.tools.find_point_permutation into find_point_to_point_mapping

* add tests for find_point_to_point_mapping

* use find_point_to_point_mapping instead of find_point_permutation

Author: majosm

meshmode/discretization/connection/direct.py

@@ -422,12 +422,17 @@ def _resample_point_pick_indices(self, to_group_index: int, ibatch_index: int,
         ibatch = self.groups[to_group_index].batches[ibatch_index]
         from_grp = self.from_discr.groups[ibatch.from_group_index]
 
-        from meshmode.mesh.tools import find_point_permutation
-        return find_point_permutation(
-                    targets=ibatch.result_unit_nodes,
-                    permutees=from_grp.unit_nodes,
+        from meshmode.mesh.tools import find_point_to_point_mapping
+        src_idx_to_tgt_idx = find_point_to_point_mapping(
+                    src_points=ibatch.result_unit_nodes,
+                    tgt_points=from_grp.unit_nodes,
                     tol_multiplier=tol_multiplier)
 
+        return (
+            src_idx_to_tgt_idx
+            if np.all(src_idx_to_tgt_idx >= 0)
+            else None)
+
     @keyed_memoize_method(lambda actx, to_group_index, ibatch_index,
             tol_multiplier=None: (to_group_index, ibatch_index, tol_multiplier))
     def _frozen_resample_point_pick_indices(self, actx: ArrayContext,

meshmode/mesh/processing.py

@@ -25,7 +25,8 @@
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass, replace
 from functools import reduce
-from typing import Any, Literal
+from typing import Literal
+from warnings import warn
 
 import numpy as np
 import numpy.linalg as la
@@ -45,7 +46,7 @@
     _FaceIDs,
     make_mesh,
 )
-from meshmode.mesh.tools import AffineMap, find_point_permutation
+from meshmode.mesh.tools import AffineMap, find_point_to_point_mapping
 
 
 __doc__ = """
@@ -588,18 +589,18 @@ def evec(i: int) -> np.ndarray:
         result[i] = 1
         return result
 
-    def unpack_single(ary: np.ndarray | None) -> np.ndarray:
+    def unpack_single(ary: np.ndarray | None) -> int:
         assert ary is not None
         item, = ary
         return item
 
-    base_vertex_index = unpack_single(find_point_permutation(
-             targets=-np.ones(grp.dim),
-             permutees=grp.vertex_unit_coordinates().T))
+    base_vertex_index = unpack_single(find_point_to_point_mapping(
+             src_points=-np.ones(grp.dim).reshape(-1, 1),
+             tgt_points=grp.vertex_unit_coordinates().T))
     spanning_vertex_indices = [
-        unpack_single(find_point_permutation(
-                     targets=-np.ones(grp.dim) + 2 * evec(i),
-                     permutees=grp.vertex_unit_coordinates().T))
+        unpack_single(find_point_to_point_mapping(
+                     src_points=(-np.ones(grp.dim) + 2 * evec(i)).reshape(-1, 1),
+                     tgt_points=grp.vertex_unit_coordinates().T))
         for i in range(grp.dim)
     ]
 
@@ -747,11 +748,10 @@ def _get_tensor_product_element_flip_matrix_and_vertex_permutation(
                  unit_flip_matrix,
                  grp.vertex_unit_coordinates().T)
 
-    vertex_permutation_to = find_point_permutation(
-        targets=flipped_vertices,
-        permutees=grp.vertex_unit_coordinates().T,
-    )
-    if vertex_permutation_to is None:
+    vertex_permutation_to = find_point_to_point_mapping(
+        src_points=flipped_vertices,
+        tgt_points=grp.vertex_unit_coordinates().T)
+    if np.any(vertex_permutation_to < 0):
         raise RuntimeError("flip permutation was not found")
 
     flipped_unit_nodes = np.einsum("ij,jn->in", unit_flip_matrix, grp.unit_nodes)
@@ -1020,85 +1020,6 @@ def split_mesh_groups(
 # }}}
 
 
-# {{{ vertex matching
-
-def _match_vertices(
-        mesh: Mesh,
-        src_vertex_indices: np.ndarray,
-        tgt_vertex_indices: np.ndarray, *,
-        aff_map: AffineMap | None = None,
-        tol: float = 1e-12,
-        use_tree: bool | None = None) -> np.ndarray:
-    if mesh.vertices is None:
-        raise ValueError("Mesh must have vertices")
-
-    if aff_map is None:
-        aff_map = AffineMap()
-
-    if use_tree is None:
-        # Empirically, the tree version becomes faster at 2**13.
-        # The temporary (displacements) below at that size requires
-        # 1.6GB, which seems like a lot. Capping at 2**11 instead,
-        # which requires a more reasonable 100M.
-        use_tree = len(tgt_vertex_indices) >= 2**11
-
-    src_vertices = mesh.vertices[:, src_vertex_indices]
-    tgt_vertices = mesh.vertices[:, tgt_vertex_indices]
-
-    mapped_src_vertices = aff_map(src_vertices)
-
-    if use_tree:
-        tgt_vertex_bboxes = np.stack((
-            tgt_vertices - tol,
-            tgt_vertices + tol))
-
-        from pytools.spatial_btree import SpatialBinaryTreeBucket
-        tree = SpatialBinaryTreeBucket(
-            np.min(tgt_vertex_bboxes[0], axis=1),
-            np.max(tgt_vertex_bboxes[1], axis=1))
-        for ivertex in range(len(tgt_vertex_indices)):
-            tree.insert(ivertex, tgt_vertex_bboxes[:, :, ivertex])
-
-        matched_tgt_vertices: np.ndarray[tuple[int, ...], np.dtype[Any]] \
-            = np.full(len(src_vertex_indices), -1)
-        for ivertex in range(len(src_vertex_indices)):
-            mapped_src_vertex = mapped_src_vertices[:, ivertex]
-            matches = np.array(list(tree.generate_matches(mapped_src_vertex)))
-            match_bboxes = tgt_vertex_bboxes[:, :, matches]
-            in_bbox = np.all(
-                (mapped_src_vertex[:, np.newaxis] >= match_bboxes[0, :, :])
-                & (mapped_src_vertex[:, np.newaxis] <= match_bboxes[1, :, :]),
-                axis=0)
-            candidate_indices = matches[in_bbox]
-            if len(candidate_indices) == 0:
-                continue
-            displacements = (
-                mapped_src_vertex.reshape(-1, 1)
-                - tgt_vertices[:, candidate_indices])
-            distances_sq = np.sum(displacements**2, axis=0)
-            matched_tgt_vertices[ivertex] = (
-                tgt_vertex_indices[candidate_indices[np.argmin(distances_sq)]])
-
-    else:
-        displacements = (
-            mapped_src_vertices.reshape(mesh.dim, -1, 1)
-            - tgt_vertices.reshape(mesh.dim, 1, -1))
-        distances_sq = np.sum(displacements**2, axis=0)
-
-        vertex_indices, = np.indices((len(src_vertex_indices),))
-        min_distance_sq_indices = np.argmin(distances_sq, axis=1)
-        min_distances_sq = distances_sq[vertex_indices, min_distance_sq_indices]
-
-        matched_tgt_vertices = np.where(
-            min_distances_sq < tol**2,
-            tgt_vertex_indices[min_distance_sq_indices],
-            -1)
-
-    return matched_tgt_vertices
-
-# }}}
-
-
 # {{{ boundary face matching
 
 @dataclass(frozen=True)
@@ -1170,8 +1091,8 @@ def _get_face_vertex_indices(mesh: Mesh, face_ids: _FaceIDs) -> np.ndarray:
 
 
 def _match_boundary_faces(
-        mesh: Mesh, bdry_pair_mapping: BoundaryPairMapping, tol: float, *,
-        use_tree: bool | None = None) -> tuple[_FaceIDs, _FaceIDs]:
+        mesh: Mesh, bdry_pair_mapping: BoundaryPairMapping, tol: float,
+        ) -> tuple[_FaceIDs, _FaceIDs]:
     """
     Given a :class:`BoundaryPairMapping` *bdry_pair_mapping*, return the
     correspondence between faces of the two boundaries (expressed as a pair of
@@ -1182,8 +1103,6 @@ def _match_boundary_faces(
         whose faces are to be matched.
     :arg tol: The allowed tolerance between the transformed vertex coordinates of
         the first boundary and the vertex coordinates of the second boundary.
-    :arg use_tree: Optional argument indicating whether to use a spatial binary
-        search tree or a (quadratic) numpy algorithm when matching vertices.
     :returns: A pair of :class:`meshmode.mesh._FaceIDs`, each having a number of
         entries equal to the number of faces in the boundary, that represents the
         correspondence between the two boundaries' faces. The first element in the
@@ -1217,12 +1136,15 @@ def _match_boundary_faces(
     bdry_n_vertex_indices = np.unique(bdry_n_face_vertex_indices)
     bdry_n_vertex_indices = bdry_n_vertex_indices[bdry_n_vertex_indices >= 0]
 
-    matched_bdry_n_vertex_indices = _match_vertices(
-        mesh, bdry_m_vertex_indices, bdry_n_vertex_indices,
-        aff_map=bdry_pair_mapping.aff_map, tol=tol, use_tree=use_tree)
+    bdry_m_vertices = mesh.vertices[:, bdry_m_vertex_indices]
+    bdry_n_vertices = mesh.vertices[:, bdry_n_vertex_indices]
+
+    m_idx_to_n_idx = find_point_to_point_mapping(
+        bdry_pair_mapping.aff_map(bdry_m_vertices),
+        bdry_n_vertices)
 
     unmatched_bdry_m_vertex_indices = bdry_m_vertex_indices[
-        np.where(matched_bdry_n_vertex_indices < 0)[0]]
+        np.where(m_idx_to_n_idx < 0)[0]]
     nunmatched = len(unmatched_bdry_m_vertex_indices)
     if nunmatched > 0:
         vertices = mesh.vertices[:, unmatched_bdry_m_vertex_indices]
@@ -1235,6 +1157,9 @@ def _match_boundary_faces(
                 for i in range(min(nunmatched, 10))])
             + f"\n...\n({nunmatched-10} more omitted.)" if nunmatched > 10 else "")
 
+    matched_bdry_n_vertex_indices = bdry_n_vertex_indices[
+        m_idx_to_n_idx]
+
     from meshmode.mesh import _concatenate_face_ids
     face_ids = _concatenate_face_ids([bdry_m_face_ids, bdry_n_face_ids])
 
@@ -1291,13 +1216,15 @@ def glue_mesh_boundaries(
         coordinates of the first boundary and the vertex coordinates of the second
         boundary when attempting to match the two. Pass at most one mapping for each
         unique (order-independent) pair of boundaries.
-    :arg use_tree: Optional argument indicating whether to use a spatial binary
-        search tree or a (quadratic) numpy algorithm when matching vertices.
     """
     if any(grp.vertex_indices is None for grp in mesh.groups):
         raise ValueError(
             "gluing mesh boundaries requires 'vertex_indices' in all groups")
 
+    if use_tree is not None:
+        warn("Passing 'use_tree' is deprecated and will be removed "
+             "in Q3 2025.", DeprecationWarning, stacklevel=2)
+
     glued_btags = {
         btag
         for mapping, _ in bdry_pair_mappings_and_tols
@@ -1320,7 +1247,7 @@ def glue_mesh_boundaries(
         glued_btag_pairs.add(btag_pair)
 
     face_id_pairs_for_mapping = [
-        _match_boundary_faces(mesh, mapping, tol, use_tree=use_tree)
+        _match_boundary_faces(mesh, mapping, tol)
         for mapping, tol in bdry_pair_mappings_and_tols]
 
     facial_adjacency_groups = []