Coverage for python/lsst/atmospec/spectraction.py: 17%
215 statements
« prev ^ index » next coverage.py v6.4.1, created at 2022-07-09 09:21 +0000
« prev ^ index » next coverage.py v6.4.1, created at 2022-07-09 09:21 +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 os
23import numpy as np
24import astropy.coordinates as asCoords
25from astropy import units as u
27from spectractor import parameters
28parameters.CALLING_CODE = "LSST_DM" # this must be set IMMEDIATELY to supress colored logs
30from spectractor.config import load_config # noqa: E402
31from spectractor.extractor.images import Image, find_target, turn_image # noqa: E402
33from spectractor.extractor.dispersers import Hologram # noqa: E402
34from spectractor.extractor.extractor import (FullForwardModelFitWorkspace, # noqa: E402
35 run_ffm_minimisation, # noqa: E402
36 extract_spectrum_from_image,
37 run_spectrogram_deconvolution_psf2d)
38from spectractor.extractor.spectrum import Spectrum, calibrate_spectrum # noqa: E402
40import lsst.log as lsstLog # noqa: E402
41from lsst.daf.base import DateTime # noqa: E402
42from lsst.obs.lsst.translators.lsst import FILTER_DELIMITER # noqa: E402
45class SpectractorShim():
46 """Class for running the Spectractor code.
48 This is designed to provide an implementation of the top-level function in
49 Spectractor.spectractor.extractor.extractor.Spectractor()."""
50 TRANSPOSE = True
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)
58 self.log = lsstLog.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)
65 return
67 def overrideParameters(self, overrides):
68 """Dict of Spectractor parameters to override.
70 Default values are set in spectractor.parameters.py for use as consts.
71 This method provides a means for overriding the parameters as needed.
73 Parameters
74 ----------
75 overrides : `dict`
76 Dict of overrides to apply. Warning is logged if keys are found that do
77 not map to existing Spectractor parameters.
78 """
79 for k, v in overrides.items():
80 # NB do not use hasattr(parameters, k) here, as this is broken by
81 # the overloading of __getattr__ in parameters
82 if k in dir(parameters):
83 setattr(parameters, k, v)
84 else:
85 self.log.warn("Did not find attribute %s in parameters" % k)
86 raise RuntimeError(f"{k} not set to {v} {self.dumpParameters()}")
88 def supplementParameters(self, supplementaryItems):
89 """Dict of Spectractor parameters to add to the parameters.
91 Use this method to add entries to the parameter namespace that do not
92 already exist.
94 Parameters
95 ----------
96 supplementaryItems : `dict`
97 Dict of parameters to add. Warning is logged if keys already exist,
98 as these should be overridden rather than supplemented.
99 """
100 # NB avoid using the variable name `parameters` in this method
101 # due to scope collision
102 for k, v in supplementaryItems.items():
103 # NB do not use hasattr(parameters, k) here, as this is broken by
104 # the overloading of __getattr__ in parameters
105 if k in dir(parameters):
106 msg = ("Supplementary parameter already existed %s in parameters,"
107 " use overrideParameters() to override already existing keys instead.")
108 self.log.warn(msg, k)
109 else:
110 setattr(parameters, k, v)
112 def resetParameters(self, resetParameters):
113 """Dict of Spectractor parameters reset in the namespace.
115 Use this method assign parameters to the namespace whether they exist
116 or not.
118 Parameters
119 ----------
120 resetParameters : `dict`
121 Dict of parameters to add.
122 """
123 # NB avoid using the variable name `parameters` in this method
124 # due to scope collision
125 for k, v in resetParameters.items():
126 # NB do not use hasattr(parameters, k) here, as this is broken by
127 # the overloading of __getattr__ in parameters
128 setattr(parameters, k, v)
130 @staticmethod
131 def dumpParameters():
132 for item in dir(parameters):
133 if not item.startswith("__"):
134 print(item, getattr(parameters, item))
136 def spectractorImageFromLsstExposure(self, exp, *, target_label='', disperser_label='', filter_label=''):
137 """Construct a Spectractor Image object from LSST objects.
139 Internally we try to use functions that calculate things and return
140 them and set the values using the return rather than modifying the
141 object in place where possible. Where this is not possible the methods
142 are labeled _setSomething().
143 """
144 image = Image(file_name='', target_label=target_label, disperser_label=disperser_label,
145 filter_label=filter_label)
147 vi = exp.getInfo().getVisitInfo()
148 rotAngle = vi.getBoresightRotAngle().asDegrees()
149 # line below correct if not rotating 90 XXX remove this once resolved
150 # parameters.OBS_CAMERA_ROTATION = 180 - (rotAngle % 360)
151 parameters.OBS_CAMERA_ROTATION = 90 - (rotAngle % 360)
153 radec = vi.getBoresightRaDec()
154 image.ra = asCoords.Angle(radec.getRa().asDegrees(), unit="deg")
155 image.dec = asCoords.Angle(radec.getDec().asDegrees(), unit="deg")
156 ha = vi.getBoresightHourAngle().asDegrees()
157 image.hour_angle = asCoords.Angle(ha, unit="deg")
159 image.data = self._getImageData(exp)
160 self._setReadNoiseFromExp(image, exp, 1)
161 # xxx remove hard coding of 1 below!
162 image.gain = self._setGainFromExp(image, exp, .85) # gain required for calculating stat err
163 self._setStatErrorInImage(image, exp, useExpVariance=False)
164 # image.coord as an astropy SkyCoord - currently unused
166 self._setImageAndHeaderInfo(image, exp) # sets image attributes
168 assert image.expo is not None
169 assert image.expo != 0
170 assert image.expo > 0
172 image.convert_to_ADU_rate_units() # divides by expTime and sets units to "ADU/s"
174 image.disperser = Hologram(disperser_label, D=parameters.DISTANCE2CCD,
175 data_dir=parameters.DISPERSER_DIR, verbose=parameters.VERBOSE)
177 image.compute_parallactic_angle()
179 return image
181 @staticmethod
182 def _getFilterAndDisperserFromExp(exp):
183 filterFullName = exp.getFilter().physicalLabel
184 if FILTER_DELIMITER not in filterFullName:
185 filt = filterFullName
186 grating = exp.getInfo().getMetadata()['GRATING']
187 else:
188 filt, grating = filterFullName.split(FILTER_DELIMITER)
189 return filt, grating
191 def _setImageAndHeaderInfo(self, image, exp, useVisitInfo=True):
192 # currently set in spectractor.tools.extract_info_from_CTIO_header()
193 filt, disperser = self._getFilterAndDisperserFromExp(exp)
195 image.header.filter = filt
196 image.header.disperser_label = disperser
198 # exp time must be set in both header and in object attribute
199 image.header.expo = exp.getInfo().getVisitInfo().getExposureTime()
200 image.expo = exp.getInfo().getVisitInfo().getExposureTime()
202 image.header['LSHIFT'] = 0. # check if necessary
203 image.header['D2CCD'] = parameters.DISTANCE2CCD # necessary MFL
205 try:
206 if useVisitInfo:
207 vi = exp.getInfo().getVisitInfo()
208 image.header.airmass = vi.getBoresightAirmass() # currently returns nan for obs_ctio0m9
209 image.airmass = vi.getBoresightAirmass() # currently returns nan for obs_ctio0m9
210 # TODO: DM-33731 work out if this should be UTC or TAI.
211 image.date_obs = vi.date.toString(DateTime.UTC)
212 else:
213 md = exp.getMetadata().toDict()
214 image.header.airmass = md['AIRMASS']
215 image.airmass = md['AIRMASS']
216 image.date_obs = md['DATE']
217 except Exception:
218 self.log.warn("Failed to set AIRMASS, default value of 1 used")
219 image.header.airmass = 1.
221 return
223 def _getImageData(self, exp):
224 if self.TRANSPOSE:
225 # return exp.maskedImage.image.array.T[:, ::-1]
226 return np.rot90(exp.maskedImage.image.array, 1)
227 return exp.maskedImage.image.array
229 def _setReadNoiseFromExp(self, spectractorImage, exp, constValue=None):
230 # xxx need to implement this properly
231 if constValue is not None:
232 spectractorImage.read_out_noise = np.ones_like(spectractorImage.data) * constValue
233 else:
234 # TODO: Check with Jeremy if we want the raw read noise
235 # or the per-pixel variance. Either is doable, just need to know.
236 raise NotImplementedError("Setting noise image from exp variance not implemented")
238 def _setReadNoiseToNone(self, spectractorImage):
239 spectractorImage.read_out_noise = None
241 def _setStatErrorInImage(self, image, exp, useExpVariance=False):
242 if useExpVariance:
243 image.stat_errors = exp.maskedImage.variance.array # xxx need to deal with TRANSPOSE here
244 else:
245 image.compute_statistical_error()
247 def _setGainFromExp(self, spectractorImage, exp, constValue=None):
248 # xxx need to implement this properly
249 # Note that this is array-like and per-amplifier
250 # so could use the code from gain flats
251 if constValue:
252 return np.ones_like(spectractorImage.data) * constValue
253 return np.ones_like(spectractorImage.data)
255 def _makePath(self, dirname, plotting=True):
256 if plotting:
257 dirname = os.path.join(dirname, 'plots')
258 if not os.path.exists(dirname):
259 os.makedirs(dirname)
261 def _ensureFitsHeader(self, obj, dataDict=None):
262 if 'SIMPLE' not in obj.header:
263 obj.header.insert(0, ('SIMPLE', True))
264 # if dataDict:
265 # header = obj.header
266 # for k, v in dataDict.items():
267 # if k not in header:
268 # header[k] = v
270 @staticmethod
271 def flipImageLeftRight(image, xpos, ypos):
272 image.data = np.flip(image.data, 1)
273 xpos = image.data.shape[1] - xpos
274 return image, xpos, ypos
276 @staticmethod
277 def transposeCentroid(dmXpos, dmYpos, image):
278 # xSize, ySize = image.data.shape
279 # newX = dmXpos
280 # newY = ySize - dmYpos # image is also flipped in Y
281 # return newY, newX
283 xSize, ySize = image.data.shape
284 newX = dmYpos
285 newY = xSize - dmXpos
286 return newX, newY
288 def displayImage(self, image, centroid=None):
289 import lsst.afw.image as afwImage
290 import lsst.afw.display as afwDisp
291 disp1 = afwDisp.Display(987, open=True)
293 tempImg = afwImage.ImageF(np.zeros(image.data.shape, dtype=np.float32))
294 tempImg.array[:] = image.data
296 disp1.mtv(tempImg)
297 if centroid:
298 disp1.dot('x', centroid[0], centroid[1], size=100)
300 def setAdrParameters(self, spectrum, exp):
301 # The adr_params parameter format expected by spectractor are:
302 # [dec, hour_angle, temperature, pressure, humidity, airmass]
303 vi = exp.getInfo().getVisitInfo()
305 raDec = vi.getBoresightRaDec()
306 dec = raDec.getDec()
307 dec = asCoords.Angle(dec.asDegrees(), unit=u.deg)
309 hourAngle = vi.getBoresightHourAngle()
310 hourAngle = asCoords.Angle(hourAngle.asDegrees(), unit=u.deg)
312 weather = vi.getWeather()
313 _temperature = weather.getAirTemperature()
314 temperature = _temperature if not np.isnan(_temperature) else 10 # maybe average?
315 _pressure = weather.getAirPressure()
316 pressure = _pressure if not np.isnan(_pressure) else 732 # nominal for altitude?
317 _humidity = weather.getHumidity()
318 humidity = _humidity if not np.isnan(_humidity) else None # not a required param so no default
320 airmass = vi.getBoresightAirmass()
321 spectrum.adr_params = [dec, hourAngle, temperature, pressure, humidity, airmass]
323 def run(self, exp, xpos, ypos, target, outputRoot=None, plotting=True):
324 # run option kwargs in the original code, seems to ~always be True
325 atmospheric_lines = True
327 self.log.info('Starting SPECTRACTOR')
328 # TODO: rename _makePath _makeOutputPath
329 if outputRoot is not None: # TODO: remove post Gen3 transition
330 self._makePath(outputRoot, plotting=plotting) # early in case this fails, as processing is slow
332 # Upstream loads config file here
334 # TODO: passing exact centroids seems to be causing a serious
335 # and non-obvious problem!
336 # this needs fixing for several reasons, mostly because if we have a
337 # known good centroid then we want to skip the refitting entirely
338 xpos = int(np.round(xpos))
339 ypos = int(np.round(ypos))
341 filter_label, disperser = self._getFilterAndDisperserFromExp(exp)
342 image = self.spectractorImageFromLsstExposure(exp, target_label=target, disperser_label=disperser,
343 filter_label=filter_label)
345 if self.TRANSPOSE:
346 xpos, ypos = self.transposeCentroid(xpos, ypos, image)
348 image.target_guess = (xpos, ypos)
349 if parameters.DEBUG:
350 image.plot_image(scale='log10', target_pixcoords=image.target_guess)
351 self.log.info(f"Pixel value at centroid = {image.data[int(ypos), int(xpos)]}")
353 # XXX this needs removing or at least dealing with to not always
354 # just run! ASAP XXX
355 # if disperser == 'ronchi170lpmm':
356 # TODO: add something more robust as to whether to flip!
357 # image, xpos, ypos = self.flipImageLeftRight(image, xpos, ypos)
358 # self.displayImage(image, centroid=(xpos, ypos))
360 # Use fast mode
361 if parameters.CCD_REBIN > 1:
362 self.log.info(f'Rebinning image with rebin of {parameters.CCD_REBIN}')
363 # TODO: Fix bug here where the passed parameter isn't used!
364 image.rebin()
365 if parameters.DEBUG:
366 image.plot_image(scale='symlog', target_pixcoords=image.target_guess)
368 # image turning and target finding - use LSST code instead?
369 # and if not, at least test how the rotation code compares
370 # this part of Spectractor is certainly slow at the very least
371 if True: # TODO: change this to be an option, at least for testing vs LSST
372 self.log.info('Search for the target in the image...')
373 _ = find_target(image, image.target_guess) # sets the image.target_pixcoords
374 turn_image(image) # creates the rotated data, and sets the image.target_pixcoords_rotated
376 # Rotate the image: several methods
377 # Find the exact target position in the rotated image:
378 # several methods - but how are these controlled? MFL
379 self.log.info('Search for the target in the rotated image...')
380 _ = find_target(image, image.target_guess, rotated=True)
381 else:
382 # code path for if the image is pre-rotated by LSST code
383 raise NotImplementedError
385 # Create Spectrum object
386 spectrum = Spectrum(image=image, order=parameters.SPEC_ORDER) # XXX new in DM-33589 check SPEC_ORDER
387 self.setAdrParameters(spectrum, exp)
389 # Subtract background and bad pixels
390 w_psf1d, bgd_model_func = extract_spectrum_from_image(image, spectrum,
391 signal_width=parameters.PIXWIDTH_SIGNAL,
392 ws=(parameters.PIXDIST_BACKGROUND,
393 parameters.PIXDIST_BACKGROUND
394 + parameters.PIXWIDTH_BACKGROUND),
395 right_edge=parameters.CCD_IMSIZE)
396 spectrum.atmospheric_lines = atmospheric_lines
398 # PSF2D deconvolution
399 if parameters.SPECTRACTOR_DECONVOLUTION_PSF2D:
400 run_spectrogram_deconvolution_psf2d(spectrum, bgd_model_func=bgd_model_func)
402 # Calibrate the spectrum
403 with_adr = True
404 if parameters.OBS_OBJECT_TYPE != "STAR":
405 # XXX Check what this is set to, and how
406 # likely need to be passed through
407 with_adr = False
408 calibrate_spectrum(spectrum, with_adr=with_adr)
410 # not necessarily set during fit but required to be present for astropy
411 # fits writing to work (required to be in keeping with upstream)
412 spectrum.data_order2 = np.zeros_like(spectrum.lambdas_order2)
413 spectrum.err_order2 = np.zeros_like(spectrum.lambdas_order2)
415 # Full forward model extraction:
416 # adds transverse ADR and order 2 subtraction
417 w = None
418 if parameters.SPECTRACTOR_DECONVOLUTION_FFM:
419 w = FullForwardModelFitWorkspace(spectrum, verbose=parameters.VERBOSE, plot=True, live_fit=False,
420 amplitude_priors_method="spectrum")
421 spectrum = run_ffm_minimisation(w, method="newton", niter=2)
423 # Save the spectrum
424 self._ensureFitsHeader(spectrum) # SIMPLE is missing by default
426 # Plot the spectrum
427 parameters.DISPLAY = True
428 if parameters.VERBOSE and parameters.DISPLAY:
429 spectrum.plot_spectrum(xlim=None)
431 spectrum.chromatic_psf.table['lambdas'] = spectrum.lambdas
433 result = Spectraction()
434 result.spectrum = spectrum
435 result.image = image
436 result.w = w
438 # XXX technically this should be a pipeBase.Struct I think
439 # change it if it matters
440 return result
443class Spectraction:
444 """A simple class for holding the Spectractor outputs.
446 Will likely be updated in future to provide some simple getters to allow
447 easier access to parts of the data structure, and perhaps some convenience
448 methods for interacting with the more awkward objects (e.g. the Lines).
449 """
450 # result.spectrum = spectrum
451 # result.image = image
452 # result.w = w