Coverage for python/lsst/cell_coadds/_cell_coadd_builder.py: 36%
95 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-30 03:56 -0700
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-30 03:56 -0700
1# This file is part of cell_coadds.
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/>.
22from __future__ import annotations
24from abc import ABCMeta, abstractmethod
25from collections.abc import Iterable, Mapping
26from typing import Any, ClassVar, cast
28import lsst.geom
29import lsst.pex.config as pexConfig
30import lsst.pipe.base as pipeBase
31import lsst.pipe.base.connectionTypes as cT
32import lsst.sphgeom
33import lsst.utils
34from lsst.daf.butler import DataCoordinate, DeferredDatasetHandle
35from lsst.pipe.tasks.coaddBase import makeSkyInfo
36from lsst.skymap import CellInfo, PatchInfo
38from ._common_components import CoaddUnits, CommonComponents
39from ._identifiers import ObservationIdentifiers, PatchIdentifiers
40from ._multiple_cell_coadd import MultipleCellCoadd
41from ._single_cell_coadd import SingleCellCoadd
42from ._uniform_grid import UniformGrid
44__all__ = (
45 "MultipleCellCoaddBuilderConfig",
46 "MultipleCellCoaddBuilderConnections",
47 "MultipleCellCoaddBuilderTask",
48 "SingleCellCoaddBuilderConfig",
49 "SingleCellCoaddBuilderTask",
50 "singleCellCoaddBuilderRegistry",
51)
54class SingleCellCoaddBuilderConfig(pexConfig.Config):
55 """Configuration for a single-cell coadd builder."""
57 psf_dimensions = pexConfig.Field[int](
58 doc="Dimensions of the PSF image",
59 default=41,
60 )
63class SingleCellCoaddBuilderTask(pipeBase.Task, metaclass=ABCMeta):
64 """An abstract interface for tasks building coadds in cells.
66 `SingleCellCoaddBuilderTask` is intended to serve as an abstract interface
67 for various concrete implementation of coaddition algorithms via its `run`
68 method. `MultipleCellCoaddBuilderTask` is the corresponding pipeline task
69 that needs to be called from the pipeline. `MultipleCellCoaddBuilderTask`
70 must contain a concrete implementation inherited from
71 `SingleCellCoaddBuilderTask` as its base class and registered with the
72 `singleCellCoaddBuilderTaskRegistry`, say using ``@registerConfigurable``
73 decorator.
75 See Also
76 --------
77 MultipleCellCoaddBuilderTask
78 """
80 ConfigClass: ClassVar[type[pexConfig.Config]] = SingleCellCoaddBuilderConfig
81 _DefaultName: ClassVar[str] = "singleCellCoaddBuilder"
83 config: SingleCellCoaddBuilderConfig
85 @abstractmethod
86 def run(
87 self,
88 inputs: Mapping[ObservationIdentifiers, tuple[DeferredDatasetHandle, lsst.geom.Box2I]],
89 cellInfo: CellInfo,
90 common: CommonComponents,
91 ) -> SingleCellCoadd:
92 """Build a single-cell coadd.
94 The images passed in from `MultipleCellCoaddBuilderTask` are guaranteed
95 to completely overlap the outer bounding box of the cells. Any further
96 exclusion of images based on quality assessment or other criteria
97 should be dome in this method.
99 Parameters
100 ----------
101 inputs : `~collections.abc.Mapping` [`ObservationIdentifiers`, `tuple`\
102 [`~lsst.daf.butler.DeferredDatasetHandle`, `~lsst.geom.Box2I`]]
104 A mapping from `ObservationIdentifiers` to a tuple containing a
105 `DeferredDatasetHandle` pointing to the input image
106 (calexp or warps) and a minimal bounding box that can be read
107 without loading the entire image.
108 cellInfo : `~lsst.skymap.CellInfo`
109 An object with the following attributes:
110 - wcs: `lsst.afw.geom.SkyWcs`
111 - outer_bbox: `lsst.geom.Box2I`
112 common : `~lsst.cell_coadds.CommonComponents`
113 A dataclass object with properties that are common to the entire
114 `MultipleCellCoadd` object that the cell produced is a part of.
116 Returns
117 -------
118 single_cell_coadd : `~lsst.cell_coadds.SingleCellCoadd`
119 A single cell coadd.
120 """
121 raise NotImplementedError()
123 registry = pexConfig.makeRegistry(doc="Internal registry")
126singleCellCoaddBuilderRegistry = pexConfig.makeRegistry(doc="Registry of single cell coadd builders")
129class MultipleCellCoaddBuilderConnections(
130 pipeBase.PipelineTaskConnections,
131 dimensions=("tract", "patch", "band", "skymap"),
132 defaultTemplates={"inputCoaddName": "deep", "outputCoaddName": "deep", "warpType": "direct"},
133):
134 # Since we defer loading of images, we could take in both calexp and
135 # warps as inputs. Unless, we don't want a dependency on warps existing.
136 # The type of image will be specified in MultipleCellCoaddConfig.
137 calexps = cT.Input(
138 doc="Input exposures to be resampled and optionally PSF-matched onto a SkyMap projection/patch",
139 name="calexp",
140 storageClass="ExposureF",
141 dimensions=("instrument", "visit", "detector"),
142 deferLoad=True,
143 multiple=True,
144 )
146 skymap = cT.Input(
147 doc="Input definition of geometry/box and projection/wcs for coadded exposures",
148 name="skyMap",
149 storageClass="SkyMap",
150 dimensions=("skymap",),
151 )
153 cell_coadd = cT.Output(
154 doc="Coadded image",
155 name="{outputCoaddName}CellCoaddPickled",
156 storageClass="MultipleCellCoadd",
157 dimensions=("tract", "patch", "skymap", "band", "instrument"),
158 )
161class MultipleCellCoaddBuilderConfig(
162 pipeBase.PipelineTaskConfig, pipelineConnections=MultipleCellCoaddBuilderConnections
163):
164 """Configuration parameters for the `MultipleCellCoaddBuilderTask`."""
166 input_type = pexConfig.ChoiceField[str](
167 doc="Type of input dataset",
168 allowed={"input_warps": "Use warps", "calexps": "Use calexps"},
169 default="calexps",
170 )
172 singleCellCoaddBuilder = singleCellCoaddBuilderRegistry.makeField(
173 doc="Coaddition algorithm to use. See `SingleCellCoaddBuilderTask` for details",
174 optional=False,
175 )
178class MultipleCellCoaddBuilderTask(pipeBase.PipelineTask):
179 """Task to build cell-based coadded images.
181 This is the pipeline task that needs to be called from the pipeline. It
182 contains ``singleCellCoaddBuilder`` as a subtask which must implement a
183 concrete coaddition algorithm in its ``run`` method. The images to be
184 coadded can either be of type ``calexp`` or ``warp``, as specified in the
185 ``inputType`` configuration field. A ``skymap`` with ``cells`` tract
186 builder must be passed. For each cell to be coadded, the task will query
187 the butler for all the input images that completely overlap the cell's
188 outer bounding box and passed them to the ``run`` method of the
189 ``singleCellCoaddBuilder``.
191 See Also
192 --------
193 SingleCellCoaddBuilderTask
194 """
196 ConfigClass: ClassVar[type[pipeBase.PipelineTaskConfig]] = MultipleCellCoaddBuilderConfig
197 _DefaultName: ClassVar[str] = "multipleCellCoaddBuilder"
199 def __init__(self, **kwargs: Any):
200 super().__init__(**kwargs)
201 self.makeSubtask(name="singleCellCoaddBuilder")
202 self.singleCellCoaddBuilder: SingleCellCoaddBuilderTask
204 def runQuantum( # type: ignore[override]
205 self,
206 butlerQC: pipeBase.QuantumContext,
207 inputRefs: pipeBase.InputQuantizedConnection,
208 outputRefs: pipeBase.OutputQuantizedConnection,
209 ) -> MultipleCellCoadd:
210 # Docstring inherited.
211 self.config: MultipleCellCoaddBuilderConfig
212 inputs: dict = butlerQC.get(inputRefs)
213 skymap = inputs.pop("skymap") # skyInfo below will contain this skymap
214 # Ideally, we should do this check early on.
215 # But skymap is not available during config time.
216 if not skymap.config.tractBuilder.name == "cells":
217 raise pipeBase.InvalidQuantumError("skymap is not a cells skymap")
219 quantumDataId = butlerQC.quantum.dataId
220 if quantumDataId is None:
221 raise ValueError("quantumDataId is None")
223 skyInfo = makeSkyInfo(
224 skymap,
225 tractId=quantumDataId["tract"],
226 patchId=quantumDataId["patch"],
227 )
229 # Run the (warp and) coaddition code
230 multipleCellCoadd = self.run(
231 inputs[self.config.input_type],
232 skyInfo=skyInfo,
233 quantumDataId=quantumDataId,
234 )
236 # Persist the results via the butler
237 butlerQC.put(multipleCellCoadd, outputRefs.cell_coadd)
238 return multipleCellCoadd # The return statement may be eventually removed.
240 def run( # type: ignore[override] # TODO: Remove after DM-34696 is fixed.
241 self,
242 expList: Iterable[DeferredDatasetHandle],
243 skyInfo: pipeBase.Struct,
244 quantumDataId: DataCoordinate,
245 ) -> MultipleCellCoadd:
246 """Run coaddition algorithm for all the cells.
248 Parameters
249 ----------
250 expList : `list` [`lsst.daf.butler.DeferredDatasetHandle`]
251 An iterable of `~lsst.daf.butler.DeferredDatasetHandle` objects,
252 where the objects can be either calexp or warp images.
253 skyInfo : `~lsst.pipe.base.Struct`
254 Struct with geometric information about the patches and cells.
255 quantumDataId : `~lsst.daf.butler.DataCoordinate`
256 An immutable dataID dictionary that uniquely refers to the dataset.
258 Returns
259 -------
260 multipleCellCoadd : `MultipleCellCoadd`
261 Cell-based coadded image.
262 """
263 cellCoadds: list[SingleCellCoadd] = []
264 patchInfo: PatchInfo = skyInfo.patchInfo
265 common = CommonComponents(
266 units=CoaddUnits.nJy,
267 wcs=patchInfo.wcs,
268 band=cast(str, quantumDataId.get("band", "")),
269 identifiers=PatchIdentifiers.from_data_id(quantumDataId),
270 )
272 # This for loop could/should be parallelized, may be with asyncio?
273 for cellInfo in patchInfo:
274 # Select calexps that completely overlap with the cell
275 try:
276 bbox_list = self._select_overlaps(expList, cellInfo=cellInfo, skyInfo=skyInfo)
277 except KeyError:
278 continue # Exception handling is likely a relic, should be removed.
280 if not bbox_list:
281 continue # We should raise exception here, since the task would fail anyway eventually.
283 scc_inputs = {
284 ObservationIdentifiers.from_data_id(handle.ref.dataId): (handle, bbox)
285 for handle, bbox in zip(expList, bbox_list, strict=True)
286 }
287 if len(scc_inputs) == 0:
288 continue
290 cellCoadd = self.singleCellCoaddBuilder.run(scc_inputs, cellInfo, common)
291 cellCoadds.append(cellCoadd)
293 # grid has no notion about border or inner/outer boundaries.
294 # So we have to clip the outermost border when constructing the grid.
295 grid_bbox = patchInfo.outer_bbox.erodedBy(patchInfo.getCellBorder())
296 grid = UniformGrid.from_bbox_cell_size(grid_bbox, patchInfo.getCellInnerDimensions())
298 multipleCellCoadd = MultipleCellCoadd(
299 cellCoadds,
300 grid=grid,
301 outer_cell_size=cellInfo.outer_bbox.getDimensions(),
302 inner_bbox=None,
303 common=common,
304 psf_image_size=cellCoadds[0].psf_image.getDimensions(),
305 )
306 return multipleCellCoadd
308 @staticmethod
309 def _select_overlaps(
310 explist: Iterable[DeferredDatasetHandle],
311 cellInfo: CellInfo,
312 skyInfo: pipeBase.Struct | None,
313 ) -> Iterable[lsst.geom.Box2I]:
314 """Filter exposures for cell-based coadds.
316 This methods selects from a list of exposures/warps those images that
317 completely overlap with the cell, thus enabling edgeless coadds.
319 Parameters
320 ----------
321 explist : `list` [`~lsst.daf.butler.DeferredDatasetHandle`]
322 List of handles for exposures to be coadded
323 cellInfo : `~lsst.pipe.base.Struct` or `collections.namedtuple`
324 The cellInfo dict, must have .wcs and .outerBBox.
325 skyInfo : `~lsst.pipe.base.Struct`
326 Struct with geometric information about the patches and cells.
329 Returns
330 -------
331 overlapping_bbox : `list` [`lsst.geom.Box2I`]
332 List of bounding boxes for each image in ``explist`` that overlaps
333 the given cell.
334 """
335 cell_bbox = cellInfo.outer_bbox
336 cell_wcs = cellInfo.wcs
337 cell_corners = [cell_wcs.pixelToSky(corner.x, corner.y) for corner in cell_bbox.getCorners()]
339 overlapping_bbox = []
340 for exp in explist:
341 calexp_bbox = exp.get(component="bbox")
342 calexp_wcs = exp.get(component="wcs")
343 calexp_corners = [
344 calexp_wcs.pixelToSky(corner.x, corner.y) for corner in calexp_bbox.getCorners()
345 ]
346 skyCell = lsst.sphgeom.ConvexPolygon([corner.getVector() for corner in cell_corners])
347 skyCalexp = lsst.sphgeom.ConvexPolygon([corner.getVector() for corner in calexp_corners])
349 if skyInfo and skyInfo.tractInfo.outer_sky_polygon.contains(skyCalexp):
350 pass
352 if skyCell.isWithin(skyCalexp):
353 tiny_bbox_min_corner = calexp_wcs.skyToPixel(
354 cell_wcs.pixelToSky(cell_bbox.minX, cell_bbox.minY)
355 )
356 tiny_bbox_max_corner = calexp_wcs.skyToPixel(
357 cell_wcs.pixelToSky(cell_bbox.maxX, cell_bbox.maxY)
358 )
359 tiny_bbox = lsst.geom.Box2D(minimum=tiny_bbox_min_corner, maximum=tiny_bbox_max_corner)
360 tiny_bbox = lsst.geom.Box2I(tiny_bbox)
361 overlapping_bbox.append(tiny_bbox)
363 return overlapping_bbox