squash me

python/lsst/meas/base/forcedPhotCoadd.py

@@ -74,6 +74,12 @@ class ForcedPhotCoaddConnections(pipeBase.PipelineTaskConnections,
         storageClass="SourceCatalog",
         dimensions=("band", "skymap", "tract", "patch")
     )
+    footprintCatInBand = pipeBase.connectionTypes.Input(
+        doc="Catalog of footprints to attach to sources",
+        name="{inputCoaddName}Coadd_deblendedFlux",
+        storageClass="SourceCatalog",
+        dimensions=("band", "skymap", "tract", "patch")
+    )
     scarletModels = pipeBase.connectionTypes.Input(
         doc="Multiband scarlet models produced by the deblender",
         name="{inputCoaddName}Coadd_scarletDataModels",
@@ -95,8 +101,10 @@ class ForcedPhotCoaddConnections(pipeBase.PipelineTaskConnections,
 
     def __init__(self, *, config=None):
         super().__init__(config=config)
-        if config.footprintDatasetName != "scarletDeblendModels":
-            self.inputs -= set(("scarletModels",))
+        if config.footprintDatasetName != "ScarletModelData":
+            self.inputs.remove("scarletModels")
+        if config.footprintDatasetName != "DeblenededFlux":
+            self.inputs.remove("footprintCatInBand")
 
 
 class ForcedPhotCoaddConfig(pipeBase.PipelineTaskConfig,
@@ -132,15 +140,15 @@ class ForcedPhotCoaddConfig(pipeBase.PipelineTaskConfig,
             "Must have IDs that match those of the reference catalog."
             "If None, Footprints will be generated by transforming the reference Footprints.",
         dtype=str,
-        default="scarletDeblendModels",
+        default="ScarletModelData",
         optional=True
     )
     doConserveFlux = lsst.pex.config.Field(
         dtype=bool,
         default=True,
         doc="Whether to use the deblender models as templates to re-distribute the flux "
             "from the 'exposure' (True), or to perform measurements on the deblender model footprints. "
-            "If footprintDatasetName != 'scarletDeblendModels' then this field is ignored.")
+            "If footprintDatasetName != 'ScarletModelData' then this field is ignored.")
     doStripFootprints = lsst.pex.config.Field(
         dtype=bool,
         default=True,
@@ -221,27 +229,29 @@ def runQuantum(self, butlerQC, inputRefs, outputRefs):
         inputs = butlerQC.get(inputRefs)
 
         refCatInBand = inputs.pop('refCatInBand')
-        if "scarletModels" in inputs:
-            dataModel = inputs.pop("scarletModels")
+        if self.config.footprintDatasetName == "ScarletModelData":
+            footprintData = inputs.pop("scarletModels")
+        elif self.config.footprintDatasetName == "DeblendedFlux":
+            footprintData = inputs.pop("footprintCatIndBand")
         else:
-            dataModel = None
+            footprintData = None
         inputs['measCat'], inputs['exposureId'] = self.generateMeasCat(inputRefs.exposure.dataId,
                                                                        inputs['exposure'],
                                                                        inputs['refCat'],
                                                                        refCatInBand,
                                                                        inputs['refWcs'],
                                                                        "tract_patch",
-                                                                       dataModel)
+                                                                       footprintData)
         outputs = self.run(**inputs)
         # Strip HeavyFootprints to save space on disk
-        if self.config.footprintDatasetName == "scarletDeblendModels" and self.config.doStripFootprints:
+        if self.config.footprintDatasetName == "ScarletModelData" and self.config.doStripFootprints:
             sources = outputs.measCat
             for source in sources[sources["parent"] != 0]:
                 source.setFootprint(None)
         butlerQC.put(outputs, outputRefs)
 
     def generateMeasCat(self, exposureDataId, exposure, refCat, refCatInBand, refWcs, idPackerName,
-                        dataModel):
+                        footprintData):
         """Generate a measurement catalog for Gen3.
 
         Parameters
@@ -259,6 +269,11 @@ def generateMeasCat(self, exposureDataId, exposure, refCat, refCatInBand, refWcs
             Reference world coordinate system.
         idPackerName : `str`
             Type of ID packer to construct from the registry.
+        footprintData : `ScarletDataModel` or `lsst.afw.table.SourceCatalog`
+            Either the scarlet data models or the deblended catalog
+            containing footprints.
+            If `footprintData` is `None` then the footprints contained
+            in `refCatInBand` are used.
 
         Returns
         -------
@@ -281,21 +296,25 @@ def generateMeasCat(self, exposureDataId, exposure, refCat, refCatInBand, refWcs
                                                    idFactory=idFactory)
         # attach footprints here, as the attachFootprints method is geared for gen2
         # and is not worth modifying, as this can naturally live inside this method
-        if self.config.footprintDatasetName == "scarletDeblendModels":
+        if self.config.footprintDatasetName == "ScarletModelData":
             # Load the scarlet models
             self._attachScarletFootprints(
                 catalog=measCat,
-                dataModel=dataModel,
+                modelData=footprintData,
                 exposure=exposure,
                 band=exposureDataId["band"]
             )
         else:
+            if self.config.footprintDatasetName is None:
+                footprintCat = refCatInBand
+            else:
+                footprintCat = footprintData
             for srcRecord in measCat:
-                fpRecord = refCatInBand.find(srcRecord.getId())
+                fpRecord = footprintCat.find(srcRecord.getId())
                 if fpRecord is None:
                     raise LookupError("Cannot find Footprint for source {}; please check that {} "
                                       "IDs are compatible with reference source IDs"
-                                      .format(srcRecord.getId(), self.config.connections.refCatInBand))
+                                      .format(srcRecord.getId(), footprintCat))
                 srcRecord.setFootprint(fpRecord.getFootprint())
         return measCat, exposureIdInfo.expId
 
@@ -443,7 +462,7 @@ def attachFootprints(self, sources, refCat, exposure, refWcs, dataRef):
         self.log.info("Loading deblended footprints for sources from %s, %s",
                       self.config.footprintDatasetName, dataRef.dataId)
 
-        if self.config.footprintDatasetName == "scarletDeblendModels":
+        if self.config.footprintDatasetName == "ScarletModelData":
             # Load the scarlet models
             dataModel = dataRef.get("%sCoadd_%s" % (self.config.coaddName, self.config.footprintDatasetName),
                                     immediate=True)
@@ -460,19 +479,19 @@ def attachFootprints(self, sources, refCat, exposure, refWcs, dataRef):
                                        self.config.footprintDatasetName))
                 srcRecord.setFootprint(fpRecord.getFootprint())
 
-    def _attachScarletFootprints(self, catalog, dataModel, exposure, band):
+    def _attachScarletFootprints(self, catalog, modelData, exposure, band):
         """Attach scarlet models as HeavyFootprints
         """
         if self.config.doConserveFlux:
-            redistributeCoadd = exposure
+            redistributeImage = exposure.image
         else:
-            redistributeCoadd = None
+            redistributeImage = None
         # Attach the footprints
-        dataModel.updateCatalogFootprints(
+        modelData.updateCatalogFootprints(
             catalog=catalog,
             band=band,
             psfModel=exposure.getPsf(),
-            redistributeCoadd=redistributeCoadd,
+            redistributeImage=redistributeImage,
             removeScarletData=True,
             updateFluxColumns=False,
         )