DM-55521: update coadd processing tasks to work with lsst.images

python/lsst/pipe/tasks/deblendCoaddSourcesPipeline.py

@@ -21,17 +21,21 @@
 
 __all__ = ["DeblendCoaddSourcesMultiConfig", "DeblendCoaddSourcesMultiTask"]
 
+import dataclasses
+
 import numpy as np
 
 from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections
 import lsst.pipe.base.connectionTypes as cT
 
-from lsst.pex.config import ConfigurableField, Field
+from lsst.pex.config import ChoiceField, ConfigurableField, Field
 from lsst.meas.base import SkyMapIdGeneratorConfig
 from lsst.meas.extensions.scarlet import ScarletDeblendTask
 
 import lsst.afw.image as afwImage
 import lsst.afw.table as afwTable
+import lsst.images as imgs
+from lsst.images.cells import CellCoadd
 
 from .coaddBase import reorderRefs
 
@@ -112,12 +116,16 @@ class DeblendCoaddSourcesMultiConnections(PipelineTaskConnections,
 
     def __init__(self, *, config=None):
         super().__init__(config=config)
-        if config:
-            if config.useCellCoadds:
-                del self.coadds
-            else:
-                del self.coadds_cell
-                del self.backgrounds
+        if self.config.imageType == "future":
+            self.coadds = dataclasses.replace(self.coadds, storageClass="CellCoadd")
+            self.deconvolvedCoadds = dataclasses.replace(self.deconvolvedCoadds, storageClass="MaskedImageV2")
+            del self.coadds_cell
+            del self.backgrounds
+        elif self.config.useCellCoadds:
+            del self.coadds
+        else:
+            del self.coadds_cell
+            del self.backgrounds
 
 
 class DeblendCoaddSourcesMultiConfig(PipelineTaskConfig,
@@ -131,6 +139,25 @@ class DeblendCoaddSourcesMultiConfig(PipelineTaskConfig,
         doc="Task to deblend an images in multiple bands"
     )
     idGenerator = SkyMapIdGeneratorConfig.make_field()
+    imageType = ChoiceField(
+        "Which image type to expect for the input coadds. "
+        "This option only directly affects connection storage classes and hence 'runQuantum'; the 'run' "
+        "method behavior is determined by which type is actually passed in.",
+        allowed={
+            "legacy": (
+                "Read a lsst.cell_coadds.MultipleCellCoadd via 'coadds_cells` and restore 'background' "
+                "(if useCellCoadd) or lsst.afw.image.Exposure via `coadds` (if not useCellCoadd), and read "
+                "lsst.afw.image.Exposure via 'deconvolvedCoadds'."
+            ),
+            "future": (
+                "Read lsst.images.cells.CellCoadd via 'coadds' and lsst.images.MaskedImage via "
+                "'deconvolvedCoadds'.  The useCellCoadds options is ignored."
+            ),
+        },
+        dtype=str,
+        optional=False,
+        default="legacy",
+    )
 
 
 class DeblendCoaddSourcesMultiTask(PipelineTask):
@@ -160,17 +187,25 @@ def runQuantum(self, butlerQC, inputRefs, outputRefs):
         inputs = butlerQC.get(inputRefs)
         bands = [dRef.dataId["band"] for dRef in deconvolvedRefs]
         mergedDetections = inputs.pop("mergedDetections")
-        if self.config.useCellCoadds:
-            exposures = [mcc.stitch().asExposure() for mcc in inputs.pop("coadds_cell")]
-            backgrounds = inputs.pop("backgrounds")
-            for exposure, background in zip(exposures, backgrounds):
-                exposure.image -= background.getImage()
-            coadds = exposures
-        else:
-            coadds = inputs.pop("coadds")
+        match self.config.imageType:
+            case "legacy":
+                if self.config.useCellCoadds:
+                    exposures = [mcc.stitch().asExposure() for mcc in inputs.pop("coadds_cell")]
+                    backgrounds = inputs.pop("backgrounds")
+                    for exposure, background in zip(exposures, backgrounds):
+                        exposure.image -= background.getImage()
+                    coadds = exposures
+                    coaddRefs = inputRefs.coadds_cell
+                else:
+                    coadds = inputs.pop("coadds")
+                    coaddRefs = inputRefs.coadds
+            case "future":
+                coadds = inputs.pop("coadds")  # conversion deferred to run().
+                coaddRefs = inputRefs.coadds
+            case _:
+                raise AssertionError(f"Invalid choice {self.config.imageType!r} for imageType.")
 
         # Ensure that the coadd bands and deconvolved coadd bands match
-        coaddRefs = inputRefs.coadds_cell if self.config.useCellCoadds else inputRefs.coadds
         coaddBands = [dRef.dataId["band"] for dRef in coaddRefs]
         if bands != coaddBands:
             self.log.error("Coadd bands %s != deconvolved coadd bands %s", bands, coaddBands)
@@ -194,12 +229,47 @@ def runQuantum(self, butlerQC, inputRefs, outputRefs):
         butlerQC.put(outputs, outputRefs)
 
     def run(self, coadds, bands, mergedDetections, deconvolvedCoadds, idFactory):
+        """Deblend coadds from multiple bands together.
+
+        Parameters
+        ----------
+        coadds : `list` [`lsst.afw.image.Exposure` | \
+                `lsst.images.cells.CellCoadd`]
+            Coadds to deblend.
+        bands : `list` [`str`]
+            Names or the bands for ``coadds`` (zip-iteration compatible).
+        mergedDetections : `lsst.afw.table.SourceCatalog`
+            Input catalog of detections, already merged across bands.
+        deconvolvedCoadds : `list` [`lsst.afw.image.Exposure` | \
+                `lsst.images.MaskedImage`]
+            Deconvolved versions of ``coadds`` (zip-iteration compatible).
+        idFactory : `lsst.afw.table.IdFactory`
+            Factory used to generate output source IDs.
+
+        Returns
+        -------
+        struct : `lsst.pipe.base.Struct`
+            Unmodified outputs of the ``multibandDeblend`` subtask.
+        """
+        coadds = [c.to_legacy() if isinstance(c, CellCoadd) else c for c in coadds]
+        deconvolvedCoadds = [self._coerceDeconvolvedInput(d, c) for d, c in zip(deconvolvedCoadds, coadds)]
         sources = self._makeSourceCatalog(mergedDetections, idFactory)
         multiExposure = afwImage.MultibandExposure.fromExposures(bands, coadds)
         mDeconvolved = afwImage.MultibandExposure.fromExposures(bands, deconvolvedCoadds)
         result = self.multibandDeblend.run(multiExposure, mDeconvolved, sources)
         return result
 
+    def _coerceDeconvolvedInput(
+        self, deconvolved: afwImage.Exposure | imgs.MaskedImage, coadd: afwImage.Exposure
+    ) -> afwImage.Exposure:
+        if isinstance(deconvolved, imgs.MaskedImage):
+            deconvolved = afwImage.Exposure(
+                maskedImage=deconvolved.to_legacy(plane_map=imgs.get_legacy_deep_coadd_mask_planes()),
+                exposureInfo=coadd.getInfo(),
+                dtype=deconvolved.image.array.dtype,
+            )
+        return deconvolved
+
     def _makeSourceCatalog(self, mergedDetections, idFactory):
         # There may be gaps in the mergeDet catalog, which will cause the
         # source ids to be inconsistent. So we update the id factory

python/lsst/pipe/tasks/fit_coadd_multiband.py

@@ -34,6 +34,7 @@
 import lsst.pipe.base.connectionTypes as cT
 
 import astropy.table
+import dataclasses
 from abc import ABC, abstractmethod
 from pydantic import Field
 from pydantic.dataclasses import dataclass
@@ -210,7 +211,11 @@ def __init__(self, *, config=None):
         if config.drop_psf_connection:
             del self.models_psf
 
-        if config.use_cell_coadds:
+        if config.image_type == "future":
+            self.coadds = dataclasses.replace(self.coadds, storageClass="CellCoadd")
+            del self.coadds_cell
+            del self.backgrounds
+        elif config.use_cell_coadds:
             del self.coadds
         else:
             del self.coadds_cell
@@ -316,6 +321,23 @@ class CoaddMultibandFitBaseConfig(
         default=False,
     )
     idGenerator = SkyMapIdGeneratorConfig.make_field()
+    image_type = pexConfig.ChoiceField(
+        "Which image type to expect for the input coadd. "
+        "This option only directly affects connection storage classes and hence 'runQuantum'; the 'run' "
+        "method behavior is determined by which type is actually passed in.",
+        allowed={
+            "legacy": (
+                "Read a lsst.cell_coadds.MultipleCellCoadd via 'coadds_cell` and restore 'background' "
+                "(if use_cell_coadd) or lsst.afw.image.Exposure via `coadds` (if not use_cell_coadd)."
+            ),
+            "future": (
+                "Read lsst.images.cells.CellCoadd via the 'coadds' connection.  use_cell_coadd is ignored."
+            ),
+        },
+        dtype=str,
+        optional=False,
+        default="legacy",
+    )
 
     def get_band_sets(self):
         """Get the set of bands required by the fit_coadd_multiband subtask.
@@ -352,7 +374,7 @@ class CoaddMultibandFitBase:
     def build_catexps(self, butlerQC, inputRefs, inputs) -> list[CatalogExposureInputs]:
         id_tp = self.config.idGenerator.apply(butlerQC.quantum.dataId).catalog_id
         # This is a roundabout way of ensuring all inputs get sorted and matched
-        if self.config.use_cell_coadds:
+        if self.config.use_cell_coadds and self.config.image_type == "legacy":
             keys = ["cats_meas", "coadds_cell", "backgrounds"]
         else:
             keys = ["cats_meas", "coadds"]
@@ -365,7 +387,9 @@ def build_catexps(self, butlerQC, inputRefs, inputs) -> list[CatalogExposureInpu
             for key, (refs, objs) in input_refs_objs.items()
         }
         cats = inputs_sorted["cats_meas"]
-        if self.config.use_cell_coadds:
+        if self.config.image_type == "future":
+            exps = {data_id: coadd.to_legacy() for data_id, coadd in inputs_sorted["coadds"].items()}
+        elif self.config.use_cell_coadds:
             exps = {}
             for data_id, background in inputs_sorted["backgrounds"].items():
                 mcc = inputs_sorted["coadds_cell"][data_id]

python/lsst/pipe/tasks/fit_coadd_psf.py

@@ -31,6 +31,7 @@
 import lsst.pipe.base as pipeBase
 import lsst.pipe.base.connectionTypes as cT
 
+import dataclasses
 from abc import ABC, abstractmethod
 from pydantic.dataclasses import dataclass
 
@@ -93,7 +94,11 @@ def __init__(self, *, config=None):
         if config is None:
             return
 
-        if config.use_cell_coadds:
+        if config.image_type == "future":
+            self.coadd = dataclasses.replace(self.coadd, storageClass="CellCoadd")
+            del self.coadd_cell
+            del self.background
+        elif config.use_cell_coadds:
             del self.coadd
         else:
             del self.coadd_cell
@@ -168,6 +173,23 @@ class CoaddPsfFitConfig(
         doc="Task to fit PSF models for a single coadd",
     )
     idGenerator = SkyMapIdGeneratorConfig.make_field()
+    image_type = pexConfig.ChoiceField(
+        "Which image type to expect for the input coadd. "
+        "This option only directly affects connection storage classes and hence 'runQuantum'; the 'run' "
+        "method behavior is determined by which type is actually passed in.",
+        allowed={
+            "legacy": (
+                "Read a lsst.cell_coadds.MultipleCellCoadd via 'coadd_cell` and restore 'background' "
+                "(if use_cell_coadd) or lsst.afw.image.Exposure via `coadd` (if not use_cell_coadd)."
+            ),
+            "future": (
+                "Read lsst.images.cells.CellCoadd via the 'coadd' connection.  use_cell_coadd is ignored."
+            ),
+        },
+        dtype=str,
+        optional=False,
+        default="legacy",
+    )
 
 
 class CoaddPsfFitTask(pipeBase.PipelineTask):
@@ -191,7 +213,10 @@ def runQuantum(self, butlerQC, inputRefs, outputRefs):
         id_tp = self.config.idGenerator.apply(butlerQC.quantum.dataId).catalog_id
         dataId = inputRefs.cat_meas.dataId
 
-        if self.config.use_cell_coadds:
+        if self.config.image_type == "future":
+            coaddDataRef = inputRefs.coadd
+            exposure = inputs.pop('coadd').to_legacy()
+        elif self.config.use_cell_coadds:
             coaddDataRef = inputRefs.coadd_cell
             multiple_cell_coadd = inputs.pop('coadd_cell')
             background = inputs.pop('background')

python/lsst/pipe/tasks/healSparseMapping.py

@@ -27,6 +27,7 @@
            "ConsolidateHealSparsePropertyMapTask"]
 
 from collections import defaultdict
+import astropy.units
 import esutil
 import warnings
 import numbers
@@ -706,8 +707,9 @@ def run(self, sky_map, tract, band, coadd_dict, input_map_dict, visit_summary_di
                                  band, tract, patch)
                 continue
 
-            coadd_photo_calib = coadd_dict[patch].get(component="photoCalib")
-            coadd_zeropoint = 2.5*np.log10(coadd_photo_calib.getInstFluxAtZeroMagnitude())
+            # LSST coadds are now always in nJy, and the lsst.images formats
+            # don't even have a PhotoCalib anymore.
+            coadd_zeropoint = float((1.0 * astropy.units.nJy).to_value(astropy.units.ABmag))
 
             # Crop input_map to the inner polygon of the patch
             poly_vertices = patch_info.getInnerSkyPolygon(tract_info.getWcs()).getVertices()