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

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) 

29 

30import dataclasses 

31import itertools 

32 

33import numpy as np 

34 

35import lsst.afw.geom as afwGeom 

36import lsst.afw.image as afwImage 

37import lsst.afw.math as afwMath 

38import lsst.geom as geom 

39from lsst.afw.detection import InvalidPsfError 

40from lsst.afw.geom import SinglePolygonException, makeWcsPairTransform 

41from lsst.cell_coadds import ( 

42 CellIdentifiers, 

43 CoaddApCorrMapStacker, 

44 CoaddInputs, 

45 CoaddUnits, 

46 CommonComponents, 

47 GridContainer, 

48 MultipleCellCoadd, 

49 ObservationIdentifiers, 

50 OwnedImagePlanes, 

51 PatchIdentifiers, 

52 SingleCellCoadd, 

53 UniformGrid, 

54) 

55from lsst.daf.butler import DataCoordinate, DeferredDatasetHandle 

56from lsst.meas.algorithms import AccumulatorMeanStack 

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

58from lsst.pipe.base import ( 

59 InMemoryDatasetHandle, 

60 NoWorkFound, 

61 PipelineTask, 

62 PipelineTaskConfig, 

63 PipelineTaskConnections, 

64 Struct, 

65) 

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

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

68from lsst.pipe.tasks.interpImage import InterpImageTask 

69from lsst.pipe.tasks.scaleZeroPoint import ScaleZeroPointTask 

70from lsst.skymap import BaseSkyMap 

71 

72 

73@dataclasses.dataclass 

74class WarpInputs: 

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

76 

77 warp: DeferredDatasetHandle | InMemoryDatasetHandle 

78 """Handle for the warped exposure.""" 

79 

80 masked_fraction: DeferredDatasetHandle | InMemoryDatasetHandle | None = None 

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

82 

83 artifact_mask: DeferredDatasetHandle | InMemoryDatasetHandle | None = None 

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

85 

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

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

88 

89 @property 

90 def dataId(self) -> DataCoordinate: 

91 """DataID corresponding to the warp. 

92 

93 Returns 

94 ------- 

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

96 DataID of the warp. 

97 """ 

98 return self.warp.dataId 

99 

100 

101class AssembleCellCoaddConnections( 

102 PipelineTaskConnections, 

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

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

105): 

106 inputWarps = Input( 

107 doc="Input warps", 

108 name="{inputWarpName}Coadd_directWarp", 

109 storageClass="ExposureF", 

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

111 deferLoad=True, 

112 multiple=True, 

113 ) 

114 

115 maskedFractionWarps = Input( 

116 doc="Mask fraction warps", 

117 name="{inputWarpName}Coadd_directWarp_maskedFraction", 

118 storageClass="ImageF", 

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

120 deferLoad=True, 

121 multiple=True, 

122 ) 

123 

124 artifactMasks = Input( 

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

126 name="compare_warp_artifact_mask", 

127 storageClass="Mask", 

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

129 deferLoad=True, 

130 multiple=True, 

131 ) 

132 

133 visitSummaryList = Input( 

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

135 name="finalVisitSummary", 

136 storageClass="ExposureCatalog", 

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

138 deferLoad=True, 

139 multiple=True, 

140 ) 

141 

142 skyMap = Input( 

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

144 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, 

145 storageClass="SkyMap", 

146 dimensions=("skymap",), 

147 ) 

148 

149 multipleCellCoadd = Output( 

150 doc="Output multiple cell coadd", 

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

152 storageClass="MultipleCellCoadd", 

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

154 ) 

155 

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

157 super().__init__(config=config) 

158 

159 if not config: 

160 return 

161 

162 if config.do_calculate_weight_from_warp: 

163 del self.visitSummaryList 

164 

165 if not config.do_use_artifact_mask: 

166 del self.artifactMasks 

167 

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

169 # number of noise realizations specified in the config. 

170 for n in range(config.num_noise_realizations): 

171 noise_warps = Input( 

172 doc="Input noise warps", 

173 name=f"direct_warp_noise{n}", 

174 storageClass="MaskedImageF", 

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

176 deferLoad=True, 

177 multiple=True, 

178 ) 

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

180 

181 

182class AssembleCellCoaddConfig(PipelineTaskConfig, pipelineConnections=AssembleCellCoaddConnections): 

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

184 interpolate_coadd = ConfigurableField( 

185 target=InterpImageTask, 

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

187 ) 

188 do_scale_zero_point = Field[bool]( 

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

190 default=False, 

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

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

193 "after v29.", 

194 ) 

195 scale_zero_point = ConfigurableField( 

196 target=ScaleZeroPointTask, 

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

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

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

200 "after v29.", 

201 ) 

202 do_calculate_weight_from_warp = Field[bool]( 

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

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

205 default=False, 

206 ) 

207 do_use_artifact_mask = Field[bool]( 

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

209 default=True, 

210 ) 

211 do_coadd_inverse_aperture_corrections = Field[bool]( 

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

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

214 default=False, 

215 ) 

216 min_overlap_fraction = RangeField[float]( 

217 doc="The minimum overlap fraction required for a single (visit, detector) input to be included in a " 

218 "cell.", 

219 # A value of 1.0 corresponds to ideal, edge-free cells. 

220 # A value of 0.0 corresponds to the deep_coadd style coadds. 

221 # This has to be at least 0.5 to ensure that the an input overlaps the 

222 # cell center. Inputs will overlap fraction less than 0.25 will 

223 # definitely not overlap the cell center. 

224 default=1.0, 

225 min=0.0, 

226 max=1.0, 

227 inclusiveMin=True, 

228 inclusiveMax=True, 

229 ) 

230 bad_mask_planes = ListField[str]( 

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

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

233 ) 

234 remove_mask_planes = ListField[str]( 

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

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

237 ) 

238 calc_error_from_input_variance = Field[bool]( 

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

240 "statistic. Passed to AccumulatorMeanStack.", 

241 default=True, 

242 ) 

243 mask_propagation_thresholds = DictField[str, float]( 

244 doc=( 

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

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

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

248 "would have contributed exceeds this value." 

249 ), 

250 default={"SAT": 0.1}, 

251 ) 

252 max_maskfrac = RangeField[float]( 

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

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

255 default=0.99, 

256 min=0.0, 

257 max=1.0, 

258 inclusiveMin=True, 

259 inclusiveMax=False, 

260 ) 

261 num_noise_realizations = Field[int]( 

262 default=0, 

263 doc=( 

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

265 "This should not exceed the corresponding config parameter " 

266 "specified in `MakeDirectWarpConfig`. " 

267 ), 

268 check=lambda x: x >= 0, 

269 ) 

270 psf_warper = ConfigField( 

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

272 "warp the images.", 

273 dtype=afwMath.Warper.ConfigClass, 

274 ) 

275 psf_dimensions = Field[int]( 

276 default=35, 

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

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

279 ) 

280 

281 

282class AssembleCellCoaddTask(PipelineTask): 

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

284 

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

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

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

288 

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

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

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

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

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

294 with the same weights. 

295 

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

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

298 assumed to be spatially constant within a cell. 

299 

300 Raises 

301 ------ 

302 NoWorkFound 

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

304 RuntimeError 

305 Raised if the skymap is not cell-based. 

306 

307 Notes 

308 ----- 

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

310 especially its Config and Connections are experimental and subject to 

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

312 until it is included in the DRP pipeline. 

313 """ 

314 

315 ConfigClass = AssembleCellCoaddConfig 

316 _DefaultName = "assembleCellCoadd" 

317 

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

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

320 if self.config.do_interpolate_coadd: 

321 self.makeSubtask("interpolate_coadd") 

322 if self.config.do_scale_zero_point: 

323 self.makeSubtask("scale_zero_point") 

324 

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

326 if (warping_kernel_name := self.config.psf_warper.warpingKernelName.lower()).startswith("lanczos"): 

327 psf_padding = 2 * int(warping_kernel_name.lstrip("lanczos")) - 1 

328 self.log.debug( 

329 "Padding PSF image by %d pixels since the warping kernel is %s.", 

330 psf_padding, 

331 self.config.psf_warper.warpingKernelName, 

332 ) 

333 else: 

334 psf_padding = 10 

335 self.log.info( 

336 "Padding PSF image by %d pixels since the warping kernel is not Lanczos.", 

337 psf_padding, 

338 ) 

339 self.psf_padding = psf_padding 

340 

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

342 # Docstring inherited. 

343 if not inputRefs.inputWarps: 

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

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

346 

347 # Construct skyInfo expected by run 

348 # Do not remove skyMap from inputData in case _makeSupplementaryData 

349 # needs it 

350 skyMap = butlerQC.get(inputRefs.skyMap) 

351 

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

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

354 

355 outputDataId = butlerQC.quantum.dataId 

356 

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

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

359 

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

361 self.common = CommonComponents( 

362 units=units, 

363 wcs=skyInfo.patchInfo.wcs, 

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

365 identifiers=PatchIdentifiers.from_data_id(outputDataId), 

366 ) 

367 

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

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

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

371 

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

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

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

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

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

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

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

379 

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

381 butlerQC.put(returnStruct, outputRefs) 

382 return returnStruct 

383 

384 @staticmethod 

385 def _compute_weight(maskedImage, statsCtrl): 

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

387 

388 Parameters 

389 ---------- 

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

391 The masked image to compute the weight. 

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

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

394 

395 Returns 

396 ------- 

397 weight : `float` 

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

399 """ 

400 statObj = afwMath.makeStatistics( 

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

402 ) 

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

404 weight = 1.0 / float(meanVar) 

405 return weight 

406 

407 @staticmethod 

408 def _construct_grid(skyInfo): 

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

410 

411 Parameters 

412 ---------- 

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

414 A Struct object 

415 

416 Returns 

417 ------- 

418 grid : `~lsst.cell_coadds.UniformGrid` 

419 A UniformGrid object. 

420 """ 

421 padding = skyInfo.patchInfo.getCellBorder() 

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

423 grid = UniformGrid.from_bbox_cell_size( 

424 grid_bbox, 

425 skyInfo.patchInfo.getCellInnerDimensions(), 

426 padding=padding, 

427 ) 

428 return grid 

429 

430 def _construct_grid_container(self, skyInfo, statsCtrl): 

431 """Construct a grid of AccumulatorMeanStack instances. 

432 

433 Parameters 

434 ---------- 

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

436 A Struct object 

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

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

439 

440 Returns 

441 ------- 

442 gc : `~lsst.cell_coadds.GridContainer` 

443 A GridContainer object container one AccumulatorMeanStack per cell. 

444 """ 

445 grid = self._construct_grid(skyInfo) 

446 

447 maskMap = setRejectedMaskMapping(statsCtrl) 

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

449 thresholdDict = AccumulatorMeanStack.stats_ctrl_to_threshold_dict(statsCtrl) 

450 

451 # Initialize the grid container with AccumulatorMeanStacks 

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

453 for cellInfo in skyInfo.patchInfo: 

454 stacker = AccumulatorMeanStack( 

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

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

457 bit_mask_value=statsCtrl.getAndMask(), 

458 mask_threshold_dict=thresholdDict, 

459 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

460 compute_n_image=False, 

461 mask_map=maskMap, 

462 no_good_pixels_mask=statsCtrl.getNoGoodPixelsMask(), 

463 ) 

464 gc[cellInfo.index] = stacker 

465 

466 return gc 

467 

468 def _construct_stats_control(self): 

469 """Construct a StatisticsControl object for coadd. 

470 

471 Unlike AssembleCoaddTask or CompareWarpAssembleCoaddTask, there is 

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

473 optionally mask propagation thresholds. 

474 

475 Returns 

476 ------- 

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

478 A control object for StatisticsStack. 

479 """ 

480 statsCtrl = afwMath.StatisticsControl() 

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

482 # CompareWarpAssembleCoaddTask to get consistent weights. This is NOT 

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

484 # fallback option that is not recommended for production. 

485 statsCtrl.setNumIter(2) 

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

487 statsCtrl.setNanSafe(True) 

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

489 bit = afwImage.Mask.getMaskPlane(plane) 

490 statsCtrl.setMaskPropagationThreshold(bit, threshold) 

491 return statsCtrl 

492 

493 def _construct_ap_corr_grid_container(self, skyInfo): 

494 """Construct a grid of CoaddApCorrMapStacker instances. 

495 

496 Parameters 

497 ---------- 

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

499 A Struct object 

500 

501 Returns 

502 ------- 

503 gc : `~lsst.cell_coadds.GridContainer` 

504 A GridContainer object container one CoaddApCorrMapStacker per 

505 cell. 

506 """ 

507 grid = self._construct_grid(skyInfo) 

508 

509 # Initialize the grid container with CoaddApCorrMapStacker. 

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

511 for cellInfo in skyInfo.patchInfo: 

512 stacker = CoaddApCorrMapStacker( 

513 evaluation_point=cellInfo.inner_bbox.getCenter(), 

514 do_coadd_inverse_ap_corr=self.config.do_coadd_inverse_aperture_corrections, 

515 ) 

516 gc[cellInfo.index] = stacker 

517 

518 return gc 

519 

520 def run( 

521 self, 

522 *, 

523 inputs: dict[DataCoordinate, WarpInputs], 

524 skyInfo, 

525 visitSummaryList: list | None = None, 

526 ): 

527 for mask_plane in self.config.bad_mask_planes: 

528 afwImage.Mask.addMaskPlane(mask_plane) 

529 for mask_plane in self.config.mask_propagation_thresholds: 

530 afwImage.Mask.addMaskPlane(mask_plane) 

531 

532 statsCtrl = self._construct_stats_control() 

533 

534 warp_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

535 maskfrac_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

536 noise_stacker_gc_list = [ 

537 self._construct_grid_container(skyInfo, statsCtrl) 

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

539 ] 

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

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

542 ap_corr_stacker_gc = self._construct_ap_corr_grid_container(skyInfo) 

543 

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

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

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

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

548 # get the corresponding PSFs as well. 

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

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

551 observation_identifiers_gc = GridContainer[dict](warp_stacker_gc.shape) 

552 # Populate them. 

553 for cellInfo in skyInfo.patchInfo: 

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

555 observation_identifiers_gc[cellInfo.index] = {} 

556 cell_center_pixel = geom.Point2D(geom.Point2I(cellInfo.inner_bbox.getCenter())) 

557 cell_centers_sky[cellInfo.index] = skyInfo.wcs.pixelToSky(cell_center_pixel) 

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

559 cell_center_pixel, 

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

561 ) 

562 psf_stacker_gc[cellInfo.index] = AccumulatorMeanStack( 

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

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

565 bit_mask_value=0, 

566 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

567 compute_n_image=False, 

568 ) 

569 

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

571 if not visitSummaryList: 

572 visitSummaryList = [] 

573 visitSummaryRefDict = { 

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

575 } 

576 

577 # Keep track of the polygons corresponding to each (visit, detector). 

578 visit_polygons: dict[ObservationIdentifiers, afwGeom.Polygon] = {} 

579 

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

581 # it completely overlaps. 

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

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

584 masked_fraction_image = ( 

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

586 if warp_input.masked_fraction 

587 else None 

588 ) 

589 

590 # Pre-process the warp before coadding. 

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

592 warp.mask.addMaskPlane("CLIPPED") 

593 warp.mask.addMaskPlane("REJECTED") 

594 warp.mask.addMaskPlane("SENSOR_EDGE") 

595 warp.mask.addMaskPlane("INEXACT_PSF") 

596 

597 if artifact_mask_ref := warp_input.artifact_mask: 

598 # Apply the artifact mask to the warp. 

599 artifact_mask = artifact_mask_ref.get() 

600 assert ( 

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

602 ), "Mask dicts do not agree." 

603 warp.mask.array = artifact_mask.array 

604 del artifact_mask 

605 

606 if self.config.do_scale_zero_point: 

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

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

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

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

611 zero_point_scale_factor = imageScaler.scale 

612 self.log.debug( 

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

614 warp_input.dataId, 

615 zero_point_scale_factor, 

616 ) 

617 else: 

618 zero_point_scale_factor = 1.0 

619 if "BUNIT" not in warp.metadata: 

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

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

622 raise ValueError( 

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

624 ) 

625 

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

627 # Instead of using self.config.bad_mask_planes, we explicitly 

628 # ask statsCtrl which pixels are going to be ignored/rejected. 

629 rejected = afwImage.Mask.getPlaneBitMask( 

630 ["CLIPPED", "REJECTED"] + afwImage.Mask.interpret(statsCtrl.getAndMask()).split(",") 

631 ) 

632 

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

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

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

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

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

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

639 full_ccd_table["ccd"], 

640 0.0, 

641 ) # Mapping from detector to weight. 

642 

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

644 visitSummary = visitSummaryRef.get() 

645 for detector in full_ccd_table["ccd"]: 

646 visitSummaryRow = visitSummary.find(detector) 

647 mean_variance = visitSummaryRow["meanVar"] 

648 mean_variance *= zero_point_scale_factor**2 

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

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

651 weights[detector] = 1.0 / mean_variance 

652 del visitSummary 

653 else: 

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

655 weight = self._compute_weight(warp, statsCtrl) 

656 if not np.isfinite(weight): 

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

658 continue 

659 

660 for detector in weights: 

661 weights[detector] = weight 

662 

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

664 

665 # Create an image where each pixel value corresponds to the 

666 # detector ID that pixel comes from. 

667 detector_map = afwImage.ImageI(bbox=warp.getBBox(), initialValue=-1) 

668 for row in full_ccd_table: 

669 transform = makeWcsPairTransform(row.wcs, warp.wcs) 

670 if (src_polygon := row.validPolygon) is None: 

671 src_polygon = afwGeom.Polygon(geom.Box2D(row.getBBox())) 

672 try: 

673 dest_polygon = src_polygon.transform(transform).intersectionSingle( 

674 geom.Box2D(warp.getBBox()) 

675 ) 

676 except SinglePolygonException: 

677 continue 

678 

679 observation_identifier = ObservationIdentifiers.from_data_id( 

680 warp_input.dataId, 

681 backup_detector=row["ccd"], 

682 ) 

683 visit_polygons[observation_identifier] = dest_polygon 

684 

685 detector_map_slice = dest_polygon.createImage(detector_map.getBBox()).array > 0 

686 if not (detector_map.array[detector_map_slice] < 0).all(): 

687 self.log.warning("Multiple detectors from visit %s are overlapping", warp_input.dataId) 

688 detector_map.array[detector_map_slice] = row["ccd"] 

689 

690 if (detector_map.array < 0).all(): 

691 self.log.warning("Unable to split the warp %s into single-detector warps.", warp_input.dataId) 

692 detector_map.array[:, :] = 0 

693 

694 for cellInfo, ccd_row in itertools.product(skyInfo.patchInfo, full_ccd_table): 

695 bbox = cellInfo.outer_bbox 

696 inner_bbox = cellInfo.inner_bbox 

697 

698 overlap_fraction = (detector_map[inner_bbox].array == ccd_row["ccd"]).mean() 

699 assert -1e-4 < overlap_fraction < 1.0001, "Overlap fraction is not within [0, 1]." 

700 if (overlap_fraction < self.config.min_overlap_fraction) or (overlap_fraction <= 0.0): 

701 self.log.debug( 

702 "Skipping %s in cell %s because it had only %.3f < %.3f fractional overlap.", 

703 warp_input.dataId, 

704 cellInfo.index, 

705 overlap_fraction, 

706 self.config.min_overlap_fraction, 

707 ) 

708 continue 

709 

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

711 if not np.isfinite(weight): 

712 self.log.warn( 

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

714 ) 

715 continue 

716 

717 if weight == 0: 

718 self.log.info( 

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

720 ) 

721 continue 

722 

723 # Decide if a deep copy is necessary to apply the single 

724 # detector cuts since it involves modifying the image in-place. 

725 # If within the inner cell, there are three or more different 

726 # values that detector map takes, then there are definitely 

727 # multiple detectors (one for chip gaps, two for two detectors) 

728 deep_copy = len(set(detector_map[inner_bbox].array.ravel())) >= 3 

729 if deep_copy: 

730 single_detector_mask_array = detector_map[bbox].array != ccd_row["ccd"] 

731 

732 mi = afwImage.MaskedImageF(warp[bbox].maskedImage, deep=deep_copy) 

733 if deep_copy: 

734 mi.image.array[single_detector_mask_array] = 0.0 

735 mi.variance.array[single_detector_mask_array] = np.inf 

736 nodata_or_mask = (single_detector_mask_array) * afwImage.Mask.getPlaneBitMask("NO_DATA") 

737 mi.mask[bbox].array |= nodata_or_mask 

738 warp_stacker_gc[cellInfo.index].add_masked_image(mi, weight=weight) 

739 

740 if masked_fraction_image: 

741 mi = afwImage.ImageF(masked_fraction_image[bbox], deep=deep_copy) 

742 if deep_copy: 

743 mi.array[single_detector_mask_array] = 0.0 

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

745 

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

747 mi = afwImage.MaskedImageF(noise_warps[n][bbox], deep=deep_copy) 

748 if deep_copy: 

749 mi.image.array[single_detector_mask_array] = 0.0 

750 mi.variance.array[single_detector_mask_array] = np.inf 

751 mi.mask[bbox].array |= nodata_or_mask 

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

753 

754 # Set the defaults for PSF shape quantities. 

755 psf_shape = afwGeom.Quadrupole() 

756 psf_shape_flag = True 

757 psf_eval_point = None 

758 try: 

759 # The `if` branch is buggy. `dest_polygon` is technically 

760 # out of scope, but Python does not raise an error. 

761 # TODO: Fix this properly in DM-53479, but sweep it under 

762 # the rug for now. 

763 if overlap_fraction < 0.5: 

764 psf_eval_point = dest_polygon.intersectionSingle( 

765 geom.Box2D(inner_bbox) 

766 ).calculateCenter() 

767 else: 

768 psf_eval_point = geom.Point2D(geom.Point2I(inner_bbox.getCenter())) 

769 psf_shape = warp.psf.computeShape(psf_eval_point) 

770 psf_shape_flag = False 

771 except SinglePolygonException: 

772 self.log.info( 

773 "Unable to find the overlapping polygon between %d detector in %s and cell %s", 

774 ccd_row["ccd"], 

775 warp_input.dataId, 

776 cellInfo.index, 

777 ) 

778 except InvalidPsfError: 

779 self.log.info( 

780 "Unable to compute PSF shape from %d detector in %s at %s", 

781 ccd_row["ccd"], 

782 warp_input.dataId, 

783 psf_eval_point, 

784 ) 

785 

786 overlaps_center = detector_map[geom.Point2I(bbox.getCenter())] == ccd_row["ccd"] 

787 

788 observation_identifier = ObservationIdentifiers.from_data_id( 

789 warp_input.dataId, 

790 backup_detector=ccd_row["ccd"], 

791 ) 

792 observation_identifiers_gc[cellInfo.index][observation_identifier] = CoaddInputs( 

793 overlaps_center=overlaps_center, 

794 overlap_fraction=overlap_fraction, 

795 weight=weight, 

796 psf_shape=psf_shape, 

797 psf_shape_flag=psf_shape_flag, 

798 ) 

799 if overlaps_center is False: 

800 self.log.debug( 

801 "%s does not overlap with the center of the cell %s", 

802 warp_input.dataId, 

803 cellInfo.index, 

804 ) 

805 continue 

806 

807 # Everything below this has to do with the center of the cell 

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

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

810 

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

812 calexp_point, 

813 undistorted_psf_im.getDimensions(), 

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

815 

816 # Convert the PSF image from Image to MaskedImage and 

817 # zero-pad the image. 

818 undistorted_psf_bbox = undistorted_psf_im.getBBox() 

819 undistorted_psf_maskedImage = afwImage.MaskedImageD( 

820 undistorted_psf_bbox.dilatedBy(self.psf_padding) 

821 ) 

822 undistorted_psf_maskedImage.image[undistorted_psf_bbox].array[:, :] = undistorted_psf_im.array 

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

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

825 

826 warped_psf_maskedImage = self.psf_warper.warpImage( 

827 destWcs=skyInfo.wcs, 

828 srcImage=undistorted_psf_maskedImage, 

829 srcWcs=ccd_row.getWcs(), 

830 destBBox=psf_bbox_gc[cellInfo.index], 

831 ) 

832 

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

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

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