Coverage for python/lsst/meas/extensions/scarlet/deconvolveExposureTask.py: 26%

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/>. 

21 

22import logging 

23 

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 

32 

33from . import utils 

34 

35log = logging.getLogger(__name__) 

36 

37__all__ = [ 

38 "DeconvolveExposureTask", 

39 "DeconvolveExposureConfig", 

40 "DeconvolveExposureConnections", 

41] 

42 

43 

44def calculate_update_step(observation: scl.Observation) -> float: 

45 """Calculate the scale factor for the update step in deconvolution. 

46 

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. 

50 

51 Parameters 

52 ---------- 

53 observation : 

54 Scarlet lite Observation. 

55 

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 

66 

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 

72 

73 # Scale factor that decreases with sparsity and increases with SNR 

74 scale = min(1.0, (sparsity * np.sqrt(snr)) / 0.1) 

75 

76 return max(0.01, scale) 

77 

78 

79class DeconvolveExposureConnections( 

80 pipeBase.PipelineTaskConnections, 

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

82 defaultTemplates={"inputCoaddName": "deep"}, 

83): 

84 """Connections for DeconvolveExposureTask""" 

85 

86 coadd = cT.Input( 

87 doc="Exposure to deconvolve", 

88 name="{inputCoaddName}Coadd_calexp", 

89 storageClass="ExposureF", 

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

91 ) 

92 

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 ) 

99 

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 ) 

106 

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 ) 

113 

114 deconvolved = cT.Output( 

115 doc="Deconvolved exposure", 

116 name="deconvolved_{inputCoaddName}_coadd", 

117 storageClass="ExposureF", 

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

119 ) 

120 

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") 

126 

127 if config.useCellCoadds: 

128 del self.coadd 

129 else: 

130 del self.coadd_cell 

131 del self.background 

132 

133 

134class DeconvolveExposureConfig( 

135 pipeBase.PipelineTaskConfig, 

136 pipelineConnections=DeconvolveExposureConnections, 

137): 

138 """Configuration for DeconvolveExposureTask""" 

139 

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 ) 

165 

166 

167class DeconvolveExposureTask(pipeBase.PipelineTask): 

168 """Deconvolve an Exposure using scarlet lite.""" 

169 

170 ConfigClass = DeconvolveExposureConfig 

171 _DefaultName = "deconvolveExposure" 

172 

173 def __init__(self, initInputs=None, **kwargs): 

174 if initInputs is None: 

175 initInputs = {} 

176 super().__init__(initInputs=initInputs, **kwargs) 

177 

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

179 inputs = butlerQC.get(inputRefs) 

180 inputs['band'] = inputRefs.coadd.dataId['band'] 

181 

182 # Stitch together cell-based coadds (if necessary) 

183 if self.config.useCellCoadds: 

184 cellCoadd = inputs.pop('coadd_cell') 

185 background = inputs.pop('background') 

186 coadd = cellCoadd.stitch().asExposure() 

187 coadd.image -= background.getImage() 

188 inputs['coadd'] = coadd 

189 

190 outputs = self.run(**inputs) 

191 butlerQC.put(outputs, outputRefs) 

192 

193 def run( 

194 self, 

195 coadd: afwImage.Exposure, 

196 catalog: afwTable.SourceCatalog | None = None, 

197 band: str = 'dummy' 

198 ) -> pipeBase.Struct: 

199 """Deconvolve an Exposure 

200 

201 Parameters 

202 ---------- 

203 coadd : 

204 Coadd image to deconvolve 

205 

206 catalog : 

207 Catalog of sources detected in the merged catalog. 

208 This is used to supress noise in regions with no 

209 significant flux about the noise in the coadds. 

210 

211 band : 

212 Band of the coadd image. 

213 Since this is a single band task the band isn't really necessary 

214 but can be useful for debugging so we keep it as a parameter. 

215 

216 Returns 

217 ------- 

218 deconvolved : `pipeBase.Struct` 

219 Deconvolved exposure 

220 """ 

221 observation = self._buildObservation(coadd, catalog, band) 

222 self.bbox = coadd.getBBox() 

223 

224 # Deconvolve. 

225 # Store the loss history for debugging purposes. 

226 model, self.loss = self._deconvolve(observation, catalog) 

227 

228 # Store the model in an Exposure 

229 exposure = self._modelToExposure(model.data[0], coadd) 

230 return pipeBase.Struct(deconvolved=exposure) 

231 

232 def _buildObservation( 

233 self, 

234 coadd: afwImage.Exposure, 

235 catalog: afwTable.SourceCatalog | None = None, 

236 band: str = 'dummy' 

237 ) -> scl.Observation: 

238 """Build a scarlet lite Observation from an Exposure. 

239 

240 We don't actually use scarlet, but the optimized convolutions 

241 using scarlet data products are still useful. 

242 

243 Parameters 

244 ---------- 

245 coadd : 

246 Coadd image to deconvolve. 

247 catalog : 

248 Catalog of sources. 

249 This is used to find a location for the PSF if it cannot be 

250 generated at the center of the coadd. 

251 

252 band : 

253 Band of the coadd image. 

254 

255 """ 

256 bands = (band,) 

257 model_psf = scl.utils.integrated_circular_gaussian(sigma=0.8) 

258 

259 image = coadd.image.array 

260 psfCenter = coadd.getBBox().getCenter() 

261 if catalog is not None: 

262 psf, _, _ = utils.computeNearestPsf(coadd, catalog, band, psfCenter) 

263 if psf is None: 

264 # There were no valid locations from 

265 # which a PSF could be obtained 

266 raise pipeBase.NoWorkFound("No valid PSF could be obtained for deconvolution") 

267 psf = psf.array 

268 else: 

269 psf = coadd.getPsf().computeKernelImage(psfCenter).array 

270 

271 weights = np.ones_like(coadd.image.array) 

272 badPixelMasks = utils.defaultBadPixelMasks 

273 badPixels = coadd.mask.getPlaneBitMask(badPixelMasks) 

274 mask = coadd.mask.array & badPixels 

275 weights[mask > 0] = 0 

276 

277 observation = scl.Observation( 

278 images=image.copy()[None], 

279 variance=coadd.variance.array.copy()[None], 

280 weights=weights[None], 

281 psfs=psf[None], 

282 model_psf=model_psf[None], 

283 convolution_mode="fft", 

284 bands=bands, 

285 bbox=utils.bboxToScarletBox(coadd.getBBox()), 

286 ) 

287 return observation 

288 

289 def _deconvolve( 

290 self, 

291 observation: scl.Observation, 

292 catalog: afwTable.SourceCatalog | None = None, 

293 ) -> tuple[scl.Image, list[float]]: 

294 """Deconvolve the observed image. 

295 

296 Parameters 

297 ---------- 

298 observation : 

299 Scarlet lite Observation. 

300 catalog : 

301 Catalog of sources detected in the deconvolved image. 

302 This is used to mask the deconvolved image so that 

303 the deconvolved footprints detected downstream will always 

304 fit inside of the original footprints. 

305 """ 

306 model = observation.images.copy() 

307 loss = [] 

308 step = calculate_update_step(observation) 

309 if catalog is not None: 

310 width, height = self.bbox.getDimensions() 

311 x0, y0 = self.bbox.getMin() 

312 footprintImage = afwDet.footprintsToNumpy(catalog, shape=(height, width), xy0=(x0, y0)) 

313 for n in range(self.config.maxIter): 

314 residual = observation.images - observation.convolve(model) 

315 loss.append(-0.5 * np.sum(residual.data**2)) 

316 update = observation.convolve(residual, grad=True) 

317 update.data[:] *= step 

318 model += update 

319 model.data[model.data < 0] = 0 

320 if catalog is not None: 

321 # Ensure that the deconvolved model footprints fit 

322 # inside of the original footprints by setting regions 

323 # outside of the original footprints to zero. 

324 model.data[:] *= footprintImage 

325 

326 # Check for a diverging model 

327 if len(loss) > 1 and loss[-1] < loss[-2]: 

328 step = step / 2 

329 self.log.warning(f"Loss increased at iteration {n}, decreasing scale to {step}") 

330 

331 # Check for convergence 

332 if n > self.config.minIter and np.abs(loss[-1] - loss[-2]) < self.config.eRel * np.abs(loss[-1]): 

333 break 

334 

335 return model, loss 

336 

337 def _modelToExposure(self, model: np.ndarray, coadd: afwImage.Exposure) -> afwImage.Exposure: 

338 """Convert a scarlet lite Image to an Exposure. 

339 

340 Parameters 

341 ---------- 

342 image : 

343 Scarlet lite Image. 

344 """ 

345 image = afwImage.Image( 

346 array=model, 

347 xy0=coadd.getBBox().getMin(), 

348 deep=False, 

349 dtype=coadd.image.array.dtype, 

350 ) 

351 maskedImage = afwImage.MaskedImage( 

352 image=image, 

353 mask=coadd.mask, 

354 variance=coadd.variance, 

355 dtype=coadd.image.array.dtype, 

356 ) 

357 exposure = afwImage.Exposure( 

358 maskedImage=maskedImage, 

359 exposureInfo=coadd.getInfo(), 

360 dtype=coadd.image.array.dtype, 

361 ) 

362 return exposure