DM-31313: Add HtmIndex20 functor

python/lsst/pipe/tasks/functors.py

@@ -30,6 +30,9 @@
 
 from lsst.daf.persistence import doImport
 from lsst.daf.butler import DeferredDatasetHandle
+import lsst.geom as geom
+import lsst.sphgeom as sphgeom
+
 from .parquetTable import ParquetTable, MultilevelParquetTable
 
 
@@ -706,6 +709,36 @@ def __call__(self, catalog, **kwargs):
         return super().__call__(catalog, **kwargs)
 
 
+class HtmIndex20(Functor):
+    """Compute the level 20 HtmIndex for the catalog.
+    """
+    name = "Htm20"
+    htmLevel = 20
+    _radians = True
+
+    def __init__(self, ra, decl, **kwargs):
+        self.pixelator = sphgeom.HtmPixelization(self.htmLevel)
+        self.ra = ra
+        self.decl = decl
+        self._columns = [self.ra, self.decl]
+        super().__init__(**kwargs)
+
+    def _func(self, df):
+
+        def computePixel(row):
+            if self._radians:
+                sphPoint = geom.SpherePoint(row[self.ra],
+                                            row[self.decl],
+                                            geom.radians)
+            else:
+                sphPoint = geom.SpherePoint(row[self.ra],
+                                            row[self.decl],
+                                            geom.degrees)
+            return self.pixelator.index(sphPoint.getVector())
+
+        return df.apply(computePixel, axis=1)
+
+
 def fluxName(col):
     if not col.endswith('_instFlux'):
         col += '_instFlux'

tests/test_functors.py

@@ -46,7 +46,7 @@
                                       LocalDipoleMeanFlux, LocalDipoleMeanFluxErr,
                                       LocalDipoleDiffFlux, LocalDipoleDiffFluxErr,
                                       LocalWcs, ComputePixelScale, ConvertPixelToArcseconds,
-                                      ConvertPixelSqToArcsecondsSq, Ratio)
+                                      ConvertPixelSqToArcsecondsSq, Ratio, HtmIndex20)
 
 ROOT = os.path.abspath(os.path.dirname(__file__))
 
@@ -765,6 +765,19 @@ def testRatio(self):
                                     atol=1e-16,
                                     rtol=1e-16))
 
+    def testHtm(self):
+        """Test that HtmIndxes are created as expected.
+        """
+        parq = self.simulateMultiParquet(self.dataDict)
+        func = HtmIndex20("coord_ra", "coord_dec")
+
+        val = self._funcVal(func, parq)
+        # Test that the HtmIds come out as the ra/dec in dataDict.
+        self.assertTrue(np.all(np.equal(
+            val.values,
+            [14924528684992, 14924528689697, 14924528501716, 14924526434259,
+             14924526433879])))
+
 
 class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
     pass