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

1# This file is part of atmospec.

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.

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/>.

22import logging

23import os

24import numpy as np

25import astropy.coordinates as asCoords

26from astropy import units as u

28from spectractor import parameters

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

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

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

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)

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

46from .utils import getFilterAndDisperserFromExp # noqa: E402

49class SpectractorShim:

50 """Class for running the Spectractor code.

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

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

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)

69 if parameters.DEBUG:

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

71 dumpParameters()

73 return

75 def overrideParameters(self, overrides):

76 """Dict of Spectractor parameters to override.

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.

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.warning("Did not find attribute %s in parameters" % k)

94 raise RuntimeError(f"{k} not set to {v} {self.dumpParameters()}")

96 def supplementParameters(self, supplementaryItems):

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

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.warning(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, data_dir=parameters.DISPERSER_DIR)

231

232 image.compute_parallactic_angle()

233

234 return image

235

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

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

238 filt, disperser = getFilterAndDisperserFromExp(exp)

239

240 image.header.filter = filt

241 image.header.disperser_label = disperser

242

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

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

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

246

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

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

249

250 try:

251 if useVisitInfo:

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

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

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

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

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

257 else:

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

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

260 image.airmass = md['AIRMASS']

261 image.date_obs = md['DATE']

262 except Exception:

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

264 image.header.airmass = 1.

265

266 return

267

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

269 if trimToSquare:

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

271 else:

272 data = exp.image.array

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

274

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

276 # xxx need to implement this properly

277 if constValue is not None:

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

279 else:

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

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

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

283

284 def _setReadNoiseToNone(self, spectractorImage):

285 spectractorImage.read_out_noise = None

286

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

288 if useExpVariance:

289 image.err = exp.maskedImage.variance.array # xxx need to deal with TRANSPOSE here

290 else:

291 image.compute_statistical_error()

292

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

294 # xxx need to implement this properly

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

296 # so could use the code from gain flats

297 if constValue:

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

299 return np.ones_like(spectractorImage.data)

300

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

302 if plotting:

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

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

305 os.makedirs(dirname)

306

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

308 if 'SIMPLE' not in obj.header:

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

310 # if dataDict:

311 # header = obj.header

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

313 # if k not in header:

314 # header[k] = v

315

316 @staticmethod

317 def flipImageLeftRight(image, xpos, ypos):

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

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

320 return image, xpos, ypos

321

322 @staticmethod

323 def transposeCentroid(dmXpos, dmYpos, image):

324 xSize, ySize = image.data.shape

325 newX = dmYpos

326 newY = xSize - dmXpos

327 return newX, newY

328

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

330 import lsst.afw.image as afwImage

331 import lsst.afw.display as afwDisp

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

333

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

335 tempImg.array[:] = image.data

336

337 disp1.mtv(tempImg)

338 if centroid:

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

340

341 def setAdrParameters(self, spectrum, exp):

342 # The adr_params parameter format expected by spectractor are:

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

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

345

346 raDec = vi.getBoresightRaDec()

347 dec = raDec.getDec()

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

349

350 hourAngle = vi.getBoresightHourAngle()

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

352

353 weather = vi.getWeather()

354

355 _temperature = weather.getAirTemperature()

356 if _temperature is None or np.isnan(_temperature):

357 self.log.warning("Temperature not set, using nominal value of 10 C")

358 _temperature = 10 # nominal value

359 temperature = _temperature

360

361 _pressure = weather.getAirPressure()

362 if _pressure is not None and not np.isnan(_pressure):

363 if _pressure > 10_000:

364 _pressure /= 100 # convert from Pa to hPa

365 else:

366 self.log.warning("Pressure not set, using nominal value of 743 hPa")

367 _pressure = 743 # nominal for altitude?

368 pressure = _pressure

369 _humidity = weather.getHumidity()

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

371

372 airmass = vi.getBoresightAirmass()

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

374 spectrum.pressure = pressure

375 spectrum.humidity = humidity

376 spectrum.airmass = airmass

377 spectrum.temperature = temperature

378

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

380 outputRoot=None, plotting=True):

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

382 atmospheric_lines = True

383

384 self.log.info('Starting SPECTRACTOR')

385 # TODO: rename _makePath _makeOutputPath

386 if outputRoot is not None: # TODO: remove post Gen3 transition

387 self._makePath(outputRoot, plotting=plotting) # early in case this fails, as processing is slow

388

389 # Upstream loads config file here

390

391 filter_label, disperser = getFilterAndDisperserFromExp(exp)

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

393 disperser_label=disperser,

394 filter_label=filter_label)

395

396 if parameters.DEBUG:

397 self.debugPrintTargetCentroidValue(image)

398 title = 'Raw image with input target location'

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

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

401

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

403 # just run! ASAP XXX

404 # if disperser == 'ronchi170lpmm':

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

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

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

408

409 if parameters.CCD_REBIN > 1:

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

411 apply_rebinning_to_parameters()

412 image.rebin()

413 if parameters.DEBUG: 413 ↛ 425line 413 didn't jump to line 425 because the condition on line 413 was always true

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

415 dumpParameters()

416 self.debugPrintTargetCentroidValue(image)

417 title = 'Rebinned image with input target location'

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

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

420 self.debugPrintTargetCentroidValue(image)

421

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

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

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

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

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

427 # sets the image.target_pixcoords

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

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

430

431 # Rotate the image: several methods

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

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

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

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

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

437 else:

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

439 raise NotImplementedError

440

441 # Create Spectrum object

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

443 self.setAdrParameters(spectrum, exp)

444

445 # Subtract background and bad pixels

446 w_psf1d, bgd_model_func = extract_spectrum_from_image(image, spectrum,

447 signal_width=parameters.PIXWIDTH_SIGNAL,

448 ws=(parameters.PIXDIST_BACKGROUND,

449 parameters.PIXDIST_BACKGROUND

450 + parameters.PIXWIDTH_BACKGROUND))

451 spectrum.atmospheric_lines = atmospheric_lines

452

453 # PSF2D deconvolution

454 if parameters.SPECTRACTOR_DECONVOLUTION_PSF2D:

455 run_spectrogram_deconvolution_psf2d(spectrum, bgd_model_func=bgd_model_func)

456

457 # Calibrate the spectrum

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

459 with_adr = True

460 if parameters.OBS_OBJECT_TYPE != "STAR":

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

462 # likely need to be passed through

463 with_adr = False

464 calibrate_spectrum(spectrum, with_adr=with_adr)

465

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

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

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

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

470

471 # Full forward model extraction:

472 # adds transverse ADR and order 2 subtraction

473 ffmWorkspace = None

474 if parameters.SPECTRACTOR_DECONVOLUTION_FFM:

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

476 plot=True,

477 live_fit=False,

478 amplitude_priors_method="spectrum")

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

480

481 # Fit the atmosphere on the spectrum using uvspec binary

482 spectrumAtmosphereWorkspace = None

483 if doFitAtmosphere:

484 spectrumAtmosphereWorkspace = SpectrumFitWorkspace(spectrum,

485 fit_angstrom_exponent=True,

486 verbose=parameters.VERBOSE,

487 plot=True)

488 run_spectrum_minimisation(spectrumAtmosphereWorkspace, method="newton")

489

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

491 spectrogramAtmosphereWorkspace = None

492 if doFitAtmosphereOnSpectrogram:

493 spectrogramAtmosphereWorkspace = SpectrogramFitWorkspace(spectrum,

494 fit_angstrom_exponent=True,

495 verbose=parameters.VERBOSE,

496 plot=True)

497 run_spectrogram_minimisation(spectrogramAtmosphereWorkspace, method="newton")

498

499 # Save the spectrum

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

501

502 # Plot the spectrum

503 parameters.DISPLAY = True

504 if parameters.VERBOSE and parameters.DISPLAY:

505 spectrum.plot_spectrum(xlim=None)

506

507 result = Spectraction()

508 result.spectrum = spectrum

509 result.image = image

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

511 result.spectrumLibradtranFitParameters = (spectrumAtmosphereWorkspace.params if

512 spectrumAtmosphereWorkspace is not None else None)

513 result.spectrogramLibradtranFitParameters = (spectrogramAtmosphereWorkspace.params if

514 spectrogramAtmosphereWorkspace is not None else None)

515

516 return result

517

518

519class Spectraction:

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

521

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

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

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

525 """

526 # result.spectrum = spectrum

527 # result.image = image

528 # result.w = w