Coverage for python/lsst/pipe/tasks/fit_coadd_multiband.py: 54%
158 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-04 17:20 +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__ = [
23 "CoaddMultibandFitConfig", "CoaddMultibandFitConnections", "CoaddMultibandFitSubConfig",
24 "CoaddMultibandFitSubTask", "CoaddMultibandFitTask",
25]
27from .fit_multiband import CatalogExposure, CatalogExposureConfig
29import lsst.afw.table as afwTable
30from lsst.meas.base import SkyMapIdGeneratorConfig
31from lsst.meas.extensions.scarlet.io import updateCatalogFootprints
32import lsst.pex.config as pexConfig
33import lsst.pipe.base as pipeBase
34import lsst.pipe.base.connectionTypes as cT
36import astropy.table
37from abc import ABC, abstractmethod
38from pydantic import Field
39from pydantic.dataclasses import dataclass
40from typing import Iterable
42CoaddMultibandFitBaseTemplates = {
43 "name_coadd": "deep",
44 "name_method": "multiprofit",
45 "name_table": "objects",
46}
49@dataclass(frozen=True, kw_only=True, config=CatalogExposureConfig)
50class CatalogExposureInputs(CatalogExposure):
51 table_psf_fits: astropy.table.Table = Field(title="A table of PSF fit parameters for each source")
53 def get_catalog(self):
54 return self.catalog
57class CoaddMultibandFitInputConnections(
58 pipeBase.PipelineTaskConnections,
59 dimensions=("tract", "patch", "skymap"),
60 defaultTemplates=CoaddMultibandFitBaseTemplates,
61):
62 cat_ref = cT.Input(
63 doc="Reference multiband source catalog",
64 name="{name_coadd}Coadd_ref",
65 storageClass="SourceCatalog",
66 dimensions=("tract", "patch", "skymap"),
67 )
68 cats_meas = cT.Input(
69 doc="Deblended single-band source catalogs",
70 name="{name_coadd}Coadd_meas",
71 storageClass="SourceCatalog",
72 dimensions=("tract", "patch", "band", "skymap"),
73 multiple=True,
74 )
75 coadds = cT.Input(
76 doc="Exposures on which to run fits",
77 name="{name_coadd}Coadd_calexp",
78 storageClass="ExposureF",
79 dimensions=("tract", "patch", "band", "skymap"),
80 multiple=True,
81 )
82 coadds_cell = cT.Input(
83 doc="Cell-coadd exposures on which to run fits",
84 name="{name_coadd}CoaddCell",
85 storageClass="MultipleCellCoadd",
86 dimensions=("tract", "patch", "band", "skymap"),
87 multiple=True,
88 )
89 backgrounds = cT.Input(
90 doc="Background models to subtract from the coadds_cell",
91 name="{name_coadd}Coadd_calexp_background",
92 storageClass="Background",
93 dimensions=("tract", "patch", "band", "skymap"),
94 multiple=True,
95 )
96 models_psf = cT.Input(
97 doc="Input PSF model parameter catalog",
98 # Consider allowing independent psf fit method
99 name="{name_coadd}Coadd_psfs_{name_method}",
100 storageClass="ArrowAstropy",
101 dimensions=("tract", "patch", "band", "skymap"),
102 multiple=True,
103 deferLoad=True,
104 )
105 models_scarlet = pipeBase.connectionTypes.Input(
106 doc="Multiband scarlet models produced by the deblender",
107 name="{name_coadd}Coadd_scarletModelData",
108 storageClass="LsstScarletModelData",
109 dimensions=("tract", "patch", "skymap"),
110 )
112 def adjustQuantum(self, inputs, outputs, label, data_id):
113 """Validates the `lsst.daf.butler.DatasetRef` bands against the
114 subtask's list of bands to fit and drops unnecessary bands.
116 Parameters
117 ----------
118 inputs : `dict`
119 Dictionary whose keys are an input (regular or prerequisite)
120 connection name and whose values are a tuple of the connection
121 instance and a collection of associated `DatasetRef` objects.
122 The exact type of the nested collections is unspecified; it can be
123 assumed to be multi-pass iterable and support `len` and ``in``, but
124 it should not be mutated in place. In contrast, the outer
125 dictionaries are guaranteed to be temporary copies that are true
126 `dict` instances, and hence may be modified and even returned; this
127 is especially useful for delegating to `super` (see notes below).
128 outputs : `Mapping`
129 Mapping of output datasets, with the same structure as ``inputs``.
130 label : `str`
131 Label for this task in the pipeline (should be used in all
132 diagnostic messages).
133 data_id : `lsst.daf.butler.DataCoordinate`
134 Data ID for this quantum in the pipeline (should be used in all
135 diagnostic messages).
137 Returns
138 -------
139 adjusted_inputs : `Mapping`
140 Mapping of the same form as ``inputs`` with updated containers of
141 input `DatasetRef` objects. All inputs involving the 'band'
142 dimension are adjusted to put them in consistent order and remove
143 unneeded bands.
144 adjusted_outputs : `Mapping`
145 Mapping of updated output datasets; always empty for this task.
147 Raises
148 ------
149 lsst.pipe.base.NoWorkFound
150 Raised if there are not enough of the right bands to run the task
151 on this quantum.
152 """
153 # Check which bands are going to be fit
154 bands_fit, bands_read_only = self.config.get_band_sets()
155 bands_needed = bands_fit + [band for band in bands_read_only if band not in bands_fit]
156 bands_needed_set = set(bands_needed)
158 adjusted_inputs = {}
159 inputs_to_adjust = {}
160 bands_found = bands_needed_set
161 for connection_name, (connection, dataset_refs) in inputs.items():
162 # Datasets without bands in their dimensions should be fine
163 if 'band' in connection.dimensions: 163 ↛ 161line 163 didn't jump to line 161 because the condition on line 163 was always true
164 datasets_by_band = {dref.dataId['band']: dref for dref in dataset_refs}
165 bands_set = set(datasets_by_band.keys())
166 if self.config.allow_missing_bands:
167 if len(bands_found) == 0: 167 ↛ 168line 167 didn't jump to line 168 because the condition on line 167 was never true
168 raise pipeBase.NoWorkFound(
169 f'DatasetRefs={dataset_refs} for {connection_name=} is empty'
170 )
171 bands_found &= bands_set
172 # All configured bands are treated as necessary
173 elif not bands_needed_set.issubset(bands_set):
174 raise pipeBase.NoWorkFound(
175 f'DatasetRefs={dataset_refs} have data with bands in the'
176 f' set={set(datasets_by_band.keys())},'
177 f' which is not a superset of the required bands={bands_needed} defined by'
178 f' {self.config.__class__}.fit_coadd_multiband='
179 f'{self.config.fit_coadd_multiband._value.__class__}\'s attributes'
180 f' bands_fit={bands_fit} and bands_read_only()={bands_read_only}.'
181 f' Add the required bands={set(bands_needed).difference(datasets_by_band.keys())}.'
182 )
183 # Adjust all datasets with band dimensions to include just
184 # the needed bands, in consistent order.
185 inputs_to_adjust[connection_name] = (connection, datasets_by_band)
187 if self.config.allow_missing_bands: 187 ↛ 195line 187 didn't jump to line 195 because the condition on line 187 was always true
188 bands_needed = [band for band in bands_fit if band in bands_found] + [
189 band for band in bands_read_only if band not in bands_found
190 ]
191 if len(bands_needed) == 0:
192 raise pipeBase.NoWorkFound(
193 f'No common bands remaining for inputs {",".join(inputs_to_adjust.keys())}'
194 )
195 for connection_name, (connection, datasets_by_band) in inputs_to_adjust.items():
196 adjusted_inputs[connection_name] = (
197 connection,
198 [datasets_by_band[band] for band in bands_needed]
199 )
201 # Delegate to super for more checks.
202 inputs.update(adjusted_inputs)
203 super().adjustQuantum(inputs, outputs, label, data_id)
204 return adjusted_inputs, {}
206 def __init__(self, *, config=None):
207 super().__init__(config=config)
208 assert isinstance(config, CoaddMultibandFitBaseConfig)
210 if config.drop_psf_connection: 210 ↛ 211line 210 didn't jump to line 211 because the condition on line 210 was never true
211 del self.models_psf
213 if config.use_cell_coadds: 213 ↛ 214line 213 didn't jump to line 214 because the condition on line 213 was never true
214 del self.coadds
215 else:
216 del self.coadds_cell
217 del self.backgrounds
220class CoaddMultibandFitConnections(CoaddMultibandFitInputConnections):
221 cat_output = cT.Output(
222 doc="Output source model fit parameter catalog",
223 name="{name_coadd}Coadd_{name_table}_{name_method}",
224 storageClass="ArrowTable",
225 dimensions=("tract", "patch", "skymap"),
226 )
229class CoaddMultibandFitSubConfig(pexConfig.Config):
230 """Configuration for implementing fitter subtasks.
231 """
233 bands_fit = pexConfig.ListField[str](
234 default=[],
235 doc="list of bandpass filters to fit",
236 listCheck=lambda x: (len(x) > 0) and (len(set(x)) == len(x)),
237 )
239 @abstractmethod
240 def bands_read_only(self) -> set:
241 """Return the set of bands that the Task needs to read (e.g. for
242 defining priors) but not necessarily fit.
244 Returns
245 -------
246 The set of such bands.
247 """
248 return set()
251class CoaddMultibandFitSubTask(pipeBase.Task, ABC):
252 """Subtask interface for multiband fitting of deblended sources.
254 Parameters
255 ----------
256 **kwargs
257 Additional arguments to be passed to the `lsst.pipe.base.Task`
258 constructor.
259 """
260 ConfigClass = CoaddMultibandFitSubConfig
262 def __init__(self, **kwargs):
263 super().__init__(**kwargs)
265 @abstractmethod
266 def run(
267 self, catexps: Iterable[CatalogExposureInputs], cat_ref: afwTable.SourceCatalog
268 ) -> pipeBase.Struct:
269 """Fit models to deblended sources from multi-band inputs.
271 Parameters
272 ----------
273 catexps : `typing.List [CatalogExposureInputs]`
274 A list of catalog-exposure pairs with metadata in a given band.
275 cat_ref : `lsst.afw.table.SourceCatalog`
276 A reference source catalog to fit.
278 Returns
279 -------
280 retStruct : `lsst.pipe.base.Struct`
281 A struct with a cat_output attribute containing the output
282 measurement catalog.
284 Notes
285 -----
286 Subclasses may have further requirements on the input parameters,
287 including:
288 - Passing only one catexp per band;
289 - Catalogs containing HeavyFootprints with deblended images;
290 - Fitting only a subset of the sources.
291 If any requirements are not met, the subtask should fail as soon as
292 possible.
293 """
296class CoaddMultibandFitBaseConfig(
297 pipeBase.PipelineTaskConfig,
298 pipelineConnections=CoaddMultibandFitInputConnections,
299):
300 """Base class for multiband fitting."""
302 allow_missing_bands = pexConfig.Field[bool](
303 doc="Whether to still fit even if some bands are missing",
304 default=True,
305 )
306 drop_psf_connection = pexConfig.Field[bool](
307 doc="Whether to drop the PSF model connection, e.g. because PSF parameters are in the input catalog",
308 default=False,
309 )
310 fit_coadd_multiband = pexConfig.ConfigurableField(
311 target=CoaddMultibandFitSubTask,
312 doc="Task to fit sources using multiple bands",
313 )
314 use_cell_coadds = pexConfig.Field[bool](
315 doc="Use cell coadd images for object fitting?",
316 default=False,
317 )
318 idGenerator = SkyMapIdGeneratorConfig.make_field()
320 def get_band_sets(self):
321 """Get the set of bands required by the fit_coadd_multiband subtask.
323 Returns
324 -------
325 bands_fit : `set`
326 The set of bands that the subtask will fit.
327 bands_read_only : `set`
328 The set of bands that the subtask will only read data
329 (measurement catalog and exposure) for.
330 """
331 try:
332 bands_fit = self.fit_coadd_multiband.bands_fit
333 except AttributeError:
334 raise RuntimeError(f'{__class__}.fit_coadd_multiband must have bands_fit attribute') from None
335 bands_read_only = self.fit_coadd_multiband.bands_read_only()
336 return tuple(list({band: None for band in bands}.keys()) for bands in (bands_fit, bands_read_only))
339class CoaddMultibandFitConfig(
340 CoaddMultibandFitBaseConfig,
341 pipelineConnections=CoaddMultibandFitConnections,
342):
343 """Configuration for a CoaddMultibandFitTask."""
346class CoaddMultibandFitBase:
347 """Base class for tasks that fit or rebuild multiband models.
349 This class only implements data reconstruction.
350 """
352 def build_catexps(self, butlerQC, inputRefs, inputs) -> list[CatalogExposureInputs]:
353 id_tp = self.config.idGenerator.apply(butlerQC.quantum.dataId).catalog_id
354 # This is a roundabout way of ensuring all inputs get sorted and matched
355 if self.config.use_cell_coadds:
356 keys = ["cats_meas", "coadds_cell", "backgrounds"]
357 else:
358 keys = ["cats_meas", "coadds"]
359 has_psf_models = "models_psf" in inputs
360 if has_psf_models:
361 keys.append("models_psf")
362 input_refs_objs = {key: (getattr(inputRefs, key), inputs[key]) for key in keys}
363 inputs_sorted = {
364 key: {dRef.dataId: obj for dRef, obj in zip(refs, objs, strict=True)}
365 for key, (refs, objs) in input_refs_objs.items()
366 }
367 cats = inputs_sorted["cats_meas"]
368 if self.config.use_cell_coadds:
369 exps = {}
370 for data_id, background in inputs_sorted["backgrounds"].items():
371 mcc = inputs_sorted["coadds_cell"][data_id]
372 stitched_coadd = mcc.stitch()
373 exposure = stitched_coadd.asExposure()
374 exposure.image -= background.getImage()
375 exps[data_id] = exposure
376 else:
377 exps = inputs_sorted["coadds"]
379 # Ensure that psf models are loaded with full metadata.
380 if has_psf_models:
381 ref0 = list(inputs_sorted["models_psf"].values())[0]
382 parameters = None
383 if ref0.ref.datasetType.storageClass_name == "ArrowAstropy":
384 parameters = {"strip_astropy_meta_yaml": False}
385 models_psf = {
386 key: ref.get(parameters=parameters)
387 for key, ref in inputs_sorted["models_psf"].items()
388 }
389 else:
390 models_psf = None
392 dataIds = set(cats).union(set(exps))
393 models_scarlet = inputs["models_scarlet"]
394 catexp_dict = {}
395 dataId = None
396 for dataId in dataIds:
397 catalog = cats[dataId]
398 exposure = exps[dataId]
399 updateCatalogFootprints(
400 modelData=models_scarlet,
401 catalog=catalog,
402 band=dataId['band'],
403 imageForRedistribution=exposure,
404 removeScarletData=False,
405 updateFluxColumns=False,
406 )
407 catexp_dict[dataId['band']] = CatalogExposureInputs(
408 catalog=catalog,
409 exposure=exposure,
410 table_psf_fits=models_psf[dataId] if has_psf_models else astropy.table.Table(),
411 dataId=dataId,
412 id_tract_patch=id_tp,
413 )
414 # This shouldn't happen unless this is called with no inputs, but check anyway
415 if dataId is None:
416 raise RuntimeError(f"Did not build any catexps for {inputRefs=}")
417 catexps = []
418 for band in self.config.get_band_sets()[0]:
419 if band in catexp_dict:
420 catexp = catexp_dict[band]
421 else:
422 # Make a dummy catexp with a dataId if there's no data
423 # This should be handled by any subtasks
424 dataId_band = dataId.to_simple(minimal=True)
425 dataId_band.dataId["band"] = band
426 catexp = CatalogExposureInputs(
427 catalog=afwTable.SourceCatalog(),
428 exposure=None,
429 table_psf_fits=astropy.table.Table(),
430 dataId=dataId.from_simple(dataId_band, universe=dataId.universe),
431 id_tract_patch=id_tp,
432 )
433 catexps.append(catexp)
434 return catexps
437class CoaddMultibandFitTask(CoaddMultibandFitBase, pipeBase.PipelineTask):
438 """Fit deblended exposures in multiple bands simultaneously.
440 It is generally assumed but not enforced (except optionally by the
441 configurable `fit_coadd_multiband` subtask) that there is only one exposure
442 per band, presumably a coadd.
443 """
445 ConfigClass = CoaddMultibandFitConfig
446 _DefaultName = "coaddMultibandFit"
448 def __init__(self, initInputs, **kwargs):
449 super().__init__(initInputs=initInputs, **kwargs)
450 self.makeSubtask("fit_coadd_multiband")
452 def make_kwargs(self, butlerQC, inputRefs, inputs):
453 """Make any kwargs needed to be passed to run.
455 This method should be overloaded by subclasses that are configured to
456 use a specific subtask that needs additional arguments derived from
457 the inputs but do not otherwise need to overload runQuantum."""
458 return {}
460 def runQuantum(self, butlerQC, inputRefs, outputRefs):
461 inputs = butlerQC.get(inputRefs)
462 catexps = self.build_catexps(butlerQC, inputRefs, inputs)
463 if not self.config.allow_missing_bands and any([catexp is None for catexp in catexps]):
464 raise RuntimeError(
465 f"Got a None catexp with {self.config.allow_missing_band=}; NoWorkFound should have been"
466 f" raised earlier"
467 )
468 kwargs = self.make_kwargs(butlerQC, inputRefs, inputs)
469 outputs = self.run(catexps=catexps, cat_ref=inputs['cat_ref'], **kwargs)
470 butlerQC.put(outputs, outputRefs)
472 def run(
473 self,
474 catexps: list[CatalogExposure],
475 cat_ref: afwTable.SourceCatalog,
476 **kwargs
477 ) -> pipeBase.Struct:
478 """Fit sources from a reference catalog using data from multiple
479 exposures in the same region (patch).
481 Parameters
482 ----------
483 catexps : `typing.List [CatalogExposure]`
484 A list of catalog-exposure pairs in a given band.
485 cat_ref : `lsst.afw.table.SourceCatalog`
486 A reference source catalog to fit.
488 Returns
489 -------
490 retStruct : `lsst.pipe.base.Struct`
491 A struct with a cat_output attribute containing the output
492 measurement catalog.
494 Notes
495 -----
496 Subtasks may have further requirements; see `CoaddMultibandFitSubTask.run`.
497 """
498 cat_output = self.fit_coadd_multiband.run(catalog_multi=cat_ref, catexps=catexps, **kwargs).output
499 retStruct = pipeBase.Struct(cat_output=cat_output)
500 return retStruct