Coverage for python / lsst / meas / extensions / scarlet / deconvolveExposureTask.py: 26%
116 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-19 01:31 -0800
1# This file is part of meas_extensions_scarlet.
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/>.
22import logging
24import lsst.afw.detection as afwDet
25import lsst.afw.image as afwImage
26import lsst.afw.table as afwTable
27import lsst.pex.config as pexConfig
28import lsst.pipe.base as pipeBase
29import lsst.pipe.base.connectionTypes as cT
30import lsst.scarlet.lite as scl
31import numpy as np
33from . import utils
35log = logging.getLogger(__name__)
37__all__ = [
38 "DeconvolveExposureTask",
39 "DeconvolveExposureConfig",
40 "DeconvolveExposureConnections",
41]
44def calculate_update_step(observation: scl.Observation) -> float:
45 """Calculate the scale factor for the update step in deconvolution.
47 For most images this will be 1.0 but for images with low SNR
48 and/or high sparsity (for example LSST u-band images) the scale
49 factor will be less than 1.0.
51 Parameters
52 ----------
53 observation :
54 Scarlet lite Observation.
56 Returns
57 -------
58 scale : float
59 Scale factor for the update step.
60 """
61 # Calculate sparsity as fraction of pixels significantly above noise
62 noise_level = observation.noise_rms[0]
63 signal_mask = observation.images.data > 3*noise_level
64 signal_pixels = np.sum(signal_mask)
65 sparsity = signal_pixels / observation.images.data.size
67 if np.any(signal_mask):
68 median_signal = np.median(observation.images.data[signal_mask])
69 snr = median_signal / noise_level
70 else:
71 snr = 1.0
73 # Scale factor that decreases with sparsity and increases with SNR
74 scale = min(1.0, (sparsity * np.sqrt(snr)) / 0.1)
76 return max(0.01, scale)
79class DeconvolveExposureConnections(
80 pipeBase.PipelineTaskConnections,
81 dimensions=("tract", "patch", "skymap", "band"),
82 defaultTemplates={"inputCoaddName": "deep"},
83):
84 """Connections for DeconvolveExposureTask"""
86 coadd = cT.Input(
87 doc="Exposure to deconvolve",
88 name="{inputCoaddName}Coadd_calexp",
89 storageClass="ExposureF",
90 dimensions=("tract", "patch", "band", "skymap"),
91 )
93 coadd_cell = cT.Input(
94 doc="Exposure on which to run deblending",
95 name="{inputCoaddName}CoaddCell",
96 storageClass="MultipleCellCoadd",
97 dimensions=("tract", "patch", "band", "skymap")
98 )
100 background = cT.Input(
101 doc="Background model to subtract from the cell-based coadd",
102 name="{inputCoaddName}Coadd_calexp_background",
103 storageClass="Background",
104 dimensions=("tract", "patch", "band", "skymap")
105 )
107 catalog = cT.Input(
108 doc="Catalog of sources detected in the deconvolved image",
109 name="{inputCoaddName}Coadd_mergeDet",
110 storageClass="SourceCatalog",
111 dimensions=("tract", "patch", "skymap"),
112 )
114 deconvolved = cT.Output(
115 doc="Deconvolved exposure",
116 name="deconvolved_{inputCoaddName}_coadd",
117 storageClass="ExposureF",
118 dimensions=("tract", "patch", "band", "skymap"),
119 )
121 def __init__(self, *, config=None):
122 if not config.useFootprints:
123 # Deconvolution will not use input catalog if
124 # footprints are not used
125 self.inputs.remove("catalog")
127 if config.useCellCoadds:
128 del self.coadd
129 else:
130 del self.coadd_cell
131 del self.background
134class DeconvolveExposureConfig(
135 pipeBase.PipelineTaskConfig,
136 pipelineConnections=DeconvolveExposureConnections,
137):
138 """Configuration for DeconvolveExposureTask"""
140 maxIter = pexConfig.Field[int](
141 doc="Maximum number of iterations",
142 default=100,
143 )
144 minIter = pexConfig.Field[int](
145 doc="Minimum number of iterations",
146 default=10,
147 )
148 eRel = pexConfig.Field[float](
149 doc="Relative error threshold",
150 default=1e-3,
151 )
152 backgroundThreshold = pexConfig.Field[float](
153 default=0,
154 doc="Threshold for background subtraction. "
155 "Pixels in the fit below this threshold will be set to zero",
156 )
157 useFootprints = pexConfig.Field[bool](
158 default=True,
159 doc="Use footprints to constrain the deconvolved model",
160 )
161 useCellCoadds = pexConfig.Field[bool](
162 doc="Use cell-based coadd instead of regular coadd?",
163 default=False,
164 )
167class DeconvolveExposureTask(pipeBase.PipelineTask):
168 """Deconvolve an Exposure using scarlet lite."""
170 ConfigClass = DeconvolveExposureConfig
171 _DefaultName = "deconvolveExposure"
173 def __init__(self, initInputs=None, **kwargs):
174 if initInputs is None:
175 initInputs = {}
176 super().__init__(initInputs=initInputs, **kwargs)
178 def runQuantum(self, butlerQC, inputRefs, outputRefs):
179 inputs = butlerQC.get(inputRefs)
181 # Stitch together cell-based coadds (if necessary)
182 if self.config.useCellCoadds:
183 band = inputRefs.coadd_cell.dataId['band']
184 cellCoadd = inputs.pop('coadd_cell')
185 background = inputs.pop('background')
186 coadd = cellCoadd.stitch().asExposure()
187 coadd.image -= background.getImage()
188 else:
189 coadd = inputs.pop("coadd")
190 band = inputRefs.coadd.dataId['band']
192 catalog = inputs.pop('catalog', None)
194 assert not inputs, "runQuantum got more inputs than expected."
195 outputs = self.run(
196 coadd=coadd,
197 catalog=catalog,
198 band=band,
199 )
200 butlerQC.put(outputs, outputRefs)
202 def run(
203 self,
204 coadd: afwImage.Exposure,
205 catalog: afwTable.SourceCatalog | None = None,
206 band: str = 'dummy'
207 ) -> pipeBase.Struct:
208 """Deconvolve an Exposure
210 Parameters
211 ----------
212 coadd :
213 Coadd image to deconvolve
215 catalog :
216 Catalog of sources detected in the merged catalog.
217 This is used to supress noise in regions with no
218 significant flux about the noise in the coadds.
220 band :
221 Band of the coadd image.
222 Since this is a single band task the band isn't really necessary
223 but can be useful for debugging so we keep it as a parameter.
225 Returns
226 -------
227 deconvolved : `pipeBase.Struct`
228 Deconvolved exposure
229 """
230 observation = self._buildObservation(coadd, catalog, band)
231 self.bbox = coadd.getBBox()
233 # Deconvolve.
234 # Store the loss history for debugging purposes.
235 model, self.loss = self._deconvolve(observation, catalog)
237 # Store the model in an Exposure
238 exposure = self._modelToExposure(model.data[0], coadd)
239 return pipeBase.Struct(deconvolved=exposure)
241 def _buildObservation(
242 self,
243 coadd: afwImage.Exposure,
244 catalog: afwTable.SourceCatalog | None = None,
245 band: str = 'dummy'
246 ) -> scl.Observation:
247 """Build a scarlet lite Observation from an Exposure.
249 We don't actually use scarlet, but the optimized convolutions
250 using scarlet data products are still useful.
252 Parameters
253 ----------
254 coadd :
255 Coadd image to deconvolve.
256 catalog :
257 Catalog of sources.
258 This is used to find a location for the PSF if it cannot be
259 generated at the center of the coadd.
261 band :
262 Band of the coadd image.
264 """
265 bands = (band,)
266 model_psf = scl.utils.integrated_circular_gaussian(sigma=0.8)
268 image = coadd.image.array
269 psfCenter = coadd.getBBox().getCenter()
270 if catalog is not None:
271 psf, _, _ = utils.computeNearestPsf(coadd, catalog, band, psfCenter)
272 if psf is None:
273 # There were no valid locations from
274 # which a PSF could be obtained
275 raise pipeBase.NoWorkFound("No valid PSF could be obtained for deconvolution")
276 psf = psf.array
277 else:
278 psf = coadd.getPsf().computeKernelImage(psfCenter).array
280 weights = np.ones_like(coadd.image.array)
281 badPixelMasks = utils.defaultBadPixelMasks
282 badPixels = coadd.mask.getPlaneBitMask(badPixelMasks)
283 mask = coadd.mask.array & badPixels
284 weights[mask > 0] = 0
286 observation = scl.Observation(
287 images=image.copy()[None],
288 variance=coadd.variance.array.copy()[None],
289 weights=weights[None],
290 psfs=psf[None],
291 model_psf=model_psf[None],
292 convolution_mode="fft",
293 bands=bands,
294 bbox=utils.bboxToScarletBox(coadd.getBBox()),
295 )
296 return observation
298 def _deconvolve(
299 self,
300 observation: scl.Observation,
301 catalog: afwTable.SourceCatalog | None = None,
302 ) -> tuple[scl.Image, list[float]]:
303 """Deconvolve the observed image.
305 Parameters
306 ----------
307 observation :
308 Scarlet lite Observation.
309 catalog :
310 Catalog of sources detected in the deconvolved image.
311 This is used to mask the deconvolved image so that
312 the deconvolved footprints detected downstream will always
313 fit inside of the original footprints.
314 """
315 model = observation.images.copy()
316 loss = []
317 step = calculate_update_step(observation)
318 if catalog is not None:
319 width, height = self.bbox.getDimensions()
320 x0, y0 = self.bbox.getMin()
321 footprintImage = afwDet.footprintsToNumpy(catalog, shape=(height, width), xy0=(x0, y0))
322 for n in range(self.config.maxIter):
323 residual = observation.images - observation.convolve(model)
324 loss.append(-0.5 * np.sum(residual.data**2))
325 update = observation.convolve(residual, grad=True)
326 update.data[:] *= step
327 model += update
328 model.data[model.data < 0] = 0
329 if catalog is not None:
330 # Ensure that the deconvolved model footprints fit
331 # inside of the original footprints by setting regions
332 # outside of the original footprints to zero.
333 model.data[:] *= footprintImage
335 # Check for a diverging model
336 if len(loss) > 1 and loss[-1] < loss[-2]:
337 step = step / 2
338 self.log.warning(f"Loss increased at iteration {n}, decreasing scale to {step}")
340 # Check for convergence
341 if n > self.config.minIter and np.abs(loss[-1] - loss[-2]) < self.config.eRel * np.abs(loss[-1]):
342 break
344 return model, loss
346 def _modelToExposure(self, model: np.ndarray, coadd: afwImage.Exposure) -> afwImage.Exposure:
347 """Convert a scarlet lite Image to an Exposure.
349 Parameters
350 ----------
351 image :
352 Scarlet lite Image.
353 """
354 image = afwImage.Image(
355 array=model,
356 xy0=coadd.getBBox().getMin(),
357 deep=False,
358 dtype=coadd.image.array.dtype,
359 )
360 maskedImage = afwImage.MaskedImage(
361 image=image,
362 mask=coadd.mask,
363 variance=coadd.variance,
364 dtype=coadd.image.array.dtype,
365 )
366 exposure = afwImage.Exposure(
367 maskedImage=maskedImage,
368 exposureInfo=coadd.getInfo(),
369 dtype=coadd.image.array.dtype,
370 )
371 return exposure