Coverage for python/lsst/atmospec/extraction.py: 16%

216 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-14 10:49 +0000

1# This file is part of atmospec. 

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 numpy as np 

23from scipy.optimize import curve_fit 

24from scipy import integrate 

25from scipy.interpolate import interp1d 

26from astropy.modeling import models, fitting 

27import matplotlib.pyplot as pl 

28 

29import lsstDebug 

30import lsst.pipe.base as pipeBase 

31import lsst.pex.config as pexConfig 

32import lsst.afw.math as afwMath 

33import lsst.afw.image as afwImage 

34from .utils import getSamplePoints # , argMaxNd 

35 

36 

37__all__ = ['SpectralExtractionTask', 'SpectralExtractionTaskConfig'] 

38 

39PREVENT_RUNAWAY = False 

40 

41 

42class SpectralExtractionTaskConfig(pexConfig.Config): 

43 perRowBackground = pexConfig.Field( 

44 doc="If True, subtract backgroud per-row, else subtract for whole image", 

45 dtype=bool, 

46 default=False, 

47 ) 

48 perRowBackgroundSize = pexConfig.Field( 

49 doc="Background box size (or width, if perRowBackground)", 

50 dtype=int, 

51 default=10, 

52 ) 

53 writeResiduals = pexConfig.Field( 

54 doc="Background box size (or width, if perRowBackground)", 

55 dtype=bool, 

56 default=False, 

57 ) 

58 doSmoothBackround = pexConfig.Field( 

59 doc="Spline smooth the 1 x n background boxes?", 

60 dtype=bool, 

61 default=True, 

62 ) 

63 doSigmaClipBackground = pexConfig.Field( 

64 doc="Sigma clip the background model in addition to masking detected (and other) pixels?", 

65 dtype=bool, 

66 default=False 

67 ) 

68 nSigmaClipBackground = pexConfig.Field( 

69 doc="Number of sigma to clip to if appying sigma clipping to background model", 

70 dtype=float, 

71 default=5 

72 ) 

73 nSigmaClipBackgroundIterations = pexConfig.Field( 

74 doc="Number of iterations if appying sigma clipping to background model", 

75 dtype=int, 

76 default=4 

77 ) 

78 

79 

80class SpectralExtractionTask(pipeBase.Task): 

81 

82 ConfigClass = SpectralExtractionTaskConfig 

83 _DefaultName = "spectralExtraction" 

84 

85 def __init__(self, **kwargs): 

86 super().__init__(**kwargs) 

87 self.debug = lsstDebug.Info(__name__) 

88 

89 def initialise(self, exp, sourceCentroid, spectrumBbox, dispersionRelation): 

90 """xxx Docstring here. 

91 

92 Parameters: 

93 ----------- 

94 par : `type 

95 Description 

96 

97 """ 

98 # xxx if the rest of exp is never used, remove this 

99 # and just pass exp[spectrumBbox] 

100 self.expRaw = exp 

101 self.footprintExp = exp[spectrumBbox] 

102 self.footprintMi = self.footprintExp.maskedImage 

103 self.sourceCentroid = sourceCentroid 

104 self.spectrumBbox = spectrumBbox 

105 self.dispersionRelation = dispersionRelation 

106 self.spectrumWidth = self.spectrumBbox.getWidth() 

107 self.spectrumHeight = self.spectrumBbox.getHeight() 

108 

109 # profiles - aperture and psf 

110 self.apertureFlux = np.zeros(self.spectrumHeight) 

111 self.rowWiseMax = np.zeros(self.spectrumHeight) 

112 # self.psfSigma = np.zeros(self.spectrumHeight) 

113 # self.psfMu = np.zeros(self.spectrumHeight) 

114 # self.psfFlux = np.zeros(self.spectrumHeight) 

115 

116 # profiles - Moffat 

117 # self.moffatFlux = np.zeros(self.spectrumHeight) 

118 # self.moffatX0 = np.zeros(self.spectrumHeight) 

119 # self.moffatGamma = np.zeros(self.spectrumHeight) 

120 # self.moffatAlpha = np.zeros(self.spectrumHeight) 

121 

122 # profiles - Gauss + Moffat 

123 # self.gmIntegralGM = np.zeros(self.spectrumHeight) 

124 # self.gmIntegralG = np.zeros(self.spectrumHeight) 

125 

126 # probably delete these - from the GausMoffat fitting 

127 # self.amplitude_0 = np.zeros(self.spectrumHeight) 

128 # self.x_0_0 = np.zeros(self.spectrumHeight) 

129 # self.gamma_0 = np.zeros(self.spectrumHeight) 

130 # self.alpha_0 = np.zeros(self.spectrumHeight) 

131 # self.amplitude_1 = np.zeros(self.spectrumHeight) 

132 # self.mean_1 = np.zeros(self.spectrumHeight) 

133 # self.stddev_1 = np.zeros(self.spectrumHeight) 

134 

135 # each will be an list of length self.spectrumHeight 

136 # with each element containing all the fit parameters 

137 self.psfFitPars = [None] * self.spectrumHeight 

138 self.moffatFitPars = [None] * self.spectrumHeight 

139 self.gausMoffatFitPars = [None] * self.spectrumHeight 

140 

141 if self.debug.display: 

142 try: 

143 import lsst.afw.display as afwDisp 

144 afwDisp.setDefaultBackend(self.debug.displayBackend) 

145 # afwDisp.Display.delAllDisplays() 

146 self.disp1 = afwDisp.Display(0, open=True) 

147 self.disp1.mtv(self.expRaw[self.spectrumBbox]) 

148 self.disp1.erase() 

149 except Exception: 

150 self.log.warn('Failed to initialise debug display') 

151 self.debug.display = False 

152 

153 # xxx probably need to change this once per-spectrum background is done 

154 # xsize = self.spectrumWidth - 2*self.config.perRowBackgroundSize # 20 

155 # residuals = np.zeros([xsize, self.spectrumHeight]) 

156 

157 self.backgroundMi = self._calculateBackground(self.footprintExp.maskedImage, # xxx remove hardcoding 

158 15, smooth=self.config.doSmoothBackround) 

159 self.bgSubMi = self.footprintMi.clone() 

160 self.bgSubMi -= self.backgroundMi 

161 if self.debug.display and 'spectrumBgSub' in self.debug.displayItems: 

162 self.disp1.mtv(self.bgSubMi) 

163 

164 return 

165 

166 def _calculateBackground(self, maskedImage, nbins, ignorePlanes=['DETECTED', 'BAD', 'SAT'], smooth=True): 

167 

168 assert nbins > 0 

169 if nbins > maskedImage.getHeight() - 1: 

170 self.log.warn("More bins selected for background than pixels in image height." 

171 + f" Reducing numbers of bins from {nbins} to {maskedImage.getHeight()-1}.") 

172 nbins = maskedImage.getHeight() - 1 

173 

174 nx = 1 

175 ny = nbins 

176 sctrl = afwMath.StatisticsControl() 

177 MaskPixel = afwImage.MaskPixel 

178 sctrl.setAndMask(afwImage.Mask[MaskPixel].getPlaneBitMask(ignorePlanes)) 

179 

180 if self.config.doSigmaClipBackground: 

181 sctrl.setNumSigmaClip(self.config.nSigmaClipBackground) 

182 sctrl.setNumIter(self.config.nSigmaClipBackgroundIterations) 

183 statType = afwMath.MEANCLIP 

184 else: 

185 statType = afwMath.MEAN 

186 

187 # xxx consider adding a custom mask plane with GROW set high after 

188 # detection to allow better masking 

189 

190 bctrl = afwMath.BackgroundControl(nx, ny, sctrl, statType) 

191 bkgd = afwMath.makeBackground(maskedImage, bctrl) 

192 

193 bgImg = bkgd.getImageF(afwMath.Interpolate.CONSTANT) 

194 

195 if not smooth: 

196 return bgImg 

197 

198 # Note that nbins functions differently for scipy.interp1d than for 

199 # afwMath.BackgroundControl 

200 

201 nbins += 1 # if you want 1 bin you must specify two to getSamplePoints() because of ends 

202 

203 kind = 'cubic' if nbins >= 4 else 'linear' # note nbinbs has been incrememented 

204 

205 # bgImg is now an image the same size as the input, as a column of 

206 # 1 x n blocks, which we now interpolate to make a smooth background 

207 xs = getSamplePoints(0, bgImg.getHeight()-1, nbins, includeEndpoints=True, integers=True) 

208 vals = [bgImg[0, point, afwImage.LOCAL] for point in xs] 

209 

210 interpFunc = interp1d(xs, vals, kind=kind) 

211 xsNew = np.linspace(0, bgImg.getHeight()-1, bgImg.getHeight()) # integers over range 

212 assert xsNew[1]-xsNew[0] == 1 # easy to make off-by-one errors here 

213 

214 interpVals = interpFunc(xsNew) 

215 

216 for col in range(bgImg.getWidth()): # there must be a more pythonic way to tile these into the image 

217 bgImg.array[:, col] = interpVals 

218 

219 return bgImg 

220 

221 def getFluxBasic(self): 

222 """Docstring here.""" 

223 

224 # xxx check if this is modified and dispose of copy if not 

225 # footprint = self.exp[self.spectrumBbox]... 

226 # ...maskedImage.image.array.copy() 

227 

228 if not self.config.perRowBackground: 

229 maskedImage = self.bgSubMi 

230 pixels = np.arange(self.spectrumWidth) 

231 else: 

232 maskedImage = self.footprintMi 

233 pixels = np.arange(0 + self.config.perRowBackgroundSize, 

234 (self.spectrumWidth - self.config.perRowBackgroundSize)) 

235 

236 # footprintMi = copy.copy(self.exp[self.spectrumBbox].maskedImage) 

237 footprintArray = maskedImage.image.array 

238 

239 # delete this next xxx merlin 

240 # residuals = np.zeros([len(pixels), self.spectrumHeight]) 

241 

242 psfAmp = np.max(footprintArray[0]) 

243 psfSigma = np.std(footprintArray[0]) 

244 psfMu = np.argmax(footprintArray[0]) 

245 if abs(psfMu - self.spectrumWidth/2.) >= 10: 

246 self.log.warn('initial mu more than 10 pixels from footprint center') 

247 

248 # loop over the rows, calculating basic parameters 

249 for rowNum in range(self.spectrumHeight): # take row slices 

250 if self.config.perRowBackground: 

251 footprintSlice = self.subtractBkgd(footprintArray[rowNum], self.spectrumWidth, 

252 self.config.perRowBackgroundSize) 

253 else: 

254 footprintSlice = footprintArray[rowNum] 

255 

256 # improve this to be a window around max point in row at least 

257 self.apertureFlux[rowNum] = np.sum(footprintSlice) 

258 self.rowWiseMax[rowNum] = np.max(footprintSlice) 

259 

260 if PREVENT_RUNAWAY: 

261 # xxx values seem odd, probably need changing. 

262 # xxx Should be independent of width. Also should check if 

263 # redoing this each time is better/worse 

264 # if so then psfAmp = np.max(footprintArray[0]) 

265 if ((psfMu > .7*self.spectrumWidth) or (psfMu < 0.3*self.spectrumWidth)): 

266 # psfMu = width/2. # Augustin's method 

267 psfMu = np.argmax(footprintSlice) 

268 self.log.warn(f'psfMu was running away on row {rowNum} - reset to nominal') 

269 if ((psfSigma > 20.) or (psfSigma < 0.1)): 

270 # psfSigma = 3. # Augustin's method 

271 psfSigma = np.std(footprintSlice) 

272 self.log.warn(f'psfSigma was running away on row {rowNum} - reset to nominal') 

273 

274 initialPars = [psfAmp, psfMu, psfSigma] # use value from previous iteration 

275 

276 try: 

277 (psfAmp, psfMu, psfSigma), varMatrix = \ 

278 curve_fit(self.gauss1D, pixels, footprintSlice, p0=initialPars, 

279 bounds=(0., [100*np.max(footprintSlice), 

280 self.spectrumWidth, 2*self.spectrumWidth])) 

281 psfFlux = np.sqrt(2*np.pi) * psfSigma * psfAmp # Gaussian integral 

282 # parErr = np.sqrt(np.diag(varMatrix)) 

283 self.psfFitPars[rowNum] = (psfAmp, psfMu, psfSigma, psfFlux) 

284 

285 except (RuntimeError, ValueError) as e: 

286 self.log.warn(f'\nRuntimeError for basic 1D Gauss fit! rowNum = {rowNum}\n{e}') 

287 

288 try: 

289 fitValsMoffat = self.moffatFit(pixels, footprintSlice, psfAmp, psfMu, psfSigma) 

290 self.moffatFitPars[rowNum] = fitValsMoffat 

291 except (RuntimeError, ValueError) as e: 

292 self.log.warn(f'\n\nRuntimeError during Moffat fit! rowNum = {rowNum}\n{e}') 

293 

294 try: 

295 if rowNum == 0: # bootstrapping, hence all the noqa: F821 

296 initialPars = (np.max(footprintSlice), 

297 np.argmax(footprintSlice), 

298 0.5, 

299 0.5, 

300 self.psfFitPars[rowNum][3]/2, 

301 np.argmax(footprintSlice), 

302 0.5) 

303 elif self.psfFitPars[rowNum] is None: # basic PSF fitting failed! 

304 initialPars = (np.sum(footprintSlice), 

305 fitValsGM[3], # noqa: F821 

306 fitValsGM[4], # noqa: F821 

307 fitValsGM[5], # noqa: F821 

308 fitValsGM[6], # noqa: F821 

309 fitValsGM[7], # noqa: F821 

310 fitValsGM[8]) # noqa: F821 

311 else: 

312 initialPars = (self.psfFitPars[rowNum][3], 

313 fitValsGM[3], # noqa: F821 

314 fitValsGM[4], # noqa: F821 

315 fitValsGM[5], # noqa: F821 

316 fitValsGM[6], # noqa: F821 

317 fitValsGM[7], # noqa: F821 

318 fitValsGM[8]) # noqa: F821 

319 

320 fitValsGM = self.gaussMoffatFit(pixels, footprintSlice, initialPars) 

321 

322 self.gausMoffatFitPars[rowNum] = fitValsGM 

323 

324 except (RuntimeError, ValueError) as e: 

325 msg = f'\n\nRuntimeError during GaussMoffatFit fit! rowNum = {rowNum}\n{e}' 

326 self.log.warn(msg) # serious, and should never happen 

327 # self.gausMoffatFitPars.append(fitValsGM) 

328 

329 '''Filling in the residuals''' 

330 # fit = self.gauss1D(pixels, *coeffs) 

331 # residuals[:, rowNum] = fit-footprintSlice #currently at Moffat 

332 

333 if self.debug.plot and ('all' in self.debug.plot or 'GausMoffat' in self.debug.plot): 

334 # if((not rowNum % 10) and (rowNum < 400)): 

335 if True: 

336 print('aper : ', rowNum, self.apertureFlux[rowNum]) 

337 print(rowNum, self.psf_gauss_flux[rowNum], self.psf_gauss_psfSigma[rowNum], 

338 self.psf_gauss_psfMu[rowNum], psfAmp) 

339 # fit = Gauss1D(pixels, *coeffs) 

340 pl.xlabel('Spectrum spatial profile (pixel)') 

341 pl.ylabel('Amplitude (ADU)') 

342 pl.title('CTIO .9m - %s'%(self.spectrum.object_name)) 

343 # pl.plot(pixels, fit, label='Gauss') 

344 pl.yscale('log') 

345 pl.ylim(1., 1E6) 

346 pl.plot(pixels, footprintSlice) 

347 pl.legend() 

348 pl.show() 

349 

350 ########################################################## 

351 # after the row-by-row processing 

352 # maxRow = argMaxNd(footprintArray)[0] 

353 

354 if self.config.writeResiduals: 

355 self.log.warn('Not implemented yet') 

356 

357 return self 

358 

359 @staticmethod 

360 def subtractBkgd(slice, width, bkgd_size): 

361 # xxx remove these slices - seems totally unnecessary 

362 subFootprint = slice[bkgd_size:(-1-bkgd_size+1)] 

363 bkgd1 = slice[:bkgd_size] 

364 bkgd2 = slice[width-bkgd_size:] 

365 bkgd1 = np.median(bkgd1) 

366 bkgd2 = np.median(bkgd2) 

367 bkgd = np.mean([bkgd1, bkgd2]) 

368 subFootprint = subFootprint - bkgd 

369 return subFootprint 

370 

371 @staticmethod 

372 def gauss1D(x, *pars): 

373 amp, mu, sigma = pars 

374 return amp*np.exp(-(x-mu)**2/(2.*sigma**2)) 

375 

376 @staticmethod 

377 def moffatFit(pixels, footprint, amp, mu, sigma): 

378 pars = (amp, mu, sigma) 

379 initialMoffat = moffatModel(pars) 

380 fitter = fitting.LevMarLSQFitter() 

381 mof = fitter(initialMoffat, pixels, footprint) 

382 

383 start = mof.x_0 - 5 * mof.gamma 

384 end = mof.x_0 + 5 * mof.gamma 

385 integral = (integrate.quad(lambda pixels: mof(pixels), start, end))[0] 

386 

387 return integral, mof.x_0.value, mof.gamma.value, mof.alpha.value 

388 

389 @staticmethod 

390 def gaussMoffatFit(pixels, footprint, initialPars): 

391 model = gausMoffatModel(initialPars) 

392 

393 model.amplitude_1.min = 0. 

394 model.x_0_1.min = min(footprint)-5 

395 model.x_0_1.max = max(pixels)+5 

396 # model.amplitude_1.max = gausAmp/10. 

397 model.gamma_1.min = 1. 

398 model.gamma_1.max = 2. 

399 model.alpha_1.min = 1. 

400 model.alpha_1.max = 2. 

401 

402 model.amplitude_0.min = 0. 

403 model.mean_0.min = min(footprint)-5 

404 model.mean_0.max = max(footprint)+5 

405 model.stddev_0.min = 0.5 

406 model.stddev_0.max = 5. 

407 

408 fitter = fitting.LevMarLSQFitter() 

409 psf = fitter(model, pixels, footprint) 

410 

411 start = psf.x_0_1 - 10 * psf.stddev_0 

412 end = psf.x_0_1 + 10 * psf.stddev_0 

413 intGM = (integrate.quad(lambda pixels: psf(pixels), start, end))[0] 

414 intG = np.sqrt(2 * np.pi) * psf.stddev_0 * psf.amplitude_0 

415 

416 fitGausAmp = psf.amplitude_1.value 

417 fitX0 = psf.x_0_1.value 

418 fitGam = psf.gamma_1.value 

419 fitAlpha = psf.alpha_1.value 

420 fitMofAmp = psf.amplitude_0.value 

421 fitMofMean = psf.mean_0.value 

422 fitMofWid = psf.stddev_0.value 

423 return intGM, intG, fitGausAmp, fitX0, fitGam, fitAlpha, fitMofAmp, fitMofMean, fitMofWid 

424 

425 

426def gausMoffatModel(pars): 

427 moffatAmp, x0, gamma, alpha, gausAmp, mu, gausSigma = pars 

428 moff = models.Moffat1D(amplitude=moffatAmp, x_0=x0, gamma=gamma, alpha=alpha) 

429 gaus = models.Gaussian1D(amplitude=gausAmp, mean=mu, stddev=gausSigma) 

430 model = gaus + moff 

431 return model 

432 

433 

434def moffatModel(pars): 

435 amp, mu, sigma = pars 

436 moff = models.Moffat1D(amplitude=amp, x_0=mu, gamma=sigma) 

437 return moff