Coverage for python / lsst / drp / tasks / assemble_cell_coadd.py: 17%
353 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:55 -0800
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:55 -0800
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/>.
22from __future__ import annotations
24__all__ = (
25 "AssembleCellCoaddTask",
26 "AssembleCellCoaddConfig",
27 "ConvertMultipleCellCoaddToExposureTask",
28)
30import dataclasses
31import itertools
33import numpy as np
35import lsst.afw.geom as afwGeom
36import lsst.afw.image as afwImage
37import lsst.afw.math as afwMath
38import lsst.geom as geom
39from lsst.afw.detection import InvalidPsfError
40from lsst.afw.geom import SinglePolygonException, makeWcsPairTransform
41from lsst.cell_coadds import (
42 CellIdentifiers,
43 CoaddApCorrMapStacker,
44 CoaddInputs,
45 CoaddUnits,
46 CommonComponents,
47 GridContainer,
48 MultipleCellCoadd,
49 ObservationIdentifiers,
50 OwnedImagePlanes,
51 PatchIdentifiers,
52 SingleCellCoadd,
53 UniformGrid,
54)
55from lsst.daf.butler import DataCoordinate, DeferredDatasetHandle
56from lsst.meas.algorithms import AccumulatorMeanStack
57from lsst.pex.config import ConfigField, ConfigurableField, DictField, Field, ListField, RangeField
58from lsst.pipe.base import (
59 InMemoryDatasetHandle,
60 NoWorkFound,
61 PipelineTask,
62 PipelineTaskConfig,
63 PipelineTaskConnections,
64 Struct,
65)
66from lsst.pipe.base.connectionTypes import Input, Output
67from lsst.pipe.tasks.coaddBase import makeSkyInfo, removeMaskPlanes, setRejectedMaskMapping
68from lsst.pipe.tasks.interpImage import InterpImageTask
69from lsst.pipe.tasks.scaleZeroPoint import ScaleZeroPointTask
70from lsst.skymap import BaseSkyMap
73@dataclasses.dataclass
74class WarpInputs:
75 """Collection of associate inputs along with warps."""
77 warp: DeferredDatasetHandle | InMemoryDatasetHandle
78 """Handle for the warped exposure."""
80 masked_fraction: DeferredDatasetHandle | InMemoryDatasetHandle | None = None
81 """Handle for the masked fraction image."""
83 artifact_mask: DeferredDatasetHandle | InMemoryDatasetHandle | None = None
84 """Handle for the CompareWarp artifact mask."""
86 noise_warps: list[DeferredDatasetHandle | InMemoryDatasetHandle] = dataclasses.field(default_factory=list)
87 """List of handles for the noise warps"""
89 @property
90 def dataId(self) -> DataCoordinate:
91 """DataID corresponding to the warp.
93 Returns
94 -------
95 data_id : `~lsst.daf.butler.DataCoordinate`
96 DataID of the warp.
97 """
98 return self.warp.dataId
101class AssembleCellCoaddConnections(
102 PipelineTaskConnections,
103 dimensions=("tract", "patch", "band", "skymap"),
104 defaultTemplates={"inputWarpName": "deep", "outputCoaddSuffix": "Cell"},
105):
106 inputWarps = Input(
107 doc="Input warps",
108 name="{inputWarpName}Coadd_directWarp",
109 storageClass="ExposureF",
110 dimensions=("tract", "patch", "skymap", "visit", "instrument"),
111 deferLoad=True,
112 multiple=True,
113 )
115 maskedFractionWarps = Input(
116 doc="Mask fraction warps",
117 name="{inputWarpName}Coadd_directWarp_maskedFraction",
118 storageClass="ImageF",
119 dimensions=("tract", "patch", "skymap", "visit", "instrument"),
120 deferLoad=True,
121 multiple=True,
122 )
124 artifactMasks = Input(
125 doc="Artifact masks to be applied to the input warps",
126 name="compare_warp_artifact_mask",
127 storageClass="Mask",
128 dimensions=("tract", "patch", "skymap", "visit", "instrument"),
129 deferLoad=True,
130 multiple=True,
131 )
133 visitSummaryList = Input(
134 doc="Input visit-summary catalogs with updated calibration objects. Mainly used for coadd weights.",
135 name="finalVisitSummary",
136 storageClass="ExposureCatalog",
137 dimensions=("instrument", "visit"),
138 deferLoad=True,
139 multiple=True,
140 )
142 skyMap = Input(
143 doc="Input definition of geometry/bbox and projection/wcs. This must be cell-based.",
144 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME,
145 storageClass="SkyMap",
146 dimensions=("skymap",),
147 )
149 multipleCellCoadd = Output(
150 doc="Output multiple cell coadd",
151 name="{inputWarpName}Coadd{outputCoaddSuffix}",
152 storageClass="MultipleCellCoadd",
153 dimensions=("tract", "patch", "band", "skymap"),
154 )
156 def __init__(self, *, config=None):
157 super().__init__(config=config)
159 if not config:
160 return
162 if config.do_calculate_weight_from_warp:
163 del self.visitSummaryList
165 if not config.do_use_artifact_mask:
166 del self.artifactMasks
168 # Dynamically set input connections for noise images, depending on the
169 # number of noise realizations specified in the config.
170 for n in range(config.num_noise_realizations):
171 noise_warps = Input(
172 doc="Input noise warps",
173 name=f"direct_warp_noise{n}",
174 storageClass="MaskedImageF",
175 dimensions=("tract", "patch", "skymap", "visit", "instrument"),
176 deferLoad=True,
177 multiple=True,
178 )
179 setattr(self, f"noise{n}_warps", noise_warps)
182class AssembleCellCoaddConfig(PipelineTaskConfig, pipelineConnections=AssembleCellCoaddConnections):
183 do_interpolate_coadd = Field[bool](doc="Interpolate over pixels with NO_DATA mask set?", default=True)
184 interpolate_coadd = ConfigurableField(
185 target=InterpImageTask,
186 doc="Task to interpolate (and extrapolate) over pixels with NO_DATA mask on cell coadds",
187 )
188 do_scale_zero_point = Field[bool](
189 doc="Scale warps to a common zero point? This is not needed if they have absolute flux calibration.",
190 default=False,
191 deprecated="Now that visits are scaled to nJy it is no longer necessary or "
192 "recommended to scale the zero point, so this will be removed "
193 "after v29.",
194 )
195 scale_zero_point = ConfigurableField(
196 target=ScaleZeroPointTask,
197 doc="Task to scale warps to a common zero point",
198 deprecated="Now that visits are scaled to nJy it is no longer necessary or "
199 "recommended to scale the zero point, so this will be removed "
200 "after v29.",
201 )
202 do_calculate_weight_from_warp = Field[bool](
203 doc="Calculate coadd weight from the input warp? Otherwise, the weight is obtained from the "
204 "visitSummaryList connection. This is meant as a fallback when run outside the pipeline.",
205 default=False,
206 )
207 do_use_artifact_mask = Field[bool](
208 doc="Substitute the mask planes input warp with an alternative artifact mask?",
209 default=True,
210 )
211 do_coadd_inverse_aperture_corrections = Field[bool](
212 doc="Coadd the inverse aperture corrections for each cell? This is formally the more accurate way "
213 "but may be turned off for parity with deepCoadd.",
214 default=False,
215 )
216 min_overlap_fraction = RangeField[float](
217 doc="The minimum overlap fraction required for a single (visit, detector) input to be included in a "
218 "cell.",
219 # A value of 1.0 corresponds to ideal, edge-free cells.
220 # A value of 0.0 corresponds to the deep_coadd style coadds.
221 # This has to be at least 0.5 to ensure that the an input overlaps the
222 # cell center. Inputs will overlap fraction less than 0.25 will
223 # definitely not overlap the cell center.
224 default=1.0,
225 min=0.0,
226 max=1.0,
227 inclusiveMin=True,
228 inclusiveMax=True,
229 )
230 bad_mask_planes = ListField[str](
231 doc="Mask planes that count towards the masked fraction within a cell.",
232 default=("BAD", "NO_DATA", "SAT", "CLIPPED"),
233 )
234 remove_mask_planes = ListField[str](
235 doc="Mask planes to remove before coadding",
236 default=["NOT_DEBLENDED"],
237 )
238 calc_error_from_input_variance = Field[bool](
239 doc="Calculate coadd variance from input variance by stacking "
240 "statistic. Passed to AccumulatorMeanStack.",
241 default=True,
242 )
243 mask_propagation_thresholds = DictField[str, float](
244 doc=(
245 "Threshold (in fractional weight) of rejection at which we "
246 "propagate a mask plane to the coadd; that is, we set the mask "
247 "bit on the coadd if the fraction the rejected frames "
248 "would have contributed exceeds this value."
249 ),
250 default={"SAT": 0.1},
251 )
252 max_maskfrac = RangeField[float](
253 doc="Maximum fraction of masked pixels in a cell. This is currently "
254 "just a placeholder and is not used now",
255 default=0.99,
256 min=0.0,
257 max=1.0,
258 inclusiveMin=True,
259 inclusiveMax=False,
260 )
261 num_noise_realizations = Field[int](
262 default=0,
263 doc=(
264 "Number of noise planes to include in the coadd. "
265 "This should not exceed the corresponding config parameter "
266 "specified in `MakeDirectWarpConfig`. "
267 ),
268 check=lambda x: x >= 0,
269 )
270 psf_warper = ConfigField(
271 doc="Configuration for the warper that warps the PSFs. It must have the same configuration used to "
272 "warp the images.",
273 dtype=afwMath.Warper.ConfigClass,
274 )
275 psf_dimensions = Field[int](
276 default=35,
277 doc="Dimensions of the PSF image stamp size to be assigned to cells (must be odd).",
278 check=lambda x: (x > 0) and (x % 2 == 1),
279 )
282class AssembleCellCoaddTask(PipelineTask):
283 """Assemble a cell-based coadded image from a set of warps.
285 This task reads in the warp one at a time, and accumulates it in all the
286 cells that it completely overlaps with. This is the optimal I/O pattern but
287 this also implies that it is not possible to build one or only a few cells.
289 Each cell coadds is guaranteed to have a well-defined PSF. This is done by
290 1) excluding warps that only partially overlap a cell from that cell coadd;
291 2) interpolating bad pixels in the warps rather than excluding them;
292 3) by computing the coadd as a weighted mean of the warps without clipping;
293 4) by computing the coadd PSF as the weighted mean of the PSF of the warps
294 with the same weights.
296 The cells are (and must be) defined in the skymap, and cannot be configured
297 or redefined here. The cells are assumed to be small enough that the PSF is
298 assumed to be spatially constant within a cell.
300 Raises
301 ------
302 NoWorkFound
303 Raised if no input warps are provided, or no cells could be populated.
304 RuntimeError
305 Raised if the skymap is not cell-based.
307 Notes
308 -----
309 This is not yet a part of the standard DRP pipeline. As such, the Task and
310 especially its Config and Connections are experimental and subject to
311 change any time without a formal RFC or standard deprecation procedures
312 until it is included in the DRP pipeline.
313 """
315 ConfigClass = AssembleCellCoaddConfig
316 _DefaultName = "assembleCellCoadd"
318 def __init__(self, *args, **kwargs):
319 super().__init__(*args, **kwargs)
320 if self.config.do_interpolate_coadd:
321 self.makeSubtask("interpolate_coadd")
322 if self.config.do_scale_zero_point:
323 self.makeSubtask("scale_zero_point")
325 self.psf_warper = afwMath.Warper.fromConfig(self.config.psf_warper)
326 if (warping_kernel_name := self.config.psf_warper.warpingKernelName.lower()).startswith("lanczos"):
327 psf_padding = 2 * int(warping_kernel_name.lstrip("lanczos")) - 1
328 self.log.debug(
329 "Padding PSF image by %d pixels since the warping kernel is %s.",
330 psf_padding,
331 self.config.psf_warper.warpingKernelName,
332 )
333 else:
334 psf_padding = 10
335 self.log.info(
336 "Padding PSF image by %d pixels since the warping kernel is not Lanczos.",
337 psf_padding,
338 )
339 self.psf_padding = psf_padding
341 def runQuantum(self, butlerQC, inputRefs, outputRefs):
342 # Docstring inherited.
343 if not inputRefs.inputWarps:
344 raise NoWorkFound("No input warps provided for co-addition")
345 self.log.info("Found %d input warps", len(inputRefs.inputWarps))
347 # Construct skyInfo expected by run
348 # Do not remove skyMap from inputData in case _makeSupplementaryData
349 # needs it
350 skyMap = butlerQC.get(inputRefs.skyMap)
352 if not skyMap.config.tractBuilder.name == "cells":
353 raise RuntimeError("AssembleCellCoaddTask requires a cell-based skymap.")
355 outputDataId = butlerQC.quantum.dataId
357 skyInfo = makeSkyInfo(skyMap, tractId=outputDataId["tract"], patchId=outputDataId["patch"])
358 visitSummaryList = butlerQC.get(getattr(inputRefs, "visitSummaryList", []))
360 units = CoaddUnits.legacy if self.config.do_scale_zero_point else CoaddUnits.nJy
361 self.common = CommonComponents(
362 units=units,
363 wcs=skyInfo.patchInfo.wcs,
364 band=outputDataId.get("band", None),
365 identifiers=PatchIdentifiers.from_data_id(outputDataId),
366 )
368 inputs: dict[DataCoordinate, WarpInputs] = {}
369 for handle in butlerQC.get(inputRefs.inputWarps):
370 inputs[handle.dataId] = WarpInputs(warp=handle, noise_warps=[])
372 for ref in getattr(inputRefs, "artifactMasks", []):
373 inputs[ref.dataId].artifact_mask = butlerQC.get(ref)
374 for ref in getattr(inputRefs, "maskedFractionWarps", []):
375 inputs[ref.dataId].masked_fraction = butlerQC.get(ref)
376 for n in range(self.config.num_noise_realizations):
377 for ref in getattr(inputRefs, f"noise{n}_warps"):
378 inputs[ref.dataId].noise_warps.append(butlerQC.get(ref))
380 returnStruct = self.run(inputs=inputs, skyInfo=skyInfo, visitSummaryList=visitSummaryList)
381 butlerQC.put(returnStruct, outputRefs)
382 return returnStruct
384 @staticmethod
385 def _compute_weight(maskedImage, statsCtrl):
386 """Compute a weight for a masked image.
388 Parameters
389 ----------
390 maskedImage : `~lsst.afw.image.MaskedImage`
391 The masked image to compute the weight.
392 statsCtrl : `~lsst.afw.math.StatisticsControl`
393 A control (config-like) object for StatisticsStack.
395 Returns
396 -------
397 weight : `float`
398 Inverse of the clipped mean variance of the masked image.
399 """
400 statObj = afwMath.makeStatistics(
401 maskedImage.getVariance(), maskedImage.getMask(), afwMath.MEANCLIP, statsCtrl
402 )
403 meanVar, _ = statObj.getResult(afwMath.MEANCLIP)
404 weight = 1.0 / float(meanVar)
405 return weight
407 @staticmethod
408 def _construct_grid(skyInfo):
409 """Construct a UniformGrid object from a SkyInfo struct.
411 Parameters
412 ----------
413 skyInfo : `~lsst.pipe.base.Struct`
414 A Struct object
416 Returns
417 -------
418 grid : `~lsst.cell_coadds.UniformGrid`
419 A UniformGrid object.
420 """
421 padding = skyInfo.patchInfo.getCellBorder()
422 grid_bbox = skyInfo.patchInfo.outer_bbox.erodedBy(padding)
423 grid = UniformGrid.from_bbox_cell_size(
424 grid_bbox,
425 skyInfo.patchInfo.getCellInnerDimensions(),
426 padding=padding,
427 )
428 return grid
430 def _construct_grid_container(self, skyInfo, statsCtrl):
431 """Construct a grid of AccumulatorMeanStack instances.
433 Parameters
434 ----------
435 skyInfo : `~lsst.pipe.base.Struct`
436 A Struct object
437 statsCtrl : `~lsst.afw.math.StatisticsControl`
438 A control (config-like) object for StatisticsStack.
440 Returns
441 -------
442 gc : `~lsst.cell_coadds.GridContainer`
443 A GridContainer object container one AccumulatorMeanStack per cell.
444 """
445 grid = self._construct_grid(skyInfo)
447 maskMap = setRejectedMaskMapping(statsCtrl)
448 self.log.debug("Obtained maskMap = %s for %s", maskMap, skyInfo.patchInfo)
449 thresholdDict = AccumulatorMeanStack.stats_ctrl_to_threshold_dict(statsCtrl)
451 # Initialize the grid container with AccumulatorMeanStacks
452 gc = GridContainer[AccumulatorMeanStack](grid.shape)
453 for cellInfo in skyInfo.patchInfo:
454 stacker = AccumulatorMeanStack(
455 # The shape is for the numpy arrays, hence transposed.
456 shape=(cellInfo.outer_bbox.height, cellInfo.outer_bbox.width),
457 bit_mask_value=statsCtrl.getAndMask(),
458 mask_threshold_dict=thresholdDict,
459 calc_error_from_input_variance=self.config.calc_error_from_input_variance,
460 compute_n_image=False,
461 mask_map=maskMap,
462 no_good_pixels_mask=statsCtrl.getNoGoodPixelsMask(),
463 )
464 gc[cellInfo.index] = stacker
466 return gc
468 def _construct_stats_control(self):
469 """Construct a StatisticsControl object for coadd.
471 Unlike AssembleCoaddTask or CompareWarpAssembleCoaddTask, there is
472 very little to be configured apart from setting the mask planes and
473 optionally mask propagation thresholds.
475 Returns
476 -------
477 statsCtrl : `~lsst.afw.math.StatisticsControl`
478 A control object for StatisticsStack.
479 """
480 statsCtrl = afwMath.StatisticsControl()
481 # Hardcode the numIter parameter to the default config value set in
482 # CompareWarpAssembleCoaddTask to get consistent weights. This is NOT
483 # exposed as a config parameter, since this is only meant to be a
484 # fallback option that is not recommended for production.
485 statsCtrl.setNumIter(2)
486 statsCtrl.setAndMask(afwImage.Mask.getPlaneBitMask(self.config.bad_mask_planes))
487 statsCtrl.setNanSafe(True)
488 for plane, threshold in self.config.mask_propagation_thresholds.items():
489 bit = afwImage.Mask.getMaskPlane(plane)
490 statsCtrl.setMaskPropagationThreshold(bit, threshold)
491 return statsCtrl
493 def _construct_ap_corr_grid_container(self, skyInfo):
494 """Construct a grid of CoaddApCorrMapStacker instances.
496 Parameters
497 ----------
498 skyInfo : `~lsst.pipe.base.Struct`
499 A Struct object
501 Returns
502 -------
503 gc : `~lsst.cell_coadds.GridContainer`
504 A GridContainer object container one CoaddApCorrMapStacker per
505 cell.
506 """
507 grid = self._construct_grid(skyInfo)
509 # Initialize the grid container with CoaddApCorrMapStacker.
510 gc = GridContainer[CoaddApCorrMapStacker](grid.shape)
511 for cellInfo in skyInfo.patchInfo:
512 stacker = CoaddApCorrMapStacker(
513 evaluation_point=cellInfo.inner_bbox.getCenter(),
514 do_coadd_inverse_ap_corr=self.config.do_coadd_inverse_aperture_corrections,
515 )
516 gc[cellInfo.index] = stacker
518 return gc
520 def run(
521 self,
522 *,
523 inputs: dict[DataCoordinate, WarpInputs],
524 skyInfo,
525 visitSummaryList: list | None = None,
526 ):
527 for mask_plane in self.config.bad_mask_planes:
528 afwImage.Mask.addMaskPlane(mask_plane)
529 for mask_plane in self.config.mask_propagation_thresholds:
530 afwImage.Mask.addMaskPlane(mask_plane)
532 statsCtrl = self._construct_stats_control()
534 warp_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl)
535 maskfrac_stacker_gc = self._construct_grid_container(skyInfo, statsCtrl)
536 noise_stacker_gc_list = [
537 self._construct_grid_container(skyInfo, statsCtrl)
538 for n in range(self.config.num_noise_realizations)
539 ]
540 psf_stacker_gc = GridContainer[AccumulatorMeanStack](warp_stacker_gc.shape)
541 psf_bbox_gc = GridContainer[geom.Box2I](warp_stacker_gc.shape)
542 ap_corr_stacker_gc = self._construct_ap_corr_grid_container(skyInfo)
544 # Make a container to hold the cell centers in sky coordinates now,
545 # so we don't have to recompute them for each warp
546 # (they share a common WCS). These are needed to find the various
547 # warp + detector combinations that contributed to each cell, and later
548 # get the corresponding PSFs as well.
549 cell_centers_sky = GridContainer[geom.SpherePoint](warp_stacker_gc.shape)
550 # Make a container to hold the observation identifiers for each cell.
551 observation_identifiers_gc = GridContainer[dict](warp_stacker_gc.shape)
552 # Populate them.
553 for cellInfo in skyInfo.patchInfo:
554 # Make a list to hold the observation identifiers for each cell.
555 observation_identifiers_gc[cellInfo.index] = {}
556 cell_centers_sky[cellInfo.index] = skyInfo.wcs.pixelToSky(cellInfo.inner_bbox.getCenter())
557 psf_bbox_gc[cellInfo.index] = geom.Box2I.makeCenteredBox(
558 geom.Point2D(cellInfo.inner_bbox.getCenter()),
559 geom.Extent2I(self.config.psf_dimensions, self.config.psf_dimensions),
560 )
561 psf_stacker_gc[cellInfo.index] = AccumulatorMeanStack(
562 # The shape is for the numpy arrays, hence transposed.
563 shape=(self.config.psf_dimensions, self.config.psf_dimensions),
564 bit_mask_value=0,
565 calc_error_from_input_variance=self.config.calc_error_from_input_variance,
566 compute_n_image=False,
567 )
569 # visit_summary do not have (tract, patch, band, skymap) dimensions.
570 if not visitSummaryList:
571 visitSummaryList = []
572 visitSummaryRefDict = {
573 visitSummaryRef.dataId["visit"]: visitSummaryRef for visitSummaryRef in visitSummaryList
574 }
576 # Keep track of the polygons corresponding to each (visit, detector).
577 visit_polygons: dict[ObservationIdentifiers, afwGeom.Polygon] = {}
579 # Read in one warp at a time, and accumulate it in all the cells that
580 # it completely overlaps.
581 for _, warp_input in inputs.items():
582 warp = warp_input.warp.get(parameters={"bbox": skyInfo.bbox})
583 masked_fraction_image = (
584 warp_input.masked_fraction.get(parameters={"bbox": skyInfo.bbox})
585 if warp_input.masked_fraction
586 else None
587 )
589 # Pre-process the warp before coadding.
590 # TODO: Can we get these mask names from artifactMask?
591 warp.mask.addMaskPlane("CLIPPED")
592 warp.mask.addMaskPlane("REJECTED")
593 warp.mask.addMaskPlane("SENSOR_EDGE")
594 warp.mask.addMaskPlane("INEXACT_PSF")
596 if artifact_mask_ref := warp_input.artifact_mask:
597 # Apply the artifact mask to the warp.
598 artifact_mask = artifact_mask_ref.get()
599 assert (
600 warp.mask.getMaskPlaneDict() == artifact_mask.getMaskPlaneDict()
601 ), "Mask dicts do not agree."
602 warp.mask.array = artifact_mask.array
603 del artifact_mask
605 if self.config.do_scale_zero_point:
606 # Each Warp that goes into a coadd will typically have an
607 # independent photometric zero-point. Therefore, we must scale
608 # each Warp to set it to a common photometric zeropoint.
609 imageScaler = self.scale_zero_point.run(exposure=warp, dataRef=warp_input.warp).imageScaler
610 zero_point_scale_factor = imageScaler.scale
611 self.log.debug(
612 "Scaled the warp %s by %f to match zero points",
613 warp_input.dataId,
614 zero_point_scale_factor,
615 )
616 else:
617 zero_point_scale_factor = 1.0
618 if "BUNIT" not in warp.metadata:
619 raise ValueError(f"Warp {warp_input.dataId} has no BUNIT metadata")
620 if warp.metadata["BUNIT"] != "nJy":
621 raise ValueError(
622 f"Warp {warp_input.dataId} has BUNIT {warp.metadata['BUNIT']}, expected nJy"
623 )
625 removeMaskPlanes(warp.mask, self.config.remove_mask_planes, self.log)
626 # Instead of using self.config.bad_mask_planes, we explicitly
627 # ask statsCtrl which pixels are going to be ignored/rejected.
628 rejected = afwImage.Mask.getPlaneBitMask(
629 ["CLIPPED", "REJECTED"] + afwImage.Mask.interpret(statsCtrl.getAndMask()).split(",")
630 )
632 # Compute the weight for each CCD in the warp from the visitSummary
633 # or from the warp itself, if not provided. Computing the weight
634 # from the warp is not recommended, and in that case we compute one
635 # weight per warp and not bother with per-detector weights.
636 full_ccd_table = warp.getInfo().getCoaddInputs().ccds
637 weights: dict[int, float] = dict.fromkeys(
638 full_ccd_table["ccd"],
639 0.0,
640 ) # Mapping from detector to weight.
642 if visitSummaryRef := visitSummaryRefDict.get(warp_input.dataId["visit"]):
643 visitSummary = visitSummaryRef.get()
644 for detector in full_ccd_table["ccd"]:
645 visitSummaryRow = visitSummary.find(detector)
646 mean_variance = visitSummaryRow["meanVar"]
647 mean_variance *= zero_point_scale_factor**2
648 if warp.metadata.get("BUNIT", None) == "nJy":
649 mean_variance *= visitSummaryRow.photoCalib.getCalibrationMean() ** 2
650 weights[detector] = 1.0 / mean_variance
651 del visitSummary
652 else:
653 self.log.debug("No visit summary found for %s; using warp-based weights", warp_input.dataId)
654 weight = self._compute_weight(warp, statsCtrl)
655 if not np.isfinite(weight):
656 self.log.warn("Non-finite weight for %s: skipping", warp_input.dataId)
657 continue
659 for detector in weights:
660 weights[detector] = weight
662 noise_warps = [ref.get(parameters={"bbox": skyInfo.bbox}) for ref in warp_input.noise_warps]
664 # Create an image where each pixel value corresponds to the
665 # detector ID that pixel comes from.
666 detector_map = afwImage.ImageI(bbox=warp.getBBox(), initialValue=-1)
667 for row in full_ccd_table:
668 transform = makeWcsPairTransform(row.wcs, warp.wcs)
669 if (src_polygon := row.validPolygon) is None:
670 src_polygon = afwGeom.Polygon(geom.Box2D(row.getBBox()))
671 try:
672 dest_polygon = src_polygon.transform(transform).intersectionSingle(
673 geom.Box2D(warp.getBBox())
674 )
675 except SinglePolygonException:
676 continue
678 observation_identifier = ObservationIdentifiers.from_data_id(
679 warp_input.dataId,
680 backup_detector=row["ccd"],
681 )
682 visit_polygons[observation_identifier] = dest_polygon
684 detector_map_slice = dest_polygon.createImage(detector_map.getBBox()).array > 0
685 if not (detector_map.array[detector_map_slice] < 0).all():
686 self.log.warning("Multiple detectors from visit %s are overlapping", warp_input.dataId)
687 detector_map.array[detector_map_slice] = row["ccd"]
689 if (detector_map.array < 0).all():
690 self.log.warning("Unable to split the warp %s into single-detector warps.", warp_input.dataId)
691 detector_map.array[:, :] = 0
693 for cellInfo, ccd_row in itertools.product(skyInfo.patchInfo, full_ccd_table):
694 bbox = cellInfo.outer_bbox
695 inner_bbox = cellInfo.inner_bbox
697 overlap_fraction = (detector_map[inner_bbox].array == ccd_row["ccd"]).mean()
698 assert -1e-4 < overlap_fraction < 1.0001, "Overlap fraction is not within [0, 1]."
699 if (overlap_fraction < self.config.min_overlap_fraction) or (overlap_fraction <= 0.0):
700 self.log.debug(
701 "Skipping %s in cell %s because it had only %.3f < %.3f fractional overlap.",
702 warp_input.dataId,
703 cellInfo.index,
704 overlap_fraction,
705 self.config.min_overlap_fraction,
706 )
707 continue
709 weight = weights[ccd_row["ccd"]]
710 if not np.isfinite(weight):
711 self.log.warn(
712 "Non-finite weight for %s in cell %s: skipping", warp_input.dataId, cellInfo.index
713 )
714 continue
716 if weight == 0:
717 self.log.info(
718 "Zero weight for %s in cell %s: skipping", warp_input.dataId, cellInfo.index
719 )
720 continue
722 # Decide if a deep copy is necessary to apply the single
723 # detector cuts since it involves modifying the image in-place.
724 # If within the inner cell, there are three or more different
725 # values that detector map takes, then there are definitely
726 # multiple detectors (one for chip gaps, two for two detectors)
727 deep_copy = len(set(detector_map[inner_bbox].array.ravel())) >= 3
728 if deep_copy:
729 single_detector_mask_array = detector_map[bbox].array != ccd_row["ccd"]
731 mi = afwImage.MaskedImageF(warp[bbox].maskedImage, deep=deep_copy)
732 if deep_copy:
733 mi.image.array[single_detector_mask_array] = 0.0
734 mi.variance.array[single_detector_mask_array] = np.inf
735 nodata_or_mask = (single_detector_mask_array) * afwImage.Mask.getPlaneBitMask("NO_DATA")
736 mi.mask[bbox].array |= nodata_or_mask
737 warp_stacker_gc[cellInfo.index].add_masked_image(mi, weight=weight)
739 if masked_fraction_image:
740 mi = afwImage.ImageF(masked_fraction_image[bbox], deep=deep_copy)
741 if deep_copy:
742 mi.array[single_detector_mask_array] = 0.0
743 maskfrac_stacker_gc[cellInfo.index].add_image(masked_fraction_image[bbox], weight=weight)
745 for n in range(self.config.num_noise_realizations):
746 mi = afwImage.MaskedImageF(noise_warps[n][bbox], deep=deep_copy)
747 if deep_copy:
748 mi.image.array[single_detector_mask_array] = 0.0
749 mi.variance.array[single_detector_mask_array] = np.inf
750 mi.mask[bbox].array |= nodata_or_mask
751 noise_stacker_gc_list[n][cellInfo.index].add_masked_image(mi, weight=weight)
753 # Set the defaults for PSF shape quantities.
754 psf_shape = afwGeom.Quadrupole()
755 psf_shape_flag = True
756 psf_eval_point = None
757 try:
758 # The `if` branch is buggy. `dest_polygon` is technically
759 # out of scope, but Python does not raise an error.
760 # TODO: Fix this properly in DM-53479, but sweep it under
761 # the rug for now.
762 if overlap_fraction < 0.5:
763 psf_eval_point = dest_polygon.intersectionSingle(
764 geom.Box2D(inner_bbox)
765 ).calculateCenter()
766 else:
767 psf_eval_point = geom.Point2D(inner_bbox.getCenter())
768 psf_shape = warp.psf.computeShape(psf_eval_point)
769 psf_shape_flag = False
770 except SinglePolygonException:
771 self.log.info(
772 "Unable to find the overlapping polygon between %d detector in %s and cell %s",
773 ccd_row["ccd"],
774 warp_input.dataId,
775 cellInfo.index,
776 )
777 except InvalidPsfError:
778 self.log.info(
779 "Unable to compute PSF shape from %d detector in %s at %s",
780 ccd_row["ccd"],
781 warp_input.dataId,
782 psf_eval_point,
783 )
785 overlaps_center = detector_map[geom.Point2I(bbox.getCenter())] == ccd_row["ccd"]
787 observation_identifier = ObservationIdentifiers.from_data_id(
788 warp_input.dataId,
789 backup_detector=ccd_row["ccd"],
790 )
791 observation_identifiers_gc[cellInfo.index][observation_identifier] = CoaddInputs(
792 overlaps_center=overlaps_center,
793 overlap_fraction=overlap_fraction,
794 weight=weight,
795 psf_shape=psf_shape,
796 psf_shape_flag=psf_shape_flag,
797 )
798 if overlaps_center is False:
799 self.log.debug(
800 "%s does not overlap with the center of the cell %s",
801 warp_input.dataId,
802 cellInfo.index,
803 )
804 continue
806 # Everything below this has to do with the center of the cell
807 calexp_point = ccd_row.getWcs().skyToPixel(cell_centers_sky[cellInfo.index])
808 undistorted_psf_im = ccd_row.getPsf().computeImage(calexp_point)
810 assert undistorted_psf_im.getBBox() == geom.Box2I.makeCenteredBox(
811 calexp_point,
812 undistorted_psf_im.getDimensions(),
813 ), "PSF image does not share the coordinates of the 'calexp'"
815 # Convert the PSF image from Image to MaskedImage and
816 # zero-pad the image.
817 undistorted_psf_bbox = undistorted_psf_im.getBBox()
818 undistorted_psf_maskedImage = afwImage.MaskedImageD(
819 undistorted_psf_bbox.dilatedBy(self.psf_padding)
820 )
821 undistorted_psf_maskedImage.image[undistorted_psf_bbox].array[:, :] = undistorted_psf_im.array
822 # TODO: In DM-43585, use the variance plane value from noise.
823 undistorted_psf_maskedImage.variance += 1.0 # Set variance to 1
825 warped_psf_maskedImage = self.psf_warper.warpImage(
826 destWcs=skyInfo.wcs,
827 srcImage=undistorted_psf_maskedImage,
828 srcWcs=ccd_row.getWcs(),
829 destBBox=psf_bbox_gc[cellInfo.index],
830 )
832 # There may be NaNs in the PSF image. Set them to 0.0
833 warped_psf_maskedImage.variance.array[np.isnan(warped_psf_maskedImage.image.array)] = 1.0
834 warped_psf_maskedImage.image.array[np.isnan(warped_psf_maskedImage.image.array)] = 0.0
836 psf_stacker = psf_stacker_gc[cellInfo.index]
837 psf_stacker.add_masked_image(warped_psf_maskedImage, weight=weight)
839 if not (0.995 < (psf_normalization := warped_psf_maskedImage.image.array.sum()) < 1.005):
840 self.log.warning(
841 "PSF image for %s in %s is not normalized to 1.0, but instead %f",
842 warp_input.dataId,
843 cellInfo.index,
844 psf_normalization,
845 )
847 if (ap_corr_map := warp.getInfo().getApCorrMap()) is not None:
848 ap_corr_stacker_gc[cellInfo.index].add(ap_corr_map, weight=weight)
850 del warp
852 # Update common with the visit polygons.
853 self.common = dataclasses.replace(
854 self.common,
855 visit_polygons=visit_polygons,
856 )
858 cells: list[SingleCellCoadd] = []
859 for cellInfo in skyInfo.patchInfo:
860 if len(observation_identifiers_gc[cellInfo.index]) == 0:
861 self.log.debug("Skipping cell %s because it has no input warps", cellInfo.index)
862 continue
864 cell_masked_image = afwImage.MaskedImageF(cellInfo.outer_bbox)
865 cell_maskfrac_image = afwImage.ImageF(cellInfo.outer_bbox)
866 cell_noise_images = [
867 afwImage.MaskedImageF(cellInfo.outer_bbox) for n in range(self.config.num_noise_realizations)
868 ]
869 psf_masked_image = afwImage.MaskedImageF(psf_bbox_gc[cellInfo.index])
871 warp_stacker_gc[cellInfo.index].fill_stacked_masked_image(cell_masked_image)
872 maskfrac_stacker_gc[cellInfo.index].fill_stacked_image(cell_maskfrac_image)
873 for n in range(self.config.num_noise_realizations):
874 noise_stacker_gc_list[n][cellInfo.index].fill_stacked_masked_image(cell_noise_images[n])
875 psf_stacker_gc[cellInfo.index].fill_stacked_masked_image(psf_masked_image)
877 if ap_corr_stacker_gc[cellInfo.index].ap_corr_names:
878 ap_corr_map = ap_corr_stacker_gc[cellInfo.index].final_ap_corr_map
879 else:
880 ap_corr_map = None
882 # Post-process the coadd before converting to new data structures.
883 if np.isnan(cell_masked_image.image.array).all():
884 cell_masked_image.image.array[:, :] = 0.0
885 cell_masked_image.variance.array[:, :] = np.inf
886 elif self.config.do_interpolate_coadd:
887 self.interpolate_coadd.run(cell_masked_image, planeName="NO_DATA")
888 for noise_image in cell_noise_images:
889 self.interpolate_coadd.run(noise_image, planeName="NO_DATA")
890 # The variance must be positive; work around for DM-3201.
891 varArray = cell_masked_image.variance.array
892 with np.errstate(invalid="ignore"):
893 varArray[:] = np.where(varArray > 0, varArray, np.inf)
895 afwImage.Mask.addMaskPlane("INEXACT_PSF")
896 cell_masked_image.mask.array[
897 (cell_masked_image.mask.array & rejected) > 0
898 ] |= cell_masked_image.mask.getPlaneBitMask("INEXACT_PSF")
900 image_planes = OwnedImagePlanes.from_masked_image(
901 masked_image=cell_masked_image,
902 mask_fractions=cell_maskfrac_image,
903 noise_realizations=[noise_image.image for noise_image in cell_noise_images],
904 )
905 identifiers = CellIdentifiers(
906 cell=cellInfo.index,
907 skymap=self.common.identifiers.skymap,
908 tract=self.common.identifiers.tract,
909 patch=self.common.identifiers.patch,
910 band=self.common.identifiers.band,
911 )
913 singleCellCoadd = SingleCellCoadd(
914 outer=image_planes,
915 psf=psf_masked_image.image,
916 inner_bbox=cellInfo.inner_bbox,
917 inputs=observation_identifiers_gc[cellInfo.index],
918 common=self.common,
919 identifiers=identifiers,
920 aperture_correction_map=ap_corr_map,
921 )
922 # TODO: Attach transmission curve when they become available.
923 cells.append(singleCellCoadd)
925 if not cells:
926 raise NoWorkFound("No cells could be populated for the cell coadd.")
928 grid = self._construct_grid(skyInfo)
929 multipleCellCoadd = MultipleCellCoadd(
930 cells,
931 grid=grid,
932 outer_cell_size=cellInfo.outer_bbox.getDimensions(),
933 inner_bbox=None,
934 common=self.common,
935 psf_image_size=cells[0].psf_image.getDimensions(),
936 )
938 return Struct(
939 multipleCellCoadd=multipleCellCoadd,
940 )
943class ConvertMultipleCellCoaddToExposureConnections(
944 PipelineTaskConnections,
945 dimensions=("tract", "patch", "band", "skymap"),
946 defaultTemplates={"inputCoaddName": "deep", "inputCoaddSuffix": "Cell"},
947):
948 cellCoaddExposure = Input(
949 doc="Output coadded exposure, produced by stacking input warps",
950 name="{inputCoaddName}Coadd{inputCoaddSuffix}",
951 storageClass="MultipleCellCoadd",
952 dimensions=("tract", "patch", "skymap", "band"),
953 )
955 stitchedCoaddExposure = Output(
956 doc="Output stitched coadded exposure, produced by stacking input warps",
957 name="{inputCoaddName}Coadd{inputCoaddSuffix}_stitched",
958 storageClass="ExposureF",
959 dimensions=("tract", "patch", "skymap", "band"),
960 )
963class ConvertMultipleCellCoaddToExposureConfig(
964 PipelineTaskConfig, pipelineConnections=ConvertMultipleCellCoaddToExposureConnections
965):
966 """A trivial PipelineTaskConfig class for
967 ConvertMultipleCellCoaddToExposureTask.
968 """
971class ConvertMultipleCellCoaddToExposureTask(PipelineTask):
972 """An after burner PipelineTask that converts a cell-based coadd from
973 `MultipleCellCoadd` format to `ExposureF` format.
975 The run method stitches the cell-based coadd into contiguous exposure and
976 returns it in as an `Exposure` object. This is lossy as it preserves only
977 the pixels in the inner bounding box of the cells and discards the values
978 in the buffer region.
980 Notes
981 -----
982 This task has no configurable parameters.
983 """
985 ConfigClass = ConvertMultipleCellCoaddToExposureConfig
986 _DefaultName = "convertMultipleCellCoaddToExposure"
988 def run(self, cellCoaddExposure):
989 return Struct(
990 stitchedCoaddExposure=cellCoaddExposure.stitch().asExposure(),
991 )