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 require_artifact_mask = Field[bool]( 

282 default=True, 

283 doc="Require presence of artifact mask for each warp? Use true if using artifact rejection outputs" 

284 " from CompareWarpTask", 

285 ) 

286 

287 

288class AssembleCellCoaddTask(PipelineTask): 

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

290 

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

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

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

294 

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

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

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

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

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

300 with the same weights. 

301 

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

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

304 assumed to be spatially constant within a cell. 

305 

306 Raises 

307 ------ 

308 NoWorkFound 

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

310 RuntimeError 

311 Raised if the skymap is not cell-based. 

312 

313 Notes 

314 ----- 

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

316 especially its Config and Connections are experimental and subject to 

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

318 until it is included in the DRP pipeline. 

319 """ 

320 

321 ConfigClass = AssembleCellCoaddConfig 

322 _DefaultName = "assembleCellCoadd" 

323 

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

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

326 if self.config.do_interpolate_coadd: 

327 self.makeSubtask("interpolate_coadd") 

328 # Suppress the warning message about fallback. 

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

330 if self.config.do_scale_zero_point: 

331 self.makeSubtask("scale_zero_point") 

332 

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

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

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

336 self.log.debug( 

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

338 psf_padding, 

339 self.config.psf_warper.warpingKernelName, 

340 ) 

341 else: 

342 psf_padding = 10 

343 self.log.info( 

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

345 psf_padding, 

346 ) 

347 self.psf_padding = psf_padding 

348 

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

350 # Docstring inherited. 

351 if not inputRefs.inputWarps: 

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

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

354 

355 # Construct skyInfo expected by run 

356 # Do not remove skyMap from inputData in case _makeSupplementaryData 

357 # needs it 

358 skyMap = butlerQC.get(inputRefs.skyMap) 

359 

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

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

362 

363 outputDataId = butlerQC.quantum.dataId 

364 

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

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

367 

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

369 self.common = CommonComponents( 

370 units=units, 

371 wcs=skyInfo.patchInfo.wcs, 

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

373 identifiers=PatchIdentifiers.from_data_id(outputDataId), 

374 ) 

375 

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

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

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

379 

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

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

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

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

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

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

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

387 

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

389 butlerQC.put(returnStruct, outputRefs) 

390 return returnStruct 

391 

392 @staticmethod 

393 def _compute_weight(maskedImage, statsCtrl): 

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

395 

396 Parameters 

397 ---------- 

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

399 The masked image to compute the weight. 

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

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

402 

403 Returns 

404 ------- 

405 weight : `float` 

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

407 """ 

408 statObj = afwMath.makeStatistics( 

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

410 ) 

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

412 weight = 1.0 / float(meanVar) 

413 return weight 

414 

415 @staticmethod 

416 def _construct_grid(skyInfo): 

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

418 

419 Parameters 

420 ---------- 

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

422 A Struct object 

423 

424 Returns 

425 ------- 

426 grid : `~lsst.cell_coadds.UniformGrid` 

427 A UniformGrid object. 

428 """ 

429 padding = skyInfo.patchInfo.getCellBorder() 

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

431 grid = UniformGrid.from_bbox_cell_size( 

432 grid_bbox, 

433 skyInfo.patchInfo.getCellInnerDimensions(), 

434 padding=padding, 

435 ) 

436 return grid 

437 

438 def _construct_grid_container(self, skyInfo, statsCtrl): 

439 """Construct a grid of AccumulatorMeanStack instances. 

440 

441 Parameters 

442 ---------- 

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

444 A Struct object 

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

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

447 

448 Returns 

449 ------- 

450 gc : `~lsst.cell_coadds.GridContainer` 

451 A GridContainer object container one AccumulatorMeanStack per cell. 

452 """ 

453 grid = self._construct_grid(skyInfo) 

454 

455 maskMap = setRejectedMaskMapping(statsCtrl) 

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

457 thresholdDict = AccumulatorMeanStack.stats_ctrl_to_threshold_dict(statsCtrl) 

458 

459 # Initialize the grid container with AccumulatorMeanStacks 

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

461 for cellInfo in skyInfo.patchInfo: 

462 stacker = AccumulatorMeanStack( 

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

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

465 bit_mask_value=statsCtrl.getAndMask(), 

466 mask_threshold_dict=thresholdDict, 

467 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

468 compute_n_image=False, 

469 mask_map=maskMap, 

470 no_good_pixels_mask=statsCtrl.getNoGoodPixelsMask(), 

471 ) 

472 gc[cellInfo.index] = stacker 

473 

474 return gc 

475 

476 def _construct_stats_control(self): 

477 """Construct a StatisticsControl object for coadd. 

478 

479 Unlike AssembleCoaddTask or CompareWarpAssembleCoaddTask, there is 

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

481 optionally mask propagation thresholds. 

482 

483 Returns 

484 ------- 

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

486 A control object for StatisticsStack. 

487 """ 

488 statsCtrl = afwMath.StatisticsControl() 

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

490 # CompareWarpAssembleCoaddTask to get consistent weights. This is NOT 

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

492 # fallback option that is not recommended for production. 

493 statsCtrl.setNumIter(2) 

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

495 statsCtrl.setNanSafe(True) 

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

497 bit = afwImage.Mask.getMaskPlane(plane) 

498 statsCtrl.setMaskPropagationThreshold(bit, threshold) 

499 return statsCtrl 

500 

501 def _construct_ap_corr_grid_container(self, skyInfo): 

502 """Construct a grid of CoaddApCorrMapStacker instances. 

503 

504 Parameters 

505 ---------- 

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

507 A Struct object 

508 

509 Returns 

510 ------- 

511 gc : `~lsst.cell_coadds.GridContainer` 

512 A GridContainer object container one CoaddApCorrMapStacker per 

513 cell. 

514 """ 

515 grid = self._construct_grid(skyInfo) 

516 

517 # Initialize the grid container with CoaddApCorrMapStacker. 

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

519 for cellInfo in skyInfo.patchInfo: 

520 stacker = CoaddApCorrMapStacker( 

521 evaluation_point=cellInfo.inner_bbox.getCenter(), 

522 do_coadd_inverse_ap_corr=self.config.do_coadd_inverse_aperture_corrections, 

523 ) 

524 gc[cellInfo.index] = stacker 

525 

526 return gc 

527 

528 def run( 

529 self, 

530 *, 

531 inputs: dict[DataCoordinate, WarpInputs], 

532 skyInfo, 

533 visitSummaryList: list | None = None, 

534 ): 

535 for mask_plane in self.config.bad_mask_planes: 

536 afwImage.Mask.addMaskPlane(mask_plane) 

537 for mask_plane in self.config.mask_propagation_thresholds: 

538 afwImage.Mask.addMaskPlane(mask_plane) 

539 

540 statsCtrl = self._construct_stats_control() 

541 

542 warp_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

543 maskfrac_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

544 noise_stacker_gc_list = [ 

545 self._construct_grid_container(skyInfo, statsCtrl) 

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

547 ] 

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

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

550 ap_corr_stacker_gc = self._construct_ap_corr_grid_container(skyInfo) 

551 

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

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

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

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

556 # get the corresponding PSFs as well. 

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

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

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

560 # Populate them. 

561 for cellInfo in skyInfo.patchInfo: 

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

563 observation_identifiers_gc[cellInfo.index] = {} 

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

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

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

567 cell_center_pixel, 

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

569 ) 

570 psf_stacker_gc[cellInfo.index] = AccumulatorMeanStack( 

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

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

573 bit_mask_value=0, 

574 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

575 compute_n_image=False, 

576 ) 

577 

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

579 if not visitSummaryList: 

580 visitSummaryList = [] 

581 visitSummaryRefDict = { 

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

583 } 

584 

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

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

587 

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

589 # it completely overlaps. 

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

591 # warps that have been excluded from CompareWarp via visit 

592 # selection from SelectVisitsTasks will not have artifact masks. 

593 # Exclude them from the cell coadds too. 

594 if self.config.require_artifact_mask and warp_input.artifact_mask is None: 

595 self.log.info( 

596 "Excluding warp %s from cell coadds because it has no artifact mask", 

597 warp_input.dataId["visit"], 

598 ) 

599 continue 

600 

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

602 masked_fraction_image = ( 

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

604 if warp_input.masked_fraction 

605 else None 

606 ) 

607 

608 # Pre-process the warp before coadding. 

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

610 warp.mask.addMaskPlane("CLIPPED") 

611 warp.mask.addMaskPlane("REJECTED") 

612 warp.mask.addMaskPlane("SENSOR_EDGE") 

613 warp.mask.addMaskPlane("INEXACT_PSF") 

614 

615 if artifact_mask_ref := warp_input.artifact_mask: 

616 # Apply the artifact mask to the warp. 

617 artifact_mask = artifact_mask_ref.get() 

618 assert ( 

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

620 ), "Mask dicts do not agree." 

621 warp.mask.array = artifact_mask.array 

622 del artifact_mask 

623 

624 if self.config.do_scale_zero_point: 

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

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

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

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

629 zero_point_scale_factor = imageScaler.scale 

630 self.log.debug( 

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

632 warp_input.dataId, 

633 zero_point_scale_factor, 

634 ) 

635 else: 

636 zero_point_scale_factor = 1.0 

637 if "BUNIT" not in warp.metadata: 

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

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

640 raise ValueError( 

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

642 ) 

643 

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

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

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

647 rejected = afwImage.Mask.getPlaneBitMask( 

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

649 ) 

650 

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

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

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

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

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

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

657 full_ccd_table["ccd"], 

658 0.0, 

659 ) # Mapping from detector to weight. 

660 

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

662 visitSummary = visitSummaryRef.get() 

663 for detector in full_ccd_table["ccd"]: 

664 visitSummaryRow = visitSummary.find(detector) 

665 mean_variance = visitSummaryRow["meanVar"] 

666 mean_variance *= zero_point_scale_factor**2 

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

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

669 weights[detector] = 1.0 / mean_variance 

670 del visitSummary 

671 else: 

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

673 weight = self._compute_weight(warp, statsCtrl) 

674 if not np.isfinite(weight): 

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

676 continue 

677 

678 for detector in weights: 

679 weights[detector] = weight 

680 

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

682 

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

684 # detector ID that pixel comes from. 

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

686 for row in full_ccd_table: 

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

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

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

690 try: 

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

692 geom.Box2D(warp.getBBox()) 

693 ) 

694 except SinglePolygonException: 

695 continue 

696 

697 observation_identifier = ObservationIdentifiers.from_data_id( 

698 warp_input.dataId, 

699 backup_detector=row["ccd"], 

700 ) 

701 visit_polygons[observation_identifier] = dest_polygon 

702 

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

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

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

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

707 

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

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

710 detector_map.array[:, :] = 0 

711 

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

713 bbox = cellInfo.outer_bbox 

714 inner_bbox = cellInfo.inner_bbox 

715 

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

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

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

719 self.log.debug( 

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

721 warp_input.dataId, 

722 cellInfo.index, 

723 overlap_fraction, 

724 self.config.min_overlap_fraction, 

725 ) 

726 continue 

727 

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

729 if not np.isfinite(weight): 

730 self.log.warn( 

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

732 ) 

733 continue 

734 

735 if weight == 0: 

736 self.log.info( 

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

738 ) 

739 continue 

740 

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

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

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

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

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

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

747 if deep_copy: 

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

749 

750 mi = afwImage.MaskedImageF(warp[bbox].maskedImage, 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 nodata_or_mask = (single_detector_mask_array) * afwImage.Mask.getPlaneBitMask("NO_DATA") 

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

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

757 

758 if masked_fraction_image: 

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

760 if deep_copy: 

761 mi.array[single_detector_mask_array] = 0.0 

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

763 

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

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

766 if deep_copy: 

767 mi.image.array[single_detector_mask_array] = 0.0 

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

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

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

771 

772 # Set the defaults for PSF shape quantities. 

773 psf_shape = afwGeom.Quadrupole() 

774 psf_shape_flag = True 

775 psf_eval_point = None 

776 try: 

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

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

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

780 # the rug for now. 

781 if overlap_fraction < 0.5: 

782 psf_eval_point = dest_polygon.intersectionSingle( 

783 geom.Box2D(inner_bbox) 

784 ).calculateCenter() 

785 else: 

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

787 psf_shape = warp.psf.computeShape(psf_eval_point) 

788 psf_shape_flag = False 

789 except SinglePolygonException: 

790 self.log.info( 

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

792 ccd_row["ccd"], 

793 warp_input.dataId, 

794 cellInfo.index, 

795 ) 

796 except InvalidPsfError: 

797 self.log.info( 

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

799 ccd_row["ccd"], 

800 warp_input.dataId, 

801 psf_eval_point, 

802 ) 

803 

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

805 

806 observation_identifier = ObservationIdentifiers.from_data_id( 

807 warp_input.dataId, 

808 backup_detector=ccd_row["ccd"], 

809 ) 

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

811 overlaps_center=overlaps_center, 

812 overlap_fraction=overlap_fraction, 

813 weight=weight, 

814 psf_shape=psf_shape, 

815 psf_shape_flag=psf_shape_flag, 

816 ) 

817 if overlaps_center is False: 

818 self.log.debug( 

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

820 warp_input.dataId, 

821 cellInfo.index, 

822 ) 

823 continue 

824 

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

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

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

828 

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

830 calexp_point, 

831 undistorted_psf_im.getDimensions(), 

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

833 

834 # Convert the PSF image from Image to MaskedImage and 

835 # zero-pad the image. 

836 undistorted_psf_bbox = undistorted_psf_im.getBBox() 

837 undistorted_psf_maskedImage = afwImage.MaskedImageD( 

838 undistorted_psf_bbox.dilatedBy(self.psf_padding) 

839 ) 

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