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

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"{config.connections.inputWarpName}Coadd_directWarp_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=0.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=["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=21, 

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 

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

328 # Docstring inherited. 

329 if not inputRefs.inputWarps: 

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

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

332 

333 # Construct skyInfo expected by run 

334 # Do not remove skyMap from inputData in case _makeSupplementaryData 

335 # needs it 

336 skyMap = butlerQC.get(inputRefs.skyMap) 

337 

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

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

340 

341 outputDataId = butlerQC.quantum.dataId 

342 

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

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

345 

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

347 self.common = CommonComponents( 

348 units=units, 

349 wcs=skyInfo.patchInfo.wcs, 

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

351 identifiers=PatchIdentifiers.from_data_id(outputDataId), 

352 ) 

353 

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

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

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

357 

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

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

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

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

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

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

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

365 

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

367 butlerQC.put(returnStruct, outputRefs) 

368 return returnStruct 

369 

370 @staticmethod 

371 def _compute_weight(maskedImage, statsCtrl): 

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

373 

374 Parameters 

375 ---------- 

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

377 The masked image to compute the weight. 

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

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

380 

381 Returns 

382 ------- 

383 weight : `float` 

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

385 """ 

386 statObj = afwMath.makeStatistics( 

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

388 ) 

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

390 weight = 1.0 / float(meanVar) 

391 return weight 

392 

393 @staticmethod 

394 def _construct_grid(skyInfo): 

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

396 

397 Parameters 

398 ---------- 

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

400 A Struct object 

401 

402 Returns 

403 ------- 

404 grid : `~lsst.cell_coadds.UniformGrid` 

405 A UniformGrid object. 

406 """ 

407 padding = skyInfo.patchInfo.getCellBorder() 

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

409 grid = UniformGrid.from_bbox_cell_size( 

410 grid_bbox, 

411 skyInfo.patchInfo.getCellInnerDimensions(), 

412 padding=padding, 

413 ) 

414 return grid 

415 

416 def _construct_grid_container(self, skyInfo, statsCtrl): 

417 """Construct a grid of AccumulatorMeanStack instances. 

418 

419 Parameters 

420 ---------- 

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

422 A Struct object 

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

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

425 

426 Returns 

427 ------- 

428 gc : `~lsst.cell_coadds.GridContainer` 

429 A GridContainer object container one AccumulatorMeanStack per cell. 

430 """ 

431 grid = self._construct_grid(skyInfo) 

432 

433 maskMap = setRejectedMaskMapping(statsCtrl) 

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

435 thresholdDict = AccumulatorMeanStack.stats_ctrl_to_threshold_dict(statsCtrl) 

436 

437 # Initialize the grid container with AccumulatorMeanStacks 

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

439 for cellInfo in skyInfo.patchInfo: 

440 stacker = AccumulatorMeanStack( 

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

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

443 bit_mask_value=statsCtrl.getAndMask(), 

444 mask_threshold_dict=thresholdDict, 

445 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

446 compute_n_image=False, 

447 mask_map=maskMap, 

448 no_good_pixels_mask=statsCtrl.getNoGoodPixelsMask(), 

449 ) 

450 gc[cellInfo.index] = stacker 

451 

452 return gc 

453 

454 def _construct_stats_control(self): 

455 """Construct a StatisticsControl object for coadd. 

456 

457 Unlike AssembleCoaddTask or CompareWarpAssembleCoaddTask, there is 

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

459 optionally mask propagation thresholds. 

460 

461 Returns 

462 ------- 

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

464 A control object for StatisticsStack. 

465 """ 

466 statsCtrl = afwMath.StatisticsControl() 

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

468 # CompareWarpAssembleCoaddTask to get consistent weights. This is NOT 

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

470 # fallback option that is not recommended for production. 

471 statsCtrl.setNumIter(2) 

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

473 statsCtrl.setNanSafe(True) 

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

475 bit = afwImage.Mask.getMaskPlane(plane) 

476 statsCtrl.setMaskPropagationThreshold(bit, threshold) 

477 return statsCtrl 

478 

479 def _construct_ap_corr_grid_container(self, skyInfo): 

480 """Construct a grid of CoaddApCorrMapStacker instances. 

481 

482 Parameters 

483 ---------- 

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

485 A Struct object 

486 

487 Returns 

488 ------- 

489 gc : `~lsst.cell_coadds.GridContainer` 

490 A GridContainer object container one CoaddApCorrMapStacker per 

491 cell. 

492 """ 

493 grid = self._construct_grid(skyInfo) 

494 

495 # Initialize the grid container with CoaddApCorrMapStacker. 

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

497 for cellInfo in skyInfo.patchInfo: 

498 stacker = CoaddApCorrMapStacker( 

499 evaluation_point=cellInfo.inner_bbox.getCenter(), 

500 do_coadd_inverse_ap_corr=self.config.do_coadd_inverse_aperture_corrections, 

501 ) 

502 gc[cellInfo.index] = stacker 

503 

504 return gc 

505 

506 def run( 

507 self, 

508 *, 

509 inputs: dict[DataCoordinate, WarpInputs], 

510 skyInfo, 

511 visitSummaryList: list | None = None, 

512 ): 

513 for mask_plane in self.config.bad_mask_planes: 

514 afwImage.Mask.addMaskPlane(mask_plane) 

515 for mask_plane in self.config.mask_propagation_thresholds: 

516 afwImage.Mask.addMaskPlane(mask_plane) 

517 

518 statsCtrl = self._construct_stats_control() 

519 

520 warp_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

521 maskfrac_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

522 noise_stacker_gc_list = [ 

523 self._construct_grid_container(skyInfo, statsCtrl) 

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

525 ] 

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

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

528 ap_corr_stacker_gc = self._construct_ap_corr_grid_container(skyInfo) 

529 

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

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

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

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

534 # get the corresponding PSFs as well. 

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

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

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

538 # Populate them. 

539 for cellInfo in skyInfo.patchInfo: 

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

541 observation_identifiers_gc[cellInfo.index] = {} 

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

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

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

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

546 ) 

547 psf_stacker_gc[cellInfo.index] = AccumulatorMeanStack( 

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

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

550 bit_mask_value=0, 

551 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

552 compute_n_image=False, 

553 ) 

554 

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

556 if not visitSummaryList: 

557 visitSummaryList = [] 

558 visitSummaryRefDict = { 

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

560 } 

561 

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

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

564 

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

566 # it completely overlaps. 

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

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

569 masked_fraction_image = ( 

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

571 if warp_input.masked_fraction 

572 else None 

573 ) 

574 

575 # Pre-process the warp before coadding. 

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

577 warp.mask.addMaskPlane("CLIPPED") 

578 warp.mask.addMaskPlane("REJECTED") 

579 warp.mask.addMaskPlane("SENSOR_EDGE") 

580 warp.mask.addMaskPlane("INEXACT_PSF") 

581 

582 if artifact_mask_ref := warp_input.artifact_mask: 

583 # Apply the artifact mask to the warp. 

584 artifact_mask = artifact_mask_ref.get() 

585 assert ( 

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

587 ), "Mask dicts do not agree." 

588 warp.mask.array = artifact_mask.array 

589 del artifact_mask 

590 

591 if self.config.do_scale_zero_point: 

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

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

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

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

596 zero_point_scale_factor = imageScaler.scale 

597 self.log.debug( 

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

599 warp_input.dataId, 

600 zero_point_scale_factor, 

601 ) 

602 else: 

603 zero_point_scale_factor = 1.0 

604 if "BUNIT" not in warp.metadata: 

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

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

607 raise ValueError( 

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

609 ) 

610 

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

612 

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

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

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

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

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

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

619 full_ccd_table["ccd"], 

620 0.0, 

621 ) # Mapping from detector to weight. 

622 

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

624 visitSummary = visitSummaryRef.get() 

625 for detector in full_ccd_table["ccd"]: 

626 visitSummaryRow = visitSummary.find(detector) 

627 mean_variance = visitSummaryRow["meanVar"] 

628 mean_variance *= zero_point_scale_factor**2 

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

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

631 weights[detector] = 1.0 / mean_variance 

632 del visitSummary 

633 else: 

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

635 weight = self._compute_weight(warp, statsCtrl) 

636 if not np.isfinite(weight): 

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

638 continue 

639 

640 for detector in weights: 

641 weights[detector] = weight 

642 

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

644 

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

646 # detector ID that pixel comes from. 

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

648 for row in full_ccd_table: 

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

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

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

652 try: 

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

654 geom.Box2D(warp.getBBox()) 

655 ) 

656 except SinglePolygonException: 

657 continue 

658 

659 observation_identifier = ObservationIdentifiers.from_data_id( 

660 warp_input.dataId, 

661 backup_detector=row["ccd"], 

662 ) 

663 visit_polygons[observation_identifier] = dest_polygon 

664 

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

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

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

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

669 

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

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

672 detector_map.array[:, :] = 0 

673 

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

675 bbox = cellInfo.outer_bbox 

676 inner_bbox = cellInfo.inner_bbox 

677 

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

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

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

681 self.log.log( 

682 self.log.DEBUG if overlap_fraction == 0.0 else self.log.INFO, 

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

684 warp_input.dataId, 

685 cellInfo.index, 

686 overlap_fraction, 

687 self.config.min_overlap_fraction, 

688 ) 

689 continue 

690 

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

692 if not np.isfinite(weight): 

693 self.log.warn( 

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

695 ) 

696 continue 

697 

698 if weight == 0: 

699 self.log.info( 

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

701 ) 

702 continue 

703 

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

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

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

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

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

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

710 if deep_copy: 

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

712 

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

714 if deep_copy: 

715 mi.image.array[single_detector_mask_array] = 0.0 

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

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

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

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

720 

721 if masked_fraction_image: 

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

723 if deep_copy: 

724 mi.array[single_detector_mask_array] = 0.0 

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

726 

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

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

729 if deep_copy: 

730 mi.image.array[single_detector_mask_array] = 0.0 

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

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

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

734 

735 # Set the defaults for PSF shape quantities. 

736 psf_shape = afwGeom.Quadrupole() 

737 psf_shape_flag = True 

738 psf_eval_point = None 

739 try: 

740 if overlap_fraction < 1.0: 

741 psf_eval_point = dest_polygon.intersectionSingle( 

742 geom.Box2D(inner_bbox) 

743 ).calculateCenter() 

744 else: 

745 psf_eval_point = geom.Point2D(inner_bbox.getCenter()) 

746 psf_shape = warp.psf.computeShape(psf_eval_point) 

747 psf_shape_flag = False 

748 except SinglePolygonException: 

749 self.log.info( 

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

751 ccd_row["ccd"], 

752 warp_input.dataId, 

753 cellInfo.index, 

754 ) 

755 except InvalidPsfError: 

756 self.log.info( 

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

758 ccd_row["ccd"], 

759 warp_input.dataId, 

760 psf_eval_point, 

761 ) 

762 

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

764 

765 observation_identifier = ObservationIdentifiers.from_data_id( 

766 warp_input.dataId, 

767 backup_detector=ccd_row["ccd"], 

768 ) 

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

770 overlaps_center=overlaps_center, 

771 overlap_fraction=overlap_fraction, 

772 weight=weight, 

773 psf_shape=psf_shape, 

774 psf_shape_flag=psf_shape_flag, 

775 ) 

776 if overlaps_center is False: 

777 self.log.debug( 

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

779 warp_input.dataId, 

780 cellInfo.index, 

781 ) 

782 continue 

783 

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

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

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

787 

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

789 calexp_point, 

790 undistorted_psf_im.getDimensions(), 

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

792 

793 # Convert the PSF image from Image to MaskedImage. 

794 undistorted_psf_maskedImage = afwImage.MaskedImageD(image=undistorted_psf_im) 

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

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

797 

798 warped_psf_maskedImage = self.psf_warper.warpImage( 

799 destWcs=skyInfo.wcs, 

800 srcImage=undistorted_psf_maskedImage, 

801 srcWcs=ccd_row.getWcs(), 

802 destBBox=psf_bbox_gc[cellInfo.index], 

803 ) 

804 

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

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

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

808 

809 psf_stacker = psf_stacker_gc[cellInfo.index] 

810 psf_stacker.add_masked_image(warped_psf_maskedImage, weight=weight) 

811 

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

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

814 

815 del warp 

816 

817 # Update common with the visit polygons. 

818 self.common = dataclasses.replace( 

819 self.common, 

820 visit_polygons=visit_polygons, 

821 ) 

822 

823 cells: list[SingleCellCoadd] = [] 

824 for cellInfo in skyInfo.patchInfo: 

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

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

827 continue 

828 

829 cell_masked_image = afwImage.MaskedImageF(cellInfo.outer_bbox) 

830 cell_maskfrac_image = afwImage.ImageF(cellInfo.outer_bbox) 

831 cell_noise_images = [ 

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