Coverage for python/lsst/atmospec/spectraction.py: 15%

252 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-25 19:31 +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/>. 

21 

22import logging 

23import os 

24import numpy as np 

25import astropy.coordinates as asCoords 

26from astropy import units as u 

27 

28from spectractor import parameters 

29parameters.CALLING_CODE = "LSST_DM" # this must be set IMMEDIATELY to supress colored logs 

30 

31from spectractor.config import load_config, apply_rebinning_to_parameters # noqa: E402 

32from spectractor.extractor.images import Image, find_target, turn_image # noqa: E402 

33 

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) 

44 

45from lsst.daf.base import DateTime # noqa: E402 

46from .utils import getFilterAndDisperserFromExp # noqa: E402 

47 

48 

49class SpectractorShim: 

50 """Class for running the Spectractor code. 

51 

52 This is designed to provide an implementation of the top-level function in 

53 Spectractor.spectractor.extractor.extractor.Spectractor().""" 

54 

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) 

68 

69 if parameters.DEBUG: 

70 self.log.debug('Parameters pre-rebinning:') 

71 dumpParameters() 

72 

73 return 

74 

75 def overrideParameters(self, overrides): 

76 """Dict of Spectractor parameters to override. 

77 

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. 

80 

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()}") 

95 

96 def supplementParameters(self, supplementaryItems): 

97 """Dict of Spectractor parameters to add to the parameters. 

98 

99 Use this method to add entries to the parameter namespace that do not 

100 already exist. 

101 

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) 

119 

120 def resetParameters(self, resetParameters): 

121 """Dict of Spectractor parameters reset in the namespace. 

122 

123 Use this method assign parameters to the namespace whether they exist 

124 or not. 

125 

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) 

137 

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)) 

144 

145 def debugPrintTargetCentroidValue(self, image): 

146 """Print the positions and values of the centroid for debug purposes. 

147 

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]}") 

159 

160 def spectractorImageFromLsstExposure(self, exp, xpos, ypos, *, target_label='', 

161 disperser_label='', filter_label=''): 

162 """Construct a Spectractor Image object from LSST objects. 

163 

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(). 

168 

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' 

183 

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) 

192 

193 vi = exp.getInfo().getVisitInfo() 

194 rotAngle = vi.getBoresightRotAngle().asDegrees() 

195 parameters.OBS_CAMERA_ROTATION = 270 - (rotAngle % 360) 

196 

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") 

202 

203 image.data = self._getImageData(exp) 

204 

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) 

215 

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 

221 

222 self._setImageAndHeaderInfo(image, exp) # sets image attributes 

223 

224 assert image.expo is not None 

225 assert image.expo != 0 

226 assert image.expo > 0 

227 

228 image.convert_to_ADU_rate_units() # divides by expTime and sets units to "ADU/s" 

229 

230 image.disperser = Hologram(disperser_label, D=parameters.DISTANCE2CCD, 

231 data_dir=parameters.DISPERSER_DIR, verbose=parameters.VERBOSE) 

232 

233 image.compute_parallactic_angle() 

234 

235 return image 

236 

237 def _setImageAndHeaderInfo(self, image, exp, useVisitInfo=True): 

238 # currently set in spectractor.tools.extract_info_from_CTIO_header() 

239 filt, disperser = getFilterAndDisperserFromExp(exp) 

240 

241 image.header.filter = filt 

242 image.header.disperser_label = disperser 

243 

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() 

247 

248 image.header['LSHIFT'] = 0. # check if necessary 

249 image.header['D2CCD'] = parameters.DISTANCE2CCD # necessary MFL 

250 

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. 

266 

267 return 

268 

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[::, ::] 

275 

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") 

284 

285 def _setReadNoiseToNone(self, spectractorImage): 

286 spectractorImage.read_out_noise = None 

287 

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() 

293 

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) 

301 

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) 

307 

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 

316 

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 

322 

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 

329 

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) 

334 

335 tempImg = afwImage.ImageF(np.zeros(image.data.shape, dtype=np.float32)) 

336 tempImg.array[:] = image.data 

337 

338 disp1.mtv(tempImg) 

339 if centroid: 

340 disp1.dot('x', centroid[0], centroid[1], size=100) 

341 

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() 

346 

347 raDec = vi.getBoresightRaDec() 

348 dec = raDec.getDec() 

349 dec = asCoords.Angle(dec.asDegrees(), unit=u.deg) 

350 

351 hourAngle = vi.getBoresightHourAngle() 

352 hourAngle = asCoords.Angle(hourAngle.asDegrees(), unit=u.deg) 

353 

354 weather = vi.getWeather() 

355 

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 

361 

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 

372 

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 

379 

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 

384 

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 

389 

390 # Upstream loads config file here 

391 

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) 

396 

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)]}") 

402 

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)) 

409 

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) 

422 

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 

431 

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 

441 

442 # Create Spectrum object 

443 spectrum = Spectrum(image=image, order=parameters.SPEC_ORDER) 

444 self.setAdrParameters(spectrum, exp) 

445 

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 

454 

455 # PSF2D deconvolution 

456 if parameters.SPECTRACTOR_DECONVOLUTION_PSF2D: 

457 run_spectrogram_deconvolution_psf2d(spectrum, bgd_model_func=bgd_model_func) 

458 

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) 

467 

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) 

472 

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) 

482 

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") 

491 

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") 

500 

501 # Save the spectrum 

502 self._ensureFitsHeader(spectrum) # SIMPLE is missing by default 

503 

504 # Plot the spectrum 

505 parameters.DISPLAY = True 

506 if parameters.VERBOSE and parameters.DISPLAY: 

507 spectrum.plot_spectrum(xlim=None) 

508 

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) 

517 

518 return result 

519 

520 

521class Spectraction: 

522 """A simple class for holding the Spectractor outputs. 

523 

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