Merge pull request #1112 from lsst/tickets/DM-50431

DM-50431: Pass Exposure metadata from multiBand to propagateSourceFlags

Author: arunkannawadi

python/lsst/pipe/tasks/multiBand.py

@@ -19,7 +19,9 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-__all__ = ["DetectCoaddSourcesConfig", "DetectCoaddSourcesTask"]
+__all__ = ["DetectCoaddSourcesConfig", "DetectCoaddSourcesTask",
+           "MeasureMergedCoaddSourcesConfig", "MeasureMergedCoaddSourcesTask",
+           ]
 
 from lsst.pipe.base import (Struct, PipelineTask, PipelineTaskConfig, PipelineTaskConnections)
 import lsst.pipe.base.connectionTypes as cT
@@ -176,9 +178,13 @@ def __init__(self, schema=None, **kwargs):
     def runQuantum(self, butlerQC, inputRefs, outputRefs):
         inputs = butlerQC.get(inputRefs)
         idGenerator = self.config.idGenerator.apply(butlerQC.quantum.dataId)
-        inputs["idFactory"] = idGenerator.make_table_id_factory()
-        inputs["expId"] = idGenerator.catalog_id
-        outputs = self.run(**inputs)
+        exposure = inputs.pop("exposure")
+        assert not inputs, "runQuantum got more inputs than expected."
+        outputs = self.run(
+            exposure=exposure,
+            idFactory=idGenerator.make_table_id_factory(),
+            expId=idGenerator.catalog_id,
+        )
         butlerQC.put(outputs, outputRefs)
 
     def run(self, exposure, idFactory, expId):
@@ -191,7 +197,7 @@ def run(self, exposure, idFactory, expId):
         Parameters
         ----------
         exposure : `lsst.afw.image.Exposure`
-            Exposure on which to detect (may be backround-subtracted and scaled,
+            Exposure on which to detect (may be background-subtracted and scaled,
             depending on configuration).
         idFactory : `lsst.afw.table.IdFactory`
             IdFactory to set source identifiers.
@@ -294,6 +300,14 @@ class MeasureMergedCoaddSourcesConnections(
         multiple=True,
         deferLoad=True,
     )
+    finalVisitSummaryHandles = cT.Input(
+        doc="Final visit summary table",
+        name="finalVisitSummary",
+        storageClass="ExposureCatalog",
+        dimensions=("instrument", "visit"),
+        multiple=True,
+        deferLoad=True,
+    )
     # TODO[DM-47797]: remove this deprecated connection.
     inputCatalog = cT.Input(
         doc=("Name of the input catalog to use."
@@ -573,12 +587,15 @@ def runQuantum(self, butlerQC, inputRefs, outputRefs):
 
         # Set psfcache
         # move this to run after gen2 deprecation
-        inputs['exposure'].getPsf().setCacheCapacity(self.config.psfCache)
+        exposure = inputs.pop("exposure")
+        exposure.getPsf().setCacheCapacity(self.config.psfCache)
+
+        ccdInputs = exposure.getInfo().getCoaddInputs().ccds
+        apCorrMap = exposure.getInfo().getApCorrMap()
 
         # Get unique integer ID for IdFactory and RNG seeds; only the latter
         # should really be used as the IDs all come from the input catalog.
         idGenerator = self.config.idGenerator.apply(butlerQC.quantum.dataId)
-        inputs['exposureId'] = idGenerator.catalog_id
 
         # Transform inputCatalog
         table = afwTable.SourceTable.make(self.schema, idGenerator.make_table_id_factory())
@@ -596,7 +613,7 @@ def runQuantum(self, butlerQC, inputRefs, outputRefs):
         if self.config.doAddFootprints:
             modelData = inputs.pop('scarletModels')
             if self.config.doConserveFlux:
-                imageForRedistribution = inputs['exposure']
+                imageForRedistribution = exposure
             else:
                 imageForRedistribution = None
             updateCatalogFootprints(
@@ -609,7 +626,6 @@ def runQuantum(self, butlerQC, inputRefs, outputRefs):
             )
         table = sources.getTable()
         table.setMetadata(self.algMetadata)  # Capture algorithm metadata to write out to the source catalog.
-        inputs['sources'] = sources
 
         skyMap = inputs.pop('skyMap')
         tractNumber = catalogRef.dataId['tract']
@@ -622,29 +638,45 @@ def runQuantum(self, butlerQC, inputRefs, outputRefs):
             wcs=tractInfo.getWcs(),
             bbox=patchInfo.getOuterBBox()
         )
-        inputs['skyInfo'] = skyInfo
 
         if self.config.doPropagateFlags:
-            ccdInputs = inputs["exposure"].getInfo().getCoaddInputs().ccds
-            inputs["ccdInputs"] = ccdInputs
-
             if "sourceTableHandles" in inputs:
                 sourceTableHandles = inputs.pop("sourceTableHandles")
                 sourceTableHandleDict = {handle.dataId["visit"]: handle for handle in sourceTableHandles}
-                inputs["sourceTableHandleDict"] = sourceTableHandleDict
+            else:
+                sourceTableHandleDict = None
             if "finalizedSourceTableHandles" in inputs:
                 finalizedSourceTableHandles = inputs.pop("finalizedSourceTableHandles")
                 finalizedSourceTableHandleDict = {handle.dataId["visit"]: handle
                                                   for handle in finalizedSourceTableHandles}
-                inputs["finalizedSourceTableHandleDict"] = finalizedSourceTableHandleDict
-
-        outputs = self.run(**inputs)
+            else:
+                finalizedSourceTableHandleDict = None
+            if "finalVisitSummaryHandles" in inputs:
+                finalVisitSummaryHandles = inputs.pop("finalVisitSummaryHandles")
+                finalVisitSummaryHandleDict = {handle.dataId["visit"]: handle
+                                               for handle in finalVisitSummaryHandles}
+            else:
+                finalVisitSummaryHandleDict = None
+
+        assert not inputs, "runQuantum got more inputs than expected."
+        outputs = self.run(
+            exposure=exposure,
+            sources=sources,
+            skyInfo=skyInfo,
+            exposureId=idGenerator.catalog_id,
+            ccdInputs=ccdInputs,
+            sourceTableHandleDict=sourceTableHandleDict,
+            finalizedSourceTableHandleDict=finalizedSourceTableHandleDict,
+            finalVisitSummaryHandleDict=finalVisitSummaryHandleDict,
+            apCorrMap=apCorrMap,
+        )
         # Strip HeavyFootprints to save space on disk
         sources = outputs.outputSources
         butlerQC.put(outputs, outputRefs)
 
     def run(self, exposure, sources, skyInfo, exposureId, ccdInputs=None,
-            sourceTableHandleDict=None, finalizedSourceTableHandleDict=None):
+            sourceTableHandleDict=None, finalizedSourceTableHandleDict=None, finalVisitSummaryHandleDict=None,
+            apCorrMap=None):
         """Run measurement algorithms on the input exposure, and optionally populate the
         resulting catalog with extra information.
 
@@ -665,12 +697,16 @@ def run(self, exposure, sources, skyInfo, exposureId, ccdInputs=None,
             the coadd.
         sourceTableHandleDict : `dict` [`int`, `lsst.daf.butler.DeferredDatasetHandle`], optional
             Dict for sourceTable_visit handles (key is visit) for propagating flags.
-            These tables are derived from the ``CalibrateTask`` sources, and contain
-            astrometry and photometry flags, and optionally PSF flags.
+            These tables contain astrometry and photometry flags, and optionally PSF flags.
         finalizedSourceTableHandleDict : `dict` [`int`, `lsst.daf.butler.DeferredDatasetHandle`], optional
             Dict for finalized_src_table handles (key is visit) for propagating flags.
-            These tables are derived from ``FinalizeCalibrationTask`` and contain
-            PSF flags from the finalized PSF estimation.
+            These tables contain PSF flags from the finalized PSF estimation.
+        finalVisitSummaryHandleDict : `dict` [`int`, `lsst.daf.butler.DeferredDatasetHandle`], optional
+            Dict for visit_summary handles (key is visit) for visit-level information.
+            These tables contain the WCS information of the single-visit input images.
+        apCorrMap : `lsst.afw.image.ApCorrMap`, optional
+            Aperture correction map attached to the ``exposure``. If None, it
+            will be read from the ``exposure``.
 
         Returns
         -------
@@ -683,9 +719,11 @@ def run(self, exposure, sources, skyInfo, exposureId, ccdInputs=None,
         self.measurement.run(sources, exposure, exposureId=exposureId)
 
         if self.config.doApCorr:
+            if apCorrMap is None:
+                apCorrMap = exposure.getInfo().getApCorrMap()
             self.applyApCorr.run(
                 catalog=sources,
-                apCorrMap=exposure.getInfo().getApCorrMap()
+                apCorrMap=apCorrMap,
             )
 
         # TODO DM-11568: this contiguous check-and-copy could go away if we
@@ -705,7 +743,8 @@ def run(self, exposure, sources, skyInfo, exposureId, ccdInputs=None,
                 sources,
                 ccdInputs,
                 sourceTableHandleDict,
-                finalizedSourceTableHandleDict
+                finalizedSourceTableHandleDict,
+                finalVisitSummaryHandleDict,
             )
 
         results = Struct()

python/lsst/pipe/tasks/propagateSourceFlags.py

@@ -113,7 +113,8 @@ def __init__(self, schema, **kwargs):
             self.schema.addField(f, type="Flag", doc="Propagated from finalized sources")
 
     def run(self, coadd_object_cat, ccd_inputs,
-            source_table_handle_dict=None, finalized_source_table_handle_dict=None):
+            source_table_handle_dict=None, finalized_source_table_handle_dict=None,
+            visit_summary_handle_dict=None):
         """Propagate flags from single-frame sources to coadd objects.
 
         Flags are only propagated if a configurable percentage of the sources
@@ -126,13 +127,16 @@ def run(self, coadd_object_cat, ccd_inputs,
             Table of coadd objects.
         ccd_inputs : `lsst.afw.table.ExposureCatalog`
             Table of single-frame inputs to coadd.
-        source_table_handle_dict : `dict` [`int`: `lsst.daf.butler.DeferredDatasetHandle`]
+        source_table_handle_dict : `dict` [`int`: `lsst.daf.butler.DeferredDatasetHandle`], optional
             Dict for sourceTable_visit handles (key is visit).  May be None if
             ``config.source_flags`` has no entries.
         finalized_source_table_handle_dict : `dict` [`int`:
-                                                     `lsst.daf.butler.DeferredDatasetHandle`]
+                                                     `lsst.daf.butler.DeferredDatasetHandle`], optional
             Dict for finalized_src_table handles (key is visit).  May be None if
             ``config.finalized_source_flags`` has no entries.
+        visit_summary_handle_dict : `dict` [`int`: `lsst.daf.butler.DeferredDatasetHandle`], optional
+            Dict for visitSummary handles (key is visit). If None, using WCS
+            from the ccd_inputs will be attempted.
         """
         if len(self.config.source_flags) == 0 and len(self.config.finalized_source_flags) == 0:
             return
@@ -185,11 +189,18 @@ def run(self, coadd_object_cat, ccd_inputs,
                         continue
                     handle = handle_dict[visit]
                     tbl = handle.get(parameters={"columns": columns})
+                    if visit_summary_handle_dict is not None:
+                        visit_summary = visit_summary_handle_dict[visit].get()
 
                     # Loop over all ccd_inputs rows for this visit.
-                    for row in ccd_inputs[ccd_inputs["visit"] == visit]:
+                    for row in ccd_inputs:
+                        if row["visit"] != visit:
+                            continue
                         detector = row["ccd"]
-                        wcs = row.getWcs()
+                        if visit_summary_handle_dict is None:
+                            wcs = row.getWcs()
+                        else:
+                            wcs = visit_summary.find(detector).getWcs()
                         if wcs is None:
                             self.log.info("No WCS for visit %d detector %d, so can't match sources to "
                                           "propagate flags.  Skipping...", visit, detector)