Move loadDiaCatalogs to standalone PipelineTask

lsst · Aug 3, 2024 · cc0a050 · cc0a050
1 parent 0ff9b6e
commit cc0a050
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 3 deletions.
diff --git a/python/lsst/ap/verify/testPipeline.py b/python/lsst/ap/verify/testPipeline.py
@@ -34,7 +34,7 @@
 import lsst.afw.image as afwImage
 import lsst.afw.math as afwMath
 import lsst.afw.table as afwTable
-from lsst.ap.association import (TransformDiaSourceCatalogConfig,
+from lsst.ap.association import (LoadDiaCatalogsConfig, TransformDiaSourceCatalogConfig,
                                  DiaPipelineConfig, FilterDiaSourceCatalogConfig)
 from lsst.pipe.base import PipelineTask, Struct
 from lsst.ip.isr import IsrTaskConfig
@@ -45,6 +45,41 @@
 from lsst.meas.transiNet import RBTransiNetConfig
 
 
+class MockLoadDiaCatalogsTask(PipelineTask):
+    """A do-nothing substitute for LoadDiaCatalogsTask.
+    """
+    ConfigClass = LoadDiaCatalogsConfig
+    _DefaultName = "notLoadDiaCatalogs"
+
+    def run(self, regionTime):
+        """Produce preloaded DiaSource and DiaObject outputs with no processing.
+
+        Parameters
+        ----------
+        regionTime : `lsst.pipe.base.utils.RegionTimeInfo`
+            A serializable container for a sky region and timespan.
+
+        Returns
+        -------
+        result : `lsst.pipe.base.Struct`
+            Results struct with components.
+
+            - ``diaObjects`` : Complete set of DiaObjects covering the input
+              exposure padded by ``pixelMargin``. DataFrame is indexed by
+              the ``diaObjectId`` column. (`pandas.DataFrame`)
+            - ``diaSources`` : Complete set of DiaSources covering the input
+              exposure padded by ``pixelMargin``. DataFrame is indexed by
+              ``diaObjectId``, ``band``, ``diaSourceId`` columns.
+              (`pandas.DataFrame`)
+            - ``diaForcedSources`` : Complete set of forced photometered fluxes
+            on the past 12 months of difference images at DiaObject locations.
+        """
+        return Struct(diaObjects=pandas.DataFrame(),
+                      diaSources=pandas.DataFrame(),
+                      diaForcedSources=pandas.DataFrame(),
+                      )
+
+
 class MockIsrTask(PipelineTask):
     """A do-nothing substitute for IsrTask.
     """
@@ -600,6 +635,9 @@ def run(self,
             diffIm,
             exposure,
             template,
+            preloadedDiaObjects,
+            preloadedDiaSources,
+            preloadedDiaForcedSources,
             ccdExposureIdBits,
             band,
             idGenerator=None):
@@ -619,6 +657,12 @@ def run(self,
             ``diffIm``.
         template : `lsst.afw.image.ExposureF`
             Template exposure used to create diffIm.
+        preloadedDiaObjects : `pandas.DataFrame`
+            Previously detected DiaObjects, loaded from the APDB.
+        preloadedDiaSources : `pandas.DataFrame`
+            Previously detected DiaSources, loaded from the APDB.
+        preloadedDiaForcedSources : `pandas.DataFrame`
+            Catalog of previously detected forced DiaSources, from the APDB
         ccdExposureIdBits : `int`
             Number of bits used for a unique ``ccdVisitId``.  Deprecated in
             favor of ``idGenerator``, and ignored if that is present.  Pass

diff --git a/tests/test_testPipeline.py b/tests/test_testPipeline.py
@@ -39,7 +39,10 @@
     MockCalibrateTask, MockGetTemplateTask, \
     MockAlardLuptonSubtractTask, MockDetectAndMeasureTask, MockTransformDiaSourceCatalogTask, \
     MockRBTransiNetTask, MockDiaPipelineTask, MockFilterDiaSourceCatalogTask, \
-    MockSpatiallySampledMetricsTask
+    MockSpatiallySampledMetricsTask, MockLoadDiaCatalogsTask
+from lsst.daf.butler import Timespan
+from lsst.pipe.base.utils import RegionTimeInfo
+from lsst.sphgeom import UnitVector3d
 
 
 class MockTaskTestSuite(unittest.TestCase):
@@ -60,6 +63,7 @@ def setUpClass(cls):
 
         INSTRUMENT = "DummyCam"
         VISIT = 42
+        GROUP = '42'
         CCD = 101
         HTM = 42
         SKYMAP = "TreasureMap"
@@ -78,6 +82,7 @@ def setUpClass(cls):
         butlerTests.addDataIdValue(cls.repo, "subfilter", SUB_FILTER)
         butlerTests.addDataIdValue(cls.repo, "exposure", VISIT)
         butlerTests.addDataIdValue(cls.repo, "visit", VISIT)
+        butlerTests.addDataIdValue(cls.repo, "group", GROUP)
         butlerTests.addDataIdValue(cls.repo, "detector", CCD)
         butlerTests.addDataIdValue(cls.repo, "skymap", SKYMAP)
         butlerTests.addDataIdValue(cls.repo, "tract", TRACT)
@@ -88,6 +93,8 @@ def setUpClass(cls):
             {"instrument": INSTRUMENT, "exposure": VISIT, "detector": CCD})
         cls.visitId = cls.repo.registry.expandDataId(
             {"instrument": INSTRUMENT, "visit": VISIT, "detector": CCD})
+        cls.groupId = cls.repo.registry.expandDataId(
+            {"instrument": INSTRUMENT, "group": GROUP, "detector": CCD})
         cls.visitOnlyId = cls.repo.registry.expandDataId(
             {"instrument": INSTRUMENT, "visit": VISIT})
         cls.skymapId = cls.repo.registry.expandDataId({"skymap": SKYMAP})
@@ -99,6 +106,11 @@ def setUpClass(cls):
             {"skymap": SKYMAP, "tract": TRACT, "patch": PATCH, "band": BAND, "subfilter": SUB_FILTER})
         cls.htmId = cls.repo.registry.expandDataId({"htm7": HTM})
 
+        butlerTests.addDatasetType(cls.repo, "regionTimeInfo", cls.groupId.dimensions, "RegionTimeInfo")
+        butlerTests.addDatasetType(cls.repo, "preloaded_diaObjects", cls.groupId.dimensions, "DataFrame")
+        butlerTests.addDatasetType(cls.repo, "preloaded_diaSources", cls.groupId.dimensions, "DataFrame")
+        butlerTests.addDatasetType(cls.repo, "preloaded_diaForcedSources", cls.groupId.dimensions,
+                                   "DataFrame")
         butlerTests.addDatasetType(cls.repo, "postISRCCD", cls.exposureId.dimensions, "Exposure")
         butlerTests.addDatasetType(cls.repo, "icExp", cls.visitId.dimensions, "ExposureF")
         butlerTests.addDatasetType(cls.repo, "icSrc", cls.visitId.dimensions, "SourceCatalog")
@@ -141,6 +153,27 @@ def setUp(self):
         super().setUp()
         self.butler = butlerTests.makeTestCollection(self.repo, uniqueId=self.id())
 
+    def testMockLoadDiaCatalogsTask(self):
+        config = MockLoadDiaCatalogsTask.ConfigClass()
+        config.apdb_config_url = "testing_only"
+        task = MockLoadDiaCatalogsTask(config=config)
+        pipelineTests.assertValidInitOutput(task)
+        points = [UnitVector3d.Z(), UnitVector3d.X(), UnitVector3d.Y()]
+        region = lsst.sphgeom.ConvexPolygon(points)
+        timespan = Timespan.makeEmpty()
+        result = task.run(RegionTimeInfo(region=region, timespan=timespan))
+        pipelineTests.assertValidOutput(task, result)
+
+        self.butler.put(RegionTimeInfo(region=region, timespan=timespan), "regionTimeInfo", self.groupId)
+        quantum = pipelineTests.makeQuantum(
+            task, self.butler, self.groupId,
+            {"regionTime": self.groupId,
+             "diaObjects": self.groupId,
+             "diaSources": self.groupId,
+             "diaForcedSources": self.groupId,
+             })
+        pipelineTests.runTestQuantum(task, self.butler, quantum, mockRun=False)
+
     def testMockIsr(self):
         # Testing MockIsrTask is tricky because the real ISR has an unstable
         # interface with dozens of potential inputs, too many to pass through
@@ -347,9 +380,13 @@ def testMockDiaPipelineTask(self):
         task = MockDiaPipelineTask(config=config)
         pipelineTests.assertValidInitOutput(task)
         result = task.run(pandas.DataFrame(), pandas.DataFrame(), afwImage.ExposureF(),
-                          afwImage.ExposureF(), afwImage.ExposureF(), 42, 'k')
+                          afwImage.ExposureF(), afwImage.ExposureF(), pandas.DataFrame(),
+                          pandas.DataFrame(), pandas.DataFrame(), 'k', 42)
         pipelineTests.assertValidOutput(task, result)
 
+        self.butler.put(pandas.DataFrame(), "preloaded_diaObjects", self.groupId)
+        self.butler.put(pandas.DataFrame(), "preloaded_diaSources", self.groupId)
+        self.butler.put(pandas.DataFrame(), "preloaded_diaForcedSources", self.groupId)
         self.butler.put(pandas.DataFrame(), "deepDiff_diaSrcTable", self.visitId)
         self.butler.put(pandas.DataFrame(), "visitSsObjects", self.visitId)
         self.butler.put(afwImage.ExposureF(), "deepDiff_differenceExp", self.visitId)
@@ -362,6 +399,9 @@ def testMockDiaPipelineTask(self):
              "diffIm": self.visitId,
              "exposure": self.visitId,
              "template": self.visitId,
+             "preloadedDiaObjects": self.groupId,
+             "preloadedDiaSources": self.groupId,
+             "preloadedDiaForcedSources": self.groupId,
              "apdbMarker": self.visitId,
              "associatedDiaSources": self.visitId,
              "diaForcedSources": self.visitId,