DM-53757: Implement injecting into cell coadds

.pre-commit-config.yaml

@@ -6,14 +6,14 @@ repos:
       - id: end-of-file-fixer
       - id: trailing-whitespace
   - repo: https://github.com/psf/black
-    rev: 24.4.2
+    rev: 25.12.0
     hooks:
       - id: black
         # It is recommended to specify the latest version of Python
         # supported by your project here, or alternatively use
         # pre-commit's default_language_version, see
         # https://pre-commit.com/#top_level-default_language_version
-        language_version: python3.11
+        language_version: python3.13
   - repo: https://github.com/PyCQA/isort
     rev: 5.13.2
     hooks:

pipelines/match_injected_tract_catalog.yaml

@@ -1,4 +1,8 @@
 description: Match tract-level injected coadd catalogs to injected object table datasets.
+parameters:
+  bands_injected: ["u", "g", "r", "i", "z", "y"]
+  flux_algorithms: ["cModel", "gaap1p0", "kron", "psf", "sersic"]
+  source_type_column: "source_type"
 tasks:
   consolidateInjectedCatalogs:
     class: lsst.source.injection.utils.ConsolidateInjectedCatalogsTask
@@ -19,40 +23,53 @@ tasks:
       match_tract_catalog.columns_target_select_false: ["merge_peak_sky"]
       match_tract_catalog.match_n_finite_min: 2
       match_tract_catalog.order_ascending: true
+      output_matched_catalog: true
+
+      diff_matched_catalog.column_matched_prefix_ref: "ref_"
+      diff_matched_catalog.columns_target_coord_err: ["coord_raErr", "coord_decErr"]
+      diff_matched_catalog.columns_ref_copy: columns_ref_copy
+      diff_matched_catalog.columns_target_copy: columns_target_copy
+      diff_matched_catalog.include_unmatched: true
+
       python: |
         from lsst.pipe.tasks.match_tract_catalog_probabilistic import MatchTractCatalogProbabilisticTask
 
         config.match_tract_catalog.retarget(MatchTractCatalogProbabilisticTask)
-  compareObjectToInjected:
-    class: lsst.pipe.tasks.diff_matched_tract_catalog.DiffMatchedTractCatalogTask
-    config:
-      connections.name_input_cat_ref: injected_deep_coadd_predetection_catalog_tract
-      connections.name_input_cat_target: injected_objectTable_tract
-      column_matched_prefix_ref: "ref_"
-      # TODO: Remove as part of DM-44139
-      columns_ref_mag_to_nJy: {
-        "u_mag": "u_flux", "g_mag": "g_flux", "r_mag": "r_flux",
-        "i_mag": "i_flux", "z_mag": "z_flux", "y_mag": "y_flux",
-      }
-      columns_ref_copy: [
-        "injected_id", "source_type",
-        "u_mag", "g_mag", "r_mag", "i_mag", "z_mag", "y_mag",
-        "u_injection_flag", "g_injection_flag", "r_injection_flag",
-        "i_injection_flag", "z_injection_flag", "y_injection_flag",
-        "injection_flag", "injected_isPatchInner", "injected_isTractInner", "injected_isPrimary",
-      ]
-      columns_target_coord_err: ["coord_raErr", "coord_decErr"]
-      columns_target_copy: [
-        "objectId",
-        "u_psfFlux", "g_psfFlux", "r_psfFlux", "i_psfFlux", "z_psfFlux", "y_psfFlux",
-        "u_psfFluxErr", "g_psfFluxErr", "r_psfFluxErr", "i_psfFluxErr", "z_psfFluxErr", "y_psfFluxErr",
-        "u_psfFlux_flag", "g_psfFlux_flag", "r_psfFlux_flag",
-        "i_psfFlux_flag", "z_psfFlux_flag", "y_psfFlux_flag",
-        "u_sersicFlux", "g_sersicFlux", "r_sersicFlux", "i_sersicFlux", "z_sersicFlux", "y_sersicFlux",
-        "u_sersicFluxErr", "g_sersicFluxErr", "r_sersicFluxErr",
-        "i_sersicFluxErr", "z_sersicFluxErr", "y_sersicFluxErr",
-        "sersic_ra", "sersic_dec",
-        "patch", "detect_isDeblendedSource", "detect_isPatchInner", "detect_isPrimary", "merge_peak_sky",
-        "refExtendedness", "refSizeExtendedness",
-      ]
-      include_unmatched: true
+        columns_ref_mag_to_nJy = {}
+        columns_ref_copy = [
+          "injected_id", parameters.source_type_column, "injection_flag",
+          "injected_isPatchInner", "injected_isTractInner", "injected_isPrimary",
+        ]
+        columns_target_copy = [
+          "objectId", "patch", "detect_isDeblendedSource", "detect_isPatchInner", "detect_isPrimary",
+          "merge_peak_sky", "refExtendedness", "refSizeExtendedness",
+        ]
+        bands = parameters.bands_injected
+        for band in bands:
+          columns_ref_mag_to_nJy[f"{band}_mag"] = f"{band}_flux"
+        for band in bands:
+          for suffix in ("mag", "injection_flag"):
+              columns_ref_copy.append(f"{band}_{suffix}")
+          done_gaap = False
+          for algo in parameters.flux_algorithms:
+            is_gaap = algo.startswith("gaap")
+            if not done_gaap and is_gaap:
+              columns_target_copy.append(f"{band}_gaapFlux_flag")
+              done_gaap = True
+            suffixes = ["Flux", "FluxErr"]
+            if algo == "cModel":
+              suffixes.append("_flag")
+            elif (algo != "sersic") and not algo.startswith("gaap"):
+              suffixes.append("Flux_flag")
+            for suffix in suffixes:
+              columns_target_copy.append(f"{band}_{algo}{suffix}")
+        if "sersic" in parameters.flux_algorithms:
+          columns_target_copy.extend([
+            "sersic_ra", "sersic_dec", "sersic_raErr", "sersic_decErr",
+            "sersic_index", "sersic_indexErr", "sersic_rho", "sersic_rhoErr",
+            "sersic_reff_x", "sersic_reff_xErr", "sersic_reff_y", "sersic_reff_yErr",
+            "sersic_reff_major", "sersic_reff_minor", "sersic_theta",
+            "sersic_unknown_flag", "sersic_no_data_flag",
+          ])
+        # TODO: Remove as part of DM-44139
+        config.diff_matched_catalog.columns_ref_mag_to_nJy = columns_ref_mag_to_nJy

python/lsst/source/injection/inject_base.py

@@ -32,8 +32,13 @@
 from astropy.table import Table, hstack, vstack
 from astropy.units import Quantity, UnitConversionError
 
+from lsst.afw.detection import Psf
+from lsst.afw.geom import SkyWcs
+from lsst.afw.image import ExposureF, PhotoCalib
 from lsst.afw.image.exposure.exposureUtils import bbox_contains_sky_coords
 from lsst.geom import Point2D
+from lsst.images.cells import CellCoadd
+from lsst.images.fields import BaseField
 from lsst.pex.config import ChoiceField, Field, ListField
 from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections, Struct
 from lsst.pipe.base.connectionTypes import PrerequisiteInput
@@ -206,22 +211,31 @@ class BaseInjectTask(PipelineTask):
     _DefaultName = "baseInjectTask"
     ConfigClass = BaseInjectConfig
 
-    def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs):
+    def run(
+        self,
+        injection_catalogs: list[Table],
+        input_exposure: ExposureF | CellCoadd,
+        psf: Psf | None = None,
+        photo_calib: PhotoCalib | None = None,
+        wcs: SkyWcs | None = None,
+    ):
         """Inject sources into an image.
 
         Parameters
         ----------
         injection_catalogs : `list` [`astropy.table.Table`]
             Tract level injection catalogs that potentially cover the named
             input exposure.
-        input_exposure : `lsst.afw.image.ExposureF`
+        input_exposure : `lsst.afw.image.ExposureF` or `lsst.images.CellCoadd`
             The exposure sources will be injected into.
         psf: `lsst.meas.algorithms.ImagePsf`
-            PSF model.
+            PSF model. Only required if input_exposure is an ExposureF.
         photo_calib : `lsst.afw.image.PhotoCalib`
             Photometric calibration used to calibrate injected sources.
+            Only required if input_exposure is an ExposureF.
         wcs : `lsst.afw.geom.SkyWcs`
             WCS used to calibrate injected sources.
+            Only required if input_exposure is an ExposureF.
 
         Returns
         -------
@@ -231,14 +245,24 @@ def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs):
         """
         self.config = cast(BaseInjectConfig, self.config)
 
-        # Attach potential externally calibrated datasets to input_exposure.
-        # Keep originals so we can reset at the end.
-        original_psf = input_exposure.getPsf()
-        original_photo_calib = input_exposure.getPhotoCalib()
-        original_wcs = input_exposure.getWcs()
-        input_exposure.setPsf(psf)
-        input_exposure.setPhotoCalib(photo_calib)
-        input_exposure.setWcs(wcs)
+        if is_exposure := isinstance(input_exposure, ExposureF):
+            # Attach potential externally calibrated datasets to
+            # input_exposure. Keep originals to reset at the end.
+            original_psf = input_exposure.getPsf()
+            original_photo_calib = input_exposure.getPhotoCalib()
+            original_wcs = input_exposure.getWcs()
+            input_exposure.setPsf(psf)
+            input_exposure.setPhotoCalib(photo_calib)
+            input_exposure.setWcs(wcs)
+            bbox = input_exposure.getBBox()
+            photo_calib = input_exposure.getPhotoCalib()
+        elif isinstance(input_exposure, CellCoadd):
+            if wcs is None:
+                wcs = input_exposure.sky_projection.to_legacy()
+            bbox = input_exposure.bbox.to_legacy()
+            photo_calib = BaseField.make_legacy_photo_calib(input_exposure.unit)
+        else:
+            raise ValueError(f"Unsupported {type(input_exposure)=} for injection")
 
         # Make empty table if none supplied to support process_all_data_ids.
         if len(injection_catalogs) == 0:
@@ -274,11 +298,11 @@ def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs):
         injection_catalog = self._standardize_columns(
             injection_catalog,
             column_mapping,
-            input_exposure.getWcs().getPixelScale(input_exposure.getBBox().getCenter()).asArcseconds(),
+            wcs.getPixelScale(bbox.getCenter()).asArcseconds(),
         )
 
         # Clean the injection catalog of sources which are not injectable.
-        injection_catalog = self._clean_sources(injection_catalog, input_exposure)
+        injection_catalog = self._clean_sources(injection_catalog, input_exposure, bbox=bbox, wcs=wcs)
 
         # Injection binary flag lookup dictionary.
         binary_flags = {
@@ -379,17 +403,20 @@ def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs):
             )
 
         # Restore original input_exposure calibrated data.
-        input_exposure.setPsf(original_psf)
-        input_exposure.setPhotoCalib(original_photo_calib)
-        input_exposure.setWcs(original_wcs)
-
-        # Add injection provenance and injection flags metadata.
-        metadata = input_exposure.getMetadata()
-        input_dataset_type = self.config.connections.input_exposure.format(**self.config.connections.toDict())
-        metadata.set("INJECTED", input_dataset_type, "Initial source injection dataset type")
-        input_exposure.getInfo().setVisitInfo(input_exposure.visitInfo.copyWith(hasSimulatedContent=True))
-        for flag, value in sorted(binary_flags.items(), key=lambda item: item[1]):
-            injection_catalog.meta[flag] = value
+        if is_exposure:
+            input_exposure.setPsf(original_psf)
+            input_exposure.setPhotoCalib(original_photo_calib)
+            input_exposure.setWcs(original_wcs)
+
+            # Add injection provenance and injection flags metadata.
+            metadata = input_exposure.getMetadata()
+            input_dataset_type = self.config.connections.input_exposure.format(
+                **self.config.connections.toDict()
+            )
+            metadata.set("INJECTED", input_dataset_type, "Initial source injection dataset type")
+            input_exposure.getInfo().setVisitInfo(input_exposure.visitInfo.copyWith(hasSimulatedContent=True))
+            for flag, value in sorted(binary_flags.items(), key=lambda item: item[1]):
+                injection_catalog.meta[flag] = value
 
         output_struct = Struct(output_exposure=input_exposure, output_catalog=injection_catalog)
         return output_struct
@@ -494,7 +521,7 @@ def _standardize_columns(self, injection_catalog, column_mapping, pixel_scale):
                     pass
         return Table(injection_catalog)
 
-    def _clean_sources(self, injection_catalog, input_exposure):
+    def _clean_sources(self, injection_catalog, input_exposure, bbox=None, wcs=None):
         """Clean the injection catalog of sources which are not injectable.
 
         This method will remove sources which are not injectable for a variety
@@ -522,6 +549,12 @@ def _clean_sources(self, injection_catalog, input_exposure):
         """
         self.config = cast(BaseInjectConfig, self.config)
 
+        if wcs is None:
+            wcs = input_exposure.getWcs()
+
+        if bbox is None:
+            bbox = input_exposure.getBBox()
+
         # Exit early if there are no sources to inject.
         if len(injection_catalog) == 0:
             self.log.info("Catalog cleaning not applied to empty injection catalog.")
@@ -530,13 +563,12 @@ def _clean_sources(self, injection_catalog, input_exposure):
         sources_to_keep = np.ones(len(injection_catalog), dtype=bool)
 
         # Determine centroids and remove sources outside the padded bbox.
-        wcs = input_exposure.getWcs()
         has_sky = {"ra", "dec"} <= set(injection_catalog.columns)
         has_pixel = {"x", "y"} <= set(injection_catalog.columns)
         # Input catalog must contain either RA/Dec OR x/y.
         # If only x/y given, RA/Dec will be calculated.
         if not has_sky and has_pixel:
-            begin_x, begin_y = input_exposure.getBBox().getBegin()
+            begin_x, begin_y = bbox.getBegin()
             ras, decs = wcs.pixelToSkyArray(
                 begin_x + injection_catalog["x"].astype(float),
                 begin_y + injection_catalog["y"].astype(float),
@@ -550,7 +582,6 @@ def _clean_sources(self, injection_catalog, input_exposure):
         elif not has_sky and not has_pixel:
             self.log.warning("No spatial coordinates found in injection catalog; cannot inject any sources!")
         if has_sky:
-            bbox = input_exposure.getBBox()
             if self.config.trim_padding:
                 bbox.grow(int(self.config.trim_padding))
             is_contained = bbox_contains_sky_coords(

python/lsst/source/injection/inject_coadd.py

@@ -23,6 +23,10 @@
 
 __all__ = ["CoaddInjectConnections", "CoaddInjectConfig", "CoaddInjectTask"]
 
+import dataclasses
+from typing import cast
+
+from lsst.pex.config import Field
 from lsst.pipe.base.connectionTypes import Input, Output
 
 from .inject_base import BaseInjectConfig, BaseInjectConnections, BaseInjectTask
@@ -56,14 +60,24 @@ class CoaddInjectConnections(
         dimensions=("skymap", "tract", "patch", "band"),
     )
 
+    def __init__(self, *, config: CoaddInjectConfig | None = None):
+        super().__init__(config=config)
+
+        if not self.config:
+            return
+
+        if self.config.use_cell_coadds:
+            self.input_exposure = dataclasses.replace(self.input_exposure, storageClass="CellCoadd")
+            self.output_exposure = dataclasses.replace(self.output_exposure, storageClass="CellCoadd")
+
 
 class CoaddInjectConfig(  # type: ignore [call-arg]
     BaseInjectConfig,
     pipelineConnections=CoaddInjectConnections,
 ):
     """Coadd-level configuration for source injection tasks."""
 
-    pass
+    use_cell_coadds = Field(dtype=bool, default=False, doc="Whether to use cell coadds?")
 
 
 class CoaddInjectTask(BaseInjectTask):
@@ -75,9 +89,14 @@ class CoaddInjectTask(BaseInjectTask):
     def runQuantum(self, butler_quantum_context, input_refs, output_refs):
         inputs = butler_quantum_context.get(input_refs)
 
-        inputs["psf"] = inputs["input_exposure"].getPsf()
-        inputs["photo_calib"] = inputs["input_exposure"].getPhotoCalib()
-        inputs["wcs"] = inputs["input_exposure"].getWcs()
+        config = cast(self.ConfigClass, self.config)
+        exposure = inputs["input_exposure"]
+        if config.use_cell_coadds:
+            exposure.apply_background(None)
+        else:
+            inputs["psf"] = exposure.getPsf()
+            inputs["photo_calib"] = exposure.getPhotoCalib()
+            inputs["wcs"] = exposure.getWcs()
 
         input_keys = ["injection_catalogs", "input_exposure", "sky_map", "psf", "photo_calib", "wcs"]
         outputs = self.run(**{key: value for (key, value) in inputs.items() if key in input_keys})