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

244 statements  

« prev     ^ index     » next       coverage.py v7.3.1, created at 2023-09-20 13:26 +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 _temperature = weather.getAirTemperature() 

356 temperature = _temperature if not np.isnan(_temperature) else 10 # maybe average? 

357 _pressure = weather.getAirPressure() 

358 pressure = _pressure if not np.isnan(_pressure) else 743 # nominal for altitude? 

359 _humidity = weather.getHumidity() 

360 humidity = _humidity if not np.isnan(_humidity) else None # not a required param so no default 

361 

362 airmass = vi.getBoresightAirmass() 

363 spectrum.adr_params = [dec, hourAngle, temperature, pressure, humidity, airmass] 

364 spectrum.pressure = pressure 

365 spectrum.humidity = humidity 

366 spectrum.airmass = airmass 

367 spectrum.temperature = temperature 

368 

369 def run(self, exp, xpos, ypos, target, doFitAtmosphere, doFitAtmosphereOnSpectrogram, 

370 outputRoot=None, plotting=True): 

371 # run option kwargs in the original code, seems to ~always be True 

372 atmospheric_lines = True 

373 

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 

378 

379 # Upstream loads config file here 

380 

381 filter_label, disperser = getFilterAndDisperserFromExp(exp) 

382 image = self.spectractorImageFromLsstExposure(exp, xpos, ypos, target_label=target, 

383 disperser_label=disperser, 

384 filter_label=filter_label) 

385 

386 if parameters.DEBUG: 

387 self.debugPrintTargetCentroidValue(image) 

388 title = 'Raw image with input target location' 

389 image.plot_image(scale='symlog', target_pixcoords=image.target_guess, title=title) 

390 self.log.info(f"Pixel value at centroid = {image.data[int(xpos), int(ypos)]}") 

391 

392 # XXX this needs removing or at least dealing with to not always 

393 # just run! ASAP XXX 

394 # if disperser == 'ronchi170lpmm': 

395 # TODO: add something more robust as to whether to flip! 

396 # image, xpos, ypos = self.flipImageLeftRight(image, xpos, ypos) 

397 # self.displayImage(image, centroid=(xpos, ypos)) 

398 

399 if parameters.CCD_REBIN > 1: 

400 self.log.info(f'Rebinning image with rebin of {parameters.CCD_REBIN}') 

401 apply_rebinning_to_parameters() 

402 image.rebin() 

403 if parameters.DEBUG: 

404 self.log.info('Parameters post-rebinning:') 

405 dumpParameters() 

406 self.debugPrintTargetCentroidValue(image) 

407 title = 'Rebinned image with input target location' 

408 image.plot_image(scale='symlog', target_pixcoords=image.target_guess, title=title) 

409 self.log.debug('Post rebin:') 

410 self.debugPrintTargetCentroidValue(image) 

411 

412 # image turning and target finding - use LSST code instead? 

413 # and if not, at least test how the rotation code compares 

414 # this part of Spectractor is certainly slow at the very least 

415 if True: # TODO: change this to be an option, at least for testing vs LSST 

416 self.log.info('Search for the target in the image...') 

417 # sets the image.target_pixcoords 

418 _ = find_target(image, image.target_guess, widths=(parameters.XWINDOW, parameters.YWINDOW)) 

419 turn_image(image) # creates the rotated data, and sets the image.target_pixcoords_rotated 

420 

421 # Rotate the image: several methods 

422 # Find the exact target position in the rotated image: 

423 # several methods - but how are these controlled? MFL 

424 self.log.info('Search for the target in the rotated image...') 

425 _ = find_target(image, image.target_guess, rotated=True, 

426 widths=(parameters.XWINDOW_ROT, parameters.YWINDOW_ROT)) 

427 else: 

428 # code path for if the image is pre-rotated by LSST code 

429 raise NotImplementedError 

430 

431 # Create Spectrum object 

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

433 self.setAdrParameters(spectrum, exp) 

434 

435 # Subtract background and bad pixels 

436 w_psf1d, bgd_model_func = extract_spectrum_from_image(image, spectrum, 

437 signal_width=parameters.PIXWIDTH_SIGNAL, 

438 ws=(parameters.PIXDIST_BACKGROUND, 

439 parameters.PIXDIST_BACKGROUND 

440 + parameters.PIXWIDTH_BACKGROUND), 

441 right_edge=image.data.shape[1]) 

442 spectrum.atmospheric_lines = atmospheric_lines 

443 

444 # PSF2D deconvolution 

445 if parameters.SPECTRACTOR_DECONVOLUTION_PSF2D: 

446 run_spectrogram_deconvolution_psf2d(spectrum, bgd_model_func=bgd_model_func) 

447 

448 # Calibrate the spectrum 

449 self.log.info(f'Calibrating order {spectrum.order:d} spectrum...') 

450 with_adr = True 

451 if parameters.OBS_OBJECT_TYPE != "STAR": 

452 # XXX Check what this is set to, and how 

453 # likely need to be passed through 

454 with_adr = False 

455 calibrate_spectrum(spectrum, with_adr=with_adr) 

456 

457 # not necessarily set during fit but required to be present for astropy 

458 # fits writing to work (required to be in keeping with upstream) 

459 spectrum.data_order2 = np.zeros_like(spectrum.lambdas) 

460 spectrum.err_order2 = np.zeros_like(spectrum.lambdas) 

461 

462 # Full forward model extraction: 

463 # adds transverse ADR and order 2 subtraction 

464 ffmWorkspace = None 

465 if parameters.SPECTRACTOR_DECONVOLUTION_FFM: 

466 ffmWorkspace = FullForwardModelFitWorkspace(spectrum, verbose=parameters.VERBOSE, 

467 plot=True, 

468 live_fit=False, 

469 amplitude_priors_method="spectrum") 

470 spectrum = run_ffm_minimisation(ffmWorkspace, method="newton", niter=2) 

471 

472 # Fit the atmosphere on the spectrum using uvspec binary 

473 spectrumAtmosphereWorkspace = None 

474 if doFitAtmosphere: 

475 spectrumAtmosphereWorkspace = SpectrumFitWorkspace(spectrum, 

476 fit_angstrom_exponent=True, 

477 verbose=parameters.VERBOSE, 

478 plot=True) 

479 run_spectrum_minimisation(spectrumAtmosphereWorkspace, method="newton") 

480 

481 # Fit the atmosphere directly on the spectrogram using uvspec binary 

482 spectrogramAtmosphereWorkspace = None 

483 if doFitAtmosphereOnSpectrogram: 

484 spectrogramAtmosphereWorkspace = SpectrogramFitWorkspace(spectrum, 

485 fit_angstrom_exponent=True, 

486 verbose=parameters.VERBOSE, 

487 plot=True) 

488 run_spectrogram_minimisation(spectrogramAtmosphereWorkspace, method="newton") 

489 

490 # Save the spectrum 

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

492 

493 # Plot the spectrum 

494 parameters.DISPLAY = True 

495 if parameters.VERBOSE and parameters.DISPLAY: 

496 spectrum.plot_spectrum(xlim=None) 

497 

498 result = Spectraction() 

499 result.spectrum = spectrum 

500 result.image = image 

501 result.spectrumForwardModelFitParameters = ffmWorkspace.params if ffmWorkspace is not None else None 

502 result.spectrumLibradtranFitParameters = (spectrumAtmosphereWorkspace.params if 

503 spectrumAtmosphereWorkspace is not None else None) 

504 result.spectrogramLibradtranFitParameters = (spectrogramAtmosphereWorkspace.params if 

505 spectrogramAtmosphereWorkspace is not None else None) 

506 

507 return result 

508 

509 

510class Spectraction: 

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

512 

513 Will likely be updated in future to provide some simple getters to allow 

514 easier access to parts of the data structure, and perhaps some convenience 

515 methods for interacting with the more awkward objects (e.g. the Lines). 

516 """ 

517 # result.spectrum = spectrum 

518 # result.image = image 

519 # result.w = w