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 

32import logging 

33 

34import numpy as np 

35 

36import lsst.afw.geom as afwGeom 

37import lsst.afw.image as afwImage 

38import lsst.afw.math as afwMath 

39import lsst.geom as geom 

40from lsst.afw.detection import InvalidPsfError 

41from lsst.afw.geom import SinglePolygonException, makeWcsPairTransform 

42from lsst.cell_coadds import ( 

43 CellIdentifiers, 

44 CoaddApCorrMapStacker, 

45 CoaddInputs, 

46 CoaddUnits, 

47 CommonComponents, 

48 GridContainer, 

49 MultipleCellCoadd, 

50 ObservationIdentifiers, 

51 OwnedImagePlanes, 

52 PatchIdentifiers, 

53 SingleCellCoadd, 

54 UniformGrid, 

55) 

56from lsst.daf.butler import DataCoordinate, DeferredDatasetHandle 

57from lsst.meas.algorithms import AccumulatorMeanStack 

58from lsst.pex.config import ( 

59 ChoiceField, 

60 ConfigField, 

61 ConfigurableField, 

62 DictField, 

63 Field, 

64 ListField, 

65 RangeField, 

66) 

67from lsst.pipe.base import ( 

68 InMemoryDatasetHandle, 

69 NoWorkFound, 

70 PipelineTask, 

71 PipelineTaskConfig, 

72 PipelineTaskConnections, 

73 Struct, 

74) 

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

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

77from lsst.pipe.tasks.healSparseMapping import HealSparseInputMapTask 

78from lsst.pipe.tasks.interpImage import InterpImageTask 

79from lsst.pipe.tasks.scaleZeroPoint import ScaleZeroPointTask 

80from lsst.skymap import BaseSkyMap 

81 

82 

83@dataclasses.dataclass 

84class WarpInputs: 

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

86 

87 warp: DeferredDatasetHandle | InMemoryDatasetHandle 

88 """Handle for the warped exposure.""" 

89 

90 masked_fraction: DeferredDatasetHandle | InMemoryDatasetHandle | None = None 

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

92 

93 artifact_mask: DeferredDatasetHandle | InMemoryDatasetHandle | None = None 

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

95 

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

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

98 

99 @property 

100 def dataId(self) -> DataCoordinate: 

101 """DataID corresponding to the warp. 

102 

103 Returns 

104 ------- 

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

106 DataID of the warp. 

107 """ 

108 return self.warp.dataId 

109 

110 

111class AssembleCellCoaddConnections( 

112 PipelineTaskConnections, 

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

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

115): 

116 inputWarps = Input( 

117 doc="Input warps", 

118 name="{inputWarpName}Coadd_directWarp", 

119 storageClass="ExposureF", 

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

121 deferLoad=True, 

122 multiple=True, 

123 ) 

124 

125 maskedFractionWarps = Input( 

126 doc="Mask fraction warps", 

127 name="{inputWarpName}Coadd_directWarp_maskedFraction", 

128 storageClass="ImageF", 

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

130 deferLoad=True, 

131 multiple=True, 

132 ) 

133 

134 artifactMasks = Input( 

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

136 name="compare_warp_artifact_mask", 

137 storageClass="Mask", 

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

139 deferLoad=True, 

140 multiple=True, 

141 ) 

142 

143 visitSummaryList = Input( 

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

145 name="finalVisitSummary", 

146 storageClass="ExposureCatalog", 

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

148 deferLoad=True, 

149 multiple=True, 

150 ) 

151 

152 skyMap = Input( 

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

154 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, 

155 storageClass="SkyMap", 

156 dimensions=("skymap",), 

157 ) 

158 

159 multipleCellCoadd = Output( 

160 doc="Output multiple cell coadd", 

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

162 storageClass="MultipleCellCoadd", 

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

164 ) 

165 

166 inputMap = Output( 

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

168 name="{inputWarpName}Coadd_inputMap", 

169 storageClass="HealSparseMap", 

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

171 ) 

172 

173 config: AssembleCellCoaddConfig 

174 

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

176 super().__init__(config=config) 

177 

178 if not self.config: 

179 return 

180 

181 if self.config.do_calculate_weight_from_warp: 

182 del self.visitSummaryList 

183 

184 if not self.config.do_use_artifact_mask: 

185 del self.artifactMasks 

186 

187 if not self.config.do_input_map: 

188 del self.inputMap 

189 

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

191 # number of noise realizations specified in the config. 

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

193 noise_warps = Input( 

194 doc="Input noise warps", 

195 name=f"direct_warp_noise{n}", 

196 storageClass="MaskedImageF", 

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

198 deferLoad=True, 

199 multiple=True, 

200 ) 

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

202 

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

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

205 

206 

207class AssembleCellCoaddConfig(PipelineTaskConfig, pipelineConnections=AssembleCellCoaddConnections): 

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

209 interpolate_coadd = ConfigurableField( 

210 target=InterpImageTask, 

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

212 ) 

213 do_scale_zero_point = Field[bool]( 

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

215 default=False, 

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

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

218 "after v29.", 

219 ) 

220 scale_zero_point = ConfigurableField( 

221 target=ScaleZeroPointTask, 

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

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

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

225 "after v29.", 

226 ) 

227 do_calculate_weight_from_warp = Field[bool]( 

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

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

230 default=False, 

231 ) 

232 do_use_artifact_mask = Field[bool]( 

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

234 default=True, 

235 ) 

236 do_coadd_inverse_aperture_corrections = Field[bool]( 

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

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

239 default=False, 

240 ) 

241 min_overlap_fraction = RangeField[float]( 

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

243 "cell.", 

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

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

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

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

248 # definitely not overlap the cell center. 

249 default=1.0, 

250 min=0.0, 

251 max=1.0, 

252 inclusiveMin=True, 

253 inclusiveMax=True, 

254 ) 

255 bad_mask_planes = ListField[str]( 

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

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

258 ) 

259 remove_mask_planes = ListField[str]( 

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

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

262 ) 

263 calc_error_from_input_variance = Field[bool]( 

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

265 "statistic. Passed to AccumulatorMeanStack.", 

266 default=True, 

267 ) 

268 mask_propagation_thresholds = DictField[str, float]( 

269 doc=( 

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

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

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

273 "would have contributed exceeds this value." 

274 ), 

275 default={"SAT": 0.1}, 

276 ) 

277 max_maskfrac = RangeField[float]( 

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

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

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

281 default=0.5, 

282 min=0.0, 

283 max=1.0, 

284 inclusiveMin=True, 

285 inclusiveMax=False, 

286 ) 

287 num_noise_realizations = Field[int]( 

288 default=0, 

289 doc=( 

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

291 "This should not exceed the corresponding config parameter " 

292 "specified in `MakeDirectWarpConfig`. " 

293 ), 

294 check=lambda x: x >= 0, 

295 ) 

296 psf_warper = ConfigField( 

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

298 "warp the images.", 

299 dtype=afwMath.Warper.ConfigClass, 

300 ) 

301 psf_dimensions = Field[int]( 

302 default=35, 

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

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

305 ) 

306 require_artifact_mask = Field[bool]( 

307 default=True, 

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

309 " from CompareWarpTask", 

310 ) 

311 do_input_map = Field[bool]( 

312 default=False, 

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

314 ) 

315 input_mapper = ConfigurableField( 

316 target=HealSparseInputMapTask, 

317 doc="Input map creation subtask.", 

318 ) 

319 output_image_type = ChoiceField[str]( 

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

321 allowed={ 

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

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

324 }, 

325 optional=False, 

326 default="legacy", 

327 ) 

328 

329 

330class AssembleCellCoaddTask(PipelineTask): 

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

332 

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

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

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

336 

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

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

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

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

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

342 with the same weights. 

343 

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

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

346 assumed to be spatially constant within a cell. 

347 

348 Raises 

349 ------ 

350 NoWorkFound 

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

352 RuntimeError 

353 Raised if the skymap is not cell-based. 

354 

355 Notes 

356 ----- 

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

358 especially its Config and Connections are experimental and subject to 

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

360 until it is included in the DRP pipeline. 

361 """ 

362 

363 ConfigClass = AssembleCellCoaddConfig 

364 _DefaultName = "assembleCellCoadd" 

365 

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

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

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

369 self.makeSubtask("interpolate_coadd") 

370 # Suppress the warning message about fallback. 

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

372 if self.config.do_scale_zero_point: 

373 self.makeSubtask("scale_zero_point") 

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

375 self.makeSubtask("input_mapper") 

376 

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

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

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

380 self.log.debug( 

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

382 psf_padding, 

383 self.config.psf_warper.warpingKernelName, 

384 ) 

385 else: 

386 psf_padding = 10 

387 self.log.info( 

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

389 psf_padding, 

390 ) 

391 self.psf_padding = psf_padding 

392 

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

394 # Docstring inherited. 

395 if not inputRefs.inputWarps: 

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

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

398 

399 # Construct skyInfo expected by run 

400 # Do not remove skyMap from inputData in case _makeSupplementaryData 

401 # needs it 

402 skyMap = butlerQC.get(inputRefs.skyMap) 

403 

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

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

406 

407 outputDataId = butlerQC.quantum.dataId 

408 

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

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

411 

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

413 self.common = CommonComponents( 

414 units=units, 

415 wcs=skyInfo.patchInfo.wcs, 

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

417 identifiers=PatchIdentifiers.from_data_id(outputDataId), 

418 ) 

419 

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

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

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

423 

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

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

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

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

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

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

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

431 

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

433 butlerQC.put(returnStruct, outputRefs) 

434 return returnStruct 

435 

436 @staticmethod 

437 def _compute_weight(maskedImage, statsCtrl): 

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

439 

440 Parameters 

441 ---------- 

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

443 The masked image to compute the weight. 

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

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

446 

447 Returns 

448 ------- 

449 weight : `float` 

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

451 """ 

452 statObj = afwMath.makeStatistics( 

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

454 ) 

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

456 weight = 1.0 / float(meanVar) 

457 return weight 

458 

459 @staticmethod 

460 def _construct_grid(skyInfo): 

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

462 

463 Parameters 

464 ---------- 

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

466 A Struct object 

467 

468 Returns 

469 ------- 

470 grid : `~lsst.cell_coadds.UniformGrid` 

471 A UniformGrid object. 

472 """ 

473 padding = skyInfo.patchInfo.getCellBorder() 

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

475 grid = UniformGrid.from_bbox_cell_size( 

476 grid_bbox, 

477 skyInfo.patchInfo.getCellInnerDimensions(), 

478 padding=padding, 

479 ) 

480 return grid 

481 

482 def _construct_grid_container(self, skyInfo, statsCtrl): 

483 """Construct a grid of AccumulatorMeanStack instances. 

484 

485 Parameters 

486 ---------- 

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

488 A Struct object 

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

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

491 

492 Returns 

493 ------- 

494 gc : `~lsst.cell_coadds.GridContainer` 

495 A GridContainer object container one AccumulatorMeanStack per cell. 

496 """ 

497 grid = self._construct_grid(skyInfo) 

498 

499 maskMap = setRejectedMaskMapping(statsCtrl) 

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

501 thresholdDict = AccumulatorMeanStack.stats_ctrl_to_threshold_dict(statsCtrl) 

502 

503 # Initialize the grid container with AccumulatorMeanStacks 

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

505 for cellInfo in skyInfo.patchInfo: 

506 stacker = AccumulatorMeanStack( 

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

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

509 bit_mask_value=statsCtrl.getAndMask(), 

510 mask_threshold_dict=thresholdDict, 

511 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

512 compute_n_image=False, 

513 mask_map=maskMap, 

514 no_good_pixels_mask=statsCtrl.getNoGoodPixelsMask(), 

515 ) 

516 gc[cellInfo.index] = stacker 

517 

518 return gc 

519 

520 def _construct_stats_control(self): 

521 """Construct a StatisticsControl object for coadd. 

522 

523 Unlike AssembleCoaddTask or CompareWarpAssembleCoaddTask, there is 

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

525 optionally mask propagation thresholds. 

526 

527 Returns 

528 ------- 

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

530 A control object for StatisticsStack. 

531 """ 

532 statsCtrl = afwMath.StatisticsControl() 

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

534 # CompareWarpAssembleCoaddTask to get consistent weights. This is NOT 

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

536 # fallback option that is not recommended for production. 

537 statsCtrl.setNumIter(2) 

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

539 statsCtrl.setNanSafe(True) 

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

541 bit = afwImage.Mask.getMaskPlane(plane) 

542 statsCtrl.setMaskPropagationThreshold(bit, threshold) 

543 return statsCtrl 

544 

545 def _construct_ap_corr_grid_container(self, skyInfo): 

546 """Construct a grid of CoaddApCorrMapStacker instances. 

547 

548 Parameters 

549 ---------- 

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

551 A Struct object 

552 

553 Returns 

554 ------- 

555 gc : `~lsst.cell_coadds.GridContainer` 

556 A GridContainer object container one CoaddApCorrMapStacker per 

557 cell. 

558 """ 

559 grid = self._construct_grid(skyInfo) 

560 

561 # Initialize the grid container with CoaddApCorrMapStacker. 

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

563 for cellInfo in skyInfo.patchInfo: 

564 stacker = CoaddApCorrMapStacker( 

565 evaluation_point=cellInfo.inner_bbox.getCenter(), 

566 do_coadd_inverse_ap_corr=self.config.do_coadd_inverse_aperture_corrections, 

567 ) 

568 gc[cellInfo.index] = stacker 

569 

570 return gc 

571 

572 def run( 

573 self, 

574 *, 

575 inputs: dict[DataCoordinate, WarpInputs], 

576 skyInfo, 

577 visitSummaryList: list | None = None, 

578 ): 

579 for mask_plane in self.config.bad_mask_planes: 

580 afwImage.Mask.addMaskPlane(mask_plane) 

581 for mask_plane in self.config.mask_propagation_thresholds: 

582 afwImage.Mask.addMaskPlane(mask_plane) 

583 

584 statsCtrl = self._construct_stats_control() 

585 

586 warp_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

587 maskfrac_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl) 

588 noise_stacker_gc_list = [ 

589 self._construct_grid_container(skyInfo, statsCtrl) 

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

591 ] 

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

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

594 ap_corr_stacker_gc = self._construct_ap_corr_grid_container(skyInfo) 

595 

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

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

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

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

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

601 

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

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

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

605 # stay in fallback mode. 

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

607 

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

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

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

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

612 # get the corresponding PSFs as well. 

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

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

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

616 

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

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

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

620 visit_detectors = [] 

621 for warp_input in warp_input_list: 

622 for row in warp_input.ccds: 

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

624 

625 self.input_mapper.initialize_cell_input_map( 

626 skyInfo.patchInfo.getOuterBBox(), 

627 skyInfo.patchInfo.wcs, 

628 visit_detectors, 

629 ) 

630 

631 # Populate them. 

632 for cellInfo in skyInfo.patchInfo: 

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

634 observation_identifiers_gc[cellInfo.index] = {} 

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

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

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

638 cell_center_pixel, 

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

640 ) 

641 psf_stacker_gc[cellInfo.index] = AccumulatorMeanStack( 

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

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

644 bit_mask_value=0, 

645 calc_error_from_input_variance=self.config.calc_error_from_input_variance, 

646 compute_n_image=False, 

647 ) 

648 is_fallback_gc[cellInfo.index] = True 

649 input_map_data_gc[cellInfo.index] = [] 

650 

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

652 if not visitSummaryList: 

653 visitSummaryList = [] 

654 visitSummaryRefDict = { 

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

656 } 

657 

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

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

660 

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

662 # it completely overlaps. 

663 for warp_input in inputs.values(): 

664 # warps that have been excluded from CompareWarp via visit 

665 # selection from SelectVisitsTasks will not have artifact masks. 

666 # Exclude them from the cell coadds too. 

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

668 self.log.info( 

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

670 warp_input.dataId["visit"], 

671 ) 

672 continue 

673 

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

675 masked_fraction_image = ( 

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

677 if warp_input.masked_fraction 

678 else None 

679 ) 

680 

681 # Pre-process the warp before coadding. 

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

683 warp.mask.addMaskPlane("CLIPPED") 

684 warp.mask.addMaskPlane("REJECTED") 

685 warp.mask.addMaskPlane("SENSOR_EDGE") 

686 warp.mask.addMaskPlane("INEXACT_PSF") 

687 

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

689 # Apply the artifact mask to the warp. 

690 artifact_mask = artifact_mask_ref.get() 

691 assert ( 

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

693 ), "Mask dicts do not agree." 

694 warp.mask.array = artifact_mask.array 

695 del artifact_mask 

696 

697 if self.config.do_scale_zero_point: 

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

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

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

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

702 zero_point_scale_factor = imageScaler.scale 

703 self.log.debug( 

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

705 warp_input.dataId, 

706 zero_point_scale_factor, 

707 ) 

708 else: 

709 zero_point_scale_factor = 1.0 

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

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

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

713 raise ValueError( 

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

715 ) 

716 

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

718 to_remove = [] 

719 for plane in self.config.remove_mask_planes: 

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

721 to_remove.append(plane) 

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

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

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

725 rejected = afwImage.Mask.getPlaneBitMask( 

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

727 ) 

728 

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

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

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

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

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

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

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

736 0.0, 

737 ) # Mapping from detector to weight. 

738 

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

740 visitSummary = visitSummaryRef.get() 

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

742 visitSummaryRow = visitSummary.find(detector) 

743 mean_variance = visitSummaryRow["meanVar"] 

744 mean_variance *= zero_point_scale_factor**2 

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

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

747 weights[detector] = 1.0 / mean_variance 

748 del visitSummary 

749 else: 

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

751 weight = self._compute_weight(warp, statsCtrl) 

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

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

754 continue 

755 

756 for detector in weights: 

757 weights[detector] = weight 

758 

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

760 

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

762 # detector ID that pixel comes from. 

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

764 for row in full_ccd_table: 

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

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

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

768 try: 

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

770 geom.Box2D(warp.getBBox()) 

771 ) 

772 except SinglePolygonException: 

773 continue 

774 

775 observation_identifier = ObservationIdentifiers.from_data_id( 

776 warp_input.dataId, 

777 backup_detector=row["ccd"], 

778 ) 

779 visit_polygons[observation_identifier] = dest_polygon 

780 

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

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

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

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

785 

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

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

788 detector_map.array[:, :] = 0 

789 

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

791 bbox = cellInfo.outer_bbox 

792 inner_bbox = cellInfo.inner_bbox 

793 

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

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

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

797 self.log.debug( 

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

799 warp_input.dataId, 

800 cellInfo.index, 

801 overlap_fraction, 

802 self.config.min_overlap_fraction, 

803 ) 

804 continue 

805 

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

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

808 self.log.warning( 

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

810 ) 

811 continue 

812 

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

814 self.log.info( 

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

816 ) 

817 continue 

818 

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

820 # cell. Used to gate on max_maskfrac. 

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

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

823 unmasked_fraction = ( 

824 inner_detector_pixels & inner_unmasked_pixels 

825 ).sum() / inner_detector_pixels.sum() 

826 is_fallback = is_fallback_gc[cellInfo.index] 

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

828 if not is_fallback: 

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

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

831 # INEXACT_PSF. 

832 self.log.debug( 

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

834 warp_input.dataId, 

835 cellInfo.index, 

836 1.0 - unmasked_fraction, 

837 self.config.max_maskfrac, 

838 ) 

839 continue 

840 else: 

841 self.log.debug( 

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

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

844 warp_input.dataId, 

845 cellInfo.index, 

846 1.0 - unmasked_fraction, 

847 self.config.max_maskfrac,