Coverage for python/lsst/pipe/tasks/deblendCoaddSourcesPipeline.py: 34%
109 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-14 08:10 +0000
1# This file is part of pipe_tasks.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
22__all__ = ["DeblendCoaddSourcesMultiConfig", "DeblendCoaddSourcesMultiTask"]
24import dataclasses
26import numpy as np
28from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections
29import lsst.pipe.base.connectionTypes as cT
31from lsst.pex.config import ChoiceField, ConfigurableField, Field
32from lsst.meas.base import SkyMapIdGeneratorConfig
33from lsst.meas.extensions.scarlet import ScarletDeblendTask
35import lsst.afw.image as afwImage
36import lsst.afw.table as afwTable
37import lsst.images as imgs
38from lsst.images.cells import CellCoadd
40from .coaddBase import reorderRefs
43deblendBaseTemplates = {"inputCoaddName": "deep", "outputCoaddName": "deep"}
46class DeblendCoaddSourcesMultiConnections(PipelineTaskConnections,
47 dimensions=("tract", "patch", "skymap"),
48 defaultTemplates=deblendBaseTemplates):
49 inputSchema = cT.InitInput(
50 doc="Input schema to use in the deblend catalog",
51 name="{inputCoaddName}Coadd_mergeDet_schema",
52 storageClass="SourceCatalog"
53 )
54 peakSchema = cT.InitInput(
55 doc="Schema of the footprint peak catalogs",
56 name="{inputCoaddName}Coadd_peak_schema",
57 storageClass="PeakCatalog"
58 )
59 mergedDetections = cT.Input(
60 doc="Detection catalog merged across bands",
61 name="{inputCoaddName}Coadd_mergeDet",
62 storageClass="SourceCatalog",
63 dimensions=("tract", "patch", "skymap")
64 )
65 coadds = cT.Input(
66 doc="Exposure on which to run deblending",
67 name="{inputCoaddName}Coadd_calexp",
68 storageClass="ExposureF",
69 multiple=True,
70 dimensions=("tract", "patch", "band", "skymap")
71 )
72 coadds_cell = cT.Input(
73 doc="Exposure on which to run deblending",
74 name="{inputCoaddName}CoaddCell",
75 storageClass="MultipleCellCoadd",
76 multiple=True,
77 dimensions=("tract", "patch", "band", "skymap")
78 )
79 backgrounds = cT.Input(
80 doc="Background model to subtract from the cell-based coadd",
81 name="{inputCoaddName}Coadd_calexp_background",
82 storageClass="Background",
83 multiple=True,
84 dimensions=("tract", "patch", "band", "skymap")
85 )
86 deconvolvedCoadds = cT.Input(
87 doc="Deconvolved coadds",
88 name="deconvolved_{inputCoaddName}_coadd",
89 storageClass="ExposureF",
90 multiple=True,
91 dimensions=("tract", "patch", "band", "skymap")
92 )
93 outputSchema = cT.InitOutput(
94 doc="Output of the schema used in deblending task",
95 name="{outputCoaddName}Coadd_deblendedFlux_schema",
96 storageClass="SourceCatalog"
97 )
98 # TODO[DM-47405]: remove this deprecated connection.
99 fluxCatalogs = cT.Output(
100 doc="Flux weighted catalogs produced by multiband deblending",
101 name="{outputCoaddName}Coadd_deblendedFlux",
102 storageClass="SourceCatalog",
103 dimensions=("tract", "patch", "band", "skymap"),
104 multiple=True,
105 deprecated="Deprecated and unused; will be removed after v29."
106 )
107 # TODO[DM-47405]: remove this deprecated connection.
108 templateCatalogs = cT.Output(
109 doc="Template catalogs produced by multiband deblending",
110 name="{outputCoaddName}Coadd_deblendedModel",
111 storageClass="SourceCatalog",
112 dimensions=("tract", "patch", "band", "skymap"),
113 multiple=True,
114 deprecated="Deprecated and unused; will be removed after v29."
115 )
116 deblendedCatalog = cT.Output(
117 doc="Catalogs produced by multiband deblending",
118 name="{outputCoaddName}Coadd_deblendedCatalog",
119 storageClass="SourceCatalog",
120 dimensions=("tract", "patch", "skymap"),
121 )
122 scarletModelData = cT.Output(
123 doc="Multiband scarlet models produced by the deblender",
124 name="{outputCoaddName}Coadd_scarletModelData",
125 storageClass="LsstScarletModelData",
126 dimensions=("tract", "patch", "skymap"),
127 )
128 objectParents = cT.Output(
129 doc="Parents of the deblended objects",
130 name="object_parent_patch",
131 storageClass="SourceCatalog",
132 dimensions=("tract", "patch", "skymap"),
133 )
135 def __init__(self, *, config=None):
136 super().__init__(config=config)
137 del self.fluxCatalogs
138 del self.templateCatalogs
139 if self.config.imageType == "future":
140 self.coadds = dataclasses.replace(self.coadds, storageClass="CellCoadd")
141 self.deconvolvedCoadds = dataclasses.replace(self.deconvolvedCoadds, storageClass="MaskedImageV2")
142 del self.coadds_cell
143 del self.backgrounds
144 elif self.config.useCellCoadds:
145 del self.coadds
146 else:
147 del self.coadds_cell
148 del self.backgrounds
151class DeblendCoaddSourcesMultiConfig(PipelineTaskConfig,
152 pipelineConnections=DeblendCoaddSourcesMultiConnections):
153 useCellCoadds = Field[bool](
154 doc="Use cell-based coadds instead of regular coadds?",
155 default=False,
156 )
157 multibandDeblend = ConfigurableField(
158 target=ScarletDeblendTask,
159 doc="Task to deblend an images in multiple bands"
160 )
161 idGenerator = SkyMapIdGeneratorConfig.make_field()
162 imageType = ChoiceField(
163 "Which image type to expect for the input coadds. "
164 "This option only directly affects connection storage classes and hence 'runQuantum'; the 'run' "
165 "method behavior is determined by which type is actually passed in.",
166 allowed={
167 "legacy": (
168 "Read a lsst.cell_coadds.MultipleCellCoadd via 'coadds_cells` and restore 'background' "
169 "(if useCellCoadd) or lsst.afw.image.Exposure via `coadds` (if not useCellCoadd), and read "
170 "lsst.afw.image.Exposure via 'deconvolvedCoadds'."
171 ),
172 "future": (
173 "Read lsst.images.cells.CellCoadd via 'coadds' and lsst.images.MaskedImage via "
174 "'deconvolvedCoadds'. The useCellCoadds options is ignored."
175 ),
176 },
177 dtype=str,
178 optional=False,
179 default="legacy",
180 )
183class DeblendCoaddSourcesMultiTask(PipelineTask):
184 ConfigClass = DeblendCoaddSourcesMultiConfig
185 _DefaultName = "deblendCoaddSourcesMulti"
187 def __init__(self, initInputs, **kwargs):
188 super().__init__(initInputs=initInputs, **kwargs)
189 schema = initInputs["inputSchema"].schema
190 self.peakSchema = initInputs["peakSchema"].schema
191 self.schemaMapper = afwTable.SchemaMapper(schema)
192 self.schemaMapper.addMinimalSchema(schema)
193 self.schema = self.schemaMapper.getOutputSchema()
194 self.makeSubtask("multibandDeblend", schema=self.schema, peakSchema=self.peakSchema)
195 self.outputSchema = afwTable.SourceCatalog(self.schema)
197 def runQuantum(self, butlerQC, inputRefs, outputRefs):
198 # Obtain the list of bands, sort them (alphabetically), then reorder
199 # all input lists to match this band order.
200 # Note: sometimes deconvolution fails. If this happens then
201 # the dataIds missing from deconvolvedRefs will be removed
202 # during the process.
203 deconvolvedRefs = inputRefs.deconvolvedCoadds
204 bandOrder = [dRef.dataId["band"] for dRef in deconvolvedRefs]
205 bandOrder.sort()
206 inputRefs = reorderRefs(inputRefs, bandOrder, dataIdKey="band")
207 inputs = butlerQC.get(inputRefs)
208 bands = [dRef.dataId["band"] for dRef in deconvolvedRefs]
209 mergedDetections = inputs.pop("mergedDetections")
210 match self.config.imageType:
211 case "legacy":
212 if self.config.useCellCoadds:
213 exposures = [mcc.stitch().asExposure() for mcc in inputs.pop("coadds_cell")]
214 backgrounds = inputs.pop("backgrounds")
215 for exposure, background in zip(exposures, backgrounds):
216 exposure.image -= background.getImage()
217 coadds = exposures
218 coaddRefs = inputRefs.coadds_cell
219 else:
220 coadds = inputs.pop("coadds")
221 coaddRefs = inputRefs.coadds
222 case "future":
223 coadds = inputs.pop("coadds") # conversion deferred to run().
224 coaddRefs = inputRefs.coadds
225 case _:
226 raise AssertionError(f"Invalid choice {self.config.imageType!r} for imageType.")
228 # Ensure that the coadd bands and deconvolved coadd bands match
229 coaddBands = [dRef.dataId["band"] for dRef in coaddRefs]
230 if bands != coaddBands:
231 self.log.error("Coadd bands %s != deconvolved coadd bands %s", bands, coaddBands)
232 raise RuntimeError(
233 "Number of coadd bands and deconvolved coadd bands do not match. "
234 "This should never happen and indicates a bug in reorderRefs."
235 )
237 deconvolvedCoadds = inputs.pop("deconvolvedCoadds")
239 # Check that all inputs have been extracted correctly.
240 assert not inputs, "runQuantum got extra inputs"
242 outputs = self.run(
243 coadds=coadds,
244 bands=bands,
245 mergedDetections=mergedDetections,
246 idFactory=self.config.idGenerator.apply(butlerQC.quantum.dataId).make_table_id_factory(),
247 deconvolvedCoadds=deconvolvedCoadds,
248 )
249 butlerQC.put(outputs, outputRefs)
251 def run(self, coadds, bands, mergedDetections, deconvolvedCoadds, idFactory):
252 """Deblend coadds from multiple bands together.
254 Parameters
255 ----------
256 coadds : `list` [`lsst.afw.image.Exposure` | \
257 `lsst.images.cells.CellCoadd`]
258 Coadds to deblend.
259 bands : `list` [`str`]
260 Names or the bands for ``coadds`` (zip-iteration compatible).
261 mergedDetections : `lsst.afw.table.SourceCatalog`
262 Input catalog of detections, already merged across bands.
263 deconvolvedCoadds : `list` [`lsst.afw.image.Exposure` | \
264 `lsst.images.MaskedImage`]
265 Deconvolved versions of ``coadds`` (zip-iteration compatible).
266 idFactory : `lsst.afw.table.IdFactory`
267 Factory used to generate output source IDs.
269 Returns
270 -------
271 struct : `lsst.pipe.base.Struct`
272 Unmodified outputs of the ``multibandDeblend`` subtask.
273 """
274 coadds = [c.to_legacy() if isinstance(c, CellCoadd) else c for c in coadds]
275 deconvolvedCoadds = [self._coerceDeconvolvedInput(d, c) for d, c in zip(deconvolvedCoadds, coadds)]
276 sources = self._makeSourceCatalog(mergedDetections, idFactory)
277 multiExposure = afwImage.MultibandExposure.fromExposures(bands, coadds)
278 mDeconvolved = afwImage.MultibandExposure.fromExposures(bands, deconvolvedCoadds)
279 result = self.multibandDeblend.run(multiExposure, mDeconvolved, sources)
280 return result
282 def _coerceDeconvolvedInput(
283 self, deconvolved: afwImage.Exposure | imgs.MaskedImage, coadd: afwImage.Exposure
284 ) -> afwImage.Exposure:
285 if isinstance(deconvolved, imgs.MaskedImage):
286 deconvolved = afwImage.Exposure(
287 maskedImage=deconvolved.to_legacy(plane_map=imgs.get_legacy_deep_coadd_mask_planes()),
288 exposureInfo=coadd.getInfo(),
289 dtype=deconvolved.image.array.dtype,
290 )
291 return deconvolved
293 def _makeSourceCatalog(self, mergedDetections, idFactory):
294 # There may be gaps in the mergeDet catalog, which will cause the
295 # source ids to be inconsistent. So we update the id factory
296 # with the largest id already in the catalog.
297 maxId = np.max(mergedDetections["id"])
298 idFactory.notify(maxId)
299 table = afwTable.SourceTable.make(self.schema, idFactory)
300 sources = afwTable.SourceCatalog(table)
301 sources.extend(mergedDetections, self.schemaMapper)
302 return sources