Use KDTree to reduce memory & compute in _xy_2dhist

[schlafly]

_xy_2dhist currently computes the differences in positions between all n * m pairs of stars in the image and reference catalogs. This can be a lot as images become large. Only pairs in a box of size 2 * r actually end up contributing to the histogram, so most of these differences end up being discarded.

Adding a kdtree to prefilter the matches to only those in a circle of radius sqrt(2) * r means that most of these matches don't need to be computed. This PR adds that prefiltering.

I have only verified that tests pass and that the particular case I was trying to run in Roman now succeeds. I'm happy to run more tests as appropriate.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (f7ad5a1) 93.84% compared to head (57e0766) 95.27%.

❗ Current head 57e0766 differs from pull request most recent head a4e6bb9. Consider uploading reports for the commit a4e6bb9 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #196      +/-   ##
==========================================
+ Coverage   93.84%   95.27%   +1.43%     
==========================================
  Files          11       11              
  Lines        1982     1947      -35     
==========================================
- Hits         1860     1855       -5     
+ Misses        122       92      -30     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mcara]

The code in this PR works fine. The change log will need to be updated and some comments cleaned (I will make detailed comments separately).

I do agree with your statements above and I understand that allocated memory for very large catalogs may be large. However, I am surprised that the bottleneck is in the 2d histogram and not in xyxymatch which normally crashes for more than 5,000 source per catalog. So, even if this PR addresses the issue of 2dhist, the matching routine should fail later. Therefore, I believe that the best solution is to limit catalog size by, e.g., selecting top 100-5,000 brightest sources regardless of this PR.

Still, this PR provides a significant improvement over the current code (at least for reasonably uniform distribution of sources), and therefore it is worth it.

Here is my test:

Test set up:

import os
from datetime import datetime
import psutil
import numpy as np
from scipy import spatial

np.random.seed(1)
a = 10000 * np.random.random((20000, 2))
d = np.random.normal((3.0, 0.5), scale=0.1, size=(20000, 2))
b = a + d

# current code on master that computes distances between all pairs:
def _xy_2dhist(imgxy, refxy, r):
    process = psutil.Process(os.getpid())
    base_memory_usage = process.memory_info().rss
    runtime_begin = datetime.now()
    dx = np.subtract.outer(imgxy[:, 0], refxy[:, 0]).ravel()
    dy = np.subtract.outer(imgxy[:, 1], refxy[:, 1]).ravel()
    idx = np.where((dx < r + 0.5) & (dx >= -r - 0.5) &
                   (dy < r + 0.5) & (dy >= -r - 0.5))
    r = int(np.ceil(r))
    h = np.histogram2d(dx[idx], dy[idx], 2 * r + 1,
                       [[-r - 0.5, r + 0.5], [-r - 0.5, r + 0.5]])
    runtime_end = datetime.now()
    print(f"RUN TIME: {runtime_end - runtime_begin}")
    memory_usage = process.memory_info().rss - base_memory_usage
    print(f"Memory usage: {memory_usage / 1024**3} GB")
    return h[0].T

# Proposed code in this PR that uses KDTree:
def _xy_2dhist2(imgxy, refxy, r):
    process = psutil.Process(os.getpid())
    base_memory_usage = process.memory_info().rss
    runtime_begin = datetime.now()
    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.copy()

    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))
    h = np.histogram2d(dx[idx], dy[idx], 2 * r + 1,
                      [[-r - 0.5, r + 0.5], [-r - 0.5, r + 0.5]])
    runtime_end = datetime.now()
    print(f"RUN TIME: {runtime_end - runtime_begin}")
    memory_usage = process.memory_info().rss - base_memory_usage
    print(f"Memory usage: {memory_usage / 1024**3} GB")
    return h[0].T

Existing code results:

In [2]: h1 = _xy_2dhist(a, b, 5)
RUN TIME: 0:00:01.433775
Memory usage: 6.7071075439453125 GB

In [27]: h1 = _xy_2dhist(a, b, 500)
RUN TIME: 0:00:02.252558
Memory usage: 5.9826507568359375 GB

Proposed code results:

In [3]: h2 = _xy_2dhist2(a, b, 5)
RUN TIME: 0:00:00.047993
Memory usage: 0.0046539306640625 GB

In [28]: h2 = _xy_2dhist2(a, b, 500)
RUN TIME: 0:00:01.492740
Memory usage: 0.2578277587890625 GB

[schlafly]

Thanks for the detailed test. I can say empirically that running the roman pipeline on these files now proceeds without crashing later in xyxymatch.

Part of the confusion might stem from the fact that while the catalogs being sent to tweakwcs are large, there is a lot of junk in them, which may make things easier somehow downstream (?!). We're also trying to improve that, of course. We can also pursue just trimming down the catalogs, but I agree that this change seems like a pure improvement to me.

I'm happy to update the comments and change log of course---thanks for looking at the code closely.

[mcara]

A couple of very minor comments that may be ignored but a changelog is needed.

[mcara]

Could you please also rebase.

[schlafly]

I've rebased, added to the changelog, and addressed the two comments; I think this is good to go.

[mcara]

This is great. Thanks!