Coverage for python/lsst/atmospec/extraction.py: 16%
216 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-04-09 03:15 -0700
« prev ^ index » next coverage.py v6.5.0, created at 2023-04-09 03:15 -0700
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/>.
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
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
37__all__ = ['SpectralExtractionTask', 'SpectralExtractionTaskConfig']
39PREVENT_RUNAWAY = False
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 )
80class SpectralExtractionTask(pipeBase.Task):
82 ConfigClass = SpectralExtractionTaskConfig
83 _DefaultName = "spectralExtraction"
85 def __init__(self, **kwargs):
86 super().__init__(**kwargs)
87 self.debug = lsstDebug.Info(__name__)
89 def initialise(self, exp, sourceCentroid, spectrumBbox, dispersionRelation):
90 """xxx Docstring here.
92 Parameters:
93 -----------
94 par : `type
95 Description
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()
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)
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)
122 # profiles - Gauss + Moffat
123 # self.gmIntegralGM = np.zeros(self.spectrumHeight)
124 # self.gmIntegralG = np.zeros(self.spectrumHeight)
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)
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
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
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])
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)
164 return
166 def _calculateBackground(self, maskedImage, nbins, ignorePlanes=['DETECTED', 'BAD', 'SAT'], smooth=True):
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
174 nx = 1
175 ny = nbins
176 sctrl = afwMath.StatisticsControl()
177 MaskPixel = afwImage.MaskPixel
178 sctrl.setAndMask(afwImage.Mask[MaskPixel].getPlaneBitMask(ignorePlanes))
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
187 # xxx consider adding a custom mask plane with GROW set high after
188 # detection to allow better masking
190 bctrl = afwMath.BackgroundControl(nx, ny, sctrl, statType)
191 bkgd = afwMath.makeBackground(maskedImage, bctrl)
193 bgImg = bkgd.getImageF(afwMath.Interpolate.CONSTANT)
195 if not smooth:
196 return bgImg
198 # Note that nbins functions differently for scipy.interp1d than for
199 # afwMath.BackgroundControl
201 nbins += 1 # if you want 1 bin you must specify two to getSamplePoints() because of ends
203 kind = 'cubic' if nbins >= 4 else 'linear' # note nbinbs has been incrememented
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]
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
214 interpVals = interpFunc(xsNew)
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
219 return bgImg
221 def getFluxBasic(self):
222 """Docstring here."""
224 # xxx check if this is modified and dispose of copy if not
225 # footprint = self.exp[self.spectrumBbox]...
226 # ...maskedImage.image.array.copy()
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))
236 # footprintMi = copy.copy(self.exp[self.spectrumBbox].maskedImage)
237 footprintArray = maskedImage.image.array
239 # delete this next xxx merlin
240 # residuals = np.zeros([len(pixels), self.spectrumHeight])
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')
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]
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)
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')
274 initialPars = [psfAmp, psfMu, psfSigma] # use value from previous iteration
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)
285 except (RuntimeError, ValueError) as e:
286 self.log.warn(f'\nRuntimeError for basic 1D Gauss fit! rowNum = {rowNum}\n{e}')
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}')
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
320 fitValsGM = self.gaussMoffatFit(pixels, footprintSlice, initialPars)
322 self.gausMoffatFitPars[rowNum] = fitValsGM
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)
329 '''Filling in the residuals'''
330 # fit = self.gauss1D(pixels, *coeffs)
331 # residuals[:, rowNum] = fit-footprintSlice #currently at Moffat
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()
350 ##########################################################
351 # after the row-by-row processing
352 # maxRow = argMaxNd(footprintArray)[0]
354 if self.config.writeResiduals:
355 self.log.warn('Not implemented yet')
357 return self
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
371 @staticmethod
372 def gauss1D(x, *pars):
373 amp, mu, sigma = pars
374 return amp*np.exp(-(x-mu)**2/(2.*sigma**2))
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)
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]
387 return integral, mof.x_0.value, mof.gamma.value, mof.alpha.value
389 @staticmethod
390 def gaussMoffatFit(pixels, footprint, initialPars):
391 model = gausMoffatModel(initialPars)
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.
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.
408 fitter = fitting.LevMarLSQFitter()
409 psf = fitter(model, pixels, footprint)
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
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
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
434def moffatModel(pars):
435 amp, mu, sigma = pars
436 moff = models.Moffat1D(amplitude=amp, x_0=mu, gamma=sigma)
437 return moff