Coverage for python/lsst/atmospec/processStar.py: 18%

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 

22__all__ = ['ProcessStarTask', 'ProcessStarTaskConfig'] 

23 

24import os 

25import numpy as np 

26import matplotlib.pyplot as plt 

27from matplotlib.backends.backend_pdf import PdfPages 

28from importlib import reload 

29import time 

30 

31import lsstDebug 

32import lsst.afw.image as afwImage 

33import lsst.geom as geom 

34from lsst.ip.isr import IsrTask 

35import lsst.pex.config as pexConfig 

36import lsst.pipe.base as pipeBase 

37import lsst.pipe.base.connectionTypes as cT 

38from lsst.pipe.base.task import TaskError 

39 

40from lsst.utils import getPackageDir 

41from lsst.pipe.tasks.characterizeImage import CharacterizeImageTask 

42from lsst.meas.algorithms import LoadIndexedReferenceObjectsTask, MagnitudeLimit 

43from lsst.meas.astrom import AstrometryTask, FitAffineWcsTask 

44 

45import lsst.afw.detection as afwDetect 

46 

47from .spectraction import SpectractorShim 

48from .utils import getTargetCentroidFromWcs, getLinearStagePosition 

49 

50COMMISSIONING = False # allows illegal things for on the mountain usage. 

51 

52# TODO: 

53# Sort out read noise and gain 

54# remove dummy image totally 

55# talk to Jeremy about turning the image beforehand and giving new coords 

56# deal with not having ambient temp 

57# Gen3ification 

58# astropy warning for units on save 

59# but actually just remove all manual saves entirely, I think? 

60# Make SED persistable 

61# Move to QFM for star finding failover case 

62# Remove old cruft functions 

63# change spectractions run method to be ~all kwargs with *,... 

64 

65 

66class ProcessStarTaskConnections(pipeBase.PipelineTaskConnections, 

67 dimensions=("instrument", "visit", "detector")): 

68 inputExp = cT.Input( 

69 name="icExp", 

70 doc="Image-characterize output exposure", 

71 storageClass="ExposureF", 

72 dimensions=("instrument", "visit", "detector"), 

73 multiple=False, 

74 ) 

75 inputCentroid = cT.Input( 

76 name="atmospecCentroid", 

77 doc="The main star centroid in yaml format.", 

78 storageClass="StructuredDataDict", 

79 dimensions=("instrument", "visit", "detector"), 

80 multiple=False, 

81 ) 

82 spectractorSpectrum = cT.Output( 

83 name="spectractorSpectrum", 

84 doc="The Spectractor output spectrum.", 

85 storageClass="SpectractorSpectrum", 

86 dimensions=("instrument", "visit", "detector"), 

87 ) 

88 spectractorImage = cT.Output( 

89 name="spectractorImage", 

90 doc="The Spectractor output image.", 

91 storageClass="SpectractorImage", 

92 dimensions=("instrument", "visit", "detector"), 

93 ) 

94 spectraction = cT.Output( 

95 name="spectraction", 

96 doc="The Spectractor output image.", 

97 storageClass="Spectraction", 

98 dimensions=("instrument", "visit", "detector"), 

99 ) 

100 

101 def __init__(self, *, config=None): 

102 super().__init__(config=config) 

103 

104 

105class ProcessStarTaskConfig(pipeBase.PipelineTaskConfig, 

106 pipelineConnections=ProcessStarTaskConnections): 

107 """Configuration parameters for ProcessStarTask.""" 

108 

109 isr = pexConfig.ConfigurableField( 

110 target=IsrTask, 

111 doc="Task to perform instrumental signature removal", 

112 ) 

113 charImage = pexConfig.ConfigurableField( 

114 target=CharacterizeImageTask, 

115 doc="""Task to characterize a science exposure: 

116 - detect sources, usually at high S/N 

117 - estimate the background, which is subtracted from the image and returned as field "background" 

118 - estimate a PSF model, which is added to the exposure 

119 - interpolate over defects and cosmic rays, updating the image, variance and mask planes 

120 """, 

121 ) 

122 doWrite = pexConfig.Field( 

123 dtype=bool, 

124 doc="Write out the results?", 

125 default=True, 

126 ) 

127 mainSourceFindingMethod = pexConfig.ChoiceField( 

128 doc="Which attribute to prioritize when selecting the main source object", 

129 dtype=str, 

130 default="BRIGHTEST", 

131 allowed={ 

132 "BRIGHTEST": "Select the brightest object with roundness > roundnessCut", 

133 "ROUNDEST": "Select the roundest object with brightness > fluxCut", 

134 } 

135 ) 

136 mainStarRoundnessCut = pexConfig.Field( 

137 dtype=float, 

138 doc="Value of ellipticity above which to reject the brightest object." 

139 " Ignored if mainSourceFindingMethod == BRIGHTEST", 

140 default=0.2 

141 ) 

142 mainStarFluxCut = pexConfig.Field( 

143 dtype=float, 

144 doc="Object flux below which to reject the roundest object." 

145 " Ignored if mainSourceFindingMethod == ROUNDEST", 

146 default=1e7 

147 ) 

148 mainStarNpixMin = pexConfig.Field( 

149 dtype=int, 

150 doc="Minimum number of pixels for object detection of main star", 

151 default=10 

152 ) 

153 mainStarNsigma = pexConfig.Field( 

154 dtype=int, 

155 doc="nSigma for detection of main star", 

156 default=200 # the m=0 is very bright indeed, and we don't want to detect much spectrum 

157 ) 

158 mainStarGrow = pexConfig.Field( 

159 dtype=int, 

160 doc="Number of pixels to grow by when detecting main star. This" 

161 " encourages the spectrum to merge into one footprint, but too much" 

162 " makes everything round, compromising mainStarRoundnessCut's" 

163 " effectiveness", 

164 default=5 

165 ) 

166 mainStarGrowIsotropic = pexConfig.Field( 

167 dtype=bool, 

168 doc="Grow main star's footprint isotropically?", 

169 default=False 

170 ) 

171 aperture = pexConfig.Field( 

172 dtype=int, 

173 doc="Width of the aperture to use in pixels", 

174 default=250 

175 ) 

176 spectrumLengthPixels = pexConfig.Field( 

177 dtype=int, 

178 doc="Length of the spectrum in pixels", 

179 default=5000 

180 ) 

181 offsetFromMainStar = pexConfig.Field( 

182 dtype=int, 

183 doc="Number of pixels from the main star's centroid to start extraction", 

184 default=100 

185 ) 

186 dispersionDirection = pexConfig.ChoiceField( 

187 doc="Direction along which the image is dispersed", 

188 dtype=str, 

189 default="y", 

190 allowed={ 

191 "x": "Dispersion along the serial direction", 

192 "y": "Dispersion along the parallel direction", 

193 } 

194 ) 

195 spectralOrder = pexConfig.ChoiceField( 

196 doc="Direction along which the image is dispersed", 

197 dtype=str, 

198 default="+1", 

199 allowed={ 

200 "+1": "Use the m+1 spectrum", 

201 "-1": "Use the m-1 spectrum", 

202 "both": "Use both spectra", 

203 } 

204 ) 

205 binning = pexConfig.Field( 

206 dtype=int, 

207 doc="Bin the input image by this factor", 

208 default=4 

209 ) 

210 doFlat = pexConfig.Field( 

211 dtype=bool, 

212 doc="Flatfield the image?", 

213 default=True 

214 ) 

215 doCosmics = pexConfig.Field( 

216 dtype=bool, 

217 doc="Repair cosmic rays?", 

218 default=True 

219 ) 

220 doDisplayPlots = pexConfig.Field( 

221 dtype=bool, 

222 doc="Matplotlib show() the plots, so they show up in a notebook or X window", 

223 default=False 

224 ) 

225 doSavePlots = pexConfig.Field( 

226 dtype=bool, 

227 doc="Save matplotlib plots to output rerun?", 

228 default=False 

229 ) 

230 spectractorDebugMode = pexConfig.Field( 

231 dtype=bool, 

232 doc="Set debug mode for Spectractor", 

233 default=True 

234 ) 

235 spectractorDebugLogging = pexConfig.Field( # TODO: tie this to the task debug level? 

236 dtype=bool, 

237 doc="Set debug logging for Spectractor", 

238 default=False 

239 ) 

240 forceObjectName = pexConfig.Field( 

241 dtype=str, 

242 doc="A supplementary name for OBJECT. Will be forced to apply to ALL visits, so this should only" 

243 " ONLY be used for immediate commissioning debug purposes. All long term fixes should be" 

244 " supplied as header fix-up yaml files.", 

245 default="" 

246 ) 

247 referenceFilterOverride = pexConfig.Field( 

248 dtype=str, 

249 doc="Which filter in the reference catalog to match to?", 

250 default="phot_g_mean" 

251 ) 

252 

253 def setDefaults(self): 

254 self.isr.doWrite = False 

255 self.charImage.doWriteExposure = False 

256 

257 self.charImage.doApCorr = False 

258 self.charImage.doMeasurePsf = False 

259 self.charImage.repair.cosmicray.nCrPixelMax = 100000 

260 self.charImage.repair.doCosmicRay = False 

261 if self.charImage.doMeasurePsf: 

262 self.charImage.measurePsf.starSelector['objectSize'].signalToNoiseMin = 10.0 

263 self.charImage.measurePsf.starSelector['objectSize'].fluxMin = 5000.0 

264 self.charImage.detection.includeThresholdMultiplier = 3 

265 self.isr.overscan.fitType = 'MEDIAN_PER_ROW' 

266 

267 

268class ProcessStarTask(pipeBase.PipelineTask): 

269 """Task for the spectral extraction of single-star dispersed images. 

270 

271 For a full description of how this tasks works, see the run() method. 

272 """ 

273 

274 ConfigClass = ProcessStarTaskConfig 

275 _DefaultName = "processStar" 

276 

277 def __init__(self, *, butler=None, **kwargs): 

278 # TODO: rename psfRefObjLoader to refObjLoader 

279 super().__init__(**kwargs) 

280 self.makeSubtask("isr") 

281 self.makeSubtask("charImage", butler=butler, refObjLoader=None) 

282 

283 self.debug = lsstDebug.Info(__name__) 

284 if self.debug.enabled: 

285 self.log.info("Running with debug enabled...") 

286 # If we're displaying, test it works and save displays for later. 

287 # It's worth testing here as displays are flaky and sometimes 

288 # can't be contacted, and given processing takes a while, 

289 # it's a shame to fail late due to display issues. 

290 if self.debug.display: 

291 try: 

292 import lsst.afw.display as afwDisp 

293 afwDisp.setDefaultBackend(self.debug.displayBackend) 

294 afwDisp.Display.delAllDisplays() 

295 # pick an unlikely number to be safe xxx replace this 

296 self.disp1 = afwDisp.Display(987, open=True) 

297 

298 im = afwImage.ImageF(2, 2) 

299 im.array[:] = np.ones((2, 2)) 

300 self.disp1.mtv(im) 

301 self.disp1.erase() 

302 afwDisp.setDefaultMaskTransparency(90) 

303 except NameError: 

304 self.debug.display = False 

305 self.log.warn('Failed to setup/connect to display! Debug display has been disabled') 

306 

307 if self.debug.notHeadless: 

308 pass # other backend options can go here 

309 else: # this stop windows popping up when plotting. When headless, use 'agg' backend too 

310 plt.interactive(False) 

311 

312 self.config.validate() 

313 self.config.freeze() 

314 

315 def findObjects(self, exp, nSigma=None, grow=0): 

316 """Find the objects in a postISR exposure.""" 

317 nPixMin = self.config.mainStarNpixMin 

318 if not nSigma: 

319 nSigma = self.config.mainStarNsigma 

320 if not grow: 

321 grow = self.config.mainStarGrow 

322 isotropic = self.config.mainStarGrowIsotropic 

323 

324 threshold = afwDetect.Threshold(nSigma, afwDetect.Threshold.STDEV) 

325 footPrintSet = afwDetect.FootprintSet(exp.getMaskedImage(), threshold, "DETECTED", nPixMin) 

326 if grow > 0: 

327 footPrintSet = afwDetect.FootprintSet(footPrintSet, grow, isotropic) 

328 return footPrintSet 

329 

330 def _getEllipticity(self, shape): 

331 """Calculate the ellipticity given a quadrupole shape. 

332 

333 Parameters 

334 ---------- 

335 shape : `lsst.afw.geom.ellipses.Quadrupole` 

336 The quadrupole shape 

337 

338 Returns 

339 ------- 

340 ellipticity : `float` 

341 The magnitude of the ellipticity 

342 """ 

343 ixx = shape.getIxx() 

344 iyy = shape.getIyy() 

345 ixy = shape.getIxy() 

346 ePlus = (ixx - iyy) / (ixx + iyy) 

347 eCross = 2*ixy / (ixx + iyy) 

348 return (ePlus**2 + eCross**2)**0.5 

349 

350 def getRoundestObject(self, footPrintSet, parentExp, fluxCut=1e-15): 

351 """Get the roundest object brighter than fluxCut from a footPrintSet. 

352 

353 Parameters 

354 ---------- 

355 footPrintSet : `lsst.afw.detection.FootprintSet` 

356 The set of footprints resulting from running detection on parentExp 

357 

358 parentExp : `lsst.afw.image.exposure` 

359 The parent exposure for the footprint set. 

360 

361 fluxCut : `float` 

362 The flux, below which, sources are rejected. 

363 

364 Returns 

365 ------- 

366 source : `lsst.afw.detection.Footprint` 

367 The winning footprint from the input footPrintSet 

368 """ 

369 self.log.debug("ellipticity\tflux/1e6\tcentroid") 

370 sourceDict = {} 

371 for fp in footPrintSet.getFootprints(): 

372 shape = fp.getShape() 

373 e = self._getEllipticity(shape) 

374 flux = fp.getSpans().flatten(parentExp.image.array, parentExp.image.getXY0()).sum() 

375 self.log.debug("%.4f\t%.2f\t%s"%(e, flux/1e6, str(fp.getCentroid()))) 

376 if flux > fluxCut: 

377 sourceDict[e] = fp 

378 

379 return sourceDict[sorted(sourceDict.keys())[0]] 

380 

381 def getBrightestObject(self, footPrintSet, parentExp, roundnessCut=1e9): 

382 """Get the brightest object rounder than the cut from a footPrintSet. 

383 

384 Parameters 

385 ---------- 

386 footPrintSet : `lsst.afw.detection.FootprintSet` 

387 The set of footprints resulting from running detection on parentExp 

388 

389 parentExp : `lsst.afw.image.exposure` 

390 The parent exposure for the footprint set. 

391 

392 roundnessCut : `float` 

393 The ellipticity, above which, sources are rejected. 

394 

395 Returns 

396 ------- 

397 source : `lsst.afw.detection.Footprint` 

398 The winning footprint from the input footPrintSet 

399 """ 

400 self.log.debug("ellipticity\tflux\tcentroid") 

401 sourceDict = {} 

402 for fp in footPrintSet.getFootprints(): 

403 shape = fp.getShape() 

404 e = self._getEllipticity(shape) 

405 flux = fp.getSpans().flatten(parentExp.image.array, parentExp.image.getXY0()).sum() 

406 self.log.debug("%.4f\t%.2f\t%s"%(e, flux/1e6, str(fp.getCentroid()))) 

407 if e < roundnessCut: 

408 sourceDict[flux] = fp 

409 

410 return sourceDict[sorted(sourceDict.keys())[-1]] 

411 

412 def findMainSource(self, exp): 

413 """Return the x,y of the brightest or roundest object in an exposure. 

414 

415 Given a postISR exposure, run source detection on it, and return the 

416 centroid of the main star. Depending on the task configuration, this 

417 will either be the roundest object above a certain flux cutoff, or 

418 the brightest object which is rounder than some ellipticity cutoff. 

419 

420 Parameters 

421 ---------- 

422 exp : `afw.image.Exposure` 

423 The postISR exposure in which to find the main star 

424 

425 Returns 

426 ------- 

427 x, y : `tuple` of `float` 

428 The centroid of the main star in the image 

429 

430 Notes 

431 ----- 

432 Behavior of this method is controlled by many task config params 

433 including, for the detection stage: 

434 config.mainStarNpixMin 

435 config.mainStarNsigma 

436 config.mainStarGrow 

437 config.mainStarGrowIsotropic 

438 

439 And post-detection, for selecting the main source: 

440 config.mainSourceFindingMethod 

441 config.mainStarFluxCut 

442 config.mainStarRoundnessCut 

443 """ 

444 # TODO: probably replace all this with QFM 

445 fpSet = self.findObjects(exp) 

446 if self.config.mainSourceFindingMethod == 'ROUNDEST': 

447 source = self.getRoundestObject(fpSet, exp, fluxCut=self.config.mainStarFluxCut) 

448 elif self.config.mainSourceFindingMethod == 'BRIGHTEST': 

449 source = self.getBrightestObject(fpSet, exp, 

450 roundnessCut=self.config.mainStarRoundnessCut) 

451 else: 

452 # should be impossible as this is a choice field, but still 

453 raise RuntimeError("Invalid source finding method " 

454 f"selected: {self.config.mainSourceFindingMethod}") 

455 return source.getCentroid() 

456 

457 def updateMetadata(self, exp, **kwargs): 

458 md = exp.getMetadata() 

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

460 

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

462 airmass = vi.getBoresightAirmass() 

463 

464 md['HA'] = ha 

465 md.setComment('HA', 'Hour angle of observation start') 

466 

467 md['AIRMASS'] = airmass 

468 md.setComment('AIRMASS', 'Airmass at observation start') 

469 

470 if 'centroid' in kwargs: 

471 centroid = kwargs['centroid'] 

472 else: 

473 centroid = (None, None) 

474 

475 md['OBJECTX'] = centroid[0] 

476 md.setComment('OBJECTX', 'x pixel coordinate of object centroid') 

477 

478 md['OBJECTY'] = centroid[1] 

479 md.setComment('OBJECTY', 'y pixel coordinate of object centroid') 

480 

481 exp.setMetadata(md) 

482 

483 def runQuantum(self, butlerQC, inputRefs, outputRefs): 

484 inputs = butlerQC.get(inputRefs) 

485 

486 inputs['dataIdDict'] = inputRefs.inputExp.dataId.byName() 

487 

488 outputs = self.run(**inputs) 

489 butlerQC.put(outputs, outputRefs) 

490 

491 def run(self, *, inputExp, inputCentroid, dataIdDict): 

492 starNames = self.loadStarNames() 

493 

494 overrideDict = {'SAVE': False, 

495 'OBS_NAME': 'AUXTEL', 

496 'DEBUG': self.config.spectractorDebugMode, 

497 'DEBUG_LOGGING': self.config.spectractorDebugLogging, 

498 'DISPLAY': self.config.doDisplayPlots, 

499 'CCD_REBIN': self.config.binning, 

500 'VERBOSE': 0, 

501 # 'CCD_IMSIZE': 4000} 

502 } 

503 supplementDict = {'CALLING_CODE': 'LSST_DM', 

504 'STAR_NAMES': starNames} 

505 

506 # anything that changes between dataRefs! 

507 resetParameters = {} 

508 # TODO: look at what to do with config option doSavePlots 

509 

510 # TODO: think if this is the right place for this 

511 # probably wants to go in spectraction.py really 

512 linearStagePosition = getLinearStagePosition(inputExp) 

513 overrideDict['DISTANCE2CCD'] = linearStagePosition 

514 

515 target = inputExp.getMetadata()['OBJECT'] 

516 if self.config.forceObjectName: 

517 self.log.info(f"Forcing target name from {target} to {self.config.forceObjectName}") 

518 target = self.config.forceObjectName 

519 

520 if target in ['FlatField position', 'Park position', 'Test', 'NOTSET']: 

521 raise ValueError(f"OBJECT set to {target} - this is not a celestial object!") 

522 

523 packageDir = getPackageDir('atmospec') 

524 configFilename = os.path.join(packageDir, 'config', 'auxtel.ini') 

525 

526 spectractor = SpectractorShim(configFile=configFilename, 

527 paramOverrides=overrideDict, 

528 supplementaryParameters=supplementDict, 

529 resetParameters=resetParameters) 

530 

531 if 'astrometricMatch' in inputCentroid: 

532 centroid = inputCentroid['centroid'] 

533 else: # it's a raw tuple 

534 centroid = inputCentroid # TODO: put this support in the docstring 

535 

536 spectraction = spectractor.run(inputExp, *centroid, target) 

537 

538 self.log.info("Finished processing %s" % (dataIdDict)) 

539 

540 self.makeResultPickleable(spectraction) 

541 

542 return pipeBase.Struct(spectractorSpectrum=spectraction.spectrum, 

543 spectractorImage=spectraction.image, 

544 spectraction=spectraction) 

545 

546 def runDataRef(self, dataRef): 

547 """Run the ProcessStarTask on a ButlerDataRef for a single exposure. 

548 

549 Runs isr to get the postISR exposure from the dataRef and passes this 

550 to the run() method. 

551 

552 Parameters 

553 ---------- 

554 dataRef : `daf.persistence.butlerSubset.ButlerDataRef` 

555 Butler reference of the detector and exposure ID 

556 """ 

557 t0 = time.time() 

558 butler = dataRef.getButler() 

559 dataId = dataRef.dataId 

560 self.log.info("Processing %s" % (dataRef.dataId)) 

561 

562 if COMMISSIONING: 

563 from lsst.rapid.analysis.bestEffort import BestEffortIsr # import here because not part of DM 

564 # TODO: some evidence suggests that CR repair is *significantly* 

565 # degrading spectractor performance investigate this for the 

566 # default task config as well as ensuring that it doesn't run here 

567 # if it does turn out to be problematic. 

568 bestEffort = BestEffortIsr(butler=dataRef.getButler()) 

569 exposure = bestEffort.getExposure(dataId) 

570 else: 

571 if butler.datasetExists('postISRCCD', dataId): 

572 exposure = butler.get('postISRCCD', dataId) 

573 self.log.info("Loaded postISRCCD from disk") 

574 else: 

575 exposure = self.isr.runDataRef(dataRef).exposure 

576 self.updateMetadata(exposure) 

577 butler.put(exposure, 'postISRCCD', dataId) 

578 

579 if butler.datasetExists('icExp', dataId) and butler.datasetExists('icSrc', dataId): 

580 exposure = butler.get('icExp', dataId) 

581 icSrc = butler.get('icSrc', dataId) 

582 self.log.info("Loaded icExp and icSrc from disk") 

583 else: 

584 charRes = self.charImage.runDataRef(dataRef=dataRef, exposure=exposure, doUnpersist=False) 

585 exposure = charRes.exposure 

586 icSrc = charRes.sourceCat 

587 butler.put(exposure, 'icExp', dataId) 

588 butler.put(icSrc, 'icSrc', dataId) 

589 

590 if butler.datasetExists('calexp', dataId): 

591 exposure = butler.get('calexp', dataId) 

592 self.log.info("Loaded calexp from disk") 

593 md = exposure.getMetadata() 

594 sourceCentroid = (md['OBJECTX'], md['OBJECTY']) # set in saved md if previous fit succeeded 

595 else: 

596 astromResult = self.runAstrometry(butler, exposure, icSrc) 

597 if astromResult and astromResult.scatterOnSky.asArcseconds() < 1: 

598 target = exposure.getMetadata()['OBJECT'] 

599 sourceCentroid = getTargetCentroidFromWcs(exposure, target, logger=self.log) 

600 else: 

601 sourceCentroid = self.findMainSource(exposure) 

602 self.log.warn("Astrometric fit failed, failing over to source-finding centroid") 

603 self.log.info(f"Centroid of main star at: {sourceCentroid!r}") 

604 

605 self.updateMetadata(exposure, centroid=sourceCentroid) 

606 butler.put(exposure, 'calexp', dataId) 

607 

608 if self.debug.display and 'raw' in self.debug.displayItems: 

609 self.disp1.mtv(exposure) 

610 self.disp1.dot('x', sourceCentroid[0], sourceCentroid[1], size=100) 

611 self.log.info("Showing full postISR image") 

612 self.log.info(f"Centroid of main star at: {sourceCentroid}") 

613 self.pause() 

614 

615 outputRoot = dataRef.getUri(datasetType='spectractorOutputRoot', write=True) 

616 if not os.path.exists(outputRoot): 

617 os.makedirs(outputRoot) 

618 if not os.path.exists(outputRoot): 

619 raise RuntimeError(f"Failed to create output dir {outputRoot}") 

620 expId = dataRef.dataId['expId'] 

621 

622 result = self.runGen2(exposure, outputRoot, expId, sourceCentroid) 

623 self.log.info("Finished processing %s" % (dataRef.dataId)) 

624 

625 result.dataId = dataId 

626 self.makeResultPickleable(result) 

627 butler.put(result, 'spectraction', dataId) 

628 

629 t1 = time.time() - t0 

630 self.log.info(f'Successfully ran to completion in {t1:.1f}s for {dataId}') 

631 

632 return result 

633 

634 def makeResultPickleable(self, result): 

635 """Remove unpicklable components from the output""" 

636 result.image.target.build_sed = None 

637 result.spectrum.target.build_sed = None 

638 result.image.target.sed = None 

639 result.spectrum.disperser.load_files = None 

640 result.image.disperser.load_files = None 

641 

642 result.spectrum.disperser.N_fit = None 

643 result.spectrum.disperser.N_interp = None 

644 result.spectrum.disperser.ratio_order_2over1 = None 

645 result.spectrum.disperser.theta = None 

646 

647 def runAstrometry(self, butler, exp, icSrc): 

648 refObjLoaderConfig = LoadIndexedReferenceObjectsTask.ConfigClass() 

649 refObjLoaderConfig.ref_dataset_name = 'gaia_dr2_20191105' 

650 refObjLoaderConfig.pixelMargin = 1000 

651 refObjLoader = LoadIndexedReferenceObjectsTask(butler=butler, config=refObjLoaderConfig) 

652 

653 astromConfig = AstrometryTask.ConfigClass() 

654 astromConfig.wcsFitter.retarget(FitAffineWcsTask) 

655 astromConfig.referenceSelector.doMagLimit = True 

656 magLimit = MagnitudeLimit() 

657 magLimit.minimum = 1 

658 magLimit.maximum = 15 

659 astromConfig.referenceSelector.magLimit = magLimit 

660 astromConfig.referenceSelector.magLimit.fluxField = "phot_g_mean_flux" 

661 astromConfig.matcher.maxRotationDeg = 5.99 

662 astromConfig.matcher.maxOffsetPix = 3000 

663 astromConfig.sourceSelector['matcher'].minSnr = 10 

664 solver = AstrometryTask(config=astromConfig, refObjLoader=refObjLoader) 

665 

666 # TODO: Change this to doing this the proper way 

667 referenceFilterName = self.config.referenceFilterOverride 

668 referenceFilterLabel = afwImage.FilterLabel(physical=referenceFilterName, band=referenceFilterName) 

669 originalFilterLabel = exp.getFilter() # there's a better way of doing this with the task I think 

670 exp.setFilter(referenceFilterLabel) 

671 

672 try: 

673 astromResult = solver.run(sourceCat=icSrc, exposure=exp) 

674 exp.setFilter(originalFilterLabel) 

675 except (RuntimeError, TaskError): 

676 self.log.warn("Solver failed to run completely") 

677 exp.setFilter(originalFilterLabel) 

678 return None 

679 

680 scatter = astromResult.scatterOnSky.asArcseconds() 

681 if scatter < 1: 

682 return astromResult 

683 else: 

684 self.log.warn("Failed to find an acceptable match") 

685 return None 

686 

687 def pause(self): 

688 if self.debug.pauseOnDisplay: 

689 input("Press return to continue...") 

690 return 

691 

692 def loadStarNames(self): 

693 starNameFile = os.path.join(getPackageDir('atmospec'), 'data', 'starNames.txt') 

694 with open(starNameFile, 'r') as f: 

695 lines = f.readlines() 

696 return [line.strip() for line in lines] 

697 

698 def runGen2(self, exp, spectractorOutputRoot, expId, sourceCentroid): 

699 """Calculate the wavelength calibrated 1D spectrum from a postISRCCD. 

700 

701 An outline of the steps in the processing is as follows: 

702 * Source extraction - find the objects in image 

703 * Process sources to find the x,y of the main star 

704 

705 * Given the centroid, the dispersion direction, and the order(s), 

706 calculate the spectrum's bounding box 

707 

708 * (Rotate the image such that the dispersion direction is vertical 

709 TODO: DM-18138) 

710 

711 * Create an initial dispersion relation object from the geometry 

712 or alternative bootstrapping method 

713 

714 * Apply an initial flatfielding - TODO: DM-18141 

715 

716 * Find and interpolate over cosmics if necessary - TODO: DM-18140 

717 

718 * Perform an initial spectral extraction, depending on selected method 

719 * Fit a background model and subtract 

720 * Perform row-wise fits for extraction 

721 * TODO: DM-18136 for doing a full-spectrum fit with PSF model 

722 

723 * Given knowledge of features in the spectrum, find lines in the 

724 measured spectrum and re-fit to refine the dispersion relation 

725 * Reflatfield the image with the refined dispersion relation 

726 

727 Parameters 

728 ---------- 

729 exp : `afw.image.Exposure` 

730 The postISR exposure in which to find the main star 

731 

732 Returns 

733 ------- 

734 spectrum : `lsst.atmospec.spectrum` - TODO: DM-18133 

735 The wavelength-calibrated 1D stellar spectrum 

736 """ 

737 reload(plt) # reset matplotlib color cycles when multiprocessing 

738 pdfPath = os.path.join(spectractorOutputRoot, 'plots.pdf') 

739 starNames = self.loadStarNames() 

740 

741 if True: 

742 overrideDict = {'SAVE': False, 

743 'OBS_NAME': 'AUXTEL', 

744 'DEBUG': self.config.spectractorDebugMode, 

745 'DEBUG_LOGGING': self.config.spectractorDebugLogging, 

746 'DISPLAY': self.config.doDisplayPlots, 

747 'CCD_REBIN': self.config.binning, 

748 'VERBOSE': 0, 

749 # 'CCD_IMSIZE': 4000} 

750 } 

751 supplementDict = {'CALLING_CODE': 'LSST_DM', 

752 'STAR_NAMES': starNames} 

753 

754 # anything that changes between dataRefs! 

755 resetParameters = {} 

756 if self.config.doSavePlots: 

757 resetParameters['LSST_SAVEFIGPATH'] = spectractorOutputRoot 

758 else: 

759 overrideDict = {} 

760 supplementDict = {} 

761 

762 # TODO: think if this is the right place for this 

763 # probably wants to go in spectraction.py really 

764 linearStagePosition = getLinearStagePosition(exp) 

765 overrideDict['DISTANCE2CCD'] = linearStagePosition 

766 

767 target = exp.getMetadata()['OBJECT'] 

768 if self.config.forceObjectName: 

769 self.log.info(f"Forcing target name from {target} to {self.config.forceObjectName}") 

770 target = self.config.forceObjectName 

771 

772 if target in ['FlatField position', 'Park position', 'Test', 'NOTSET']: 

773 raise ValueError(f"OBJECT set to {target} - this is not a celestial object!") 

774 

775 packageDir = getPackageDir('atmospec') 

776 configFilename = os.path.join(packageDir, 'config', 'auxtel.ini') 

777 

778 if self.config.doDisplayPlots: # no pdfpages backend - isn't compatible with display-as-you-go 

779 spectractor = SpectractorShim(configFile=configFilename, 

780 paramOverrides=overrideDict, 

781 supplementaryParameters=supplementDict, 

782 resetParameters=resetParameters) 

783 result = spectractor.run(exp, *sourceCentroid, target, spectractorOutputRoot) 

784 else: 

785 try: # need a try here so that the context manager always exits cleanly so plots always written 

786 with PdfPages(pdfPath) as pdf: # TODO: Doesn't the try need to be inside the with?! 

787 resetParameters['PdfPages'] = pdf 

788 spectractor = SpectractorShim(configFile=configFilename, 

789 paramOverrides=overrideDict, 

790 supplementaryParameters=supplementDict, 

791 resetParameters=resetParameters) 

792 

793 result = spectractor.run(exp, *sourceCentroid, target, spectractorOutputRoot) 

794 except Exception as e: 

795 self.log.warn(f"Caught exception {e}, passing here so pdf can be written to {pdfPath}") 

796 result = None 

797 return result 

798 

799 def flatfield(self, exp, disp): 

800 """Placeholder for wavelength dependent flatfielding: TODO: DM-18141 

801 

802 Will probably need a dataRef, as it will need to be retrieving flats 

803 over a range. Also, it will be somewhat complex, so probably needs 

804 moving to its own task""" 

805 self.log.warn("Flatfielding not yet implemented") 

806 return exp 

807 

808 def repairCosmics(self, exp, disp): 

809 self.log.warn("Cosmic ray repair not yet implemented") 

810 return exp 

811 

812 def measureSpectrum(self, exp, sourceCentroid, spectrumBBox, dispersionRelation): 

813 """Perform the spectral extraction, given a source location and exp.""" 

814 

815 self.extraction.initialise(exp, sourceCentroid, spectrumBBox, dispersionRelation) 

816 

817 # xxx this method currently doesn't return an object - fix this 

818 spectrum = self.extraction.getFluxBasic() 

819 

820 return spectrum 

821 

822 def calcSpectrumBBox(self, exp, centroid, aperture, order='+1'): 

823 """Calculate the bbox for the spectrum, given the centroid. 

824 

825 XXX Longer explanation here, inc. parameters 

826 TODO: Add support for order = "both" 

827 """ 

828 extent = self.config.spectrumLengthPixels 

829 halfWidth = aperture//2 

830 translate_y = self.config.offsetFromMainStar 

831 sourceX = centroid[0] 

832 sourceY = centroid[1] 

833 

834 if(order == '-1'): 

835 translate_y = - extent - self.config.offsetFromMainStar 

836 

837 xStart = sourceX - halfWidth 

838 xEnd = sourceX + halfWidth - 1 

839 yStart = sourceY + translate_y 

840 yEnd = yStart + extent - 1 

841 

842 xEnd = min(xEnd, exp.getWidth()-1) 

843 yEnd = min(yEnd, exp.getHeight()-1)