Coverage for python/lsst/pipe/tasks/deblendCoaddSourcesPipeline.py: 33%
105 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-09-29 10:07 +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 deblendedCatalog = cT.Output(
99 doc="Catalogs produced by multiband deblending",
100 name="{outputCoaddName}Coadd_deblendedCatalog",
101 storageClass="SourceCatalog",
102 dimensions=("tract", "patch", "skymap"),
103 )
104 scarletModelData = cT.Output(
105 doc="Multiband scarlet models produced by the deblender",
106 name="{outputCoaddName}Coadd_scarletModelData",
107 storageClass="LsstScarletModelData",
108 dimensions=("tract", "patch", "skymap"),
109 )
110 objectParents = cT.Output(
111 doc="Parents of the deblended objects",
112 name="object_parent_patch",
113 storageClass="SourceCatalog",
114 dimensions=("tract", "patch", "skymap"),
115 )
117 def __init__(self, *, config=None):
118 super().__init__(config=config)
119 if self.config.imageType == "future":
120 self.coadds = dataclasses.replace(self.coadds, storageClass="CellCoadd")
121 self.deconvolvedCoadds = dataclasses.replace(self.deconvolvedCoadds, storageClass="MaskedImageV2")
122 del self.coadds_cell
123 del self.backgrounds
124 elif self.config.useCellCoadds:
125 del self.coadds
126 else:
127 del self.coadds_cell
128 del self.backgrounds
131class DeblendCoaddSourcesMultiConfig(PipelineTaskConfig,
132 pipelineConnections=DeblendCoaddSourcesMultiConnections):
133 useCellCoadds = Field[bool](
134 doc="Use cell-based coadds instead of regular coadds?",
135 default=False,
136 )
137 multibandDeblend = ConfigurableField(
138 target=ScarletDeblendTask,
139 doc="Task to deblend an images in multiple bands"
140 )
141 idGenerator = SkyMapIdGeneratorConfig.make_field()
142 imageType = ChoiceField(
143 "Which image type to expect for the input coadds. "
144 "This option only directly affects connection storage classes and hence 'runQuantum'; the 'run' "
145 "method behavior is determined by which type is actually passed in.",
146 allowed={
147 "legacy": (
148 "Read a lsst.cell_coadds.MultipleCellCoadd via 'coadds_cells` and restore 'background' "
149 "(if useCellCoadd) or lsst.afw.image.Exposure via `coadds` (if not useCellCoadd), and read "
150 "lsst.afw.image.Exposure via 'deconvolvedCoadds'."
151 ),
152 "future": (
153 "Read lsst.images.cells.CellCoadd via 'coadds' and lsst.images.MaskedImage via "
154 "'deconvolvedCoadds'. The useCellCoadds options is ignored."
155 ),
156 },
157 dtype=str,
158 optional=False,
159 default="legacy",
160 )
163class DeblendCoaddSourcesMultiTask(PipelineTask):
164 ConfigClass = DeblendCoaddSourcesMultiConfig
165 _DefaultName = "deblendCoaddSourcesMulti"
167 def __init__(self, initInputs, **kwargs):
168 super().__init__(initInputs=initInputs, **kwargs)
169 schema = initInputs["inputSchema"].schema
170 self.peakSchema = initInputs["peakSchema"].schema
171 self.schemaMapper = afwTable.SchemaMapper(schema)
172 self.schemaMapper.addMinimalSchema(schema)
173 self.schema = self.schemaMapper.getOutputSchema()
174 self.makeSubtask("multibandDeblend", schema=self.schema, peakSchema=self.peakSchema)
175 self.outputSchema = afwTable.SourceCatalog(self.schema)
177 def runQuantum(self, butlerQC, inputRefs, outputRefs):
178 # Obtain the list of bands, sort them (alphabetically), then reorder
179 # all input lists to match this band order.
180 # Note: sometimes deconvolution fails. If this happens then
181 # the dataIds missing from deconvolvedRefs will be removed
182 # during the process.
183 deconvolvedRefs = inputRefs.deconvolvedCoadds
184 bandOrder = [dRef.dataId["band"] for dRef in deconvolvedRefs]
185 bandOrder.sort()
186 inputRefs = reorderRefs(inputRefs, bandOrder, dataIdKey="band")
187 inputs = butlerQC.get(inputRefs)
188 bands = [dRef.dataId["band"] for dRef in deconvolvedRefs]
189 mergedDetections = inputs.pop("mergedDetections")
190 match self.config.imageType:
191 case "legacy":
192 if self.config.useCellCoadds:
193 exposures = [mcc.stitch().asExposure() for mcc in inputs.pop("coadds_cell")]
194 backgrounds = inputs.pop("backgrounds")
195 for exposure, background in zip(exposures, backgrounds):
196 exposure.image -= background.getImage()
197 coadds = exposures
198 coaddRefs = inputRefs.coadds_cell
199 else:
200 coadds = inputs.pop("coadds")
201 coaddRefs = inputRefs.coadds
202 case "future":
203 coadds = inputs.pop("coadds") # conversion deferred to run().
204 coaddRefs = inputRefs.coadds
205 case _:
206 raise AssertionError(f"Invalid choice {self.config.imageType!r} for imageType.")
208 # Ensure that the coadd bands and deconvolved coadd bands match
209 coaddBands = [dRef.dataId["band"] for dRef in coaddRefs]
210 if bands != coaddBands:
211 self.log.error("Coadd bands %s != deconvolved coadd bands %s", bands, coaddBands)
212 raise RuntimeError(
213 "Number of coadd bands and deconvolved coadd bands do not match. "
214 "This should never happen and indicates a bug in reorderRefs."
215 )
217 deconvolvedCoadds = inputs.pop("deconvolvedCoadds")
219 # Check that all inputs have been extracted correctly.
220 assert not inputs, "runQuantum got extra inputs"
222 outputs = self.run(
223 coadds=coadds,
224 bands=bands,
225 mergedDetections=mergedDetections,
226 idFactory=self.config.idGenerator.apply(butlerQC.quantum.dataId).make_table_id_factory(),
227 deconvolvedCoadds=deconvolvedCoadds,
228 )
229 butlerQC.put(outputs, outputRefs)
231 def run(self, coadds, bands, mergedDetections, deconvolvedCoadds, idFactory):
232 """Deblend coadds from multiple bands together.
234 Parameters
235 ----------
236 coadds : `list` [`lsst.afw.image.Exposure` | \
237 `lsst.images.cells.CellCoadd`]
238 Coadds to deblend.
239 bands : `list` [`str`]
240 Names or the bands for ``coadds`` (zip-iteration compatible).
241 mergedDetections : `lsst.afw.table.SourceCatalog`
242 Input catalog of detections, already merged across bands.
243 deconvolvedCoadds : `list` [`lsst.afw.image.Exposure` | \
244 `lsst.images.MaskedImage`]
245 Deconvolved versions of ``coadds`` (zip-iteration compatible).
246 idFactory : `lsst.afw.table.IdFactory`
247 Factory used to generate output source IDs.
249 Returns
250 -------
251 struct : `lsst.pipe.base.Struct`
252 Unmodified outputs of the ``multibandDeblend`` subtask.
253 """
254 coadds = [c.to_legacy() if isinstance(c, CellCoadd) else c for c in coadds]
255 deconvolvedCoadds = [self._coerceDeconvolvedInput(d, c) for d, c in zip(deconvolvedCoadds, coadds)]
256 sources = self._makeSourceCatalog(mergedDetections, idFactory)
257 multiExposure = afwImage.MultibandExposure.fromExposures(bands, coadds)
258 mDeconvolved = afwImage.MultibandExposure.fromExposures(bands, deconvolvedCoadds)
259 result = self.multibandDeblend.run(multiExposure, mDeconvolved, sources)
260 return result
262 def _coerceDeconvolvedInput(
263 self, deconvolved: afwImage.Exposure | imgs.MaskedImage, coadd: afwImage.Exposure
264 ) -> afwImage.Exposure:
265 if isinstance(deconvolved, imgs.MaskedImage):
266 deconvolved = afwImage.Exposure(
267 maskedImage=deconvolved.to_legacy(plane_map=imgs.get_legacy_deep_coadd_mask_planes()),
268 exposureInfo=coadd.getInfo(),
269 dtype=deconvolved.image.array.dtype,
270 )
271 return deconvolved
273 def _makeSourceCatalog(self, mergedDetections, idFactory):
274 # There may be gaps in the mergeDet catalog, which will cause the
275 # source ids to be inconsistent. So we update the id factory
276 # with the largest id already in the catalog.
277 maxId = np.max(mergedDetections["id"])
278 idFactory.notify(maxId)
279 table = afwTable.SourceTable.make(self.schema, idFactory)
280 sources = afwTable.SourceCatalog(table)
281 sources.extend(mergedDetections, self.schemaMapper)
282 return sources