Coverage for python/lsst/pipe/tasks/fit_coadd_psf.py: 41%
82 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-09-30 11:39 +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 "CatalogExposurePsf", "CoaddPsfFitConfig", "CoaddPsfFitConnections",
24 "CoaddPsfFitSubConfig", "CoaddPsfFitSubTask", "CoaddPsfFitTask",
25]
27from .fit_multiband import CatalogExposure, CatalogExposureConfig
28from lsst.geom import Point2D
29from lsst.meas.base import SkyMapIdGeneratorConfig
30import lsst.pex.config as pexConfig
31import lsst.pipe.base as pipeBase
32import lsst.pipe.base.connectionTypes as cT
34import dataclasses
35from abc import ABC, abstractmethod
36from pydantic.dataclasses import dataclass
39@dataclass(frozen=True, kw_only=True, config=CatalogExposureConfig)
40class CatalogExposurePsf(CatalogExposure):
41 def get_catalog(self):
42 return self.catalog
44 def get_psf_image(self, source):
45 """Return the PSF image for this object."""
46 center = Point2D(round(source.getX()), round(source.getY()))
47 return self.exposure.getPsf().computeKernelImage(center).array
50CoaddPsfFitBaseTemplates = {
51 "name_coadd": "deep",
52 "name_output_method": "multiprofit",
53}
56class CoaddPsfFitConnections(
57 pipeBase.PipelineTaskConnections,
58 dimensions=("tract", "patch", "band", "skymap"),
59 defaultTemplates=CoaddPsfFitBaseTemplates,
60):
61 coadd = cT.Input(
62 doc="Coadd image to fit a PSF model to",
63 name="{name_coadd}Coadd_calexp",
64 storageClass="ExposureF",
65 dimensions=("tract", "patch", "band", "skymap"),
66 )
67 coadd_cell = cT.Input(
68 doc="Cell-coadd image to fit a PSF model to",
69 name="{name_coadd}CoaddCell",
70 storageClass="MultipleCellCoadd",
71 dimensions=("tract", "patch", "band", "skymap"),
72 )
73 background = cT.Input(
74 doc="Background model to subtract from the coadd_cell",
75 name="{name_coadd}Coadd_calexp_background",
76 storageClass="Background",
77 dimensions=("tract", "patch", "band", "skymap"),
78 )
79 cat_meas = cT.Input(
80 doc="Deblended single-band source catalog",
81 name="{name_coadd}Coadd_meas",
82 storageClass="SourceCatalog",
83 dimensions=("tract", "patch", "band", "skymap"),
84 )
85 cat_output = cT.Output(
86 doc="Output PSF fit parameter catalog",
87 name="{name_coadd}Coadd_psfs_{name_output_method}",
88 storageClass="ArrowTable",
89 dimensions=("tract", "patch", "band", "skymap"),
90 )
92 def __init__(self, *, config=None):
93 super().__init__(config=config)
94 if config is None:
95 return
97 if config.image_type == "future":
98 self.coadd = dataclasses.replace(self.coadd, storageClass="CellCoadd")
99 del self.coadd_cell
100 del self.background
101 elif config.use_cell_coadds:
102 del self.coadd
103 else:
104 del self.coadd_cell
105 del self.background
108class CoaddPsfFitSubConfig(pexConfig.Config):
109 """Base config class for the CoaddPsfFitTask.
111 Implementing classes may add any necessary attributes.
112 """
115class CoaddPsfFitSubTask(pipeBase.Task, ABC):
116 """Interface for CoaddPsfFitTask subtasks to fit PSFs.
118 Parameters
119 ----------
120 **kwargs
121 Additional arguments to be passed to the `lsst.pipe.base.Task`
122 constructor.
123 """
124 ConfigClass = CoaddPsfFitSubConfig
126 def __init__(self, **kwargs):
127 super().__init__(**kwargs)
129 @abstractmethod
130 def run(
131 self, catexp: CatalogExposurePsf
132 ) -> pipeBase.Struct:
133 """Fit PSF images at locations of sources in a single exposure.
135 Parameters
136 ----------
137 catexp : `CatalogExposurePsf`
138 An exposure to fit a model PSF at the position of all
139 sources in the corresponding catalog.
141 Returns
142 -------
143 retStruct : `lsst.pipe.base.Struct`
144 A struct with a cat_output attribute containing the output
145 measurement catalog.
147 Notes
148 -----
149 Subclasses may have further requirements on the input parameters,
150 including:
151 - Passing only one catexp per band;
152 - Catalogs containing HeavyFootprints with deblended images;
153 - Fitting only a subset of the sources.
154 If any requirements are not met, the subtask should fail as soon as
155 possible.
156 """
157 raise NotImplementedError()
160class CoaddPsfFitConfig(
161 pipeBase.PipelineTaskConfig,
162 pipelineConnections=CoaddPsfFitConnections,
163):
164 """Configure a CoaddPsfFitTask, including a configurable fitting subtask.
165 """
166 use_cell_coadds = pexConfig.Field(
167 dtype=bool,
168 default=False,
169 doc="Use cell coadd images for PSF fitting",
170 )
171 fit_coadd_psf = pexConfig.ConfigurableField(
172 target=CoaddPsfFitSubTask,
173 doc="Task to fit PSF models for a single coadd",
174 )
175 idGenerator = SkyMapIdGeneratorConfig.make_field()
176 image_type = pexConfig.ChoiceField(
177 "Which image type to expect for the input coadd. "
178 "This option only directly affects connection storage classes and hence 'runQuantum'; the 'run' "
179 "method behavior is determined by which type is actually passed in.",
180 allowed={
181 "legacy": (
182 "Read a lsst.cell_coadds.MultipleCellCoadd via 'coadd_cell` and restore 'background' "
183 "(if use_cell_coadd) or lsst.afw.image.Exposure via `coadd` (if not use_cell_coadd)."
184 ),
185 "future": (
186 "Read lsst.images.cells.CellCoadd via the 'coadd' connection. use_cell_coadd is ignored."
187 ),
188 },
189 dtype=str,
190 optional=False,
191 default="legacy",
192 )
195class CoaddPsfFitTask(pipeBase.PipelineTask):
196 """Fit a PSF model at the location of sources in a coadd.
198 This task is intended to fit only a single PSF model at the
199 centroid of all of the sources in a single coadd exposure.
200 Subtasks may choose to filter which sources they fit,
201 and may output whatever columns they desire in addition to
202 the minimum of 'id'.
203 """
204 ConfigClass = CoaddPsfFitConfig
205 _DefaultName = "CoaddPsfFit"
207 def __init__(self, initInputs, **kwargs):
208 super().__init__(initInputs=initInputs, **kwargs)
209 self.makeSubtask("fit_coadd_psf")
211 def runQuantum(self, butlerQC, inputRefs, outputRefs):
212 inputs = butlerQC.get(inputRefs)
213 id_tp = self.config.idGenerator.apply(butlerQC.quantum.dataId).catalog_id
214 dataId = inputRefs.cat_meas.dataId
216 if self.config.image_type == "future":
217 coaddDataRef = inputRefs.coadd
218 exposure = inputs.pop('coadd').to_legacy()
219 elif self.config.use_cell_coadds:
220 coaddDataRef = inputRefs.coadd_cell
221 multiple_cell_coadd = inputs.pop('coadd_cell')
222 background = inputs.pop('background')
223 exposure = multiple_cell_coadd.stitch().asExposure()
224 exposure.image -= background.getImage()
225 else:
226 coaddDataRef = inputRefs.coadd
227 exposure = inputs.pop('coadd')
229 for dataRef in (coaddDataRef,):
230 if dataRef.dataId != dataId:
231 raise RuntimeError(f'{dataRef=}.dataId != {inputRefs.cat_meas.dataId=}')
233 catalog = inputs.pop('cat_meas')
234 catexp = CatalogExposurePsf(
235 catalog=catalog, exposure=exposure, dataId=dataId, id_tract_patch=id_tp,
236 )
237 assert not inputs, "runQuantum got more inputs than expected"
238 outputs = self.run(catexp=catexp)
239 butlerQC.put(outputs, outputRefs)
241 def run(self, catexp: CatalogExposurePsf) -> pipeBase.Struct:
242 """Fit a PSF model at the location of sources in a coadd.
244 Parameters
245 ----------
246 catexp : `typing.List [CatalogExposurePsf]`
247 A list of catalog-exposure pairs in a given band.
249 Returns
250 -------
251 retStruct : `lsst.pipe.base.Struct`
252 A struct with a cat_output attribute containing the output
253 measurement catalog.
255 Notes
256 -----
257 Subtasks may have further requirements; see `CoaddPsfFitSubTask.run`.
258 """
259 cat_output = self.fit_coadd_psf.run(catexp).output
260 retStruct = pipeBase.Struct(cat_output=cat_output)
261 return retStruct