track inputs for CoaddPsf

python/lsst/drp/tasks/assemble_cell_coadd.py

@@ -37,11 +37,12 @@
     SingleCellCoadd,
     UniformGrid,
 )
-from lsst.meas.algorithms import AccumulatorMeanStack
-from lsst.pex.config import ConfigurableField, Field, ListField, RangeField
+from lsst.meas.algorithms import AccumulatorMeanStack, CoaddPsf, CoaddPsfConfig
+from lsst.pex.config import ConfigField, ConfigurableField, Field, ListField, RangeField
 from lsst.pipe.base import NoWorkFound, PipelineTask, PipelineTaskConfig, PipelineTaskConnections, Struct
 from lsst.pipe.base.connectionTypes import Input, Output
 from lsst.pipe.tasks.coaddBase import makeSkyInfo
+from lsst.pipe.tasks.coaddInputRecorder import CoaddInputRecorderTask
 from lsst.pipe.tasks.interpImage import InterpImageTask
 from lsst.pipe.tasks.scaleZeroPoint import ScaleZeroPointTask
 from lsst.skymap import BaseSkyMap
@@ -106,6 +107,15 @@ class AssembleCellCoaddConfig(PipelineTaskConfig, pipelineConnections=AssembleCe
         inclusiveMin=True,
         inclusiveMax=False,
     )
+    # The following config options are specific to the CoaddPsf.
+    coaddPsf = ConfigField(
+        doc="Configuration for CoaddPsf",
+        dtype=CoaddPsfConfig,
+    )
+    inputRecorder = ConfigurableField(
+        doc="Subtask that helps fill CoaddInputs catalogs added to the final Exposure",
+        target=CoaddInputRecorderTask,
+    )
 
 
 class AssembleCellCoaddTask(PipelineTask):
@@ -114,6 +124,7 @@ class AssembleCellCoaddTask(PipelineTask):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.makeSubtask("inputRecorder")
         self.makeSubtask("interpImage")
         self.makeSubtask("scaleZeroPoint")
 
@@ -199,7 +210,14 @@ def run(self, inputWarps, skyMap, skyInfo, **kwargs):
         statsCtrl = self._construct_stats_control()
 
         gc = self._construct_grid_container(skyInfo, statsCtrl)
-
+        coadd_inputs_gc = GridContainer(gc.shape)
+        for cellInfo in skyInfo.patchInfo:
+            coadd_inputs = self.inputRecorder.makeCoaddInputs()
+            # Reserve the absolute maximum of how many ccds, visits
+            # we could potentially have.
+            coadd_inputs.ccds.reserve(len(inputWarps))
+            coadd_inputs.visits.reserve(len(inputWarps))
+            coadd_inputs_gc[cellInfo.index] = coadd_inputs
         # Read in one warp at a time, and accumulate it in all the cells that
         # it completely overlaps.
 
@@ -236,6 +254,9 @@ def run(self, inputWarps, skyMap, skyInfo, **kwargs):
 
                 stacker.add_masked_image(mi, weight=weight)
 
+                coadd_inputs = coadd_inputs_gc[cellInfo.index]
+                self.inputRecorder.addVisitToCoadd(coadd_inputs, warp[bbox], weight)
+
             # del warp
 
         cells: list[SingleCellCoadd] = []
@@ -254,6 +275,14 @@ def run(self, inputWarps, skyMap, skyInfo, **kwargs):
 
             # This is where brightObjectMasks need to be set.
 
+            coadd_inputs = coadd_inputs_gc[cellInfo.index]
+            coadd_inputs.ccds.sort()
+            coadd_inputs.visits.sort()
+            cell_coadd_psf = CoaddPsf(
+                coadd_inputs.ccds, cell_exposure.getWcs(), self.config.coaddPsf.makeControl()
+            )
+            cell_exposure.setPsf(cell_coadd_psf)
+
             image_planes = OwnedImagePlanes.from_exposure(cell_exposure)
             identifiers = CellIdentifiers(
                 cell=cellInfo.index,
@@ -262,10 +291,10 @@ def run(self, inputWarps, skyMap, skyInfo, **kwargs):
                 patch=self.common.identifiers.patch,
                 band=self.common.identifiers.band,
             )
-            cell_coadd_psf = cell_exposure.getPsf()
+
             singleCellCoadd = SingleCellCoadd(
                 outer=image_planes,
-                psf=warp.getPsf().computeKernelImage(warp.getPsf().getAveragePosition()),
+                psf=cell_coadd_psf.computeKernelImage(cell_coadd_psf.getAveragePosition()),
                 inner_bbox=cellInfo.inner_bbox,  # TODO
                 inputs=None,
                 common=self.common,