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

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.healSparseMapping import HealSparseInputMapTask 

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 inputMap = Output( 

158 doc="Output healsparse map of input images", 

159 name="{inputWarpName}Coadd_inputMap", 

160 storageClass="HealSparseMap", 

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

162 ) 

163 

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

165 super().__init__(config=config) 

166 

167 if not config: 

168 return 

169 

170 if config.do_calculate_weight_from_warp: 

171 del self.visitSummaryList 

172 

173 if not config.do_use_artifact_mask: 

174 del self.artifactMasks 

175 

176 if not config.do_input_map: 

177 del self.inputMap 

178 

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

180 # number of noise realizations specified in the config. 

181 for n in range(config.num_noise_realizations): 

182 noise_warps = Input( 

183 doc="Input noise warps", 

184 name=f"direct_warp_noise{n}", 

185 storageClass="MaskedImageF", 

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

187 deferLoad=True, 

188 multiple=True, 

189 ) 

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

191 

192 

193class AssembleCellCoaddConfig(PipelineTaskConfig, pipelineConnections=AssembleCellCoaddConnections): 

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

195 interpolate_coadd = ConfigurableField( 

196 target=InterpImageTask, 

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

198 ) 

199 do_scale_zero_point = Field[bool]( 

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

201 default=False, 

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

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

204 "after v29.", 

205 ) 

206 scale_zero_point = ConfigurableField( 

207 target=ScaleZeroPointTask, 

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

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

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

211 "after v29.", 

212 ) 

213 do_calculate_weight_from_warp = Field[bool]( 

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

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

216 default=False, 

217 ) 

218 do_use_artifact_mask = Field[bool]( 

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

220 default=True, 

221 ) 

222 do_coadd_inverse_aperture_corrections = Field[bool]( 

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

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

225 default=False, 

226 ) 

227 min_overlap_fraction = RangeField[float]( 

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

229 "cell.", 

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

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

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

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

234 # definitely not overlap the cell center. 

235 default=1.0, 

236 min=0.0, 

237 max=1.0, 

238 inclusiveMin=True, 

239 inclusiveMax=True, 

240 ) 

241 bad_mask_planes = ListField[str]( 

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

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

244 ) 

245 remove_mask_planes = ListField[str]( 

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

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

248 ) 

249 calc_error_from_input_variance = Field[bool]( 

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

251 "statistic. Passed to AccumulatorMeanStack.", 

252 default=True, 

253 ) 

254 mask_propagation_thresholds = DictField[str, float]( 

255 doc=( 

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

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

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

259 "would have contributed exceeds this value." 

260 ), 

261 default={"SAT": 0.1}, 

262 ) 

263 max_maskfrac = RangeField[float]( 

264 doc="Maximum fraction of masked pixels in a cell for a given warp. " 

265 "Warps exceeding this threshold are excluded from the science coadd, " 

266 "PSF, aperture corrections, and input maps.", 

267 default=0.5, 

268 min=0.0, 

269 max=1.0, 

270 inclusiveMin=True, 

271 inclusiveMax=False, 

272 ) 

273 num_noise_realizations = Field[int]( 

274 default=0, 

275 doc=( 

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

277 "This should not exceed the corresponding config parameter " 

278 "specified in `MakeDirectWarpConfig`. " 

279 ), 

280 check=lambda x: x >= 0, 

281 ) 

282 psf_warper = ConfigField( 

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

284 "warp the images.", 

285 dtype=afwMath.Warper.ConfigClass, 

286 ) 

287 psf_dimensions = Field[int]( 

288 default=35, 

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

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

291 ) 

292 require_artifact_mask = Field[bool]( 

293 default=True, 

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

295 " from CompareWarpTask", 

296 ) 

297 do_input_map = Field[bool]( 

298 default=False, 

299 doc="Create a bitwise map of coadd inputs.", 

300 ) 

301 input_mapper = ConfigurableField( 

302 target=HealSparseInputMapTask, 

303 doc="Input map creation subtask.", 

304 ) 

305 

306 

307class AssembleCellCoaddTask(PipelineTask): 

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

309 

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

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

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

313 

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

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

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

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

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

319 with the same weights. 

320 

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

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

323 assumed to be spatially constant within a cell. 

324 

325 Raises 

326 ------ 

327 NoWorkFound 

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

329 RuntimeError 

330 Raised if the skymap is not cell-based. 

331 

332 Notes 

333 ----- 

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

335 especially its Config and Connections are experimental and subject to 

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

337 until it is included in the DRP pipeline. 

338 """ 

339 

340 ConfigClass = AssembleCellCoaddConfig 

341 _DefaultName = "assembleCellCoadd" 

342 

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

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

345 if self.config.do_interpolate_coadd: 

346 self.makeSubtask("interpolate_coadd") 

347 if self.config.do_scale_zero_point: 

348 self.makeSubtask("scale_zero_point") 

349 if self.config.do_input_map: 

350 self.makeSubtask("input_mapper") 

351 

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

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

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

355 self.log.debug( 

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

357 psf_padding, 

358 self.config.psf_warper.warpingKernelName, 

359 ) 

360 else: 

361 psf_padding = 10 

362 self.log.info( 

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

364 psf_padding, 

365 ) 

366 self.psf_padding = psf_padding 

367 

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

369 # Docstring inherited. 

370 if not inputRefs.inputWarps: 

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

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

373 

374 # Construct skyInfo expected by run 

375 # Do not remove skyMap from inputData in case _makeSupplementaryData 

376 # needs it 

377 skyMap = butlerQC.get(inputRefs.skyMap) 

378 

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

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

381 

382 outputDataId = butlerQC.quantum.dataId 

383 

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

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

386 

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

388 self.common = CommonComponents( 

389 units=units, 

390 wcs=skyInfo.patchInfo.wcs, 

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

392 identifiers=PatchIdentifiers.from_data_id(outputDataId), 

393 ) 

394 

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

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

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

398 

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

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

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

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

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

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

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

406 

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

408 butlerQC.put(returnStruct, outputRefs) 

409 return returnStruct 

410 

411 @staticmethod 

412 def _compute_weight(maskedImage, statsCtrl): 

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

414 

415 Parameters 

416 ---------- 

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

418 The masked image to compute the weight. 

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

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

421 

422 Returns 

423 ------- 

424 weight : `float` 

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

426 """ 

427 statObj = afwMath.makeStatistics( 

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

429 ) 

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

431 weight = 1.0 / float(meanVar) 

432 return weight 

433 

434 @staticmethod 

435 def _construct_grid(skyInfo): 

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

437 

438 Parameters 

439 ---------- 

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

441 A Struct object 

442 

443 Returns 

444 ------- 

445 grid : `~lsst.cell_coadds.UniformGrid` 

446 A UniformGrid object. 

447 """ 

448 padding = skyInfo.patchInfo.getCellBorder() 

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

450 grid = UniformGrid.from_bbox_cell_size( 

451 grid_bbox, 

452 skyInfo.patchInfo.getCellInnerDimensions(), 

453 padding=padding, 

454 ) 

455 return grid 

456 

457 def _construct_grid_container(self, skyInfo, statsCtrl): 

458 """Construct a grid of AccumulatorMeanStack instances. 

459 

460 Parameters 

461 ---------- 

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

463 A Struct object 

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

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

466 

467 Returns 

468 ------- 

469 gc : `~lsst.cell_coadds.GridContainer` 

470 A GridContainer object container one AccumulatorMeanStack per cell. 

471 """ 

472 grid = self._construct_grid(skyInfo) 

473 

474 maskMap = setRejectedMaskMapping(statsCtrl) 

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

476 thresholdDict = AccumulatorMeanStack.stats_ctrl_to_threshold_dict(statsCtrl) 

477 

478 # Initialize the grid container with AccumulatorMeanStacks 

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

480 for cellInfo in skyInfo.patchInfo: 

481 stacker = AccumulatorMeanStack( 

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

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

484 bit_mask_value=statsCtrl.getAndMask(), 

485 mask_threshold_dict=thresholdDict, 

486 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

487 compute_n_image=False, 

488 mask_map=maskMap, 

489 no_good_pixels_mask=statsCtrl.getNoGoodPixelsMask(), 

490 ) 

491 gc[cellInfo.index] = stacker 

492 

493 return gc 

494 

495 def _construct_stats_control(self): 

496 """Construct a StatisticsControl object for coadd. 

497 

498 Unlike AssembleCoaddTask or CompareWarpAssembleCoaddTask, there is 

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

500 optionally mask propagation thresholds. 

501 

502 Returns 

503 ------- 

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

505 A control object for StatisticsStack. 

506 """ 

507 statsCtrl = afwMath.StatisticsControl() 

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

509 # CompareWarpAssembleCoaddTask to get consistent weights. This is NOT 

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

511 # fallback option that is not recommended for production. 

512 statsCtrl.setNumIter(2) 

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

514 statsCtrl.setNanSafe(True) 

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

516 bit = afwImage.Mask.getMaskPlane(plane) 

517 statsCtrl.setMaskPropagationThreshold(bit, threshold) 

518 return statsCtrl 

519 

520 def _construct_ap_corr_grid_container(self, skyInfo): 

521 """Construct a grid of CoaddApCorrMapStacker instances. 

522 

523 Parameters 

524 ---------- 

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

526 A Struct object 

527 

528 Returns 

529 ------- 

530 gc : `~lsst.cell_coadds.GridContainer` 

531 A GridContainer object container one CoaddApCorrMapStacker per 

532 cell. 

533 """ 

534 grid = self._construct_grid(skyInfo) 

535 

536 # Initialize the grid container with CoaddApCorrMapStacker. 

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

538 for cellInfo in skyInfo.patchInfo: 

539 stacker = CoaddApCorrMapStacker( 

540 evaluation_point=cellInfo.inner_bbox.getCenter(), 

541 do_coadd_inverse_ap_corr=self.config.do_coadd_inverse_aperture_corrections, 

542 ) 

543 gc[cellInfo.index] = stacker 

544 

545 return gc 

546 

547 def run( 

548 self, 

549 *, 

550 inputs: dict[DataCoordinate, WarpInputs], 

551 skyInfo, 

552 visitSummaryList: list | None = None, 

553 ): 

554 for mask_plane in self.config.bad_mask_planes: 

555 afwImage.Mask.addMaskPlane(mask_plane) 

556 for mask_plane in self.config.mask_propagation_thresholds: 

557 afwImage.Mask.addMaskPlane(mask_plane) 

558 

559 statsCtrl = self._construct_stats_control() 

560 

561 warp_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

562 maskfrac_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

563 noise_stacker_gc_list = [ 

564 self._construct_grid_container(skyInfo, statsCtrl) 

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

566 ] 

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

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

569 ap_corr_stacker_gc = self._construct_ap_corr_grid_container(skyInfo) 

570 

571 # A cell is in "fallback" mode if it does not yet have any warps that 

572 # pass the per-detector cuts; in that mode, we accumulate warps 

573 # regardless of that cut, but clear the accumulators and start over if 

574 # we later see data that does pass the per-detector cuts. 

575 is_fallback_gc = GridContainer[bool](warp_stacker_gc.shape) 

576 

577 # We accumulate the information to pass to the Healsparse input-map 

578 # accumulator instead of calling it directly, so we can do that only 

579 # after we've accumulated all warps and hence know which cells will 

580 # stay in fallback mode. 

581 input_map_data_gc = GridContainer[list](warp_stacker_gc.shape) 

582 

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

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

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

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

587 # get the corresponding PSFs as well. 

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

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

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

591 

592 if self.config.do_input_map: 

593 # We need to know all the visit + detector pairs in the inputs. 

594 warp_input_list = [warp_ref.warp.get(component="coaddInputs") for warp_ref in inputs.values()] 

595 visit_detectors = [] 

596 for warp_input in warp_input_list: 

597 for row in warp_input.ccds: 

598 visit_detectors.append((int(row["visit"]), int(row["ccd"]))) 

599 

600 self.input_mapper.initialize_cell_input_map( 

601 skyInfo.patchInfo.getOuterBBox(), 

602 skyInfo.patchInfo.wcs, 

603 visit_detectors, 

604 ) 

605 

606 # Populate them. 

607 for cellInfo in skyInfo.patchInfo: 

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

609 observation_identifiers_gc[cellInfo.index] = {} 

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

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

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

613 cell_center_pixel, 

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

615 ) 

616 psf_stacker_gc[cellInfo.index] = AccumulatorMeanStack( 

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

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

619 bit_mask_value=0, 

620 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

621 compute_n_image=False, 

622 ) 

623 is_fallback_gc[cellInfo.index] = True 

624 input_map_data_gc[cellInfo.index] = [] 

625 

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

627 if not visitSummaryList: 

628 visitSummaryList = [] 

629 visitSummaryRefDict = { 

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

631 } 

632 

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

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

635 

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

637 # it completely overlaps. 

638 for warp_input in inputs.values(): 

639 # warps that have been excluded from CompareWarp via visit 

640 # selection from SelectVisitsTasks will not have artifact masks. 

641 # Exclude them from the cell coadds too. 

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

643 self.log.info( 

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

645 warp_input.dataId["visit"], 

646 ) 

647 continue 

648 

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

650 masked_fraction_image = ( 

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

652 if warp_input.masked_fraction 

653 else None 

654 ) 

655 

656 # Pre-process the warp before coadding. 

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

658 warp.mask.addMaskPlane("CLIPPED") 

659 warp.mask.addMaskPlane("REJECTED") 

660 warp.mask.addMaskPlane("SENSOR_EDGE") 

661 warp.mask.addMaskPlane("INEXACT_PSF") 

662 

663 if artifact_mask_ref := warp_input.artifact_mask: 

664 # Apply the artifact mask to the warp. 

665 artifact_mask = artifact_mask_ref.get() 

666 assert ( 

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

668 ), "Mask dicts do not agree." 

669 warp.mask.array = artifact_mask.array 

670 del artifact_mask 

671 

672 if self.config.do_scale_zero_point: 

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

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

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

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

677 zero_point_scale_factor = imageScaler.scale 

678 self.log.debug( 

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

680 warp_input.dataId, 

681 zero_point_scale_factor, 

682 ) 

683 else: 

684 zero_point_scale_factor = 1.0 

685 if "BUNIT" not in warp.metadata: 

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

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

688 raise ValueError( 

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

690 ) 

691 

692 # Only try to remove maks planes that have been registered. 

693 to_remove = [] 

694 for plane in self.config.remove_mask_planes: 

695 if plane in warp.mask.getMaskPlaneDict(): 

696 to_remove.append(plane) 

697 removeMaskPlanes(warp.mask, to_remove, self.log) 

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

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

700 rejected = afwImage.Mask.getPlaneBitMask( 

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

702 ) 

703 

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

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

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

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

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

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

710 full_ccd_table["ccd"].tolist(), 

711 0.0, 

712 ) # Mapping from detector to weight. 

713 

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

715 visitSummary = visitSummaryRef.get() 

716 for detector in full_ccd_table["ccd"].tolist(): 

717 visitSummaryRow = visitSummary.find(detector) 

718 mean_variance = visitSummaryRow["meanVar"] 

719 mean_variance *= zero_point_scale_factor**2 

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

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

722 weights[detector] = 1.0 / mean_variance 

723 del visitSummary 

724 else: 

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

726 weight = self._compute_weight(warp, statsCtrl) 

727 if not np.isfinite(weight): 

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

729 continue 

730 

731 for detector in weights: 

732 weights[detector] = weight 

733 

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

735 

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

737 # detector ID that pixel comes from. 

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

739 for row in full_ccd_table: 

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

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

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

743 try: 

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

745 geom.Box2D(warp.getBBox()) 

746 ) 

747 except SinglePolygonException: 

748 continue 

749 

750 observation_identifier = ObservationIdentifiers.from_data_id( 

751 warp_input.dataId, 

752 backup_detector=row["ccd"], 

753 ) 

754 visit_polygons[observation_identifier] = dest_polygon 

755 

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

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

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

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

760 

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

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

763 detector_map.array[:, :] = 0 

764 

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

766 bbox = cellInfo.outer_bbox 

767 inner_bbox = cellInfo.inner_bbox 

768 

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

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

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

772 self.log.debug( 

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

774 warp_input.dataId, 

775 cellInfo.index, 

776 overlap_fraction, 

777 self.config.min_overlap_fraction, 

778 ) 

779 continue 

780 

781 weight = weights[int(ccd_row["ccd"])] 

782 if not np.isfinite(weight): 

783 self.log.warning( 

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

785 ) 

786 continue 

787 

788 if weight == 0: 

789 self.log.info( 

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

791 ) 

792 continue 

793 

794 # Compute the unmasked fraction for this detector in the inner 

795 # cell. Used to gate on max_maskfrac. 

796 inner_detector_pixels = detector_map[inner_bbox].array == ccd_row["ccd"] 

797 inner_unmasked_pixels = (warp[inner_bbox].mask.array & rejected) == 0 

798 unmasked_fraction = ( 

799 inner_detector_pixels & inner_unmasked_pixels 

800 ).sum() / inner_detector_pixels.sum() 

801 is_fallback = is_fallback_gc[cellInfo.index] 

802 if unmasked_fraction <= max(1.0 - self.config.max_maskfrac, 0.0): 

803 if not is_fallback: 

804 # We already have good data in this cell, so we don't 

805 # want this heavily masked warp - it will add too much 

806 # INEXACT_PSF. 

807 self.log.debug( 

808 "Skipping %s in cell %s: masked fraction %.3f exceeds threshold %.3f", 

809 warp_input.dataId, 

810 cellInfo.index, 

811 1.0 - unmasked_fraction, 

812 self.config.max_maskfrac, 

813 ) 

814 continue 

815 else: 

816 self.log.debug( 

817 "Including %s in cell %s only as potential fallback: " 

818 "masked fraction %.3f exceeds threshold %.3f", 

819 warp_input.dataId, 

820 cellInfo.index, 

821 1.0 - unmasked_fraction, 

822 self.config.max_maskfrac, 

823 ) 

824 elif is_fallback: 

825 # This is the first good data we've gotten for this cell; 

826 # wipe out the fallback coadd we've been accumulating so 

827 # far, so we can start fresh. 

828 warp_stacker_gc[cellInfo.index].reset() 

829 maskfrac_stacker_gc[cellInfo.index].reset() 

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

831 noise_stacker_gc_list[n][cellInfo.index].reset() 

832 psf_stacker_gc[cellInfo.index].reset() 

833 ap_corr_stacker_gc[cellInfo.index].reset() 

834 observation_identifiers_gc[cellInfo.index].clear() 

835 input_map_data_gc[cellInfo.index].clear() 

836 is_fallback_gc[cellInfo.index] = False 

837 

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

839 if not overlaps_center: 

840 self.log.debug( 

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

842 warp_input.dataId, 

843 cellInfo.index, 

844 )