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

228 statements  

« prev     ^ index     » next       coverage.py v7.2.5, created at 2023-05-13 03:20 -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/>. 

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 

41 

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

43from .utils import getFilterAndDisperserFromExp # noqa: E402 

44 

45 

46class SpectractorShim: 

47 """Class for running the Spectractor code. 

48 

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

50 Spectractor.spectractor.extractor.extractor.Spectractor().""" 

51 

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) 

65 

66 if parameters.DEBUG: 

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

68 dumpParameters() 

69 

70 return 

71 

72 def overrideParameters(self, overrides): 

73 """Dict of Spectractor parameters to override. 

74 

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. 

77 

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

92 

93 def supplementParameters(self, supplementaryItems): 

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

95 

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

97 already exist. 

98 

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) 

116 

117 def resetParameters(self, resetParameters): 

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

119 

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

121 or not. 

122 

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) 

134 

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

141 

142 def debugPrintTargetCentroidValue(self, image): 

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

144 

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

156 

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

158 disperser_label='', filter_label=''): 

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

160 

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

165 

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' 

180 

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) 

189 

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

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

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

193 

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

199 

200 image.data = self._getImageData(exp) 

201 

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) 

212 

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 

218 

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

220 

221 assert image.expo is not None 

222 assert image.expo != 0 

223 assert image.expo > 0 

224 

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

226 

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

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

229 

230 image.compute_parallactic_angle() 

231 

232 return image 

233 

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

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

236 filt, disperser = getFilterAndDisperserFromExp(exp) 

237 

238 image.header.filter = filt 

239 image.header.disperser_label = disperser 

240 

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

244 

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

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

247 

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. 

263 

264 return 

265 

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

272 

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

281 

282 def _setReadNoiseToNone(self, spectractorImage): 

283 spectractorImage.read_out_noise = None 

284 

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

290 

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) 

298 

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) 

304 

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 

313 

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 

319 

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 

326 

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) 

331 

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

333 tempImg.array[:] = image.data 

334 

335 disp1.mtv(tempImg) 

336 if centroid: 

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

338 

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

343 

344 raDec = vi.getBoresightRaDec() 

345 dec = raDec.getDec() 

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

347 

348 hourAngle = vi.getBoresightHourAngle() 

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

350 

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 

358 

359 airmass = vi.getBoresightAirmass() 

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

361 

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 

365 

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 

370 

371 # Upstream loads config file here 

372 

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) 

377 

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

383 

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

390 

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) 

403 

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 

412 

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 

422 

423 # Create Spectrum object 

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

425 self.setAdrParameters(spectrum, exp) 

426 

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 

435 

436 # PSF2D deconvolution 

437 if parameters.SPECTRACTOR_DECONVOLUTION_PSF2D: 

438 run_spectrogram_deconvolution_psf2d(spectrum, bgd_model_func=bgd_model_func) 

439 

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) 

448 

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) 

453 

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) 

461 

462 # Save the spectrum 

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

464 

465 # Plot the spectrum 

466 parameters.DISPLAY = True 

467 if parameters.VERBOSE and parameters.DISPLAY: 

468 spectrum.plot_spectrum(xlim=None) 

469 

470 result = Spectraction() 

471 result.spectrum = spectrum 

472 result.image = image 

473 result.w = w 

474 

475 # XXX technically this should be a pipeBase.Struct I think 

476 # change it if it matters 

477 return result 

478 

479 

480class Spectraction: 

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

482 

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