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

215 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-12-16 11:13 +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 # 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 run_spectrogram_deconvolution_psf2d) 

39from spectractor.extractor.spectrum import Spectrum, calibrate_spectrum # noqa: E402 

40 

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

42from lsst.obs.lsst.translators.lsst import FILTER_DELIMITER # noqa: E402 

43 

44 

45class SpectractorShim: 

46 """Class for running the Spectractor code. 

47 

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

49 Spectractor.spectractor.extractor.extractor.Spectractor().""" 

50 TRANSPOSE = True 

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) 

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 return 

66 

67 def overrideParameters(self, overrides): 

68 """Dict of Spectractor parameters to override. 

69 

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. 

72 

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

87 

88 def supplementParameters(self, supplementaryItems): 

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

90 

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

92 already exist. 

93 

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) 

111 

112 def resetParameters(self, resetParameters): 

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

114 

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

116 or not. 

117 

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) 

129 

130 @staticmethod 

131 def dumpParameters(): 

132 for item in dir(parameters): 

133 if not item.startswith("__"): 

134 print(item, getattr(parameters, item)) 

135 

136 def spectractorImageFromLsstExposure(self, exp, *, target_label='', disperser_label='', filter_label=''): 

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

138 

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) 

146 

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) 

152 

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

158 

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 

165 

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

167 

168 assert image.expo is not None 

169 assert image.expo != 0 

170 assert image.expo > 0 

171 

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

173 

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

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

176 

177 image.compute_parallactic_angle() 

178 

179 return image 

180 

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 

190 

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) 

194 

195 image.header.filter = filt 

196 image.header.disperser_label = disperser 

197 

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

201 

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

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

204 

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. 

220 

221 return 

222 

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 

228 

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

237 

238 def _setReadNoiseToNone(self, spectractorImage): 

239 spectractorImage.read_out_noise = None 

240 

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

246 

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) 

254 

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) 

260 

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 

269 

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 

275 

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 

282 

283 xSize, ySize = image.data.shape 

284 newX = dmYpos 

285 newY = xSize - dmXpos 

286 return newX, newY 

287 

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) 

292 

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

294 tempImg.array[:] = image.data 

295 

296 disp1.mtv(tempImg) 

297 if centroid: 

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

299 

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

304 

305 raDec = vi.getBoresightRaDec() 

306 dec = raDec.getDec() 

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

308 

309 hourAngle = vi.getBoresightHourAngle() 

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

311 

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 

319 

320 airmass = vi.getBoresightAirmass() 

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

322 

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 

326 

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 

331 

332 # Upstream loads config file here 

333 

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

340 

341 filter_label, disperser = self._getFilterAndDisperserFromExp(exp) 

342 image = self.spectractorImageFromLsstExposure(exp, target_label=target, disperser_label=disperser, 

343 filter_label=filter_label) 

344 

345 if self.TRANSPOSE: 

346 xpos, ypos = self.transposeCentroid(xpos, ypos, image) 

347 

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

352 

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

359 

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) 

367 

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 

375 

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 

384 

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) 

388 

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 

397 

398 # PSF2D deconvolution 

399 if parameters.SPECTRACTOR_DECONVOLUTION_PSF2D: 

400 run_spectrogram_deconvolution_psf2d(spectrum, bgd_model_func=bgd_model_func) 

401 

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) 

409 

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) 

414 

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) 

422 

423 # Save the spectrum 

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

425 

426 # Plot the spectrum 

427 parameters.DISPLAY = True 

428 if parameters.VERBOSE and parameters.DISPLAY: 

429 spectrum.plot_spectrum(xlim=None) 

430 

431 spectrum.chromatic_psf.table['lambdas'] = spectrum.lambdas 

432 

433 result = Spectraction() 

434 result.spectrum = spectrum 

435 result.image = image 

436 result.w = w 

437 

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

439 # change it if it matters 

440 return result 

441 

442 

443class Spectraction: 

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

445 

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