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

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 ( 

58 ChoiceField, 

59 ConfigField, 

60 ConfigurableField, 

61 DictField, 

62 Field, 

63 ListField, 

64 RangeField, 

65) 

66from lsst.pipe.base import ( 

67 InMemoryDatasetHandle, 

68 NoWorkFound, 

69 PipelineTask, 

70 PipelineTaskConfig, 

71 PipelineTaskConnections, 

72 Struct, 

73) 

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

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

76from lsst.pipe.tasks.healSparseMapping import HealSparseInputMapTask 

77from lsst.pipe.tasks.interpImage import InterpImageTask 

78from lsst.pipe.tasks.scaleZeroPoint import ScaleZeroPointTask 

79from lsst.skymap import BaseSkyMap 

80 

81 

82@dataclasses.dataclass 

83class WarpInputs: 

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

85 

86 warp: DeferredDatasetHandle | InMemoryDatasetHandle 

87 """Handle for the warped exposure.""" 

88 

89 masked_fraction: DeferredDatasetHandle | InMemoryDatasetHandle | None = None 

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

91 

92 artifact_mask: DeferredDatasetHandle | InMemoryDatasetHandle | None = None 

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

94 

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

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

97 

98 @property 

99 def dataId(self) -> DataCoordinate: 

100 """DataID corresponding to the warp. 

101 

102 Returns 

103 ------- 

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

105 DataID of the warp. 

106 """ 

107 return self.warp.dataId 

108 

109 

110class AssembleCellCoaddConnections( 

111 PipelineTaskConnections, 

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

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

114): 

115 inputWarps = Input( 

116 doc="Input warps", 

117 name="{inputWarpName}Coadd_directWarp", 

118 storageClass="ExposureF", 

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

120 deferLoad=True, 

121 multiple=True, 

122 ) 

123 

124 maskedFractionWarps = Input( 

125 doc="Mask fraction warps", 

126 name="{inputWarpName}Coadd_directWarp_maskedFraction", 

127 storageClass="ImageF", 

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

129 deferLoad=True, 

130 multiple=True, 

131 ) 

132 

133 artifactMasks = Input( 

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

135 name="compare_warp_artifact_mask", 

136 storageClass="Mask", 

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

138 deferLoad=True, 

139 multiple=True, 

140 ) 

141 

142 visitSummaryList = Input( 

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

144 name="finalVisitSummary", 

145 storageClass="ExposureCatalog", 

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

147 deferLoad=True, 

148 multiple=True, 

149 ) 

150 

151 skyMap = Input( 

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

153 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, 

154 storageClass="SkyMap", 

155 dimensions=("skymap",), 

156 ) 

157 

158 multipleCellCoadd = Output( 

159 doc="Output multiple cell coadd", 

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

161 storageClass="MultipleCellCoadd", 

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

163 ) 

164 

165 inputMap = Output( 

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

167 name="{inputWarpName}Coadd_inputMap", 

168 storageClass="HealSparseMap", 

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

170 ) 

171 

172 config: AssembleCellCoaddConfig 

173 

174 def __init__(self, *, config: AssembleCellCoaddConfig | None = None): 

175 super().__init__(config=config) 

176 

177 if not self.config: 

178 return 

179 

180 if self.config.do_calculate_weight_from_warp: 

181 del self.visitSummaryList 

182 

183 if not self.config.do_use_artifact_mask: 

184 del self.artifactMasks 

185 

186 if not self.config.do_input_map: 

187 del self.inputMap 

188 

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

190 # number of noise realizations specified in the config. 

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

192 noise_warps = Input( 

193 doc="Input noise warps", 

194 name=f"direct_warp_noise{n}", 

195 storageClass="MaskedImageF", 

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

197 deferLoad=True, 

198 multiple=True, 

199 ) 

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

201 

202 if self.config.output_image_type == "future": 

203 self.multipleCellCoadd = dataclasses.replace(self.multipleCellCoadd, storageClass="CellCoadd") 

204 

205 

206class AssembleCellCoaddConfig(PipelineTaskConfig, pipelineConnections=AssembleCellCoaddConnections): 

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

208 interpolate_coadd = ConfigurableField( 

209 target=InterpImageTask, 

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

211 ) 

212 do_scale_zero_point = Field[bool]( 

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

214 default=False, 

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

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

217 "after v29.", 

218 ) 

219 scale_zero_point = ConfigurableField( 

220 target=ScaleZeroPointTask, 

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

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

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

224 "after v29.", 

225 ) 

226 do_calculate_weight_from_warp = Field[bool]( 

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

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

229 default=False, 

230 ) 

231 do_use_artifact_mask = Field[bool]( 

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

233 default=True, 

234 ) 

235 do_coadd_inverse_aperture_corrections = Field[bool]( 

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

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

238 default=False, 

239 ) 

240 min_overlap_fraction = RangeField[float]( 

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

242 "cell.", 

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

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

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

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

247 # definitely not overlap the cell center. 

248 default=1.0, 

249 min=0.0, 

250 max=1.0, 

251 inclusiveMin=True, 

252 inclusiveMax=True, 

253 ) 

254 bad_mask_planes = ListField[str]( 

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

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

257 ) 

258 remove_mask_planes = ListField[str]( 

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

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

261 ) 

262 calc_error_from_input_variance = Field[bool]( 

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

264 "statistic. Passed to AccumulatorMeanStack.", 

265 default=True, 

266 ) 

267 mask_propagation_thresholds = DictField[str, float]( 

268 doc=( 

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

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

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

272 "would have contributed exceeds this value." 

273 ), 

274 default={"SAT": 0.1}, 

275 ) 

276 max_maskfrac = RangeField[float]( 

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

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

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

280 default=0.5, 

281 min=0.0, 

282 max=1.0, 

283 inclusiveMin=True, 

284 inclusiveMax=False, 

285 ) 

286 num_noise_realizations = Field[int]( 

287 default=0, 

288 doc=( 

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

290 "This should not exceed the corresponding config parameter " 

291 "specified in `MakeDirectWarpConfig`. " 

292 ), 

293 check=lambda x: x >= 0, 

294 ) 

295 psf_warper = ConfigField( 

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

297 "warp the images.", 

298 dtype=afwMath.Warper.ConfigClass, 

299 ) 

300 psf_dimensions = Field[int]( 

301 default=35, 

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

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

304 ) 

305 require_artifact_mask = Field[bool]( 

306 default=True, 

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

308 " from CompareWarpTask", 

309 ) 

310 do_input_map = Field[bool]( 

311 default=False, 

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

313 ) 

314 input_mapper = ConfigurableField( 

315 target=HealSparseInputMapTask, 

316 doc="Input map creation subtask.", 

317 ) 

318 output_image_type = ChoiceField[str]( 

319 "Which image type to use for the output coadd.", 

320 allowed={ 

321 "legacy": "Write as a lsst.cell_Coadds.MultipleCellCoadd.", 

322 "future": "Write as a lsst.images.cells.CellCoadd.", 

323 }, 

324 optional=False, 

325 default="legacy", 

326 ) 

327 

328 

329class AssembleCellCoaddTask(PipelineTask): 

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

331 

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

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

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

335 

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

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

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

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

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

341 with the same weights. 

342 

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

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

345 assumed to be spatially constant within a cell. 

346 

347 Raises 

348 ------ 

349 NoWorkFound 

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

351 RuntimeError 

352 Raised if the skymap is not cell-based. 

353 

354 Notes 

355 ----- 

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

357 especially its Config and Connections are experimental and subject to 

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

359 until it is included in the DRP pipeline. 

360 """ 

361 

362 ConfigClass = AssembleCellCoaddConfig 

363 _DefaultName = "assembleCellCoadd" 

364 

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

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

367 if self.config.do_interpolate_coadd: 367 ↛ 369line 367 didn't jump to line 369 because the condition on line 367 was always true

368 self.makeSubtask("interpolate_coadd") 

369 if self.config.do_scale_zero_point: 

370 self.makeSubtask("scale_zero_point") 

371 if self.config.do_input_map: 371 ↛ 374line 371 didn't jump to line 374 because the condition on line 371 was always true

372 self.makeSubtask("input_mapper") 

373 

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

375 if (warping_kernel_name := self.config.psf_warper.warpingKernelName.lower()).startswith("lanczos"): 375 ↛ 383line 375 didn't jump to line 383 because the condition on line 375 was always true

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

377 self.log.debug( 

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

379 psf_padding, 

380 self.config.psf_warper.warpingKernelName, 

381 ) 

382 else: 

383 psf_padding = 10 

384 self.log.info( 

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

386 psf_padding, 

387 ) 

388 self.psf_padding = psf_padding 

389 

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

391 # Docstring inherited. 

392 if not inputRefs.inputWarps: 

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

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

395 

396 # Construct skyInfo expected by run 

397 # Do not remove skyMap from inputData in case _makeSupplementaryData 

398 # needs it 

399 skyMap = butlerQC.get(inputRefs.skyMap) 

400 

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

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

403 

404 outputDataId = butlerQC.quantum.dataId 

405 

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

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

408 

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

410 self.common = CommonComponents( 

411 units=units, 

412 wcs=skyInfo.patchInfo.wcs, 

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

414 identifiers=PatchIdentifiers.from_data_id(outputDataId), 

415 ) 

416 

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

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

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

420 

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

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

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

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

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

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

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

428 

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

430 butlerQC.put(returnStruct, outputRefs) 

431 return returnStruct 

432 

433 @staticmethod 

434 def _compute_weight(maskedImage, statsCtrl): 

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

436 

437 Parameters 

438 ---------- 

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

440 The masked image to compute the weight. 

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

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

443 

444 Returns 

445 ------- 

446 weight : `float` 

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

448 """ 

449 statObj = afwMath.makeStatistics( 

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

451 ) 

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

453 weight = 1.0 / float(meanVar) 

454 return weight 

455 

456 @staticmethod 

457 def _construct_grid(skyInfo): 

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

459 

460 Parameters 

461 ---------- 

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

463 A Struct object 

464 

465 Returns 

466 ------- 

467 grid : `~lsst.cell_coadds.UniformGrid` 

468 A UniformGrid object. 

469 """ 

470 padding = skyInfo.patchInfo.getCellBorder() 

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

472 grid = UniformGrid.from_bbox_cell_size( 

473 grid_bbox, 

474 skyInfo.patchInfo.getCellInnerDimensions(), 

475 padding=padding, 

476 ) 

477 return grid 

478 

479 def _construct_grid_container(self, skyInfo, statsCtrl): 

480 """Construct a grid of AccumulatorMeanStack instances. 

481 

482 Parameters 

483 ---------- 

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

485 A Struct object 

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

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

488 

489 Returns 

490 ------- 

491 gc : `~lsst.cell_coadds.GridContainer` 

492 A GridContainer object container one AccumulatorMeanStack per cell. 

493 """ 

494 grid = self._construct_grid(skyInfo) 

495 

496 maskMap = setRejectedMaskMapping(statsCtrl) 

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

498 thresholdDict = AccumulatorMeanStack.stats_ctrl_to_threshold_dict(statsCtrl) 

499 

500 # Initialize the grid container with AccumulatorMeanStacks 

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

502 for cellInfo in skyInfo.patchInfo: 

503 stacker = AccumulatorMeanStack( 

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

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

506 bit_mask_value=statsCtrl.getAndMask(), 

507 mask_threshold_dict=thresholdDict, 

508 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

509 compute_n_image=False, 

510 mask_map=maskMap, 

511 no_good_pixels_mask=statsCtrl.getNoGoodPixelsMask(), 

512 ) 

513 gc[cellInfo.index] = stacker 

514 

515 return gc 

516 

517 def _construct_stats_control(self): 

518 """Construct a StatisticsControl object for coadd. 

519 

520 Unlike AssembleCoaddTask or CompareWarpAssembleCoaddTask, there is 

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

522 optionally mask propagation thresholds. 

523 

524 Returns 

525 ------- 

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

527 A control object for StatisticsStack. 

528 """ 

529 statsCtrl = afwMath.StatisticsControl() 

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

531 # CompareWarpAssembleCoaddTask to get consistent weights. This is NOT 

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

533 # fallback option that is not recommended for production. 

534 statsCtrl.setNumIter(2) 

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

536 statsCtrl.setNanSafe(True) 

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

538 bit = afwImage.Mask.getMaskPlane(plane) 

539 statsCtrl.setMaskPropagationThreshold(bit, threshold) 

540 return statsCtrl 

541 

542 def _construct_ap_corr_grid_container(self, skyInfo): 

543 """Construct a grid of CoaddApCorrMapStacker instances. 

544 

545 Parameters 

546 ---------- 

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

548 A Struct object 

549 

550 Returns 

551 ------- 

552 gc : `~lsst.cell_coadds.GridContainer` 

553 A GridContainer object container one CoaddApCorrMapStacker per 

554 cell. 

555 """ 

556 grid = self._construct_grid(skyInfo) 

557 

558 # Initialize the grid container with CoaddApCorrMapStacker. 

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

560 for cellInfo in skyInfo.patchInfo: 

561 stacker = CoaddApCorrMapStacker( 

562 evaluation_point=cellInfo.inner_bbox.getCenter(), 

563 do_coadd_inverse_ap_corr=self.config.do_coadd_inverse_aperture_corrections, 

564 ) 

565 gc[cellInfo.index] = stacker 

566 

567 return gc 

568 

569 def run( 

570 self, 

571 *, 

572 inputs: dict[DataCoordinate, WarpInputs], 

573 skyInfo, 

574 visitSummaryList: list | None = None, 

575 ): 

576 for mask_plane in self.config.bad_mask_planes: 

577 afwImage.Mask.addMaskPlane(mask_plane) 

578 for mask_plane in self.config.mask_propagation_thresholds: 

579 afwImage.Mask.addMaskPlane(mask_plane) 

580 

581 statsCtrl = self._construct_stats_control() 

582 

583 warp_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

584 maskfrac_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

585 noise_stacker_gc_list = [ 

586 self._construct_grid_container(skyInfo, statsCtrl) 

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

588 ] 

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

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

591 ap_corr_stacker_gc = self._construct_ap_corr_grid_container(skyInfo) 

592 

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

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

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

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

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

598 

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

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

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

602 # stay in fallback mode. 

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

604 

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

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

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

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

609 # get the corresponding PSFs as well. 

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

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

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

613 

614 if self.config.do_input_map: 614 ↛ 629line 614 didn't jump to line 629 because the condition on line 614 was always true

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

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

617 visit_detectors = [] 

618 for warp_input in warp_input_list: 

619 for row in warp_input.ccds: 

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

621 

622 self.input_mapper.initialize_cell_input_map( 

623 skyInfo.patchInfo.getOuterBBox(), 

624 skyInfo.patchInfo.wcs, 

625 visit_detectors, 

626 ) 

627 

628 # Populate them. 

629 for cellInfo in skyInfo.patchInfo: 

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

631 observation_identifiers_gc[cellInfo.index] = {} 

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

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

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

635 cell_center_pixel, 

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

637 ) 

638 psf_stacker_gc[cellInfo.index] = AccumulatorMeanStack( 

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

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

641 bit_mask_value=0, 

642 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

643 compute_n_image=False, 

644 ) 

645 is_fallback_gc[cellInfo.index] = True 

646 input_map_data_gc[cellInfo.index] = [] 

647 

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

649 if not visitSummaryList: 

650 visitSummaryList = [] 

651 visitSummaryRefDict = { 

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

653 } 

654 

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

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

657 

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

659 # it completely overlaps. 

660 for warp_input in inputs.values(): 

661 # warps that have been excluded from CompareWarp via visit 

662 # selection from SelectVisitsTasks will not have artifact masks. 

663 # Exclude them from the cell coadds too. 

664 if self.config.require_artifact_mask and warp_input.artifact_mask is None: 664 ↛ 665line 664 didn't jump to line 665 because the condition on line 664 was never true

665 self.log.info( 

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

667 warp_input.dataId["visit"], 

668 ) 

669 continue 

670 

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

672 masked_fraction_image = ( 

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

674 if warp_input.masked_fraction 

675 else None 

676 ) 

677 

678 # Pre-process the warp before coadding. 

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

680 warp.mask.addMaskPlane("CLIPPED") 

681 warp.mask.addMaskPlane("REJECTED") 

682 warp.mask.addMaskPlane("SENSOR_EDGE") 

683 warp.mask.addMaskPlane("INEXACT_PSF") 

684 

685 if artifact_mask_ref := warp_input.artifact_mask: 685 ↛ 687line 685 didn't jump to line 687 because the condition on line 685 was never true

686 # Apply the artifact mask to the warp. 

687 artifact_mask = artifact_mask_ref.get() 

688 assert ( 

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

690 ), "Mask dicts do not agree." 

691 warp.mask.array = artifact_mask.array 

692 del artifact_mask 

693 

694 if self.config.do_scale_zero_point: 

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

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

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

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

699 zero_point_scale_factor = imageScaler.scale 

700 self.log.debug( 

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

702 warp_input.dataId, 

703 zero_point_scale_factor, 

704 ) 

705 else: 

706 zero_point_scale_factor = 1.0 

707 if "BUNIT" not in warp.metadata: 707 ↛ 708line 707 didn't jump to line 708 because the condition on line 707 was never true

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

709 if warp.metadata["BUNIT"] != "nJy": 709 ↛ 710line 709 didn't jump to line 710 because the condition on line 709 was never true

710 raise ValueError( 

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

712 ) 

713 

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

715 to_remove = [] 

716 for plane in self.config.remove_mask_planes: 

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

718 to_remove.append(plane) 

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

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

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

722 rejected = afwImage.Mask.getPlaneBitMask( 

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

724 ) 

725 

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

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

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

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

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

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

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

733 0.0, 

734 ) # Mapping from detector to weight. 

735 

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

737 visitSummary = visitSummaryRef.get() 

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

739 visitSummaryRow = visitSummary.find(detector) 

740 mean_variance = visitSummaryRow["meanVar"] 

741 mean_variance *= zero_point_scale_factor**2 

742 if warp.metadata.get("BUNIT", None) == "nJy": 742 ↛ 744line 742 didn't jump to line 744 because the condition on line 742 was always true

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

744 weights[detector] = 1.0 / mean_variance 

745 del visitSummary 

746 else: 

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

748 weight = self._compute_weight(warp, statsCtrl) 

749 if not np.isfinite(weight): 749 ↛ 750line 749 didn't jump to line 750 because the condition on line 749 was never true

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

751 continue 

752 

753 for detector in weights: 

754 weights[detector] = weight 

755 

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

757 

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

759 # detector ID that pixel comes from. 

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

761 for row in full_ccd_table: 

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

763 if (src_polygon := row.validPolygon) is None: 763 ↛ 764line 763 didn't jump to line 764 because the condition on line 763 was never true

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

765 try: 

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

767 geom.Box2D(warp.getBBox()) 

768 ) 

769 except SinglePolygonException: 

770 continue 

771 

772 observation_identifier = ObservationIdentifiers.from_data_id( 

773 warp_input.dataId, 

774 backup_detector=row["ccd"], 

775 ) 

776 visit_polygons[observation_identifier] = dest_polygon 

777 

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

779 if not (detector_map.array[detector_map_slice] < 0).all(): 779 ↛ 780line 779 didn't jump to line 780 because the condition on line 779 was never true

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

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

782 

783 if (detector_map.array < 0).all(): 783 ↛ 784line 783 didn't jump to line 784 because the condition on line 783 was never true

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

785 detector_map.array[:, :] = 0 

786 

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

788 bbox = cellInfo.outer_bbox 

789 inner_bbox = cellInfo.inner_bbox 

790 

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

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

793 if (overlap_fraction < self.config.min_overlap_fraction) or (overlap_fraction <= 0.0): 793 ↛ 794line 793 didn't jump to line 794 because the condition on line 793 was never true

794 self.log.debug( 

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

796 warp_input.dataId, 

797 cellInfo.index, 

798 overlap_fraction, 

799 self.config.min_overlap_fraction, 

800 ) 

801 continue 

802 

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

804 if not np.isfinite(weight): 804 ↛ 805line 804 didn't jump to line 805 because the condition on line 804 was never true

805 self.log.warning( 

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

807 ) 

808 continue 

809 

810 if weight == 0: 810 ↛ 811line 810 didn't jump to line 811 because the condition on line 810 was never true

811 self.log.info( 

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

813 ) 

814 continue 

815 

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

817 # cell. Used to gate on max_maskfrac. 

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

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

820 unmasked_fraction = ( 

821 inner_detector_pixels & inner_unmasked_pixels 

822 ).sum() / inner_detector_pixels.sum() 

823 is_fallback = is_fallback_gc[cellInfo.index] 

824 if unmasked_fraction <= max(1.0 - self.config.max_maskfrac, 0.0): 824 ↛ 825line 824 didn't jump to line 825 because the condition on line 824 was never true

825 if not is_fallback: 

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

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

828 # INEXACT_PSF. 

829 self.log.debug( 

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

831 warp_input.dataId, 

832 cellInfo.index, 

833 1.0 - unmasked_fraction, 

834 self.config.max_maskfrac, 

835 ) 

836 continue 

837 else: 

838 self.log.debug( 

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

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

841 warp_input.dataId, 

842 cellInfo.index, 

843 1.0 - unmasked_fraction, 

844 self.config.max_maskfrac, 

845 ) 

846 elif is_fallback: 

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

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