Coverage for python/lsst/drp/tasks/assemble_cell_coadd.py: 20%

1# This file is part of drp_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/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ( 

25 "AssembleCellCoaddTask", 

26 "AssembleCellCoaddConfig", 

27 "ConvertMultipleCellCoaddToExposureTask", 

28 "EmptyCellCoaddError", 

29) 

30 

31import dataclasses 

32 

33import numpy as np 

34 

35import lsst.afw.image as afwImage 

36import lsst.afw.math as afwMath 

37import lsst.geom as geom 

38from lsst.cell_coadds import ( 

39 CellIdentifiers, 

40 CoaddApCorrMapStacker, 

41 CoaddUnits, 

42 CommonComponents, 

43 GridContainer, 

44 MultipleCellCoadd, 

45 ObservationIdentifiers, 

46 OwnedImagePlanes, 

47 PatchIdentifiers, 

48 SingleCellCoadd, 

49 UniformGrid, 

50) 

51from lsst.daf.butler import DataCoordinate, DeferredDatasetHandle 

52from lsst.meas.algorithms import AccumulatorMeanStack 

53from lsst.pex.config import ConfigField, ConfigurableField, DictField, Field, ListField, RangeField 

54from lsst.pipe.base import ( 

55 AlgorithmError, 

56 InMemoryDatasetHandle, 

57 NoWorkFound, 

58 PipelineTask, 

59 PipelineTaskConfig, 

60 PipelineTaskConnections, 

61 Struct, 

62) 

63from lsst.pipe.base.connectionTypes import Input, Output 

64from lsst.pipe.tasks.coaddBase import makeSkyInfo, removeMaskPlanes, setRejectedMaskMapping 

65from lsst.pipe.tasks.interpImage import InterpImageTask 

66from lsst.pipe.tasks.scaleZeroPoint import ScaleZeroPointTask 

67from lsst.skymap import BaseSkyMap 

68 

69 

70class EmptyCellCoaddError(AlgorithmError): 

71 """Raised if no cells could be populated.""" 

72 

73 def __init__(self): 

74 msg = "No cells could be populated for the cell coadd." 

75 super().__init__(msg) 

76 

77 @property 

78 def metadata(self) -> dict: 

79 """There is no metadata associated with this error.""" 

80 return {} 

81 

82 

83@dataclasses.dataclass 

84class WarpInputs: 

85 """Collection of associate inputs along with warps.""" 

86 

87 warp: DeferredDatasetHandle | InMemoryDatasetHandle 

88 """Handle for the warped exposure.""" 

89 

90 masked_fraction: DeferredDatasetHandle | InMemoryDatasetHandle | None = None 

91 """Handle for the masked fraction image.""" 

92 

93 artifact_mask: DeferredDatasetHandle | InMemoryDatasetHandle | None = None 

94 """Handle for the CompareWarp artifact mask.""" 

95 

96 noise_warps: list[DeferredDatasetHandle | InMemoryDatasetHandle] = dataclasses.field(default_factory=list) 

97 """List of handles for the noise warps""" 

98 

99 @property 

100 def dataId(self) -> DataCoordinate: 

101 """DataID corresponding to the warp. 

102 

103 Returns 

104 ------- 

105 data_id : `~lsst.daf.butler.DataCoordinate` 

106 DataID of the warp. 

107 """ 

108 return self.warp.dataId 

109 

110 

111class AssembleCellCoaddConnections( 

112 PipelineTaskConnections, 

113 dimensions=("tract", "patch", "band", "skymap"), 

114 defaultTemplates={"inputWarpName": "deep", "outputCoaddSuffix": "Cell"}, 

115): 

116 inputWarps = Input( 

117 doc="Input warps", 

118 name="{inputWarpName}Coadd_directWarp", 

119 storageClass="ExposureF", 

120 dimensions=("tract", "patch", "skymap", "visit", "instrument"), 

121 deferLoad=True, 

122 multiple=True, 

123 ) 

124 

125 maskedFractionWarps = Input( 

126 doc="Mask fraction warps", 

127 name="{inputWarpName}Coadd_directWarp_maskedFraction", 

128 storageClass="ImageF", 

129 dimensions=("tract", "patch", "skymap", "visit", "instrument"), 

130 deferLoad=True, 

131 multiple=True, 

132 ) 

133 

134 artifactMasks = Input( 

135 doc="Artifact masks to be applied to the input warps", 

136 name="compare_warp_artifact_mask", 

137 storageClass="Mask", 

138 dimensions=("tract", "patch", "skymap", "visit", "instrument"), 

139 deferLoad=True, 

140 multiple=True, 

141 ) 

142 

143 visitSummaryList = Input( 

144 doc="Input visit-summary catalogs with updated calibration objects. Mainly used for coadd weights.", 

145 name="finalVisitSummary", 

146 storageClass="ExposureCatalog", 

147 dimensions=("instrument", "visit"), 

148 deferLoad=True, 

149 multiple=True, 

150 ) 

151 

152 skyMap = Input( 

153 doc="Input definition of geometry/bbox and projection/wcs. This must be cell-based.", 

154 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, 

155 storageClass="SkyMap", 

156 dimensions=("skymap",), 

157 ) 

158 

159 multipleCellCoadd = Output( 

160 doc="Output multiple cell coadd", 

161 name="{inputWarpName}Coadd{outputCoaddSuffix}", 

162 storageClass="MultipleCellCoadd", 

163 dimensions=("tract", "patch", "band", "skymap"), 

164 ) 

165 

166 def __init__(self, *, config=None): 

167 super().__init__(config=config) 

168 

169 if not config: 

170 return 

171 

172 if config.do_calculate_weight_from_warp: 

173 del self.visitSummaryList 

174 

175 if not config.do_use_artifact_mask: 

176 del self.artifactMasks 

177 

178 # Dynamically set input connections for noise images, depending on the 

179 # number of noise realizations specified in the config. 

180 for n in range(config.num_noise_realizations): 

181 noise_warps = Input( 

182 doc="Input noise warps", 

183 name=f"{config.connections.inputWarpName}Coadd_directWarp_noise{n}", 

184 storageClass="MaskedImageF", 

185 dimensions=("tract", "patch", "skymap", "visit", "instrument"), 

186 deferLoad=True, 

187 multiple=True, 

188 ) 

189 setattr(self, f"noise{n}_warps", noise_warps) 

190 

191 

192class AssembleCellCoaddConfig(PipelineTaskConfig, pipelineConnections=AssembleCellCoaddConnections): 

193 do_interpolate_coadd = Field[bool](doc="Interpolate over pixels with NO_DATA mask set?", default=True) 

194 interpolate_coadd = ConfigurableField( 

195 target=InterpImageTask, 

196 doc="Task to interpolate (and extrapolate) over pixels with NO_DATA mask on cell coadds", 

197 ) 

198 do_scale_zero_point = Field[bool]( 

199 doc="Scale warps to a common zero point? This is not needed if they have absolute flux calibration.", 

200 default=False, 

201 deprecated="Now that visits are scaled to nJy it is no longer necessary or " 

202 "recommended to scale the zero point, so this will be removed " 

203 "after v29.", 

204 ) 

205 scale_zero_point = ConfigurableField( 

206 target=ScaleZeroPointTask, 

207 doc="Task to scale warps to a common zero point", 

208 deprecated="Now that visits are scaled to nJy it is no longer necessary or " 

209 "recommended to scale the zero point, so this will be removed " 

210 "after v29.", 

211 ) 

212 do_calculate_weight_from_warp = Field[bool]( 

213 doc="Calculate coadd weight from the input warp? Otherwise, the weight is obtained from the " 

214 "visitSummaryList connection. This is meant as a fallback when run outside the pipeline.", 

215 default=False, 

216 ) 

217 do_use_artifact_mask = Field[bool]( 

218 doc="Substitute the mask planes input warp with an alternative artifact mask?", 

219 default=True, 

220 ) 

221 do_coadd_inverse_aperture_corrections = Field[bool]( 

222 doc="Coadd the inverse aperture corrections for each cell? This is formally the more accurate way " 

223 "but may be turned off for parity with deepCoadd.", 

224 default=False, 

225 ) 

226 bad_mask_planes = ListField[str]( 

227 doc="Mask planes that count towards the masked fraction within a cell.", 

228 default=("BAD", "NO_DATA", "SAT", "CLIPPED"), 

229 ) 

230 remove_mask_planes = ListField[str]( 

231 doc="Mask planes to remove before coadding", 

232 default=["NOT_DEBLENDED", "EDGE"], 

233 ) 

234 calc_error_from_input_variance = Field[bool]( 

235 doc="Calculate coadd variance from input variance by stacking " 

236 "statistic. Passed to AccumulatorMeanStack.", 

237 default=True, 

238 ) 

239 mask_propagation_thresholds = DictField[str, float]( 

240 doc=( 

241 "Threshold (in fractional weight) of rejection at which we " 

242 "propagate a mask plane to the coadd; that is, we set the mask " 

243 "bit on the coadd if the fraction the rejected frames " 

244 "would have contributed exceeds this value." 

245 ), 

246 default={"SAT": 0.1}, 

247 ) 

248 max_maskfrac = RangeField[float]( 

249 doc="Maximum fraction of masked pixels in a cell. This is currently " 

250 "just a placeholder and is not used now", 

251 default=0.99, 

252 min=0.0, 

253 max=1.0, 

254 inclusiveMin=True, 

255 inclusiveMax=False, 

256 ) 

257 num_noise_realizations = Field[int]( 

258 default=0, 

259 doc=( 

260 "Number of noise planes to include in the coadd. " 

261 "This should not exceed the corresponding config parameter " 

262 "specified in `MakeDirectWarpConfig`. " 

263 ), 

264 check=lambda x: x >= 0, 

265 ) 

266 psf_warper = ConfigField( 

267 doc="Configuration for the warper that warps the PSFs. It must have the same configuration used to " 

268 "warp the images.", 

269 dtype=afwMath.Warper.ConfigClass, 

270 ) 

271 psf_dimensions = Field[int]( 

272 default=21, 

273 doc="Dimensions of the PSF image stamp size to be assigned to cells (must be odd).", 

274 check=lambda x: (x > 0) and (x % 2 == 1), 

275 ) 

276 

277 

278class AssembleCellCoaddTask(PipelineTask): 

279 """Assemble a cell-based coadded image from a set of warps. 

280 

281 This task reads in the warp one at a time, and accumulates it in all the 

282 cells that it completely overlaps with. This is the optimal I/O pattern but 

283 this also implies that it is not possible to build one or only a few cells. 

284 

285 Each cell coadds is guaranteed to have a well-defined PSF. This is done by 

286 1) excluding warps that only partially overlap a cell from that cell coadd; 

287 2) interpolating bad pixels in the warps rather than excluding them; 

288 3) by computing the coadd as a weighted mean of the warps without clipping; 

289 4) by computing the coadd PSF as the weighted mean of the PSF of the warps 

290 with the same weights. 

291 

292 The cells are (and must be) defined in the skymap, and cannot be configured 

293 or redefined here. The cells are assumed to be small enough that the PSF is 

294 assumed to be spatially constant within a cell. 

295 

296 Raises 

297 ------ 

298 NoWorkFound 

299 Raised if no input warps are provided, or no cells could be populated. 

300 RuntimeError 

301 Raised if the skymap is not cell-based. 

302 

303 Notes 

304 ----- 

305 This is not yet a part of the standard DRP pipeline. As such, the Task and 

306 especially its Config and Connections are experimental and subject to 

307 change any time without a formal RFC or standard deprecation procedures 

308 until it is included in the DRP pipeline. 

309 """ 

310 

311 ConfigClass = AssembleCellCoaddConfig 

312 _DefaultName = "assembleCellCoadd" 

313 

314 def __init__(self, *args, **kwargs): 

315 super().__init__(*args, **kwargs) 

316 if self.config.do_interpolate_coadd: 

317 self.makeSubtask("interpolate_coadd") 

318 if self.config.do_scale_zero_point: 

319 self.makeSubtask("scale_zero_point") 

320 

321 self.psf_warper = afwMath.Warper.fromConfig(self.config.psf_warper) 

322 

323 def runQuantum(self, butlerQC, inputRefs, outputRefs): 

324 # Docstring inherited. 

325 if not inputRefs.inputWarps: 

326 raise NoWorkFound("No input warps provided for co-addition") 

327 self.log.info("Found %d input warps", len(inputRefs.inputWarps)) 

328 

329 # Construct skyInfo expected by run 

330 # Do not remove skyMap from inputData in case _makeSupplementaryData 

331 # needs it 

332 skyMap = butlerQC.get(inputRefs.skyMap) 

333 

334 if not skyMap.config.tractBuilder.name == "cells": 

335 raise RuntimeError("AssembleCellCoaddTask requires a cell-based skymap.") 

336 

337 outputDataId = butlerQC.quantum.dataId 

338 

339 skyInfo = makeSkyInfo(skyMap, tractId=outputDataId["tract"], patchId=outputDataId["patch"]) 

340 visitSummaryList = butlerQC.get(getattr(inputRefs, "visitSummaryList", [])) 

341 

342 units = CoaddUnits.legacy if self.config.do_scale_zero_point else CoaddUnits.nJy 

343 self.common = CommonComponents( 

344 units=units, 

345 wcs=skyInfo.patchInfo.wcs, 

346 band=outputDataId.get("band", None), 

347 identifiers=PatchIdentifiers.from_data_id(outputDataId), 

348 ) 

349 

350 inputs: dict[DataCoordinate, WarpInputs] = {} 

351 for handle in butlerQC.get(inputRefs.inputWarps): 

352 inputs[handle.dataId] = WarpInputs(warp=handle, noise_warps=[]) 

353 

354 for ref in getattr(inputRefs, "artifactMasks", []): 

355 inputs[ref.dataId].artifact_mask = butlerQC.get(ref) 

356 for ref in getattr(inputRefs, "maskedFractionWarps", []): 

357 inputs[ref.dataId].masked_fraction = butlerQC.get(ref) 

358 for n in range(self.config.num_noise_realizations): 

359 for ref in getattr(inputRefs, f"noise{n}_warps"): 

360 inputs[ref.dataId].noise_warps.append(butlerQC.get(ref)) 

361 

362 returnStruct = self.run(inputs=inputs, skyInfo=skyInfo, visitSummaryList=visitSummaryList) 

363 butlerQC.put(returnStruct, outputRefs) 

364 return returnStruct 

365 

366 @staticmethod 

367 def _compute_weight(maskedImage, statsCtrl): 

368 """Compute a weight for a masked image. 

369 

370 Parameters 

371 ---------- 

372 maskedImage : `~lsst.afw.image.MaskedImage` 

373 The masked image to compute the weight. 

374 statsCtrl : `~lsst.afw.math.StatisticsControl` 

375 A control (config-like) object for StatisticsStack. 

376 

377 Returns 

378 ------- 

379 weight : `float` 

380 Inverse of the clipped mean variance of the masked image. 

381 """ 

382 statObj = afwMath.makeStatistics( 

383 maskedImage.getVariance(), maskedImage.getMask(), afwMath.MEANCLIP, statsCtrl 

384 ) 

385 meanVar, _ = statObj.getResult(afwMath.MEANCLIP) 

386 weight = 1.0 / float(meanVar) 

387 return weight 

388 

389 @staticmethod 

390 def _construct_grid(skyInfo): 

391 """Construct a UniformGrid object from a SkyInfo struct. 

392 

393 Parameters 

394 ---------- 

395 skyInfo : `~lsst.pipe.base.Struct` 

396 A Struct object 

397 

398 Returns 

399 ------- 

400 grid : `~lsst.cell_coadds.UniformGrid` 

401 A UniformGrid object. 

402 """ 

403 padding = skyInfo.patchInfo.getCellBorder() 

404 grid_bbox = skyInfo.patchInfo.outer_bbox.erodedBy(padding) 

405 grid = UniformGrid.from_bbox_cell_size( 

406 grid_bbox, 

407 skyInfo.patchInfo.getCellInnerDimensions(), 

408 padding=padding, 

409 ) 

410 return grid 

411 

412 def _construct_grid_container(self, skyInfo, statsCtrl): 

413 """Construct a grid of AccumulatorMeanStack instances. 

414 

415 Parameters 

416 ---------- 

417 skyInfo : `~lsst.pipe.base.Struct` 

418 A Struct object 

419 statsCtrl : `~lsst.afw.math.StatisticsControl` 

420 A control (config-like) object for StatisticsStack. 

421 

422 Returns 

423 ------- 

424 gc : `~lsst.cell_coadds.GridContainer` 

425 A GridContainer object container one AccumulatorMeanStack per cell. 

426 """ 

427 grid = self._construct_grid(skyInfo) 

428 

429 maskMap = setRejectedMaskMapping(statsCtrl) 

430 self.log.debug("Obtained maskMap = %s for %s", maskMap, skyInfo.patchInfo) 

431 thresholdDict = AccumulatorMeanStack.stats_ctrl_to_threshold_dict(statsCtrl) 

432 

433 # Initialize the grid container with AccumulatorMeanStacks 

434 gc = GridContainer[AccumulatorMeanStack](grid.shape) 

435 for cellInfo in skyInfo.patchInfo: 

436 stacker = AccumulatorMeanStack( 

437 # The shape is for the numpy arrays, hence transposed. 

438 shape=(cellInfo.outer_bbox.height, cellInfo.outer_bbox.width), 

439 bit_mask_value=statsCtrl.getAndMask(), 

440 mask_threshold_dict=thresholdDict, 

441 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

442 compute_n_image=False, 

443 mask_map=maskMap, 

444 no_good_pixels_mask=statsCtrl.getNoGoodPixelsMask(), 

445 ) 

446 gc[cellInfo.index] = stacker 

447 

448 return gc 

449 

450 def _construct_stats_control(self): 

451 """Construct a StatisticsControl object for coadd. 

452 

453 Unlike AssembleCoaddTask or CompareWarpAssembleCoaddTask, there is 

454 very little to be configured apart from setting the mask planes and 

455 optionally mask propagation thresholds. 

456 

457 Returns 

458 ------- 

459 statsCtrl : `~lsst.afw.math.StatisticsControl` 

460 A control object for StatisticsStack. 

461 """ 

462 statsCtrl = afwMath.StatisticsControl() 

463 # Hardcode the numIter parameter to the default config value set in 

464 # CompareWarpAssembleCoaddTask to get consistent weights. This is NOT 

465 # exposed as a config parameter, since this is only meant to be a 

466 # fallback option that is not recommended for production. 

467 statsCtrl.setNumIter(2) 

468 statsCtrl.setAndMask(afwImage.Mask.getPlaneBitMask(self.config.bad_mask_planes)) 

469 statsCtrl.setNanSafe(True) 

470 for plane, threshold in self.config.mask_propagation_thresholds.items(): 

471 bit = afwImage.Mask.getMaskPlane(plane) 

472 statsCtrl.setMaskPropagationThreshold(bit, threshold) 

473 return statsCtrl 

474 

475 def _construct_ap_corr_grid_container(self, skyInfo): 

476 """Construct a grid of CoaddApCorrMapStacker instances. 

477 

478 Parameters 

479 ---------- 

480 skyInfo : `~lsst.pipe.base.Struct` 

481 A Struct object 

482 

483 Returns 

484 ------- 

485 gc : `~lsst.cell_coadds.GridContainer` 

486 A GridContainer object container one CoaddApCorrMapStacker per 

487 cell. 

488 """ 

489 grid = self._construct_grid(skyInfo) 

490 

491 # Initialize the grid container with CoaddApCorrMapStacker. 

492 gc = GridContainer[CoaddApCorrMapStacker](grid.shape) 

493 for cellInfo in skyInfo.patchInfo: 

494 stacker = CoaddApCorrMapStacker( 

495 evaluation_point=cellInfo.inner_bbox.getCenter(), 

496 do_coadd_inverse_ap_corr=self.config.do_coadd_inverse_aperture_corrections, 

497 ) 

498 gc[cellInfo.index] = stacker 

499 

500 return gc 

501 

502 def run( 

503 self, 

504 *, 

505 inputs: dict[DataCoordinate, WarpInputs], 

506 skyInfo, 

507 visitSummaryList: list | None = None, 

508 ): 

509 for mask_plane in self.config.bad_mask_planes: 

510 afwImage.Mask.addMaskPlane(mask_plane) 

511 for mask_plane in self.config.mask_propagation_thresholds: 

512 afwImage.Mask.addMaskPlane(mask_plane) 

513 

514 statsCtrl = self._construct_stats_control() 

515 

516 warp_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

517 maskfrac_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

518 noise_stacker_gc_list = [ 

519 self._construct_grid_container(skyInfo, statsCtrl) 

520 for n in range(self.config.num_noise_realizations) 

521 ] 

522 psf_stacker_gc = GridContainer[AccumulatorMeanStack](warp_stacker_gc.shape) 

523 psf_bbox_gc = GridContainer[geom.Box2I](warp_stacker_gc.shape) 

524 ap_corr_stacker_gc = self._construct_ap_corr_grid_container(skyInfo) 

525 

526 # Make a container to hold the cell centers in sky coordinates now, 

527 # so we don't have to recompute them for each warp 

528 # (they share a common WCS). These are needed to find the various 

529 # warp + detector combinations that contributed to each cell, and later 

530 # get the corresponding PSFs as well. 

531 cell_centers_sky = GridContainer[geom.SpherePoint](warp_stacker_gc.shape) 

532 # Make a container to hold the observation identifiers for each cell. 

533 observation_identifiers_gc = GridContainer[list](warp_stacker_gc.shape) 

534 # Populate them. 

535 for cellInfo in skyInfo.patchInfo: 

536 # Make a list to hold the observation identifiers for each cell. 

537 observation_identifiers_gc[cellInfo.index] = [] 

538 cell_centers_sky[cellInfo.index] = skyInfo.wcs.pixelToSky(cellInfo.inner_bbox.getCenter()) 

539 psf_bbox_gc[cellInfo.index] = geom.Box2I.makeCenteredBox( 

540 geom.Point2D(cellInfo.inner_bbox.getCenter()), 

541 geom.Extent2I(self.config.psf_dimensions, self.config.psf_dimensions), 

542 ) 

543 psf_stacker_gc[cellInfo.index] = AccumulatorMeanStack( 

544 # The shape is for the numpy arrays, hence transposed. 

545 shape=(self.config.psf_dimensions, self.config.psf_dimensions), 

546 bit_mask_value=0, 

547 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

548 compute_n_image=False, 

549 ) 

550 

551 # visit_summary do not have (tract, patch, band, skymap) dimensions. 

552 if not visitSummaryList: 

553 visitSummaryList = [] 

554 visitSummaryRefDict = { 

555 visitSummaryRef.dataId["visit"]: visitSummaryRef for visitSummaryRef in visitSummaryList 

556 } 

557 

558 # Read in one warp at a time, and accumulate it in all the cells that 

559 # it completely overlaps. 

560 for _, warp_input in inputs.items(): 

561 warp = warp_input.warp.get(parameters={"bbox": skyInfo.bbox}) 

562 masked_fraction_image = ( 

563 warp_input.masked_fraction.get(parameters={"bbox": skyInfo.bbox}) 

564 if warp_input.masked_fraction 

565 else None 

566 ) 

567 

568 # Pre-process the warp before coadding. 

569 # TODO: Can we get these mask names from artifactMask? 

570 warp.mask.addMaskPlane("CLIPPED") 

571 warp.mask.addMaskPlane("REJECTED") 

572 warp.mask.addMaskPlane("SENSOR_EDGE") 

573 warp.mask.addMaskPlane("INEXACT_PSF") 

574 

575 if artifact_mask_ref := warp_input.artifact_mask: 

576 # Apply the artifact mask to the warp. 

577 artifact_mask = artifact_mask_ref.get() 

578 assert ( 

579 warp.mask.getMaskPlaneDict() == artifact_mask.getMaskPlaneDict() 

580 ), "Mask dicts do not agree." 

581 warp.mask.array = artifact_mask.array 

582 del artifact_mask 

583 

584 if self.config.do_scale_zero_point: 

585 # Each Warp that goes into a coadd will typically have an 

586 # independent photometric zero-point. Therefore, we must scale 

587 # each Warp to set it to a common photometric zeropoint. 

588 imageScaler = self.scale_zero_point.run(exposure=warp, dataRef=warp_input.warp).imageScaler 

589 zero_point_scale_factor = imageScaler.scale 

590 self.log.debug( 

591 "Scaled the warp %s by %f to match zero points", 

592 warp_input.dataId, 

593 zero_point_scale_factor, 

594 ) 

595 else: 

596 zero_point_scale_factor = 1.0 

597 if "BUNIT" not in warp.metadata: 

598 raise ValueError(f"Warp {warp_input.dataId} has no BUNIT metadata") 

599 if warp.metadata["BUNIT"] != "nJy": 

600 raise ValueError( 

601 f"Warp {warp_input.dataId} has BUNIT {warp.metadata['BUNIT']}, expected nJy" 

602 ) 

603 

604 # Coadd the warp onto the cells it completely overlaps. 

605 edge = afwImage.Mask.getPlaneBitMask(["NO_DATA", "SENSOR_EDGE"]) 

606 reject = afwImage.Mask.getPlaneBitMask(["CLIPPED", "REJECTED"]) 

607 removeMaskPlanes(warp.mask, self.config.remove_mask_planes, self.log) 

608 

609 # Compute the weight for each CCD in the warp from the visitSummary 

610 # or from the warp itself, if not provided. Computing the weight 

611 # from the warp is not recommended, and in that case we compute one 

612 # weight per warp and not bother with per-detector weights. 

613 full_ccd_table = warp.getInfo().getCoaddInputs().ccds 

614 weights: dict[int, float] = dict.fromkeys( 

615 full_ccd_table["ccd"], 

616 0.0, 

617 ) # Mapping from detector to weight. 

618 

619 if visitSummaryRef := visitSummaryRefDict.get(warp_input.dataId["visit"]): 

620 visitSummary = visitSummaryRef.get() 

621 for detector in full_ccd_table["ccd"]: 

622 visitSummaryRow = visitSummary.find(detector) 

623 mean_variance = visitSummaryRow["meanVar"] 

624 mean_variance *= zero_point_scale_factor**2 

625 if warp.metadata.get("BUNIT", None) == "nJy": 

626 mean_variance *= visitSummaryRow.photoCalib.getCalibrationMean() ** 2 

627 weights[detector] = 1.0 / mean_variance 

628 del visitSummary 

629 else: 

630 self.log.debug("No visit summary found for %s; using warp-based weights", warp_input.dataId) 

631 weight = self._compute_weight(warp, statsCtrl) 

632 if not np.isfinite(weight): 

633 self.log.warn("Non-finite weight for %s: skipping", warp_input.dataId) 

634 continue 

635 

636 for detector in weights: 

637 weights[detector] = weight 

638 

639 noise_warps = [ref.get(parameters={"bbox": skyInfo.bbox}) for ref in warp_input.noise_warps] 

640 

641 for cellInfo in skyInfo.patchInfo: 

642 bbox = cellInfo.outer_bbox 

643 inner_bbox = cellInfo.inner_bbox 

644 mi = warp[bbox].maskedImage 

645 

646 if (mi.mask[inner_bbox].array & edge).any(): 

647 self.log.debug( 

648 "Skipping %s in cell %s because it has a pixel with SENSOR_EDGE or NO_DATA bit set", 

649 warp_input.dataId, 

650 cellInfo.index, 

651 ) 

652 continue 

653 

654 if (mi.mask[inner_bbox].array & reject).any(): 

655 self.log.debug( 

656 "Skipping %s in cell %s because it has a pixel with CLIPPED or REJECTED bit set", 

657 warp_input.dataId, 

658 cellInfo.index, 

659 ) 

660 continue 

661 

662 # Find the CCD that contributed to this cell. 

663 if len(warp.getInfo().getCoaddInputs().ccds) == 1: 

664 # If there is only one, don't bother with a WCS look up. 

665 ccd_row = full_ccd_table[0] 

666 else: 

667 ccd_table = full_ccd_table.subsetContaining(cell_centers_sky[cellInfo.index]) 

668 

669 if len(ccd_table) == 0: 

670 # This condition rarely occurs in test runs, if ever. 

671 # But the QG generated upfront in campaign management 

672 # land routinely has extra quanta that should be 

673 # dropped during runtime. These cases arise when 

674 # the tasks upstream didn't process. 

675 # See DM-52306 for example. 

676 self.log.debug("No CCD found for %s in cell %s", warp_input.dataId, cellInfo.index) 

677 continue 

678 

679 assert len(ccd_table) == 1, "More than one CCD from a warp found within a cell." 

680 ccd_row = ccd_table[0] 

681 

682 weight = weights[ccd_row["ccd"]] 

683 if not np.isfinite(weight): 

684 self.log.warn( 

685 "Non-finite weight for %s in cell %s: skipping", warp_input.dataId, cellInfo.index 

686 ) 

687 continue 

688 

689 if weight == 0: 

690 self.log.info( 

691 "Zero weight for %s in cell %s: skipping", warp_input.dataId, cellInfo.index 

692 ) 

693 continue 

694 

695 observation_identifier = ObservationIdentifiers.from_data_id( 

696 warp_input.dataId, 

697 backup_detector=ccd_row["ccd"], 

698 ) 

699 observation_identifiers_gc[cellInfo.index].append(observation_identifier) 

700 

701 stacker = warp_stacker_gc[cellInfo.index] 

702 stacker.add_masked_image(mi, weight=weight) 

703 if masked_fraction_image: 

704 maskfrac_stacker_gc[cellInfo.index].add_image(masked_fraction_image[bbox], weight=weight) 

705 

706 for n in range(self.config.num_noise_realizations): 

707 mi = noise_warps[n][bbox] 

708 noise_stacker_gc_list[n][cellInfo.index].add_masked_image(mi, weight=weight) 

709 

710 calexp_point = ccd_row.getWcs().skyToPixel(cell_centers_sky[cellInfo.index]) 

711 undistorted_psf_im = ccd_row.getPsf().computeImage(calexp_point) 

712 

713 assert undistorted_psf_im.getBBox() == geom.Box2I.makeCenteredBox( 

714 calexp_point, 

715 undistorted_psf_im.getDimensions(), 

716 ), "PSF image does not share the coordinates of the 'calexp'" 

717 

718 # Convert the PSF image from Image to MaskedImage. 

719 undistorted_psf_maskedImage = afwImage.MaskedImageD(image=undistorted_psf_im) 

720 # TODO: In DM-43585, use the variance plane value from noise. 

721 undistorted_psf_maskedImage.variance += 1.0 # Set variance to 1 

722 

723 warped_psf_maskedImage = self.psf_warper.warpImage( 

724 destWcs=skyInfo.wcs, 

725 srcImage=undistorted_psf_maskedImage, 

726 srcWcs=ccd_row.getWcs(), 

727 destBBox=psf_bbox_gc[cellInfo.index], 

728 ) 

729 

730 # There may be NaNs in the PSF image. Set them to 0.0 

731 warped_psf_maskedImage.variance.array[np.isnan(warped_psf_maskedImage.image.array)] = 1.0 

732 warped_psf_maskedImage.image.array[np.isnan(warped_psf_maskedImage.image.array)] = 0.0 

733 

734 psf_stacker = psf_stacker_gc[cellInfo.index] 

735 psf_stacker.add_masked_image(warped_psf_maskedImage, weight=weight) 

736 

737 if (ap_corr_map := warp.getInfo().getApCorrMap()) is not None: 

738 ap_corr_stacker_gc[cellInfo.index].add(ap_corr_map, weight=weight) 

739 

740 del warp 

741 

742 cells: list[SingleCellCoadd] = [] 

743 for cellInfo in skyInfo.patchInfo: 

744 if len(observation_identifiers_gc[cellInfo.index]) == 0: 

745 self.log.debug("Skipping cell %s because it has no input warps", cellInfo.index) 

746 continue 

747 

748 cell_masked_image = afwImage.MaskedImageF(cellInfo.outer_bbox) 

749 cell_maskfrac_image = afwImage.ImageF(cellInfo.outer_bbox) 

750 cell_noise_images = [ 

751 afwImage.MaskedImageF(cellInfo.outer_bbox) for n in range(self.config.num_noise_realizations) 

752 ] 

753 psf_masked_image = afwImage.MaskedImageF(psf_bbox_gc[cellInfo.index]) 

754 

755 warp_stacker_gc[cellInfo.index].fill_stacked_masked_image(cell_masked_image) 

756 maskfrac_stacker_gc[cellInfo.index].fill_stacked_image(cell_maskfrac_image) 

757 for n in range(self.config.num_noise_realizations): 

758 noise_stacker_gc_list[n][cellInfo.index].fill_stacked_masked_image(cell_noise_images[n]) 

759 psf_stacker_gc[cellInfo.index].fill_stacked_masked_image(psf_masked_image) 

760 

761 if ap_corr_stacker_gc[cellInfo.index].ap_corr_names: 

762 ap_corr_map = ap_corr_stacker_gc[cellInfo.index].final_ap_corr_map 

763 else: 

764 ap_corr_map = None 

765 

766 # Post-process the coadd before converting to new data structures. 

767 if self.config.do_interpolate_coadd: 

768 self.interpolate_coadd.run(cell_masked_image, planeName="NO_DATA") 

769 for noise_image in cell_noise_images: 

770 self.interpolate_coadd.run(noise_image, planeName="NO_DATA") 

771 # The variance must be positive; work around for DM-3201. 

772 varArray = cell_masked_image.variance.array 

773 with np.errstate(invalid="ignore"): 

774 varArray[:] = np.where(varArray > 0, varArray, np.inf) 

775 

776 image_planes = OwnedImagePlanes.from_masked_image( 

777 masked_image=cell_masked_image, 

778 mask_fractions=cell_maskfrac_image, 

779 noise_realizations=[noise_image.image for noise_image in cell_noise_images], 

780 ) 

781 identifiers = CellIdentifiers( 

782 cell=cellInfo.index, 

783 skymap=self.common.identifiers.skymap, 

784 tract=self.common.identifiers.tract, 

785 patch=self.common.identifiers.patch, 

786 band=self.common.identifiers.band, 

787 ) 

788 

789 singleCellCoadd = SingleCellCoadd( 

790 outer=image_planes, 

791 psf=psf_masked_image.image, 

792 inner_bbox=cellInfo.inner_bbox, 

793 inputs=observation_identifiers_gc[cellInfo.index], 

794 common=self.common, 

795 identifiers=identifiers, 

796 aperture_correction_map=ap_corr_map, 

797 ) 

798 # TODO: Attach transmission curve when they become available. 

799 cells.append(singleCellCoadd) 

800 

801 if not cells: 

802 raise EmptyCellCoaddError() 

803 

804 grid = self._construct_grid(skyInfo) 

805 multipleCellCoadd = MultipleCellCoadd( 

806 cells, 

807 grid=grid, 

808 outer_cell_size=cellInfo.outer_bbox.getDimensions(), 

809 inner_bbox=None, 

810 common=self.common, 

811 psf_image_size=cells[0].psf_image.getDimensions(), 

812 ) 

813 

814 return Struct( 

815 multipleCellCoadd=multipleCellCoadd, 

816 ) 

817 

818 

819class ConvertMultipleCellCoaddToExposureConnections( 

820 PipelineTaskConnections, 

821 dimensions=("tract", "patch", "band", "skymap"), 

822 defaultTemplates={"inputCoaddName": "deep", "inputCoaddSuffix": "Cell"}, 

823): 

824 cellCoaddExposure = Input( 

825 doc="Output coadded exposure, produced by stacking input warps", 

826 name="{inputCoaddName}Coadd{inputCoaddSuffix}", 

827 storageClass="MultipleCellCoadd", 

828 dimensions=("tract", "patch", "skymap", "band"), 

829 ) 

830 

831 stitchedCoaddExposure = Output( 

832 doc="Output stitched coadded exposure, produced by stacking input warps", 

833 name="{inputCoaddName}Coadd{inputCoaddSuffix}_stitched", 

834 storageClass="ExposureF", 

835 dimensions=("tract", "patch", "skymap", "band"), 

836 ) 

837 

838 

839class ConvertMultipleCellCoaddToExposureConfig(