Coverage for python/lsst/atmospec/spectraction.py: 15%
252 statements
« prev ^ index » next coverage.py v7.3.2, created at 2023-11-02 12:39 +0000
« prev ^ index » next coverage.py v7.3.2, created at 2023-11-02 12:39 +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/>.
22import logging
23import os
24import numpy as np
25import astropy.coordinates as asCoords
26from astropy import units as u
28from spectractor import parameters
29parameters.CALLING_CODE = "LSST_DM" # this must be set IMMEDIATELY to supress colored logs
31from spectractor.config import load_config, apply_rebinning_to_parameters # noqa: E402
32from spectractor.extractor.images import Image, find_target, turn_image # noqa: E402
34from spectractor.extractor.dispersers import Hologram # noqa: E402
35from spectractor.extractor.extractor import (FullForwardModelFitWorkspace, # noqa: E402
36 run_ffm_minimisation, # noqa: E402
37 extract_spectrum_from_image,
38 dumpParameters,
39 run_spectrogram_deconvolution_psf2d)
40from spectractor.extractor.spectrum import Spectrum, calibrate_spectrum # noqa: E402
41from spectractor.fit.fit_spectrum import SpectrumFitWorkspace, run_spectrum_minimisation # noqa: E402
42from spectractor.fit.fit_spectrogram import (SpectrogramFitWorkspace, # noqa: E402
43 run_spectrogram_minimisation)
45from lsst.daf.base import DateTime # noqa: E402
46from .utils import getFilterAndDisperserFromExp # noqa: E402
49class SpectractorShim:
50 """Class for running the Spectractor code.
52 This is designed to provide an implementation of the top-level function in
53 Spectractor.spectractor.extractor.extractor.Spectractor()."""
55 # leading * for kwargs only in constructor
56 def __init__(self, *, configFile=None, paramOverrides=None, supplementaryParameters=None,
57 resetParameters=None):
58 if configFile:
59 print(f"Loading config from {configFile}")
60 load_config(configFile, rebin=False)
61 self.log = logging.getLogger(__name__)
62 if paramOverrides is not None:
63 self.overrideParameters(paramOverrides)
64 if supplementaryParameters is not None:
65 self.supplementParameters(supplementaryParameters)
66 if resetParameters is not None:
67 self.resetParameters(resetParameters)
69 if parameters.DEBUG:
70 self.log.debug('Parameters pre-rebinning:')
71 dumpParameters()
73 return
75 def overrideParameters(self, overrides):
76 """Dict of Spectractor parameters to override.
78 Default values are set in spectractor.parameters.py for use as consts.
79 This method provides a means for overriding the parameters as needed.
81 Parameters
82 ----------
83 overrides : `dict`
84 Dict of overrides to apply. Warning is logged if keys are found
85 that do not map to existing Spectractor parameters.
86 """
87 for k, v in overrides.items():
88 # NB do not use hasattr(parameters, k) here, as this is broken by
89 # the overloading of __getattr__ in parameters
90 if k in dir(parameters):
91 setattr(parameters, k, v)
92 else:
93 self.log.warn("Did not find attribute %s in parameters" % k)
94 raise RuntimeError(f"{k} not set to {v} {self.dumpParameters()}")
96 def supplementParameters(self, supplementaryItems):
97 """Dict of Spectractor parameters to add to the parameters.
99 Use this method to add entries to the parameter namespace that do not
100 already exist.
102 Parameters
103 ----------
104 supplementaryItems : `dict`
105 Dict of parameters to add. Warning is logged if keys already exist,
106 as these should be overridden rather than supplemented.
107 """
108 # NB avoid using the variable name `parameters` in this method
109 # due to scope collision
110 for k, v in supplementaryItems.items():
111 # NB do not use hasattr(parameters, k) here, as this is broken by
112 # the overloading of __getattr__ in parameters
113 if k in dir(parameters):
114 msg = ("Supplementary parameter already existed %s in parameters,"
115 " use overrideParameters() to override already existing keys instead.")
116 self.log.warn(msg, k)
117 else:
118 setattr(parameters, k, v)
120 def resetParameters(self, resetParameters):
121 """Dict of Spectractor parameters reset in the namespace.
123 Use this method assign parameters to the namespace whether they exist
124 or not.
126 Parameters
127 ----------
128 resetParameters : `dict`
129 Dict of parameters to add.
130 """
131 # NB avoid using the variable name `parameters` in this method
132 # due to scope collision
133 for k, v in resetParameters.items():
134 # NB do not use hasattr(parameters, k) here, as this is broken by
135 # the overloading of __getattr__ in parameters
136 setattr(parameters, k, v)
138 @staticmethod
139 def dumpParameters():
140 """Print all the values in Spectractor's parameters module."""
141 for item in dir(parameters):
142 if not item.startswith("__"):
143 print(item, getattr(parameters, item))
145 def debugPrintTargetCentroidValue(self, image):
146 """Print the positions and values of the centroid for debug purposes.
148 Parameters
149 ----------
150 image : `spectractor.extractor.images.Image`
151 The image.
152 """
153 x, y = image.target_guess
154 self.log.debug(f"Image shape = {image.data.shape}")
155 self.log.debug(f"x, y = {x}, {y}")
156 x = int(np.round(x))
157 y = int(np.round(y))
158 self.log.debug(f"Value at {x}, {y} = {image.data[y, x]}")
160 def spectractorImageFromLsstExposure(self, exp, xpos, ypos, *, target_label='',
161 disperser_label='', filter_label=''):
162 """Construct a Spectractor Image object from LSST objects.
164 Internally we try to use functions that calculate things and return
165 them and set the values using the return rather than modifying the
166 object in place where possible. Where this is not possible the methods
167 are labeled _setSomething().
169 Parameters
170 ----------
171 exp : `lsst.afw.image.Exposure`
172 The exposure to construct the image from.
173 xpos : `float`
174 The x position of the star's centroid.
175 ypos : `float`
176 The y position of the star's centroid.
177 target_label : `str`, optional
178 The name of the object, e.g. HD12345.
179 disperser_label : `str`, optional
180 The name of the dispersed, e.g. 'holo_003'
181 filter_label : `str`, optional
182 The name of the filter, e.g. 'SDSSi'
184 Returns
185 -------
186 image : `spectractor.extractor.images.Image`
187 The image.
188 """
189 # make a blank image, with the filter/disperser set
190 image = Image(file_name='', target_label=target_label, disperser_label=disperser_label,
191 filter_label=filter_label)
193 vi = exp.getInfo().getVisitInfo()
194 rotAngle = vi.getBoresightRotAngle().asDegrees()
195 parameters.OBS_CAMERA_ROTATION = 270 - (rotAngle % 360)
197 radec = vi.getBoresightRaDec()
198 image.ra = asCoords.Angle(radec.getRa().asDegrees(), unit="deg")
199 image.dec = asCoords.Angle(radec.getDec().asDegrees(), unit="deg")
200 ha = vi.getBoresightHourAngle().asDegrees()
201 image.hour_angle = asCoords.Angle(ha, unit="deg")
203 image.data = self._getImageData(exp)
205 def _translateCentroid(dmXpos, dmYpos):
206 # this function was necessary when we were sometimes transposing
207 # and sometimes not. If we decide to always/never transpose then
208 # this function can just be removed.
209 newX = dmYpos
210 newY = dmXpos
211 return newX, newY
212 image.target_guess = _translateCentroid(xpos, ypos)
213 if parameters.DEBUG:
214 self.debugPrintTargetCentroidValue(image)
216 self._setReadNoiseFromExp(image, exp, 1)
217 # xxx remove hard coding of 1 below!
218 image.gain = self._setGainFromExp(image, exp, .85) # gain required for calculating stat err
219 self._setStatErrorInImage(image, exp, useExpVariance=False)
220 # image.coord as an astropy SkyCoord - currently unused
222 self._setImageAndHeaderInfo(image, exp) # sets image attributes
224 assert image.expo is not None
225 assert image.expo != 0
226 assert image.expo > 0
228 image.convert_to_ADU_rate_units() # divides by expTime and sets units to "ADU/s"
230 image.disperser = Hologram(disperser_label, D=parameters.DISTANCE2CCD,
231 data_dir=parameters.DISPERSER_DIR, verbose=parameters.VERBOSE)
233 image.compute_parallactic_angle()
235 return image
237 def _setImageAndHeaderInfo(self, image, exp, useVisitInfo=True):
238 # currently set in spectractor.tools.extract_info_from_CTIO_header()
239 filt, disperser = getFilterAndDisperserFromExp(exp)
241 image.header.filter = filt
242 image.header.disperser_label = disperser
244 # exp time must be set in both header and in object attribute
245 image.header.expo = exp.getInfo().getVisitInfo().getExposureTime()
246 image.expo = exp.getInfo().getVisitInfo().getExposureTime()
248 image.header['LSHIFT'] = 0. # check if necessary
249 image.header['D2CCD'] = parameters.DISTANCE2CCD # necessary MFL
251 try:
252 if useVisitInfo:
253 vi = exp.getInfo().getVisitInfo()
254 image.header.airmass = vi.getBoresightAirmass() # currently returns nan for obs_ctio0m9
255 image.airmass = vi.getBoresightAirmass() # currently returns nan for obs_ctio0m9
256 # TODO: DM-33731 work out if this should be UTC or TAI.
257 image.date_obs = vi.date.toString(DateTime.UTC)
258 else:
259 md = exp.getMetadata().toDict()
260 image.header.airmass = md['AIRMASS']
261 image.airmass = md['AIRMASS']
262 image.date_obs = md['DATE']
263 except Exception:
264 self.log.warn("Failed to set AIRMASS, default value of 1 used")
265 image.header.airmass = 1.
267 return
269 def _getImageData(self, exp, trimToSquare=False):
270 if trimToSquare:
271 data = exp.image.array[0:4000, 0:4000]
272 else:
273 data = exp.image.array
274 return data.T[::, ::]
276 def _setReadNoiseFromExp(self, spectractorImage, exp, constValue=None):
277 # xxx need to implement this properly
278 if constValue is not None:
279 spectractorImage.read_out_noise = np.ones_like(spectractorImage.data) * constValue
280 else:
281 # TODO: Check with Jeremy if we want the raw read noise
282 # or the per-pixel variance. Either is doable, just need to know.
283 raise NotImplementedError("Setting noise image from exp variance not implemented")
285 def _setReadNoiseToNone(self, spectractorImage):
286 spectractorImage.read_out_noise = None
288 def _setStatErrorInImage(self, image, exp, useExpVariance=False):
289 if useExpVariance:
290 image.stat_errors = exp.maskedImage.variance.array # xxx need to deal with TRANSPOSE here
291 else:
292 image.compute_statistical_error()
294 def _setGainFromExp(self, spectractorImage, exp, constValue=None):
295 # xxx need to implement this properly
296 # Note that this is array-like and per-amplifier
297 # so could use the code from gain flats
298 if constValue:
299 return np.ones_like(spectractorImage.data) * constValue
300 return np.ones_like(spectractorImage.data)
302 def _makePath(self, dirname, plotting=True):
303 if plotting:
304 dirname = os.path.join(dirname, 'plots')
305 if not os.path.exists(dirname):
306 os.makedirs(dirname)
308 def _ensureFitsHeader(self, obj, dataDict=None):
309 if 'SIMPLE' not in obj.header:
310 obj.header.insert(0, ('SIMPLE', True))
311 # if dataDict:
312 # header = obj.header
313 # for k, v in dataDict.items():
314 # if k not in header:
315 # header[k] = v
317 @staticmethod
318 def flipImageLeftRight(image, xpos, ypos):
319 image.data = np.flip(image.data, 1)
320 xpos = image.data.shape[1] - xpos
321 return image, xpos, ypos
323 @staticmethod
324 def transposeCentroid(dmXpos, dmYpos, image):
325 xSize, ySize = image.data.shape
326 newX = dmYpos
327 newY = xSize - dmXpos
328 return newX, newY
330 def displayImage(self, image, centroid=None):
331 import lsst.afw.image as afwImage
332 import lsst.afw.display as afwDisp
333 disp1 = afwDisp.Display(987, open=True)
335 tempImg = afwImage.ImageF(np.zeros(image.data.shape, dtype=np.float32))
336 tempImg.array[:] = image.data
338 disp1.mtv(tempImg)
339 if centroid:
340 disp1.dot('x', centroid[0], centroid[1], size=100)
342 def setAdrParameters(self, spectrum, exp):
343 # The adr_params parameter format expected by spectractor are:
344 # [dec, hour_angle, temperature, pressure, humidity, airmass]
345 vi = exp.getInfo().getVisitInfo()
347 raDec = vi.getBoresightRaDec()
348 dec = raDec.getDec()
349 dec = asCoords.Angle(dec.asDegrees(), unit=u.deg)
351 hourAngle = vi.getBoresightHourAngle()
352 hourAngle = asCoords.Angle(hourAngle.asDegrees(), unit=u.deg)
354 weather = vi.getWeather()
356 _temperature = weather.getAirTemperature()
357 if _temperature is None or np.isnan(_temperature):
358 self.log.warning("Temperature not set, using nominal value of 10 C")
359 _temperature = 10 # nominal value
360 temperature = _temperature
362 _pressure = weather.getAirPressure()
363 if _pressure is not None and not np.isnan(_pressure):
364 if _pressure > 10_000:
365 _pressure /= 100 # convert from Pa to hPa
366 else:
367 self.log.warning("Pressure not set, using nominal value of 743 hPa")
368 _pressure = 743 # nominal for altitude?
369 pressure = _pressure
370 _humidity = weather.getHumidity()
371 humidity = _humidity if not np.isnan(_humidity) else None # not a required param so no default
373 airmass = vi.getBoresightAirmass()
374 spectrum.adr_params = [dec, hourAngle, temperature, pressure, humidity, airmass]
375 spectrum.pressure = pressure
376 spectrum.humidity = humidity
377 spectrum.airmass = airmass
378 spectrum.temperature = temperature
380 def run(self, exp, xpos, ypos, target, doFitAtmosphere, doFitAtmosphereOnSpectrogram,
381 outputRoot=None, plotting=True):
382 # run option kwargs in the original code, seems to ~always be True
383 atmospheric_lines = True
385 self.log.info('Starting SPECTRACTOR')
386 # TODO: rename _makePath _makeOutputPath
387 if outputRoot is not None: # TODO: remove post Gen3 transition
388 self._makePath(outputRoot, plotting=plotting) # early in case this fails, as processing is slow
390 # Upstream loads config file here
392 filter_label, disperser = getFilterAndDisperserFromExp(exp)
393 image = self.spectractorImageFromLsstExposure(exp, xpos, ypos, target_label=target,
394 disperser_label=disperser,
395 filter_label=filter_label)
397 if parameters.DEBUG:
398 self.debugPrintTargetCentroidValue(image)
399 title = 'Raw image with input target location'
400 image.plot_image(scale='symlog', target_pixcoords=image.target_guess, title=title)
401 self.log.info(f"Pixel value at centroid = {image.data[int(xpos), int(ypos)]}")
403 # XXX this needs removing or at least dealing with to not always
404 # just run! ASAP XXX
405 # if disperser == 'ronchi170lpmm':
406 # TODO: add something more robust as to whether to flip!
407 # image, xpos, ypos = self.flipImageLeftRight(image, xpos, ypos)
408 # self.displayImage(image, centroid=(xpos, ypos))
410 if parameters.CCD_REBIN > 1:
411 self.log.info(f'Rebinning image with rebin of {parameters.CCD_REBIN}')
412 apply_rebinning_to_parameters()
413 image.rebin()
414 if parameters.DEBUG:
415 self.log.info('Parameters post-rebinning:')
416 dumpParameters()
417 self.debugPrintTargetCentroidValue(image)
418 title = 'Rebinned image with input target location'
419 image.plot_image(scale='symlog', target_pixcoords=image.target_guess, title=title)
420 self.log.debug('Post rebin:')
421 self.debugPrintTargetCentroidValue(image)
423 # image turning and target finding - use LSST code instead?
424 # and if not, at least test how the rotation code compares
425 # this part of Spectractor is certainly slow at the very least
426 if True: # TODO: change this to be an option, at least for testing vs LSST
427 self.log.info('Search for the target in the image...')
428 # sets the image.target_pixcoords
429 _ = find_target(image, image.target_guess, widths=(parameters.XWINDOW, parameters.YWINDOW))
430 turn_image(image) # creates the rotated data, and sets the image.target_pixcoords_rotated
432 # Rotate the image: several methods
433 # Find the exact target position in the rotated image:
434 # several methods - but how are these controlled? MFL
435 self.log.info('Search for the target in the rotated image...')
436 _ = find_target(image, image.target_guess, rotated=True,
437 widths=(parameters.XWINDOW_ROT, parameters.YWINDOW_ROT))
438 else:
439 # code path for if the image is pre-rotated by LSST code
440 raise NotImplementedError
442 # Create Spectrum object
443 spectrum = Spectrum(image=image, order=parameters.SPEC_ORDER)
444 self.setAdrParameters(spectrum, exp)
446 # Subtract background and bad pixels
447 w_psf1d, bgd_model_func = extract_spectrum_from_image(image, spectrum,
448 signal_width=parameters.PIXWIDTH_SIGNAL,
449 ws=(parameters.PIXDIST_BACKGROUND,
450 parameters.PIXDIST_BACKGROUND
451 + parameters.PIXWIDTH_BACKGROUND),
452 right_edge=image.data.shape[1])
453 spectrum.atmospheric_lines = atmospheric_lines
455 # PSF2D deconvolution
456 if parameters.SPECTRACTOR_DECONVOLUTION_PSF2D:
457 run_spectrogram_deconvolution_psf2d(spectrum, bgd_model_func=bgd_model_func)
459 # Calibrate the spectrum
460 self.log.info(f'Calibrating order {spectrum.order:d} spectrum...')
461 with_adr = True
462 if parameters.OBS_OBJECT_TYPE != "STAR":
463 # XXX Check what this is set to, and how
464 # likely need to be passed through
465 with_adr = False
466 calibrate_spectrum(spectrum, with_adr=with_adr)
468 # not necessarily set during fit but required to be present for astropy
469 # fits writing to work (required to be in keeping with upstream)
470 spectrum.data_order2 = np.zeros_like(spectrum.lambdas)
471 spectrum.err_order2 = np.zeros_like(spectrum.lambdas)
473 # Full forward model extraction:
474 # adds transverse ADR and order 2 subtraction
475 ffmWorkspace = None
476 if parameters.SPECTRACTOR_DECONVOLUTION_FFM:
477 ffmWorkspace = FullForwardModelFitWorkspace(spectrum, verbose=parameters.VERBOSE,
478 plot=True,
479 live_fit=False,
480 amplitude_priors_method="spectrum")
481 spectrum = run_ffm_minimisation(ffmWorkspace, method="newton", niter=2)
483 # Fit the atmosphere on the spectrum using uvspec binary
484 spectrumAtmosphereWorkspace = None
485 if doFitAtmosphere:
486 spectrumAtmosphereWorkspace = SpectrumFitWorkspace(spectrum,
487 fit_angstrom_exponent=True,
488 verbose=parameters.VERBOSE,
489 plot=True)
490 run_spectrum_minimisation(spectrumAtmosphereWorkspace, method="newton")
492 # Fit the atmosphere directly on the spectrogram using uvspec binary
493 spectrogramAtmosphereWorkspace = None
494 if doFitAtmosphereOnSpectrogram:
495 spectrogramAtmosphereWorkspace = SpectrogramFitWorkspace(spectrum,
496 fit_angstrom_exponent=True,
497 verbose=parameters.VERBOSE,
498 plot=True)
499 run_spectrogram_minimisation(spectrogramAtmosphereWorkspace, method="newton")
501 # Save the spectrum
502 self._ensureFitsHeader(spectrum) # SIMPLE is missing by default
504 # Plot the spectrum
505 parameters.DISPLAY = True
506 if parameters.VERBOSE and parameters.DISPLAY:
507 spectrum.plot_spectrum(xlim=None)
509 result = Spectraction()
510 result.spectrum = spectrum
511 result.image = image
512 result.spectrumForwardModelFitParameters = ffmWorkspace.params if ffmWorkspace is not None else None
513 result.spectrumLibradtranFitParameters = (spectrumAtmosphereWorkspace.params if
514 spectrumAtmosphereWorkspace is not None else None)
515 result.spectrogramLibradtranFitParameters = (spectrogramAtmosphereWorkspace.params if
516 spectrogramAtmosphereWorkspace is not None else None)
518 return result
521class Spectraction:
522 """A simple class for holding the Spectractor outputs.
524 Will likely be updated in future to provide some simple getters to allow
525 easier access to parts of the data structure, and perhaps some convenience
526 methods for interacting with the more awkward objects (e.g. the Lines).
527 """
528 # result.spectrum = spectrum
529 # result.image = image
530 # result.w = w