Coverage for python/lsst/atmospec/spectraction.py: 16%
228 statements
« prev ^ index » next coverage.py v7.3.0, created at 2023-08-15 09:33 +0000
« prev ^ index » next coverage.py v7.3.0, created at 2023-08-15 09:33 +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
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 # make a blank image, with the filter/disperser set
187 image = Image(file_name='', target_label=target_label, disperser_label=disperser_label,
188 filter_label=filter_label)
190 vi = exp.getInfo().getVisitInfo()
191 rotAngle = vi.getBoresightRotAngle().asDegrees()
192 parameters.OBS_CAMERA_ROTATION = 270 - (rotAngle % 360)
194 radec = vi.getBoresightRaDec()
195 image.ra = asCoords.Angle(radec.getRa().asDegrees(), unit="deg")
196 image.dec = asCoords.Angle(radec.getDec().asDegrees(), unit="deg")
197 ha = vi.getBoresightHourAngle().asDegrees()
198 image.hour_angle = asCoords.Angle(ha, unit="deg")
200 image.data = self._getImageData(exp)
202 def _translateCentroid(dmXpos, dmYpos):
203 # this function was necessary when we were sometimes transposing
204 # and sometimes not. If we decide to always/never transpose then
205 # this function can just be removed.
206 newX = dmYpos
207 newY = dmXpos
208 return newX, newY
209 image.target_guess = _translateCentroid(xpos, ypos)
210 if parameters.DEBUG:
211 self.debugPrintTargetCentroidValue(image)
213 self._setReadNoiseFromExp(image, exp, 1)
214 # xxx remove hard coding of 1 below!
215 image.gain = self._setGainFromExp(image, exp, .85) # gain required for calculating stat err
216 self._setStatErrorInImage(image, exp, useExpVariance=False)
217 # image.coord as an astropy SkyCoord - currently unused
219 self._setImageAndHeaderInfo(image, exp) # sets image attributes
221 assert image.expo is not None
222 assert image.expo != 0
223 assert image.expo > 0
225 image.convert_to_ADU_rate_units() # divides by expTime and sets units to "ADU/s"
227 image.disperser = Hologram(disperser_label, D=parameters.DISTANCE2CCD,
228 data_dir=parameters.DISPERSER_DIR, verbose=parameters.VERBOSE)
230 image.compute_parallactic_angle()
232 return image
234 def _setImageAndHeaderInfo(self, image, exp, useVisitInfo=True):
235 # currently set in spectractor.tools.extract_info_from_CTIO_header()
236 filt, disperser = getFilterAndDisperserFromExp(exp)
238 image.header.filter = filt
239 image.header.disperser_label = disperser
241 # exp time must be set in both header and in object attribute
242 image.header.expo = exp.getInfo().getVisitInfo().getExposureTime()
243 image.expo = exp.getInfo().getVisitInfo().getExposureTime()
245 image.header['LSHIFT'] = 0. # check if necessary
246 image.header['D2CCD'] = parameters.DISTANCE2CCD # necessary MFL
248 try:
249 if useVisitInfo:
250 vi = exp.getInfo().getVisitInfo()
251 image.header.airmass = vi.getBoresightAirmass() # currently returns nan for obs_ctio0m9
252 image.airmass = vi.getBoresightAirmass() # currently returns nan for obs_ctio0m9
253 # TODO: DM-33731 work out if this should be UTC or TAI.
254 image.date_obs = vi.date.toString(DateTime.UTC)
255 else:
256 md = exp.getMetadata().toDict()
257 image.header.airmass = md['AIRMASS']
258 image.airmass = md['AIRMASS']
259 image.date_obs = md['DATE']
260 except Exception:
261 self.log.warn("Failed to set AIRMASS, default value of 1 used")
262 image.header.airmass = 1.
264 return
266 def _getImageData(self, exp, trimToSquare=False):
267 if trimToSquare:
268 data = exp.image.array[0:4000, 0:4000]
269 else:
270 data = exp.image.array
271 return data.T[::, ::]
273 def _setReadNoiseFromExp(self, spectractorImage, exp, constValue=None):
274 # xxx need to implement this properly
275 if constValue is not None:
276 spectractorImage.read_out_noise = np.ones_like(spectractorImage.data) * constValue
277 else:
278 # TODO: Check with Jeremy if we want the raw read noise
279 # or the per-pixel variance. Either is doable, just need to know.
280 raise NotImplementedError("Setting noise image from exp variance not implemented")
282 def _setReadNoiseToNone(self, spectractorImage):
283 spectractorImage.read_out_noise = None
285 def _setStatErrorInImage(self, image, exp, useExpVariance=False):
286 if useExpVariance:
287 image.stat_errors = exp.maskedImage.variance.array # xxx need to deal with TRANSPOSE here
288 else:
289 image.compute_statistical_error()
291 def _setGainFromExp(self, spectractorImage, exp, constValue=None):
292 # xxx need to implement this properly
293 # Note that this is array-like and per-amplifier
294 # so could use the code from gain flats
295 if constValue:
296 return np.ones_like(spectractorImage.data) * constValue
297 return np.ones_like(spectractorImage.data)
299 def _makePath(self, dirname, plotting=True):
300 if plotting:
301 dirname = os.path.join(dirname, 'plots')
302 if not os.path.exists(dirname):
303 os.makedirs(dirname)
305 def _ensureFitsHeader(self, obj, dataDict=None):
306 if 'SIMPLE' not in obj.header:
307 obj.header.insert(0, ('SIMPLE', True))
308 # if dataDict:
309 # header = obj.header
310 # for k, v in dataDict.items():
311 # if k not in header:
312 # header[k] = v
314 @staticmethod
315 def flipImageLeftRight(image, xpos, ypos):
316 image.data = np.flip(image.data, 1)
317 xpos = image.data.shape[1] - xpos
318 return image, xpos, ypos
320 @staticmethod
321 def transposeCentroid(dmXpos, dmYpos, image):
322 xSize, ySize = image.data.shape
323 newX = dmYpos
324 newY = xSize - dmXpos
325 return newX, newY
327 def displayImage(self, image, centroid=None):
328 import lsst.afw.image as afwImage
329 import lsst.afw.display as afwDisp
330 disp1 = afwDisp.Display(987, open=True)
332 tempImg = afwImage.ImageF(np.zeros(image.data.shape, dtype=np.float32))
333 tempImg.array[:] = image.data
335 disp1.mtv(tempImg)
336 if centroid:
337 disp1.dot('x', centroid[0], centroid[1], size=100)
339 def setAdrParameters(self, spectrum, exp):
340 # The adr_params parameter format expected by spectractor are:
341 # [dec, hour_angle, temperature, pressure, humidity, airmass]
342 vi = exp.getInfo().getVisitInfo()
344 raDec = vi.getBoresightRaDec()
345 dec = raDec.getDec()
346 dec = asCoords.Angle(dec.asDegrees(), unit=u.deg)
348 hourAngle = vi.getBoresightHourAngle()
349 hourAngle = asCoords.Angle(hourAngle.asDegrees(), unit=u.deg)
351 weather = vi.getWeather()
352 _temperature = weather.getAirTemperature()
353 temperature = _temperature if not np.isnan(_temperature) else 10 # maybe average?
354 _pressure = weather.getAirPressure()
355 pressure = _pressure if not np.isnan(_pressure) else 743 # nominal for altitude?
356 _humidity = weather.getHumidity()
357 humidity = _humidity if not np.isnan(_humidity) else None # not a required param so no default
359 airmass = vi.getBoresightAirmass()
360 spectrum.adr_params = [dec, hourAngle, temperature, pressure, humidity, airmass]
362 def run(self, exp, xpos, ypos, target, outputRoot=None, plotting=True):
363 # run option kwargs in the original code, seems to ~always be True
364 atmospheric_lines = True
366 self.log.info('Starting SPECTRACTOR')
367 # TODO: rename _makePath _makeOutputPath
368 if outputRoot is not None: # TODO: remove post Gen3 transition
369 self._makePath(outputRoot, plotting=plotting) # early in case this fails, as processing is slow
371 # Upstream loads config file here
373 filter_label, disperser = getFilterAndDisperserFromExp(exp)
374 image = self.spectractorImageFromLsstExposure(exp, xpos, ypos, target_label=target,
375 disperser_label=disperser,
376 filter_label=filter_label)
378 if parameters.DEBUG:
379 self.debugPrintTargetCentroidValue(image)
380 title = 'Raw image with input target location'
381 image.plot_image(scale='symlog', target_pixcoords=image.target_guess, title=title)
382 self.log.info(f"Pixel value at centroid = {image.data[int(xpos), int(ypos)]}")
384 # XXX this needs removing or at least dealing with to not always
385 # just run! ASAP XXX
386 # if disperser == 'ronchi170lpmm':
387 # TODO: add something more robust as to whether to flip!
388 # image, xpos, ypos = self.flipImageLeftRight(image, xpos, ypos)
389 # self.displayImage(image, centroid=(xpos, ypos))
391 if parameters.CCD_REBIN > 1:
392 self.log.info(f'Rebinning image with rebin of {parameters.CCD_REBIN}')
393 apply_rebinning_to_parameters()
394 image.rebin()
395 if parameters.DEBUG:
396 self.log.info('Parameters post-rebinning:')
397 dumpParameters()
398 self.debugPrintTargetCentroidValue(image)
399 title = 'Rebinned image with input target location'
400 image.plot_image(scale='symlog', target_pixcoords=image.target_guess, title=title)
401 self.log.debug('Post rebin:')
402 self.debugPrintTargetCentroidValue(image)
404 # image turning and target finding - use LSST code instead?
405 # and if not, at least test how the rotation code compares
406 # this part of Spectractor is certainly slow at the very least
407 if True: # TODO: change this to be an option, at least for testing vs LSST
408 self.log.info('Search for the target in the image...')
409 # sets the image.target_pixcoords
410 _ = find_target(image, image.target_guess, widths=(parameters.XWINDOW, parameters.YWINDOW))
411 turn_image(image) # creates the rotated data, and sets the image.target_pixcoords_rotated
413 # Rotate the image: several methods
414 # Find the exact target position in the rotated image:
415 # several methods - but how are these controlled? MFL
416 self.log.info('Search for the target in the rotated image...')
417 _ = find_target(image, image.target_guess, rotated=True,
418 widths=(parameters.XWINDOW_ROT, parameters.YWINDOW_ROT))
419 else:
420 # code path for if the image is pre-rotated by LSST code
421 raise NotImplementedError
423 # Create Spectrum object
424 spectrum = Spectrum(image=image, order=parameters.SPEC_ORDER)
425 self.setAdrParameters(spectrum, exp)
427 # Subtract background and bad pixels
428 w_psf1d, bgd_model_func = extract_spectrum_from_image(image, spectrum,
429 signal_width=parameters.PIXWIDTH_SIGNAL,
430 ws=(parameters.PIXDIST_BACKGROUND,
431 parameters.PIXDIST_BACKGROUND
432 + parameters.PIXWIDTH_BACKGROUND),
433 right_edge=image.data.shape[1])
434 spectrum.atmospheric_lines = atmospheric_lines
436 # PSF2D deconvolution
437 if parameters.SPECTRACTOR_DECONVOLUTION_PSF2D:
438 run_spectrogram_deconvolution_psf2d(spectrum, bgd_model_func=bgd_model_func)
440 # Calibrate the spectrum
441 self.log.info(f'Calibrating order {spectrum.order:d} spectrum...')
442 with_adr = True
443 if parameters.OBS_OBJECT_TYPE != "STAR":
444 # XXX Check what this is set to, and how
445 # likely need to be passed through
446 with_adr = False
447 calibrate_spectrum(spectrum, with_adr=with_adr)
449 # not necessarily set during fit but required to be present for astropy
450 # fits writing to work (required to be in keeping with upstream)
451 spectrum.data_order2 = np.zeros_like(spectrum.lambdas)
452 spectrum.err_order2 = np.zeros_like(spectrum.lambdas)
454 # Full forward model extraction:
455 # adds transverse ADR and order 2 subtraction
456 w = None
457 if parameters.SPECTRACTOR_DECONVOLUTION_FFM:
458 w = FullForwardModelFitWorkspace(spectrum, verbose=parameters.VERBOSE, plot=True, live_fit=False,
459 amplitude_priors_method="spectrum")
460 spectrum = run_ffm_minimisation(w, method="newton", niter=2)
462 # Save the spectrum
463 self._ensureFitsHeader(spectrum) # SIMPLE is missing by default
465 # Plot the spectrum
466 parameters.DISPLAY = True
467 if parameters.VERBOSE and parameters.DISPLAY:
468 spectrum.plot_spectrum(xlim=None)
470 result = Spectraction()
471 result.spectrum = spectrum
472 result.image = image
473 result.w = w
475 # XXX technically this should be a pipeBase.Struct I think
476 # change it if it matters
477 return result
480class Spectraction:
481 """A simple class for holding the Spectractor outputs.
483 Will likely be updated in future to provide some simple getters to allow
484 easier access to parts of the data structure, and perhaps some convenience
485 methods for interacting with the more awkward objects (e.g. the Lines).
486 """
487 # result.spectrum = spectrum
488 # result.image = image
489 # result.w = w