diff --git a/python/lsst/skymap/baseSkyMap.py b/python/lsst/skymap/baseSkyMap.py index b82c43d..36c0197 100755 --- a/python/lsst/skymap/baseSkyMap.py +++ b/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] diff --git a/tests/test_discreteSkyMap.py b/tests/test_discreteSkyMap.py index f183e31..8980526 100644 --- a/tests/test_discreteSkyMap.py +++ b/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 diff --git a/tests/test_ringsSkyMap.py b/tests/test_ringsSkyMap.py index 4a41414..6cd4a75 100644 --- a/tests/test_ringsSkyMap.py +++ b/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