Coverage for python/lsst/cell_coadds/_cell_coadd_builder.py: 36%

95 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-24 10:24 +0000

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/>. 

21 

22from __future__ import annotations 

23 

24from abc import ABCMeta, abstractmethod 

25from collections.abc import Iterable, Mapping 

26from typing import Any, ClassVar 

27 

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 

37 

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 

43 

44__all__ = ( 

45 "MultipleCellCoaddBuilderConfig", 

46 "MultipleCellCoaddBuilderConnections", 

47 "MultipleCellCoaddBuilderTask", 

48 "SingleCellCoaddBuilderConfig", 

49 "SingleCellCoaddBuilderTask", 

50 "singleCellCoaddBuilderRegistry", 

51) 

52 

53 

54class SingleCellCoaddBuilderConfig(pexConfig.Config): 

55 """Configuration for a single-cell coadd builder.""" 

56 

57 psf_dimensions = pexConfig.Field[int]( 

58 doc="Dimensions of the PSF image", 

59 default=41, 

60 ) 

61 

62 

63class SingleCellCoaddBuilderTask(pipeBase.Task, metaclass=ABCMeta): 

64 """An abstract interface for tasks building coadds in cells. 

65 

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. 

74 

75 See Also 

76 -------- 

77 MultipleCellCoaddBuilderTask 

78 """ 

79 

80 ConfigClass: ClassVar[type[pexConfig.Config]] = SingleCellCoaddBuilderConfig 

81 _DefaultName: ClassVar[str] = "singleCellCoaddBuilder" 

82 

83 config: SingleCellCoaddBuilderConfig 

84 

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. 

93 

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. 

98 

99 Parameters 

100 ---------- 

101 inputs : `~collections.abc.Mapping` [`ObservationIdentifiers`, `tuple`\ 

102 [`~lsst.daf.butler.DeferredDatasetHandle`, `~lsst.geom.Box2I`]] 

103 

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. 

115 

116 Returns 

117 ------- 

118 single_cell_coadd : `~lsst.cell_coadds.SingleCellCoadd` 

119 A single cell coadd. 

120 """ 

121 raise NotImplementedError() 

122 

123 registry = pexConfig.makeRegistry(doc="Internal registry") 

124 

125 

126singleCellCoaddBuilderRegistry = pexConfig.makeRegistry(doc="Registry of single cell coadd builders") 

127 

128 

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 ) 

145 

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 ) 

152 

153 cell_coadd = cT.Output( 

154 doc="Coadded image", 

155 name="{outputCoaddName}CellCoaddPickled", 

156 storageClass="MultipleCellCoadd", 

157 dimensions=("tract", "patch", "skymap", "band", "instrument"), 

158 ) 

159 

160 

161class MultipleCellCoaddBuilderConfig( 

162 pipeBase.PipelineTaskConfig, pipelineConnections=MultipleCellCoaddBuilderConnections 

163): 

164 """Configuration parameters for the `MultipleCellCoaddBuilderTask`.""" 

165 

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 ) 

171 

172 singleCellCoaddBuilder = singleCellCoaddBuilderRegistry.makeField( 

173 doc="Coaddition algorithm to use. See `SingleCellCoaddBuilderTask` for details", 

174 optional=False, 

175 ) 

176 

177 

178class MultipleCellCoaddBuilderTask(pipeBase.PipelineTask): 

179 """Task to build cell-based coadded images. 

180 

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``. 

190 

191 See Also 

192 -------- 

193 SingleCellCoaddBuilderTask 

194 """ 

195 

196 ConfigClass: ClassVar[type[pipeBase.PipelineTaskConfig]] = MultipleCellCoaddBuilderConfig 

197 _DefaultName: ClassVar[str] = "multipleCellCoaddBuilder" 

198 

199 def __init__(self, **kwargs: Any): 

200 super().__init__(**kwargs) 

201 self.makeSubtask(name="singleCellCoaddBuilder") 

202 self.singleCellCoaddBuilder: SingleCellCoaddBuilderTask 

203 

204 def runQuantum( # type: ignore[override] 

205 self, 

206 butlerQC: pipeBase.ButlerQuantumContext, 

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") 

218 

219 quantumDataId = butlerQC.quantum.dataId 

220 if quantumDataId is None: 

221 raise ValueError("quantumDataId is None") 

222 

223 skyInfo = makeSkyInfo( 

224 skymap, 

225 tractId=quantumDataId["tract"], 

226 patchId=quantumDataId["patch"], 

227 ) 

228 

229 # Run the (warp and) coaddition code 

230 multipleCellCoadd = self.run( 

231 inputs[self.config.input_type], 

232 skyInfo=skyInfo, 

233 quantumDataId=quantumDataId, 

234 ) 

235 

236 # Persist the results via the butler 

237 butlerQC.put(multipleCellCoadd, outputRefs.cell_coadd) 

238 return multipleCellCoadd # The return statement may be eventually removed. 

239 

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. 

247 

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. 

257 

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=quantumDataId.get("band", None), 

269 identifiers=PatchIdentifiers.from_data_id(quantumDataId), 

270 ) 

271 

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. 

279 

280 if not bbox_list: 

281 continue # We should raise exception here, since the task would fail anyway eventually. 

282 

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 

289 

290 cellCoadd = self.singleCellCoaddBuilder.run(scc_inputs, cellInfo, common) 

291 cellCoadds.append(cellCoadd) 

292 

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()) 

297 

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 

307 

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. 

315 

316 This methods selects from a list of exposures/warps those images that 

317 completely overlap with the cell, thus enabling edgeless coadds. 

318 

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. 

327 

328 

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()] 

338 

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]) 

348 

349 if skyInfo and skyInfo.tractInfo.outer_sky_polygon.contains(skyCalexp): 

350 pass 

351 

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) 

362 

363 return overlapping_bbox