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

233 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-04-07 02:24 +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 

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 # TODO: passing exact centroids seems to be causing a serious 

187 # and non-obvious problem! 

188 # this needs fixing for several reasons, mostly because if we have a 

189 # known good centroid then we want to skip the refitting entirely 

190 xpos = int(np.round(xpos)) 

191 ypos = int(np.round(ypos)) 

192 self.log.debug(f'DM value at centroid: {exp.image.array[ypos, xpos]}\n') 

193 

194 # make a blank image, with the filter/disperser set 

195 image = Image(file_name='', target_label=target_label, disperser_label=disperser_label, 

196 filter_label=filter_label) 

197 

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

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

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

201 

202 radec = vi.getBoresightRaDec() 

203 image.ra = asCoords.Angle(radec.getRa().asDegrees(), unit="deg") 

204 image.dec = asCoords.Angle(radec.getDec().asDegrees(), unit="deg") 

205 ha = vi.getBoresightHourAngle().asDegrees() 

206 image.hour_angle = asCoords.Angle(ha, unit="deg") 

207 

208 image.data = self._getImageData(exp) 

209 

210 def _translateCentroid(dmXpos, dmYpos): 

211 # this function was necessary when we were sometimes transposing 

212 # and sometimes not. If we decide to always/never transpose then 

213 # this function can just be removed. 

214 newX = dmYpos 

215 newY = dmXpos 

216 return newX, newY 

217 image.target_guess = _translateCentroid(xpos, ypos) 

218 if parameters.DEBUG: 

219 self.debugPrintTargetCentroidValue(image) 

220 

221 self._setReadNoiseFromExp(image, exp, 1) 

222 # xxx remove hard coding of 1 below! 

223 image.gain = self._setGainFromExp(image, exp, .85) # gain required for calculating stat err 

224 self._setStatErrorInImage(image, exp, useExpVariance=False) 

225 # image.coord as an astropy SkyCoord - currently unused 

226 

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

228 

229 assert image.expo is not None 

230 assert image.expo != 0 

231 assert image.expo > 0 

232 

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

234 

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

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

237 

238 image.compute_parallactic_angle() 

239 

240 return image 

241 

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

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

244 filt, disperser = getFilterAndDisperserFromExp(exp) 

245 

246 image.header.filter = filt 

247 image.header.disperser_label = disperser 

248 

249 # exp time must be set in both header and in object attribute 

250 image.header.expo = exp.getInfo().getVisitInfo().getExposureTime() 

251 image.expo = exp.getInfo().getVisitInfo().getExposureTime() 

252 

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

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

255 

256 try: 

257 if useVisitInfo: 

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

259 image.header.airmass = vi.getBoresightAirmass() # currently returns nan for obs_ctio0m9 

260 image.airmass = vi.getBoresightAirmass() # currently returns nan for obs_ctio0m9 

261 # TODO: DM-33731 work out if this should be UTC or TAI. 

262 image.date_obs = vi.date.toString(DateTime.UTC) 

263 else: 

264 md = exp.getMetadata().toDict() 

265 image.header.airmass = md['AIRMASS'] 

266 image.airmass = md['AIRMASS'] 

267 image.date_obs = md['DATE'] 

268 except Exception: 

269 self.log.warn("Failed to set AIRMASS, default value of 1 used") 

270 image.header.airmass = 1. 

271 

272 return 

273 

274 def _getImageData(self, exp, trimToSquare=False): 

275 if trimToSquare: 

276 data = exp.image.array[0:4000, 0:4000] 

277 else: 

278 data = exp.image.array 

279 return data.T[::, ::] 

280 

281 def _setReadNoiseFromExp(self, spectractorImage, exp, constValue=None): 

282 # xxx need to implement this properly 

283 if constValue is not None: 

284 spectractorImage.read_out_noise = np.ones_like(spectractorImage.data) * constValue 

285 else: 

286 # TODO: Check with Jeremy if we want the raw read noise 

287 # or the per-pixel variance. Either is doable, just need to know. 

288 raise NotImplementedError("Setting noise image from exp variance not implemented") 

289 

290 def _setReadNoiseToNone(self, spectractorImage): 

291 spectractorImage.read_out_noise = None 

292 

293 def _setStatErrorInImage(self, image, exp, useExpVariance=False): 

294 if useExpVariance: 

295 image.stat_errors = exp.maskedImage.variance.array # xxx need to deal with TRANSPOSE here 

296 else: 

297 image.compute_statistical_error() 

298 

299 def _setGainFromExp(self, spectractorImage, exp, constValue=None): 

300 # xxx need to implement this properly 

301 # Note that this is array-like and per-amplifier 

302 # so could use the code from gain flats 

303 if constValue: 

304 return np.ones_like(spectractorImage.data) * constValue 

305 return np.ones_like(spectractorImage.data) 

306 

307 def _makePath(self, dirname, plotting=True): 

308 if plotting: 

309 dirname = os.path.join(dirname, 'plots') 

310 if not os.path.exists(dirname): 

311 os.makedirs(dirname) 

312 

313 def _ensureFitsHeader(self, obj, dataDict=None): 

314 if 'SIMPLE' not in obj.header: 

315 obj.header.insert(0, ('SIMPLE', True)) 

316 # if dataDict: 

317 # header = obj.header 

318 # for k, v in dataDict.items(): 

319 # if k not in header: 

320 # header[k] = v 

321 

322 @staticmethod 

323 def flipImageLeftRight(image, xpos, ypos): 

324 image.data = np.flip(image.data, 1) 

325 xpos = image.data.shape[1] - xpos 

326 return image, xpos, ypos 

327 

328 @staticmethod 

329 def transposeCentroid(dmXpos, dmYpos, image): 

330 xSize, ySize = image.data.shape 

331 newX = dmYpos 

332 newY = xSize - dmXpos 

333 return newX, newY 

334 

335 def displayImage(self, image, centroid=None): 

336 import lsst.afw.image as afwImage 

337 import lsst.afw.display as afwDisp 

338 disp1 = afwDisp.Display(987, open=True) 

339 

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

341 tempImg.array[:] = image.data 

342 

343 disp1.mtv(tempImg) 

344 if centroid: 

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

346 

347 def setAdrParameters(self, spectrum, exp): 

348 # The adr_params parameter format expected by spectractor are: 

349 # [dec, hour_angle, temperature, pressure, humidity, airmass] 

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

351 

352 raDec = vi.getBoresightRaDec() 

353 dec = raDec.getDec() 

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

355 

356 hourAngle = vi.getBoresightHourAngle() 

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

358 

359 weather = vi.getWeather() 

360 _temperature = weather.getAirTemperature() 

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

362 _pressure = weather.getAirPressure() 

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

364 _humidity = weather.getHumidity() 

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

366 

367 airmass = vi.getBoresightAirmass() 

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

369 

370 def run(self, exp, xpos, ypos, target, 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 # TODO: DM-38264: 

382 # passing exact centroids seems to be causing a serious 

383 # and non-obvious problem! this needs fixing for several reasons, 

384 # mostly because if we have a known good centroid then we want to skip 

385 # the refitting entirely 

386 xpos = int(np.round(xpos)) 

387 ypos = int(np.round(ypos)) 

388 

389 filter_label, disperser = getFilterAndDisperserFromExp(exp) 

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

391 disperser_label=disperser, 

392 filter_label=filter_label) 

393 

394 if parameters.DEBUG: 

395 self.debugPrintTargetCentroidValue(image) 

396 title = 'Raw image with input target location' 

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

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

399 

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

401 # just run! ASAP XXX 

402 # if disperser == 'ronchi170lpmm': 

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

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

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

406 

407 if parameters.CCD_REBIN > 1: 

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

409 apply_rebinning_to_parameters() 

410 image.rebin() 

411 if parameters.DEBUG: 

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

413 dumpParameters() 

414 self.debugPrintTargetCentroidValue(image) 

415 title = 'Rebinned image with input target location' 

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

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

418 self.debugPrintTargetCentroidValue(image) 

419 

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

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

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

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

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

425 # sets the image.target_pixcoords 

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

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

428 

429 # Rotate the image: several methods 

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

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

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

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

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

435 else: 

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

437 raise NotImplementedError 

438 

439 # Create Spectrum object 

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

441 self.setAdrParameters(spectrum, exp) 

442 

443 # Subtract background and bad pixels 

444 w_psf1d, bgd_model_func = extract_spectrum_from_image(image, spectrum, 

445 signal_width=parameters.PIXWIDTH_SIGNAL, 

446 ws=(parameters.PIXDIST_BACKGROUND, 

447 parameters.PIXDIST_BACKGROUND 

448 + parameters.PIXWIDTH_BACKGROUND), 

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

450 spectrum.atmospheric_lines = atmospheric_lines 

451 

452 # PSF2D deconvolution 

453 if parameters.SPECTRACTOR_DECONVOLUTION_PSF2D: 

454 run_spectrogram_deconvolution_psf2d(spectrum, bgd_model_func=bgd_model_func) 

455 

456 # Calibrate the spectrum 

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

458 with_adr = True 

459 if parameters.OBS_OBJECT_TYPE != "STAR": 

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

461 # likely need to be passed through 

462 with_adr = False 

463 calibrate_spectrum(spectrum, with_adr=with_adr) 

464 

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

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

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

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

469 

470 # Full forward model extraction: 

471 # adds transverse ADR and order 2 subtraction 

472 w = None 

473 if parameters.SPECTRACTOR_DECONVOLUTION_FFM: 

474 w = FullForwardModelFitWorkspace(spectrum, verbose=parameters.VERBOSE, plot=True, live_fit=False, 

475 amplitude_priors_method="spectrum") 

476 spectrum = run_ffm_minimisation(w, method="newton", niter=2) 

477 

478 # Save the spectrum 

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

480 

481 # Plot the spectrum 

482 parameters.DISPLAY = True 

483 if parameters.VERBOSE and parameters.DISPLAY: 

484 spectrum.plot_spectrum(xlim=None) 

485 

486 result = Spectraction() 

487 result.spectrum = spectrum 

488 result.image = image 

489 result.w = w 

490 

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

492 # change it if it matters 

493 return result 

494 

495 

496class Spectraction: 

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

498 

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

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

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

502 """ 

503 # result.spectrum = spectrum 

504 # result.image = image 

505 # result.w = w