Coverage for python/lsst/atmospec/spectraction.py: 14%
233 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-04-12 03:16 -0700
« prev ^ index » next coverage.py v6.5.0, created at 2023-04-12 03:16 -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 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
42from lsst.daf.base import DateTime # noqa: E402
43from .utils import getFilterAndDisperserFromExp # noqa: E402
46class SpectractorShim:
47 """Class for running the Spectractor code.
49 This is designed to provide an implementation of the top-level function in
50 Spectractor.spectractor.extractor.extractor.Spectractor()."""
52 # leading * for kwargs only in constructor
53 def __init__(self, *, configFile=None, paramOverrides=None, supplementaryParameters=None,
54 resetParameters=None):
55 if configFile:
56 print(f"Loading config from {configFile}")
57 load_config(configFile, rebin=False)
58 self.log = logging.getLogger(__name__)
59 if paramOverrides is not None:
60 self.overrideParameters(paramOverrides)
61 if supplementaryParameters is not None:
62 self.supplementParameters(supplementaryParameters)
63 if resetParameters is not None:
64 self.resetParameters(resetParameters)
66 if parameters.DEBUG:
67 self.log.debug('Parameters pre-rebinning:')
68 dumpParameters()
70 return
72 def overrideParameters(self, overrides):
73 """Dict of Spectractor parameters to override.
75 Default values are set in spectractor.parameters.py for use as consts.
76 This method provides a means for overriding the parameters as needed.
78 Parameters
79 ----------
80 overrides : `dict`
81 Dict of overrides to apply. Warning is logged if keys are found
82 that do not map to existing Spectractor parameters.
83 """
84 for k, v in overrides.items():
85 # NB do not use hasattr(parameters, k) here, as this is broken by
86 # the overloading of __getattr__ in parameters
87 if k in dir(parameters):
88 setattr(parameters, k, v)
89 else:
90 self.log.warn("Did not find attribute %s in parameters" % k)
91 raise RuntimeError(f"{k} not set to {v} {self.dumpParameters()}")
93 def supplementParameters(self, supplementaryItems):
94 """Dict of Spectractor parameters to add to the parameters.
96 Use this method to add entries to the parameter namespace that do not
97 already exist.
99 Parameters
100 ----------
101 supplementaryItems : `dict`
102 Dict of parameters to add. Warning is logged if keys already exist,
103 as these should be overridden rather than supplemented.
104 """
105 # NB avoid using the variable name `parameters` in this method
106 # due to scope collision
107 for k, v in supplementaryItems.items():
108 # NB do not use hasattr(parameters, k) here, as this is broken by
109 # the overloading of __getattr__ in parameters
110 if k in dir(parameters):
111 msg = ("Supplementary parameter already existed %s in parameters,"
112 " use overrideParameters() to override already existing keys instead.")
113 self.log.warn(msg, k)
114 else:
115 setattr(parameters, k, v)
117 def resetParameters(self, resetParameters):
118 """Dict of Spectractor parameters reset in the namespace.
120 Use this method assign parameters to the namespace whether they exist
121 or not.
123 Parameters
124 ----------
125 resetParameters : `dict`
126 Dict of parameters to add.
127 """
128 # NB avoid using the variable name `parameters` in this method
129 # due to scope collision
130 for k, v in resetParameters.items():
131 # NB do not use hasattr(parameters, k) here, as this is broken by
132 # the overloading of __getattr__ in parameters
133 setattr(parameters, k, v)
135 @staticmethod
136 def dumpParameters():
137 """Print all the values in Spectractor's parameters module."""
138 for item in dir(parameters):
139 if not item.startswith("__"):
140 print(item, getattr(parameters, item))
142 def debugPrintTargetCentroidValue(self, image):
143 """Print the positions and values of the centroid for debug purposes.
145 Parameters
146 ----------
147 image : `spectractor.extractor.images.Image`
148 The image.
149 """
150 x, y = image.target_guess
151 self.log.debug(f"Image shape = {image.data.shape}")
152 self.log.debug(f"x, y = {x}, {y}")
153 x = int(np.round(x))
154 y = int(np.round(y))
155 self.log.debug(f"Value at {x}, {y} = {image.data[y, x]}")
157 def spectractorImageFromLsstExposure(self, exp, xpos, ypos, *, target_label='',
158 disperser_label='', filter_label=''):
159 """Construct a Spectractor Image object from LSST objects.
161 Internally we try to use functions that calculate things and return
162 them and set the values using the return rather than modifying the
163 object in place where possible. Where this is not possible the methods
164 are labeled _setSomething().
166 Parameters
167 ----------
168 exp : `lsst.afw.image.Exposure`
169 The exposure to construct the image from.
170 xpos : `float`
171 The x position of the star's centroid.
172 ypos : `float`
173 The y position of the star's centroid.
174 target_label : `str`, optional
175 The name of the object, e.g. HD12345.
176 disperser_label : `str`, optional
177 The name of the dispersed, e.g. 'holo_003'
178 filter_label : `str`, optional
179 The name of the filter, e.g. 'SDSSi'
181 Returns
182 -------
183 image : `spectractor.extractor.images.Image`
184 The image.
185 """
186 # TODO: passing exact centroids seems to be causing a serious
187 # and non-obvious problem!
188 # this needs fixing for several reasons, mostly because if we have a
189 # known good centroid then we want to skip the refitting entirely
190 xpos = int(np.round(xpos))
191 ypos = int(np.round(ypos))
192 self.log.debug(f'DM value at centroid: {exp.image.array[ypos, xpos]}\n')
194 # make a blank image, with the filter/disperser set
195 image = Image(file_name='', target_label=target_label, disperser_label=disperser_label,
196 filter_label=filter_label)
198 vi = exp.getInfo().getVisitInfo()
199 rotAngle = vi.getBoresightRotAngle().asDegrees()
200 parameters.OBS_CAMERA_ROTATION = 270 - (rotAngle % 360)
202 radec = vi.getBoresightRaDec()
203 image.ra = asCoords.Angle(radec.getRa().asDegrees(), unit="deg")
204 image.dec = asCoords.Angle(radec.getDec().asDegrees(), unit="deg")
205 ha = vi.getBoresightHourAngle().asDegrees()
206 image.hour_angle = asCoords.Angle(ha, unit="deg")
208 image.data = self._getImageData(exp)
210 def _translateCentroid(dmXpos, dmYpos):
211 # this function was necessary when we were sometimes transposing
212 # and sometimes not. If we decide to always/never transpose then
213 # this function can just be removed.
214 newX = dmYpos
215 newY = dmXpos
216 return newX, newY
217 image.target_guess = _translateCentroid(xpos, ypos)
218 if parameters.DEBUG:
219 self.debugPrintTargetCentroidValue(image)
221 self._setReadNoiseFromExp(image, exp, 1)
222 # xxx remove hard coding of 1 below!
223 image.gain = self._setGainFromExp(image, exp, .85) # gain required for calculating stat err
224 self._setStatErrorInImage(image, exp, useExpVariance=False)
225 # image.coord as an astropy SkyCoord - currently unused
227 self._setImageAndHeaderInfo(image, exp) # sets image attributes
229 assert image.expo is not None
230 assert image.expo != 0
231 assert image.expo > 0
233 image.convert_to_ADU_rate_units() # divides by expTime and sets units to "ADU/s"
235 image.disperser = Hologram(disperser_label, D=parameters.DISTANCE2CCD,
236 data_dir=parameters.DISPERSER_DIR, verbose=parameters.VERBOSE)
238 image.compute_parallactic_angle()
240 return image
242 def _setImageAndHeaderInfo(self, image, exp, useVisitInfo=True):
243 # currently set in spectractor.tools.extract_info_from_CTIO_header()
244 filt, disperser = getFilterAndDisperserFromExp(exp)
246 image.header.filter = filt
247 image.header.disperser_label = disperser
249 # exp time must be set in both header and in object attribute
250 image.header.expo = exp.getInfo().getVisitInfo().getExposureTime()
251 image.expo = exp.getInfo().getVisitInfo().getExposureTime()
253 image.header['LSHIFT'] = 0. # check if necessary
254 image.header['D2CCD'] = parameters.DISTANCE2CCD # necessary MFL
256 try:
257 if useVisitInfo:
258 vi = exp.getInfo().getVisitInfo()
259 image.header.airmass = vi.getBoresightAirmass() # currently returns nan for obs_ctio0m9
260 image.airmass = vi.getBoresightAirmass() # currently returns nan for obs_ctio0m9
261 # TODO: DM-33731 work out if this should be UTC or TAI.
262 image.date_obs = vi.date.toString(DateTime.UTC)
263 else:
264 md = exp.getMetadata().toDict()
265 image.header.airmass = md['AIRMASS']
266 image.airmass = md['AIRMASS']
267 image.date_obs = md['DATE']
268 except Exception:
269 self.log.warn("Failed to set AIRMASS, default value of 1 used")
270 image.header.airmass = 1.
272 return
274 def _getImageData(self, exp, trimToSquare=False):
275 if trimToSquare:
276 data = exp.image.array[0:4000, 0:4000]
277 else:
278 data = exp.image.array
279 return data.T[::, ::]
281 def _setReadNoiseFromExp(self, spectractorImage, exp, constValue=None):
282 # xxx need to implement this properly
283 if constValue is not None:
284 spectractorImage.read_out_noise = np.ones_like(spectractorImage.data) * constValue
285 else:
286 # TODO: Check with Jeremy if we want the raw read noise
287 # or the per-pixel variance. Either is doable, just need to know.
288 raise NotImplementedError("Setting noise image from exp variance not implemented")
290 def _setReadNoiseToNone(self, spectractorImage):
291 spectractorImage.read_out_noise = None
293 def _setStatErrorInImage(self, image, exp, useExpVariance=False):
294 if useExpVariance:
295 image.stat_errors = exp.maskedImage.variance.array # xxx need to deal with TRANSPOSE here
296 else:
297 image.compute_statistical_error()
299 def _setGainFromExp(self, spectractorImage, exp, constValue=None):
300 # xxx need to implement this properly
301 # Note that this is array-like and per-amplifier
302 # so could use the code from gain flats
303 if constValue:
304 return np.ones_like(spectractorImage.data) * constValue
305 return np.ones_like(spectractorImage.data)
307 def _makePath(self, dirname, plotting=True):
308 if plotting:
309 dirname = os.path.join(dirname, 'plots')
310 if not os.path.exists(dirname):
311 os.makedirs(dirname)
313 def _ensureFitsHeader(self, obj, dataDict=None):
314 if 'SIMPLE' not in obj.header:
315 obj.header.insert(0, ('SIMPLE', True))
316 # if dataDict:
317 # header = obj.header
318 # for k, v in dataDict.items():
319 # if k not in header:
320 # header[k] = v
322 @staticmethod
323 def flipImageLeftRight(image, xpos, ypos):
324 image.data = np.flip(image.data, 1)
325 xpos = image.data.shape[1] - xpos
326 return image, xpos, ypos
328 @staticmethod
329 def transposeCentroid(dmXpos, dmYpos, image):
330 xSize, ySize = image.data.shape
331 newX = dmYpos
332 newY = xSize - dmXpos
333 return newX, newY
335 def displayImage(self, image, centroid=None):
336 import lsst.afw.image as afwImage
337 import lsst.afw.display as afwDisp
338 disp1 = afwDisp.Display(987, open=True)
340 tempImg = afwImage.ImageF(np.zeros(image.data.shape, dtype=np.float32))
341 tempImg.array[:] = image.data
343 disp1.mtv(tempImg)
344 if centroid:
345 disp1.dot('x', centroid[0], centroid[1], size=100)
347 def setAdrParameters(self, spectrum, exp):
348 # The adr_params parameter format expected by spectractor are:
349 # [dec, hour_angle, temperature, pressure, humidity, airmass]
350 vi = exp.getInfo().getVisitInfo()
352 raDec = vi.getBoresightRaDec()
353 dec = raDec.getDec()
354 dec = asCoords.Angle(dec.asDegrees(), unit=u.deg)
356 hourAngle = vi.getBoresightHourAngle()
357 hourAngle = asCoords.Angle(hourAngle.asDegrees(), unit=u.deg)
359 weather = vi.getWeather()
360 _temperature = weather.getAirTemperature()
361 temperature = _temperature if not np.isnan(_temperature) else 10 # maybe average?
362 _pressure = weather.getAirPressure()
363 pressure = _pressure if not np.isnan(_pressure) else 743 # nominal for altitude?
364 _humidity = weather.getHumidity()
365 humidity = _humidity if not np.isnan(_humidity) else None # not a required param so no default
367 airmass = vi.getBoresightAirmass()
368 spectrum.adr_params = [dec, hourAngle, temperature, pressure, humidity, airmass]
370 def run(self, exp, xpos, ypos, target, outputRoot=None, plotting=True):
371 # run option kwargs in the original code, seems to ~always be True
372 atmospheric_lines = True
374 self.log.info('Starting SPECTRACTOR')
375 # TODO: rename _makePath _makeOutputPath
376 if outputRoot is not None: # TODO: remove post Gen3 transition
377 self._makePath(outputRoot, plotting=plotting) # early in case this fails, as processing is slow
379 # Upstream loads config file here
381 # TODO: DM-38264:
382 # passing exact centroids seems to be causing a serious
383 # and non-obvious problem! this needs fixing for several reasons,
384 # mostly because if we have a known good centroid then we want to skip
385 # the refitting entirely
386 xpos = int(np.round(xpos))
387 ypos = int(np.round(ypos))
389 filter_label, disperser = getFilterAndDisperserFromExp(exp)
390 image = self.spectractorImageFromLsstExposure(exp, xpos, ypos, target_label=target,
391 disperser_label=disperser,
392 filter_label=filter_label)
394 if parameters.DEBUG:
395 self.debugPrintTargetCentroidValue(image)
396 title = 'Raw image with input target location'
397 image.plot_image(scale='symlog', target_pixcoords=image.target_guess, title=title)
398 self.log.info(f"Pixel value at centroid = {image.data[int(xpos), int(ypos)]}")
400 # XXX this needs removing or at least dealing with to not always
401 # just run! ASAP XXX
402 # if disperser == 'ronchi170lpmm':
403 # TODO: add something more robust as to whether to flip!
404 # image, xpos, ypos = self.flipImageLeftRight(image, xpos, ypos)
405 # self.displayImage(image, centroid=(xpos, ypos))
407 if parameters.CCD_REBIN > 1:
408 self.log.info(f'Rebinning image with rebin of {parameters.CCD_REBIN}')
409 apply_rebinning_to_parameters()
410 image.rebin()
411 if parameters.DEBUG:
412 self.log.info('Parameters post-rebinning:')
413 dumpParameters()
414 self.debugPrintTargetCentroidValue(image)
415 title = 'Rebinned image with input target location'
416 image.plot_image(scale='symlog', target_pixcoords=image.target_guess, title=title)
417 self.log.debug('Post rebin:')
418 self.debugPrintTargetCentroidValue(image)
420 # image turning and target finding - use LSST code instead?
421 # and if not, at least test how the rotation code compares
422 # this part of Spectractor is certainly slow at the very least
423 if True: # TODO: change this to be an option, at least for testing vs LSST
424 self.log.info('Search for the target in the image...')
425 # sets the image.target_pixcoords
426 _ = find_target(image, image.target_guess, widths=(parameters.XWINDOW, parameters.YWINDOW))
427 turn_image(image) # creates the rotated data, and sets the image.target_pixcoords_rotated
429 # Rotate the image: several methods
430 # Find the exact target position in the rotated image:
431 # several methods - but how are these controlled? MFL
432 self.log.info('Search for the target in the rotated image...')
433 _ = find_target(image, image.target_guess, rotated=True,
434 widths=(parameters.XWINDOW_ROT, parameters.YWINDOW_ROT))
435 else:
436 # code path for if the image is pre-rotated by LSST code
437 raise NotImplementedError
439 # Create Spectrum object
440 spectrum = Spectrum(image=image, order=parameters.SPEC_ORDER)
441 self.setAdrParameters(spectrum, exp)
443 # Subtract background and bad pixels
444 w_psf1d, bgd_model_func = extract_spectrum_from_image(image, spectrum,
445 signal_width=parameters.PIXWIDTH_SIGNAL,
446 ws=(parameters.PIXDIST_BACKGROUND,
447 parameters.PIXDIST_BACKGROUND
448 + parameters.PIXWIDTH_BACKGROUND),
449 right_edge=image.data.shape[1])
450 spectrum.atmospheric_lines = atmospheric_lines
452 # PSF2D deconvolution
453 if parameters.SPECTRACTOR_DECONVOLUTION_PSF2D:
454 run_spectrogram_deconvolution_psf2d(spectrum, bgd_model_func=bgd_model_func)
456 # Calibrate the spectrum
457 self.log.info(f'Calibrating order {spectrum.order:d} spectrum...')
458 with_adr = True
459 if parameters.OBS_OBJECT_TYPE != "STAR":
460 # XXX Check what this is set to, and how
461 # likely need to be passed through
462 with_adr = False
463 calibrate_spectrum(spectrum, with_adr=with_adr)
465 # not necessarily set during fit but required to be present for astropy
466 # fits writing to work (required to be in keeping with upstream)
467 spectrum.data_order2 = np.zeros_like(spectrum.lambdas)
468 spectrum.err_order2 = np.zeros_like(spectrum.lambdas)
470 # Full forward model extraction:
471 # adds transverse ADR and order 2 subtraction
472 w = None
473 if parameters.SPECTRACTOR_DECONVOLUTION_FFM:
474 w = FullForwardModelFitWorkspace(spectrum, verbose=parameters.VERBOSE, plot=True, live_fit=False,
475 amplitude_priors_method="spectrum")
476 spectrum = run_ffm_minimisation(w, method="newton", niter=2)
478 # Save the spectrum
479 self._ensureFitsHeader(spectrum) # SIMPLE is missing by default
481 # Plot the spectrum
482 parameters.DISPLAY = True
483 if parameters.VERBOSE and parameters.DISPLAY:
484 spectrum.plot_spectrum(xlim=None)
486 result = Spectraction()
487 result.spectrum = spectrum
488 result.image = image
489 result.w = w
491 # XXX technically this should be a pipeBase.Struct I think
492 # change it if it matters
493 return result
496class Spectraction:
497 """A simple class for holding the Spectractor outputs.
499 Will likely be updated in future to provide some simple getters to allow
500 easier access to parts of the data structure, and perhaps some convenience
501 methods for interacting with the more awkward objects (e.g. the Lines).
502 """
503 # result.spectrum = spectrum
504 # result.image = image
505 # result.w = w