Coverage for python/lsst/meas/extensions/scarlet/deconvolveExposureTask.py: 23%
137 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-06 08:34 +0000
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
32from deprecated.sphinx import deprecated
34from . import utils
36log = logging.getLogger(__name__)
38__all__ = [
39 "DeconvolveExposureTask",
40 "DeconvolveExposureConfig",
41 "DeconvolveExposureConnections",
42]
45def calculateUpdateStep(
46 observation: scl.Observation,
47 minScale: float = 0.01,
48 defaultScale: float = 0.1,
49) -> float:
50 """Calculate the scale factor for the update step in deconvolution.
52 For most images this will be 1.0 but for images with low SNR
53 and/or high sparsity (for example LSST u-band images) the scale
54 factor will be less than 1.0.
56 Parameters
57 ----------
58 observation :
59 Scarlet lite Observation.
61 minScale :
62 Minimum allowed scale factor.
64 defaultScale :
65 Default scale factor to return if noise level is non-finite.
67 Returns
68 -------
69 scale : float
70 Scale factor for the update step.
71 """
72 # Calculate sparsity as fraction of unmasked pixels significantly
73 # above noise. Pixels with zero weight (border, NO_DATA, BAD) are
74 # excluded from both numerator and denominator so heavily masked
75 # inputs are not biased toward a small step.
76 noiseLevel = observation.noise_rms[0]
77 # Guard against non-finite or non-positive noise levels
78 if noiseLevel <= 0 or not np.isfinite(noiseLevel):
79 return defaultScale
80 image = observation.images.data[0]
81 validMask = observation.weights.data[0] > 0
82 validPixels = np.sum(validMask)
83 if validPixels == 0:
84 return defaultScale
85 signalMask = (image > 3*noiseLevel) & validMask
86 signalPixels = np.sum(signalMask)
87 sparsity = signalPixels / validPixels
89 if np.any(signalMask):
90 medianSignal = np.median(image[signalMask])
91 snr = medianSignal / noiseLevel
92 else:
93 snr = 1.0
95 # Scale factor that decreases with sparsity and increases with SNR
96 scale = min(1.0, (sparsity * np.sqrt(snr)) / 0.1)
98 return max(minScale, scale)
101@deprecated(
102 reason=(
103 "Use `calculateUpdateStep` instead; the snake_case name is kept "
104 "as a shim. Will be removed after v31."
105 ),
106 version="v30.0",
107 category=FutureWarning,
108)
109def calculate_update_step(
110 observation: scl.Observation,
111 min_scale: float = 0.01,
112 default_scale: float = 0.1,
113) -> float:
114 """Deprecated snake_case alias for `calculateUpdateStep`."""
115 return calculateUpdateStep(
116 observation, minScale=min_scale, defaultScale=default_scale,
117 )
120class DeconvolveExposureConnections(
121 pipeBase.PipelineTaskConnections,
122 dimensions=("tract", "patch", "skymap", "band"),
123 defaultTemplates={"inputCoaddName": "deep"},
124):
125 """Connections for DeconvolveExposureTask"""
127 coadd = cT.Input(
128 doc="Exposure to deconvolve",
129 name="{inputCoaddName}Coadd_calexp",
130 storageClass="ExposureF",
131 dimensions=("tract", "patch", "band", "skymap"),
132 )
134 coadd_cell = cT.Input(
135 doc="Exposure on which to run deblending",
136 name="{inputCoaddName}CoaddCell",
137 storageClass="MultipleCellCoadd",
138 dimensions=("tract", "patch", "band", "skymap")
139 )
141 background = cT.Input(
142 doc="Background model to subtract from the cell-based coadd",
143 name="{inputCoaddName}Coadd_calexp_background",
144 storageClass="Background",
145 dimensions=("tract", "patch", "band", "skymap")
146 )
148 catalog = cT.Input(
149 doc="Catalog of sources detected in the deconvolved image",
150 name="{inputCoaddName}Coadd_mergeDet",
151 storageClass="SourceCatalog",
152 dimensions=("tract", "patch", "skymap"),
153 )
155 deconvolved = cT.Output(
156 doc="Deconvolved exposure",
157 name="deconvolved_{inputCoaddName}_coadd",
158 storageClass="ExposureF",
159 dimensions=("tract", "patch", "band", "skymap"),
160 )
162 def __init__(self, *, config=None):
163 if not config.useFootprints:
164 # Deconvolution will not use input catalog if
165 # footprints are not used
166 self.inputs.remove("catalog")
168 if config.useCellCoadds:
169 del self.coadd
170 else:
171 del self.coadd_cell
172 del self.background
175class DeconvolveExposureConfig(
176 pipeBase.PipelineTaskConfig,
177 pipelineConnections=DeconvolveExposureConnections,
178):
179 """Configuration for DeconvolveExposureTask"""
181 maxIter = pexConfig.Field[int](
182 doc="Maximum number of iterations",
183 default=100,
184 )
185 minIter = pexConfig.Field[int](
186 doc="Minimum number of iterations",
187 default=10,
188 )
189 eRel = pexConfig.Field[float](
190 doc="Relative error threshold",
191 default=1e-3,
192 )
193 backgroundThreshold = pexConfig.Field[float](
194 default=0,
195 doc="Threshold for background subtraction. "
196 "Pixels in the fit below this threshold will be set to zero",
197 )
198 useFootprints = pexConfig.Field[bool](
199 default=True,
200 doc="Use footprints to constrain the deconvolved model",
201 )
202 useCellCoadds = pexConfig.Field[bool](
203 doc="Use cell-based coadd instead of regular coadd?",
204 default=False,
205 )
208class DeconvolveExposureTask(pipeBase.PipelineTask):
209 """Deconvolve an Exposure using scarlet lite."""
211 ConfigClass = DeconvolveExposureConfig
212 _DefaultName = "deconvolveExposure"
214 def __init__(self, initInputs=None, **kwargs):
215 if initInputs is None:
216 initInputs = {}
217 super().__init__(initInputs=initInputs, **kwargs)
219 def runQuantum(self, butlerQC, inputRefs, outputRefs):
220 inputs = butlerQC.get(inputRefs)
222 # Stitch together cell-based coadds (if necessary)
223 if self.config.useCellCoadds:
224 band = inputRefs.coadd_cell.dataId['band']
225 cellCoadd = inputs.pop('coadd_cell')
226 background = inputs.pop('background')
227 coadd = cellCoadd.stitch().asExposure()
228 coadd.image -= background.getImage()
229 else:
230 coadd = inputs.pop("coadd")
231 band = inputRefs.coadd.dataId['band']
233 catalog = inputs.pop('catalog', None)
235 assert not inputs, "runQuantum got more inputs than expected."
236 outputs = self.run(
237 coadd=coadd,
238 catalog=catalog,
239 band=band,
240 )
241 butlerQC.put(outputs, outputRefs)
243 def run(
244 self,
245 coadd: afwImage.Exposure,
246 catalog: afwTable.SourceCatalog | None = None,
247 band: str = 'dummy'
248 ) -> pipeBase.Struct:
249 """Deconvolve an Exposure
251 Parameters
252 ----------
253 coadd :
254 Coadd image to deconvolve
256 catalog :
257 Catalog of sources detected in the merged catalog.
258 This is used to supress noise in regions with no
259 significant flux about the noise in the coadds.
261 band :
262 Band of the coadd image.
263 Since this is a single band task the band isn't really necessary
264 but can be useful for debugging so we keep it as a parameter.
266 Returns
267 -------
268 deconvolved : `pipeBase.Struct`
269 Deconvolved exposure
270 """
271 observation = self._buildObservation(coadd, catalog, band)
273 # Build the per-pixel footprint mask from the catalog, if one
274 # was supplied, so the deconvolution loop only needs to know
275 # about the mask itself rather than how it was derived.
276 if catalog is not None:
277 bbox = coadd.getBBox()
278 width, height = bbox.getDimensions()
279 x0, y0 = bbox.getMin()
280 footprintImage = afwDet.footprintsToNumpy(
281 catalog, shape=(height, width), xy0=(x0, y0)
282 )
283 else:
284 footprintImage = None
286 model, loss = self._deconvolve(observation, footprintImage=footprintImage)
288 exposure = self._modelToExposure(model.data[0], coadd)
289 return pipeBase.Struct(deconvolved=exposure, loss=loss)
291 def _buildObservation(
292 self,
293 coadd: afwImage.Exposure,
294 catalog: afwTable.SourceCatalog | None = None,
295 band: str = 'dummy'
296 ) -> scl.Observation:
297 """Build a scarlet lite Observation from an Exposure.
299 We don't actually use scarlet, but the optimized convolutions
300 using scarlet data products are still useful.
302 Parameters
303 ----------
304 coadd :
305 Coadd image to deconvolve.
306 catalog :
307 Catalog of sources.
308 This is used to find a location for the PSF if it cannot be
309 generated at the center of the coadd.
311 band :
312 Band of the coadd image.
314 """
315 bands = (band,)
316 model_psf = scl.utils.integrated_circular_gaussian(sigma=0.8)
318 # Give zero weight to non-finite pixels
319 weights = np.ones_like(coadd.image.array)
320 weights[~np.isfinite(coadd.image.array)] = 0
322 image = coadd.image.array.copy()
323 # Set non-finite pixels to zero
324 image[~np.isfinite(image)] = 0.0
325 psfCenter = coadd.getBBox().getCenter()
326 if catalog is not None:
327 psf, _, _ = utils.computeNearestPsf(coadd, catalog, band, psfCenter)
328 if psf is None:
329 # There were no valid locations from
330 # which a PSF could be obtained
331 raise pipeBase.NoWorkFound("No valid PSF could be obtained for deconvolution")
332 psf = psf.array
333 else:
334 psf = coadd.getPsf().computeKernelImage(psfCenter).array
336 badPixelMasks = utils.defaultBadPixelMasks
337 badPixels = coadd.mask.getPlaneBitMask(badPixelMasks)
338 mask = coadd.mask.array & badPixels
339 weights[mask > 0] = 0
341 observation = scl.Observation(
342 images=image[None],
343 variance=coadd.variance.array.copy()[None],
344 weights=weights[None],
345 psfs=psf[None],
346 model_psf=model_psf[None],
347 convolution_mode="fft",
348 bands=bands,
349 bbox=utils.bboxToScarletBox(coadd.getBBox()),
350 )
351 return observation
353 def _deconvolve(
354 self,
355 observation: scl.Observation,
356 footprintImage: np.ndarray | None = None,
357 ) -> tuple[scl.Image, list[float]]:
358 """Deconvolve the observed image.
360 Parameters
361 ----------
362 observation :
363 Scarlet lite Observation.
364 footprintImage :
365 Per-pixel mask matching ``observation.images.shape[1:]``.
366 When supplied, the deconvolved model is multiplied by this
367 mask after each iteration so the recovered footprints stay
368 inside the input footprints.
369 """
370 model = observation.images.copy()
371 loss = []
372 step = calculateUpdateStep(observation)
373 for n in range(self.config.maxIter):
374 # cache=True reuses the FFT plan across iterations; the
375 # image shape is stable inside the loop so this is a free
376 # speedup at zero correctness cost.
377 residual = observation.images - observation.convolve(model, cache=True)
378 if np.all(~np.isfinite(residual.data)):
379 self.log.warning(f"Residual is non-finite at iteration {n}, stopping deconvolution")
380 loss.append(-np.inf)
381 break
382 loss.append(-0.5 * np.nansum(residual.data**2))
383 update = observation.convolve(residual, grad=True, cache=True)
384 update.data[:] *= step
385 model += update
386 model.data[(model.data < 0) | ~np.isfinite(model.data)] = 0
387 if footprintImage is not None:
388 model.data[:] *= footprintImage
390 # Check for a diverging model
391 if len(loss) > 1 and loss[-1] < loss[-2]:
392 step = step / 2
393 self.log.warning(f"Loss increased at iteration {n}, decreasing scale to {step}")
395 # Check for convergence
396 if n > self.config.minIter and np.abs(loss[-1] - loss[-2]) < self.config.eRel * np.abs(loss[-1]):
397 break
399 return model, loss
401 def _modelToExposure(self, model: np.ndarray, coadd: afwImage.Exposure) -> afwImage.Exposure:
402 """Convert a deconvolved image array to an Exposure.
404 The output exposure's mask is a deep copy of the input coadd's
405 mask, and its variance plane is fresh and filled with ``inf``.
406 Convolution-then-deconvolution alters the per-pixel noise
407 covariance, so the input coadd's variance no longer describes
408 the deconvolved pixel values; the infinite variance signals
409 "no information about the noise here" and naturally zero-weights
410 these pixels under any inverse-variance scheme. Downstream
411 consumers that need a variance plane must supply their own.
413 Parameters
414 ----------
415 model :
416 Deconvolved image array.
417 coadd :
418 Input coadd exposure; its image dtype, bbox, ``ExposureInfo``,
419 and mask contents are reused.
420 """
421 image = afwImage.Image(
422 array=model,
423 xy0=coadd.getBBox().getMin(),
424 deep=False,
425 dtype=coadd.image.array.dtype,
426 )
427 # Deep-copy the mask and build a fresh inf-filled variance so
428 # the output exposure doesn't alias the input coadd's planes.
429 # The variance is deliberately invalidated because the input's
430 # variance does not describe the deconvolved pixel values.
431 mask = coadd.mask.clone()
432 variance = coadd.variance.Factory(coadd.variance.getBBox())
433 variance.array[:] = np.inf
434 maskedImage = afwImage.MaskedImage(
435 image=image,
436 mask=mask,
437 variance=variance,
438 dtype=coadd.image.array.dtype,
439 )
440 exposure = afwImage.Exposure(
441 maskedImage=maskedImage,
442 exposureInfo=coadd.getInfo(),
443 dtype=coadd.image.array.dtype,
444 )
445 return exposure