Enforce using object indices safely.

lsst · Jul 23, 2019 · 4e78b85 · 4e78b85
1 parent a2de75c
commit 4e78b85
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/python/lsst/ap/association/association.py b/python/lsst/ap/association/association.py
@@ -607,9 +607,12 @@ def score(self, dia_objects, dia_sources, max_dist):
 
         vectors = self._radec_to_xyz(dia_sources)
 
-        dists, obj_idxs = spatial_tree.query(vectors)
-        scores = np.where(dists < max_dist_rad, dists, np.inf)
-        obj_ids = dia_objects.index[obj_idxs]
+        scores, obj_idxs = spatial_tree.query(
+            vectors,
+            distance_upper_bound=max_dist_rad)
+        matched_src_idxs = np.argwhere(np.isfinite(scores))
+        obj_ids[matched_src_idxs] = dia_objects.index[
+            obj_idxs[matched_src_idxs]]
 
         return pipeBase.Struct(
             scores=scores,