Merge pull request spacetelescope#196 from schlafly/kdtree

Use KDTree to reduce memory & compute in _xy_2dhist

CHANGELOG.rst

@@ -15,6 +15,10 @@ Release Notes
   using input catalogs that failed to align in the expanded reference
   catalog. [#195]
 
+- Reduce memory & compute needed by _xy_2dhist by pruning distant
+  pairs with a kdtree.  This is a purely internal change that does not
+  affect the results of the algorithm.  [#196]
+
 
 0.8.5 (30-November-2023)
 ========================

tweakwcs/matchutils.py

@@ -16,6 +16,7 @@
 from astropy.utils.exceptions import AstropyDeprecationWarning
 
 from stsci.stimage import xyxymatch
+from scipy import spatial
 
 from . import __version__  # noqa: F401
 
@@ -277,10 +278,19 @@ def __call__(self, refcat, imcat, tp_pscale=1.0, tp_units=None, **kwargs):
 
 
 def _xy_2dhist(imgxy, refxy, r):
-    # This code replaces the C version (arrxyzero) from carrutils.c
-    # It is about 5-8 times slower than the C version.
-    dx = np.subtract.outer(imgxy[:, 0], refxy[:, 0]).ravel()
-    dy = np.subtract.outer(imgxy[:, 1], refxy[:, 1]).ravel()
+    # trim to only pairs within (r+0.5) * np.sqrt(2) using a kdtree
+    # to avoid computing differences for many widely separated pairs.
+    kdtree = spatial.KDTree(refxy)
+    neighbors = kdtree.query_ball_point(imgxy, (r + 0.5) * np.sqrt(2))
+    lens = [len(n) for n in neighbors]
+    mi = np.repeat(np.arange(imgxy.shape[0]), lens)
+    if len(mi) > 0:
+        mr = np.concatenate([n for n in neighbors if len(n) > 0])
+    else:
+        mr = mi
+
+    dx = imgxy[mi, 0] - refxy[mr, 0]
+    dy = imgxy[mi, 1] - refxy[mr, 1]
     idx = np.where((dx < r + 0.5) & (dx >= -r - 0.5) &
                    (dy < r + 0.5) & (dy >= -r - 0.5))
     r = int(np.ceil(r))