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 

32import logging 

33 

34import numpy as np 

35 

36import lsst.afw.geom as afwGeom 

37import lsst.afw.image as afwImage 

38import lsst.afw.math as afwMath 

39import lsst.geom as geom 

40from lsst.afw.detection import InvalidPsfError 

41from lsst.afw.geom import SinglePolygonException, makeWcsPairTransform 

42from lsst.cell_coadds import ( 

43 CellIdentifiers, 

44 CoaddApCorrMapStacker, 

45 CoaddInputs, 

46 CoaddUnits, 

47 CommonComponents, 

48 GridContainer, 

49 MultipleCellCoadd, 

50 ObservationIdentifiers, 

51 OwnedImagePlanes, 

52 PatchIdentifiers, 

53 SingleCellCoadd, 

54 UniformGrid, 

55) 

56from lsst.daf.butler import DataCoordinate, DeferredDatasetHandle 

57from lsst.meas.algorithms import AccumulatorMeanStack 

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

59from lsst.pipe.base import ( 

60 InMemoryDatasetHandle, 

61 NoWorkFound, 

62 PipelineTask, 

63 PipelineTaskConfig, 

64 PipelineTaskConnections, 

65 Struct, 

66) 

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

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

69from lsst.pipe.tasks.interpImage import InterpImageTask 

70from lsst.pipe.tasks.scaleZeroPoint import ScaleZeroPointTask 

71from lsst.skymap import BaseSkyMap 

72 

73 

74@dataclasses.dataclass 

75class WarpInputs: 

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

77 

78 warp: DeferredDatasetHandle | InMemoryDatasetHandle 

79 """Handle for the warped exposure.""" 

80 

81 masked_fraction: DeferredDatasetHandle | InMemoryDatasetHandle | None = None 

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

83 

84 artifact_mask: DeferredDatasetHandle | InMemoryDatasetHandle | None = None 

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

86 

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

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

89 

90 @property 

91 def dataId(self) -> DataCoordinate: 

92 """DataID corresponding to the warp. 

93 

94 Returns 

95 ------- 

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

97 DataID of the warp. 

98 """ 

99 return self.warp.dataId 

100 

101 

102class AssembleCellCoaddConnections( 

103 PipelineTaskConnections, 

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

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

106): 

107 inputWarps = Input( 

108 doc="Input warps", 

109 name="{inputWarpName}Coadd_directWarp", 

110 storageClass="ExposureF", 

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

112 deferLoad=True, 

113 multiple=True, 

114 ) 

115 

116 maskedFractionWarps = Input( 

117 doc="Mask fraction warps", 

118 name="{inputWarpName}Coadd_directWarp_maskedFraction", 

119 storageClass="ImageF", 

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

121 deferLoad=True, 

122 multiple=True, 

123 ) 

124 

125 artifactMasks = Input( 

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

127 name="compare_warp_artifact_mask", 

128 storageClass="Mask", 

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

130 deferLoad=True, 

131 multiple=True, 

132 ) 

133 

134 visitSummaryList = Input( 

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

136 name="finalVisitSummary", 

137 storageClass="ExposureCatalog", 

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

139 deferLoad=True, 

140 multiple=True, 

141 ) 

142 

143 skyMap = Input( 

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

145 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, 

146 storageClass="SkyMap", 

147 dimensions=("skymap",), 

148 ) 

149 

150 multipleCellCoadd = Output( 

151 doc="Output multiple cell coadd", 

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

153 storageClass="MultipleCellCoadd", 

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

155 ) 

156 

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

158 super().__init__(config=config) 

159 

160 if not config: 

161 return 

162 

163 if config.do_calculate_weight_from_warp: 

164 del self.visitSummaryList 

165 

166 if not config.do_use_artifact_mask: 

167 del self.artifactMasks 

168 

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

170 # number of noise realizations specified in the config. 

171 for n in range(config.num_noise_realizations): 

172 noise_warps = Input( 

173 doc="Input noise warps", 

174 name=f"direct_warp_noise{n}", 

175 storageClass="MaskedImageF", 

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

177 deferLoad=True, 

178 multiple=True, 

179 ) 

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

181 

182 

183class AssembleCellCoaddConfig(PipelineTaskConfig, pipelineConnections=AssembleCellCoaddConnections): 

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

185 interpolate_coadd = ConfigurableField( 

186 target=InterpImageTask, 

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

188 ) 

189 do_scale_zero_point = Field[bool]( 

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

191 default=False, 

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

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

194 "after v29.", 

195 ) 

196 scale_zero_point = ConfigurableField( 

197 target=ScaleZeroPointTask, 

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

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

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

201 "after v29.", 

202 ) 

203 do_calculate_weight_from_warp = Field[bool]( 

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

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

206 default=False, 

207 ) 

208 do_use_artifact_mask = Field[bool]( 

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

210 default=True, 

211 ) 

212 do_coadd_inverse_aperture_corrections = Field[bool]( 

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

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

215 default=False, 

216 ) 

217 min_overlap_fraction = RangeField[float]( 

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

219 "cell.", 

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

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

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

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

224 # definitely not overlap the cell center. 

225 default=1.0, 

226 min=0.0, 

227 max=1.0, 

228 inclusiveMin=True, 

229 inclusiveMax=True, 

230 ) 

231 bad_mask_planes = ListField[str]( 

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

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

234 ) 

235 remove_mask_planes = ListField[str]( 

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

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

238 ) 

239 calc_error_from_input_variance = Field[bool]( 

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

241 "statistic. Passed to AccumulatorMeanStack.", 

242 default=True, 

243 ) 

244 mask_propagation_thresholds = DictField[str, float]( 

245 doc=( 

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

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

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

249 "would have contributed exceeds this value." 

250 ), 

251 default={"SAT": 0.1}, 

252 ) 

253 max_maskfrac = RangeField[float]( 

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

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

256 default=0.99, 

257 min=0.0, 

258 max=1.0, 

259 inclusiveMin=True, 

260 inclusiveMax=False, 

261 ) 

262 num_noise_realizations = Field[int]( 

263 default=0, 

264 doc=( 

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

266 "This should not exceed the corresponding config parameter " 

267 "specified in `MakeDirectWarpConfig`. " 

268 ), 

269 check=lambda x: x >= 0, 

270 ) 

271 psf_warper = ConfigField( 

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

273 "warp the images.", 

274 dtype=afwMath.Warper.ConfigClass, 

275 ) 

276 psf_dimensions = Field[int]( 

277 default=35, 

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

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

280 ) 

281 

282 

283class AssembleCellCoaddTask(PipelineTask): 

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

285 

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

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

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

289 

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

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

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

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

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

295 with the same weights. 

296 

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

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

299 assumed to be spatially constant within a cell. 

300 

301 Raises 

302 ------ 

303 NoWorkFound 

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

305 RuntimeError 

306 Raised if the skymap is not cell-based. 

307 

308 Notes 

309 ----- 

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

311 especially its Config and Connections are experimental and subject to 

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

313 until it is included in the DRP pipeline. 

314 """ 

315 

316 ConfigClass = AssembleCellCoaddConfig 

317 _DefaultName = "assembleCellCoadd" 

318 

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

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

321 if self.config.do_interpolate_coadd: 

322 self.makeSubtask("interpolate_coadd") 

323 # Suppress the warning message about fallback. 

324 self.interpolate_coadd.log.setLevel(logging.ERROR) 

325 if self.config.do_scale_zero_point: 

326 self.makeSubtask("scale_zero_point") 

327 

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

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

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

331 self.log.debug( 

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

333 psf_padding, 

334 self.config.psf_warper.warpingKernelName, 

335 ) 

336 else: 

337 psf_padding = 10 

338 self.log.info( 

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

340 psf_padding, 

341 ) 

342 self.psf_padding = psf_padding 

343 

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

345 # Docstring inherited. 

346 if not inputRefs.inputWarps: 

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

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

349 

350 # Construct skyInfo expected by run 

351 # Do not remove skyMap from inputData in case _makeSupplementaryData 

352 # needs it 

353 skyMap = butlerQC.get(inputRefs.skyMap) 

354 

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

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

357 

358 outputDataId = butlerQC.quantum.dataId 

359 

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

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

362 

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

364 self.common = CommonComponents( 

365 units=units, 

366 wcs=skyInfo.patchInfo.wcs, 

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

368 identifiers=PatchIdentifiers.from_data_id(outputDataId), 

369 ) 

370 

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

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

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

374 

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

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

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

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

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

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

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

382 

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

384 butlerQC.put(returnStruct, outputRefs) 

385 return returnStruct 

386 

387 @staticmethod 

388 def _compute_weight(maskedImage, statsCtrl): 

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

390 

391 Parameters 

392 ---------- 

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

394 The masked image to compute the weight. 

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

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

397 

398 Returns 

399 ------- 

400 weight : `float` 

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

402 """ 

403 statObj = afwMath.makeStatistics( 

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

405 ) 

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

407 weight = 1.0 / float(meanVar) 

408 return weight 

409 

410 @staticmethod 

411 def _construct_grid(skyInfo): 

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

413 

414 Parameters 

415 ---------- 

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

417 A Struct object 

418 

419 Returns 

420 ------- 

421 grid : `~lsst.cell_coadds.UniformGrid` 

422 A UniformGrid object. 

423 """ 

424 padding = skyInfo.patchInfo.getCellBorder() 

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

426 grid = UniformGrid.from_bbox_cell_size( 

427 grid_bbox, 

428 skyInfo.patchInfo.getCellInnerDimensions(), 

429 padding=padding, 

430 ) 

431 return grid 

432 

433 def _construct_grid_container(self, skyInfo, statsCtrl): 

434 """Construct a grid of AccumulatorMeanStack instances. 

435 

436 Parameters 

437 ---------- 

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

439 A Struct object 

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

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

442 

443 Returns 

444 ------- 

445 gc : `~lsst.cell_coadds.GridContainer` 

446 A GridContainer object container one AccumulatorMeanStack per cell. 

447 """ 

448 grid = self._construct_grid(skyInfo) 

449 

450 maskMap = setRejectedMaskMapping(statsCtrl) 

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

452 thresholdDict = AccumulatorMeanStack.stats_ctrl_to_threshold_dict(statsCtrl) 

453 

454 # Initialize the grid container with AccumulatorMeanStacks 

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

456 for cellInfo in skyInfo.patchInfo: 

457 stacker = AccumulatorMeanStack( 

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

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

460 bit_mask_value=statsCtrl.getAndMask(), 

461 mask_threshold_dict=thresholdDict, 

462 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

463 compute_n_image=False, 

464 mask_map=maskMap, 

465 no_good_pixels_mask=statsCtrl.getNoGoodPixelsMask(), 

466 ) 

467 gc[cellInfo.index] = stacker 

468 

469 return gc 

470 

471 def _construct_stats_control(self): 

472 """Construct a StatisticsControl object for coadd. 

473 

474 Unlike AssembleCoaddTask or CompareWarpAssembleCoaddTask, there is 

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

476 optionally mask propagation thresholds. 

477 

478 Returns 

479 ------- 

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

481 A control object for StatisticsStack. 

482 """ 

483 statsCtrl = afwMath.StatisticsControl() 

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

485 # CompareWarpAssembleCoaddTask to get consistent weights. This is NOT 

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

487 # fallback option that is not recommended for production. 

488 statsCtrl.setNumIter(2) 

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

490 statsCtrl.setNanSafe(True) 

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

492 bit = afwImage.Mask.getMaskPlane(plane) 

493 statsCtrl.setMaskPropagationThreshold(bit, threshold) 

494 return statsCtrl 

495 

496 def _construct_ap_corr_grid_container(self, skyInfo): 

497 """Construct a grid of CoaddApCorrMapStacker instances. 

498 

499 Parameters 

500 ---------- 

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

502 A Struct object 

503 

504 Returns 

505 ------- 

506 gc : `~lsst.cell_coadds.GridContainer` 

507 A GridContainer object container one CoaddApCorrMapStacker per 

508 cell. 

509 """ 

510 grid = self._construct_grid(skyInfo) 

511 

512 # Initialize the grid container with CoaddApCorrMapStacker. 

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

514 for cellInfo in skyInfo.patchInfo: 

515 stacker = CoaddApCorrMapStacker( 

516 evaluation_point=cellInfo.inner_bbox.getCenter(), 

517 do_coadd_inverse_ap_corr=self.config.do_coadd_inverse_aperture_corrections, 

518 ) 

519 gc[cellInfo.index] = stacker 

520 

521 return gc 

522 

523 def run( 

524 self, 

525 *, 

526 inputs: dict[DataCoordinate, WarpInputs], 

527 skyInfo, 

528 visitSummaryList: list | None = None, 

529 ): 

530 for mask_plane in self.config.bad_mask_planes: 

531 afwImage.Mask.addMaskPlane(mask_plane) 

532 for mask_plane in self.config.mask_propagation_thresholds: 

533 afwImage.Mask.addMaskPlane(mask_plane) 

534 

535 statsCtrl = self._construct_stats_control() 

536 

537 warp_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

538 maskfrac_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

539 noise_stacker_gc_list = [ 

540 self._construct_grid_container(skyInfo, statsCtrl) 

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

542 ] 

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

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

545 ap_corr_stacker_gc = self._construct_ap_corr_grid_container(skyInfo) 

546 

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

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

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

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

551 # get the corresponding PSFs as well. 

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

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

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

555 # Populate them. 

556 for cellInfo in skyInfo.patchInfo: 

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

558 observation_identifiers_gc[cellInfo.index] = {} 

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

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

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

562 cell_center_pixel, 

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

564 ) 

565 psf_stacker_gc[cellInfo.index] = AccumulatorMeanStack( 

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

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

568 bit_mask_value=0, 

569 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

570 compute_n_image=False, 

571 ) 

572 

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

574 if not visitSummaryList: 

575 visitSummaryList = [] 

576 visitSummaryRefDict = { 

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

578 } 

579 

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

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

582 

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

584 # it completely overlaps. 

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

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

587 masked_fraction_image = ( 

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

589 if warp_input.masked_fraction 

590 else None 

591 ) 

592 

593 # Pre-process the warp before coadding. 

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

595 warp.mask.addMaskPlane("CLIPPED") 

596 warp.mask.addMaskPlane("REJECTED") 

597 warp.mask.addMaskPlane("SENSOR_EDGE") 

598 warp.mask.addMaskPlane("INEXACT_PSF") 

599 

600 if artifact_mask_ref := warp_input.artifact_mask: 

601 # Apply the artifact mask to the warp. 

602 artifact_mask = artifact_mask_ref.get() 

603 assert ( 

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

605 ), "Mask dicts do not agree." 

606 warp.mask.array = artifact_mask.array 

607 del artifact_mask 

608 

609 if self.config.do_scale_zero_point: 

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

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

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

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

614 zero_point_scale_factor = imageScaler.scale 

615 self.log.debug( 

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

617 warp_input.dataId, 

618 zero_point_scale_factor, 

619 ) 

620 else: 

621 zero_point_scale_factor = 1.0 

622 if "BUNIT" not in warp.metadata: 

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

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

625 raise ValueError( 

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

627 ) 

628 

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

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

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

632 rejected = afwImage.Mask.getPlaneBitMask( 

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

634 ) 

635 

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

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

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

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

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

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

642 full_ccd_table["ccd"], 

643 0.0, 

644 ) # Mapping from detector to weight. 

645 

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

647 visitSummary = visitSummaryRef.get() 

648 for detector in full_ccd_table["ccd"]: 

649 visitSummaryRow = visitSummary.find(detector) 

650 mean_variance = visitSummaryRow["meanVar"] 

651 mean_variance *= zero_point_scale_factor**2 

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

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

654 weights[detector] = 1.0 / mean_variance 

655 del visitSummary 

656 else: 

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

658 weight = self._compute_weight(warp, statsCtrl) 

659 if not np.isfinite(weight): 

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

661 continue 

662 

663 for detector in weights: 

664 weights[detector] = weight 

665 

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

667 

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

669 # detector ID that pixel comes from. 

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

671 for row in full_ccd_table: 

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

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

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

675 try: 

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

677 geom.Box2D(warp.getBBox()) 

678 ) 

679 except SinglePolygonException: 

680 continue 

681 

682 observation_identifier = ObservationIdentifiers.from_data_id( 

683 warp_input.dataId, 

684 backup_detector=row["ccd"], 

685 ) 

686 visit_polygons[observation_identifier] = dest_polygon 

687 

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

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

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

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

692 

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

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

695 detector_map.array[:, :] = 0 

696 

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

698 bbox = cellInfo.outer_bbox 

699 inner_bbox = cellInfo.inner_bbox 

700 

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

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

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

704 self.log.debug( 

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

706 warp_input.dataId, 

707 cellInfo.index, 

708 overlap_fraction, 

709 self.config.min_overlap_fraction, 

710 ) 

711 continue 

712 

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

714 if not np.isfinite(weight): 

715 self.log.warn( 

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

717 ) 

718 continue 

719 

720 if weight == 0: 

721 self.log.info( 

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

723 ) 

724 continue 

725 

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

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

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

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

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

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

732 if deep_copy: 

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

734 

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

736 if deep_copy: 

737 mi.image.array[single_detector_mask_array] = 0.0 

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

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

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

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

742 

743 if masked_fraction_image: 

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

745 if deep_copy: 

746 mi.array[single_detector_mask_array] = 0.0 

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

748 

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

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

751 if deep_copy: 

752 mi.image.array[single_detector_mask_array] = 0.0 

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

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

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

756 

757 # Set the defaults for PSF shape quantities. 

758 psf_shape = afwGeom.Quadrupole() 

759 psf_shape_flag = True 

760 psf_eval_point = None 

761 try: 

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

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

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

765 # the rug for now. 

766 if overlap_fraction < 0.5: 

767 psf_eval_point = dest_polygon.intersectionSingle( 

768 geom.Box2D(inner_bbox) 

769 ).calculateCenter() 

770 else: 

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

772 psf_shape = warp.psf.computeShape(psf_eval_point) 

773 psf_shape_flag = False 

774 except SinglePolygonException: 

775 self.log.info( 

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

777 ccd_row["ccd"], 

778 warp_input.dataId, 

779 cellInfo.index, 

780 ) 

781 except InvalidPsfError: 

782 self.log.info( 

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

784 ccd_row["ccd"], 

785 warp_input.dataId, 

786 psf_eval_point, 

787 ) 

788 

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

790 

791 observation_identifier = ObservationIdentifiers.from_data_id( 

792 warp_input.dataId, 

793 backup_detector=ccd_row["ccd"], 

794 ) 

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

796 overlaps_center=overlaps_center, 

797 overlap_fraction=overlap_fraction, 

798 weight=weight, 

799 psf_shape=psf_shape, 

800 psf_shape_flag=psf_shape_flag, 

801 ) 

802 if overlaps_center is False: 

803 self.log.debug( 

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

805 warp_input.dataId, 

806 cellInfo.index, 

807 ) 

808 continue 

809 

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

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

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

813 

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

815 calexp_point, 

816 undistorted_psf_im.getDimensions(), 

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

818 

819 # Convert the PSF image from Image to MaskedImage and 

820 # zero-pad the image. 

821 undistorted_psf_bbox = undistorted_psf_im.getBBox() 

822 undistorted_psf_maskedImage = afwImage.MaskedImageD( 

823 undistorted_psf_bbox.dilatedBy(self.psf_padding) 

824 ) 

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

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

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

828 

829 warped_psf_maskedImage = self.psf_warper.warpImage( 

830 destWcs=skyInfo.wcs, 

831 srcImage=undistorted_psf_maskedImage, 

832 srcWcs=ccd_row.getWcs(), 

833 destBBox=psf_bbox_gc[cellInfo.index], 

834 ) 

835 

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

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