DM-55440: Add vectorized tract/patch id lookup function

python/lsst/skymap/baseSkyMap.py

@@ -296,6 +296,71 @@ def findClosestTractPatchList(self, coordList):
                 retList.append((tractInfo, patchList))
         return retList
 
+    def findTractIdPatchIdArray(self, ra, dec, degrees=False):
+        """Find array of tract IDs and patch IDs with vectorized operations
+        (where supported).
+
+        If a given sky map does not support vectorized operations, then
+        loops will be called.
+
+        Parameters
+        ----------
+        ra : `numpy.ndarray`
+            Array of Right Ascension.  Units are radians unless
+            degrees=True.
+        dec : `numpy.ndarray`
+            Array of Declination.  Units are radians unless
+            degrees=True.
+        degrees : `bool`, optional
+            Input ra, dec arrays are degrees if `True`.
+
+        Returns
+        -------
+        tractId : `numpy.ndarray`
+            Array of tract IDs.
+        patchId : `numpy.ndarray`
+            Array of sequential patch IDs.  -1 if there is no appropriate
+            patch.
+
+        Notes
+        -----
+        - If coord is equidistant between multiple sky tract centers then one
+          is arbitrarily chosen.
+
+        .. warning::
+
+           If tracts do not cover the whole sky then the returned tract may not
+           include the given ra/dec.
+        """
+        from scipy.ndimage import value_indices
+
+        _ra = np.atleast_1d(ra)
+        _dec = np.atleast_1d(dec)
+
+        # This will be vectorized if possible.
+        tractIds = self.findTractIdArray(_ra.ravel(), _dec.ravel(), degrees=degrees)
+
+        # This will be vectorized within each patch.
+        patchIds = np.zeros(len(tractIds), dtype=np.int32) - 1
+
+        inds = value_indices(tractIds)
+        for tractId, (indices,) in inds.items():
+            tract = self[tractId]
+            wcs = tract.wcs
+            x, y = wcs.skyToPixelArray(_ra[indices], _dec[indices], degrees=degrees)
+            xInd = np.floor(x).astype(np.int64) // tract.patch_inner_dimensions[0]
+            yInd = np.floor(y).astype(np.int64) // tract.patch_inner_dimensions[1]
+            nx, ny = tract.num_patches
+            patchIds[indices] = nx * yInd + xInd
+
+            # Check for overflows
+            bad = (xInd < 0) | (xInd >= nx) | (yInd < 0) | (yInd >= ny)
+            if bad.sum() > 0:
+                tractIds[indices[bad]] = -1
+                patchIds[indices[bad]] = -1
+
+        return (tractIds, patchIds)
+
     def __getitem__(self, ind):
         return self._tractInfoList[ind]
 

tests/test_discreteSkyMap.py

@@ -1,3 +1,4 @@
+import numpy as np
 import unittest
 
 import lsst.utils.tests
@@ -65,6 +66,41 @@ def testCompare(self):
             skyMap = self.getSkyMap(config=config)
             self.assertNotEqual(skyMap, defaultSkyMap)
 
+    def testFindTractIdPatchIdArray(self):
+        np.random.seed(12345)
+
+        skymap = self.getSkyMap()
+
+        ras = np.random.uniform(low=0.0, high=360.0, size=1000)
+        decs = np.random.uniform(low=-90.0, high=90.0, size=1000)
+
+        # Make sure we have points that are just off a tract.
+        wcs = skymap[0].wcs
+        ras2, decs2 = wcs.pixelToSkyArray([-0.5, 10.0], [10.0, -0.5], degrees=True)
+        ras = np.concatenate((ras, ras2))
+        decs = np.concatenate((decs, decs2))
+
+        coords = [lsst.geom.SpherePoint(ra*lsst.geom.degrees, dec*lsst.geom.degrees)
+                  for ra, dec in zip(ras, decs)]
+
+        tractIds = [skymap.findTract(coord).getId() for coord in coords]
+        patchIds = np.zeros(len(tractIds), dtype=np.int32)
+        for i, tractId in enumerate(tractIds):
+            tract = skymap[tractId]
+            try:
+                patch = tract.findPatch(coords[i])
+                index = patch.sequential_index
+            except LookupError:
+                # Not in a tract at all.
+                tractIds[i] = -1
+                index = -1
+            patchIds[i] = index
+
+        tractIds2, patchIds2 = skymap.findTractIdPatchIdArray(ras, decs, degrees=True)
+
+        np.testing.assert_array_equal(tractIds2, tractIds)
+        np.testing.assert_array_equal(patchIds2, patchIds)
+
 
 class MemoryTester(lsst.utils.tests.MemoryTestCase):
     pass

tests/test_ringsSkyMap.py

@@ -169,6 +169,28 @@ def testFindTractIdArray(self):
 
         np.testing.assert_array_equal(indexes2, indexes)
 
+    def testFindTractIdPatchIdArray(self):
+        """Test findTractIdPatchIdArray."""
+        np.random.seed(12345)
+
+        ras = np.random.uniform(low=0.0, high=360.0, size=1000)
+        decs = np.random.uniform(low=-90.0, high=90.0, size=1000)
+
+        coords = [lsst.geom.SpherePoint(ra*lsst.geom.degrees, dec*lsst.geom.degrees)
+                  for ra, dec in zip(ras, decs)]
+
+        tractIds = [self.skymap.findTract(coord).getId() for coord in coords]
+        patchIds = np.zeros(len(tractIds), dtype=np.int32)
+        for i, tractId in enumerate(tractIds):
+            tract = self.skymap[tractId]
+            patch = tract.findPatch(coords[i])
+            patchIds[i] = patch.sequential_index
+
+        tractIds2, patchIds2 = self.skymap.findTractIdPatchIdArray(ras, decs, degrees=True)
+
+        np.testing.assert_array_equal(tractIds2, tractIds)
+        np.testing.assert_array_equal(patchIds2, patchIds)
+
     def getFirstTractLastRingCoord(self):
         """Return the coordinates of the first tract in the last ring