Coverage for python/lsst/summit/extras/animation.py: 13%

229 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-12 03:05 -0700

1# This file is part of summit_extras. 

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 gc 

23import logging 

24import math 

25import os 

26import shutil 

27import subprocess 

28import uuid 

29from typing import Any 

30 

31import matplotlib.pyplot as plt 

32 

33import lsst.afw.display as afwDisplay 

34import lsst.afw.image as afwImage 

35import lsst.afw.math as afwMath 

36import lsst.daf.butler as dafButler 

37import lsst.meas.algorithms as measAlg 

38from lsst.atmospec.utils import airMassFromRawMetadata, getTargetCentroidFromWcs 

39from lsst.pipe.tasks.quickFrameMeasurement import QuickFrameMeasurementTask, QuickFrameMeasurementTaskConfig 

40from lsst.summit.utils.butlerUtils import ( 

41 datasetExists, 

42 getDayObs, 

43 getExpRecordFromDataId, 

44 getLatissOnSkyDataIds, 

45 getSeqNum, 

46 makeDefaultLatissButler, 

47 updateDataIdOrDataCord, 

48) 

49from lsst.summit.utils.utils import dayObsIntToString, setupLogging 

50 

51logger = logging.getLogger("lsst.summit.extras.animation") 

52setupLogging() 

53 

54 

55class Animator: 

56 """Animate the list of dataIds in the order in which they are specified 

57 for the data product specified.""" 

58 

59 def __init__( 

60 self, 

61 butler: dafButler.Butler, 

62 dataIdList: list[dict], 

63 outputPath: str, 

64 outputFilename: str, 

65 *, 

66 remakePngs: bool = False, 

67 clobberVideoAndGif: bool = False, 

68 keepIntermediateGif: bool = False, 

69 smoothImages: bool = True, 

70 plotObjectCentroids: bool = True, 

71 useQfmForCentroids: bool = False, 

72 dataProductToPlot: str = "calexp", 

73 debug: bool = False, 

74 ): 

75 self.butler = butler 

76 self.dataIdList = dataIdList 

77 self.outputPath = outputPath 

78 self.outputFilename = os.path.join(outputPath, outputFilename) 

79 if not self.outputFilename.endswith(".mp4"): 

80 self.outputFilename += ".mp4" 

81 self.pngPath = os.path.join(outputPath, "pngs/") 

82 

83 self.remakePngs = remakePngs 

84 self.clobberVideoAndGif = clobberVideoAndGif 

85 self.keepIntermediateGif = keepIntermediateGif 

86 self.smoothImages = smoothImages 

87 self.plotObjectCentroids = plotObjectCentroids 

88 self.useQfmForCentroids = useQfmForCentroids 

89 self.dataProductToPlot = dataProductToPlot 

90 self.debug = debug 

91 

92 # zfilled at the start as animation is alphabetical 

93 # if you're doing more than 1e6 files you've got bigger problems 

94 self.toAnimateTemplate = "%06d-%s-%s.png" 

95 self.basicTemplate = "%s-%s.png" 

96 

97 qfmTaskConfig = QuickFrameMeasurementTaskConfig() 

98 self.qfmTask = QuickFrameMeasurementTask(config=qfmTaskConfig) 

99 

100 afwDisplay.setDefaultBackend("matplotlib") 

101 self.fig = plt.figure(figsize=(15, 15)) 

102 self.disp = afwDisplay.Display(self.fig) 

103 self.disp.setImageColormap("gray") 

104 self.disp.scale("asinh", "zscale") 

105 

106 self.pngsToMakeDataIds: list[dict] = [] 

107 

108 self.preRun() # sets the above list 

109 

110 @staticmethod 

111 def _strDataId(dataId: dict) -> str: 

112 """Make a dataId into a string suitable for use as a filename. 

113 

114 Parameters 

115 ---------- 

116 dataId : `dict` 

117 The data id. 

118 

119 Returns 

120 ------- 

121 strId : `str` 

122 The data id as a string. 

123 """ 

124 if (dayObs := getDayObs(dataId)) and (seqNum := getSeqNum(dataId)): # nicely ordered if easy 

125 return f"{dayObsIntToString(dayObs)}-{seqNum:05d}" 

126 

127 # General case (and yeah, I should probably learn regex someday) 

128 dIdStr = str(dataId) 

129 dIdStr = dIdStr.replace(" ", "") 

130 dIdStr = dIdStr.replace("{", "") 

131 dIdStr = dIdStr.replace("}", "") 

132 dIdStr = dIdStr.replace("'", "") 

133 dIdStr = dIdStr.replace(":", "-") 

134 dIdStr = dIdStr.replace(",", "-") 

135 return dIdStr 

136 

137 def dataIdToFilename(self, dataId: dict, includeNumber: bool = False, imNum: int | None = None) -> str: 

138 """Convert dataId to filename. 

139 

140 Returns a full path+filename by default. if includeNumber then 

141 returns just the filename for use in temporary dir for animation.""" 

142 if includeNumber: 

143 assert imNum is not None 

144 

145 dIdStr = self._strDataId(dataId) 

146 

147 if includeNumber: # for use in temp dir, so not full path 

148 filename = self.toAnimateTemplate % (imNum, dIdStr, self.dataProductToPlot) 

149 return os.path.join(filename) 

150 else: 

151 filename = self.basicTemplate % (dIdStr, self.dataProductToPlot) 

152 return os.path.join(self.pngPath, filename) 

153 

154 def exists(self, obj: Any) -> bool: 

155 if isinstance(obj, str): 

156 return os.path.exists(obj) 

157 raise RuntimeError("Other type checks not yet implemented") 

158 

159 def preRun(self) -> None: 

160 # check the paths work 

161 if not os.path.exists(self.pngPath): 

162 os.makedirs(self.pngPath) 

163 assert os.path.exists(self.pngPath), f"Failed to create output dir: {self.pngPath}" 

164 

165 if self.exists(self.outputFilename): 

166 if self.clobberVideoAndGif: 

167 os.remove(self.outputFilename) 

168 else: 

169 raise RuntimeError(f"Output file {self.outputFilename} exists and clobber==False") 

170 

171 # make list of found & missing files 

172 dIdsWithPngs = [d for d in self.dataIdList if self.exists(self.dataIdToFilename(d))] 

173 dIdsWithoutPngs = [d for d in self.dataIdList if d not in dIdsWithPngs] 

174 if self.debug: 

175 logger.info(f"dIdsWithPngs = {dIdsWithPngs}") 

176 logger.info(f"dIdsWithoutPngs = {dIdsWithoutPngs}") 

177 

178 # check the datasets exist for the pngs which need remaking 

179 missingData = [ 

180 d 

181 for d in dIdsWithoutPngs 

182 if not datasetExists(self.butler, self.dataProductToPlot, d, detector=0) 

183 ] 

184 

185 logger.info(f"Of the provided {len(self.dataIdList)} dataIds:") 

186 logger.info(f"{len(dIdsWithPngs)} existing pngs were found") 

187 logger.info(f"{len(dIdsWithoutPngs)} do not yet exist") 

188 

189 if missingData: 

190 for dId in missingData: 

191 msg = f"Failed to find {self.dataProductToPlot} for {dId}" 

192 logger.warning(msg) 

193 self.dataIdList.remove(dId) 

194 logger.info( 

195 f"Of the {len(dIdsWithoutPngs)} dataIds without pngs, {len(missingData)}" 

196 " did not have the corresponding dataset existing" 

197 ) 

198 

199 if self.remakePngs: 

200 self.pngsToMakeDataIds = [d for d in self.dataIdList if d not in missingData] 

201 else: 

202 self.pngsToMakeDataIds = [d for d in dIdsWithoutPngs if d not in missingData] 

203 

204 msg = f"So {len(self.pngsToMakeDataIds)} will be made" 

205 if self.remakePngs and len(dIdsWithPngs) > 0: 

206 msg += " because remakePngs=True" 

207 logger.info(msg) 

208 

209 def run(self) -> str | None: 

210 # make the missing pngs 

211 if self.pngsToMakeDataIds: 

212 logger.info("Creating necessary pngs...") 

213 for i, dataId in enumerate(self.pngsToMakeDataIds): 

214 logger.info(f"Making png for file {i+1} of {len(self.pngsToMakeDataIds)}") 

215 self.makePng(dataId, self.dataIdToFilename(dataId)) 

216 

217 # stage files in temp dir with numbers prepended to filenames 

218 if not self.dataIdList: 

219 logger.warning("No files to animate - nothing to do") 

220 return None 

221 

222 logger.info("Copying files to ordered temp dir...") 

223 pngFilesOriginal = [self.dataIdToFilename(d) for d in self.dataIdList] 

224 for filename in pngFilesOriginal: # these must all now exist, but let's assert just in case 

225 assert self.exists(filename) 

226 tempDir = os.path.join(self.pngPath, f"{uuid.uuid1()}/"[0:8]) 

227 os.makedirs(tempDir) 

228 pngFileList = [] # list of number-prepended files in the temp dir 

229 for i, dId in enumerate(self.dataIdList): 

230 srcFile = self.dataIdToFilename(dId) 

231 destFile = os.path.join(tempDir, self.dataIdToFilename(dId, includeNumber=True, imNum=i)) 

232 shutil.copy(srcFile, destFile) 

233 pngFileList.append(destFile) 

234 

235 # # create gif in temp dir 

236 # outputGifFilename = os.path.join(tempDir, 'animation.gif') 

237 # self.pngsToGif(pngFileList, outputGifFilename) 

238 

239 # # gif turn into mp4, optionally keep gif by moving up to output dir 

240 # logger.info('Turning gif into mp4...') 

241 # outputMp4Filename = self.outputFilename 

242 # self.gifToMp4(outputGifFilename, outputMp4Filename) 

243 

244 # self.tidyUp(tempDir) 

245 # logger.info('Finished!') 

246 

247 # create gif in temp dir 

248 

249 logger.info("Making mp4 of pngs...") 

250 self.pngsToMp4(tempDir, self.outputFilename, 10, verbose=False) 

251 self.tidyUp(tempDir) 

252 logger.info(f"Finished! Output at {self.outputFilename}") 

253 return self.outputFilename 

254 

255 def _titleFromExp(self, exp: afwImage.Exposure, dataId: dict) -> str: 

256 expRecord = getExpRecordFromDataId(self.butler, dataId) 

257 obj = expRecord.target_name 

258 expTime = expRecord.exposure_time 

259 filterCompound = expRecord.physical_filter 

260 filt, grating = filterCompound.split("~") 

261 rawMd = self.butler.get("raw.metadata", dataId) 

262 airmass = airMassFromRawMetadata(rawMd) # XXX this could be improved a lot 

263 if not airmass: 

264 airmass = -1 

265 dayObs = dayObsIntToString(getDayObs(dataId)) 

266 timestamp = expRecord.timespan.begin.to_datetime().strftime("%H:%M:%S") # no microseconds 

267 ms = expRecord.timespan.begin.to_datetime().strftime("%f") # always 6 chars long, 000000 if zero 

268 timestamp += f".{ms[0:2]}" 

269 title = f"seqNum {getSeqNum(dataId)} - {dayObs} {timestamp}TAI - " 

270 title += f"Object: {obj} expTime: {expTime}s Filter: {filt} Grating: {grating} Airmass: {airmass:.3f}" 

271 return title 

272 

273 def getStarPixCoord( 

274 self, exp: Any, doMotionCorrection: bool = True, useQfm: bool = False 

275 ) -> tuple[float, float] | None: 

276 target = exp.visitInfo.object 

277 

278 if self.useQfmForCentroids: 

279 try: 

280 result = self.qfmTask.run(exp) 

281 pixCoord = result.brightestObjCentroid 

282 expId = exp.info.id 

283 logger.info(f"expId {expId} has centroid {pixCoord}") 

284 except Exception: 

285 return None 

286 else: 

287 pixCoord = getTargetCentroidFromWcs(exp, target, doMotionCorrection=doMotionCorrection) 

288 return pixCoord 

289 

290 def makePng(self, dataId: dict, saveFilename: str) -> None: 

291 if self.exists(saveFilename) and not self.remakePngs: # should not be possible due to prerun 

292 assert False, f"Almost overwrote {saveFilename} - how is this possible?" 

293 

294 if self.debug: 

295 logger.info(f"Creating {saveFilename}") 

296 

297 self.fig.clear() 

298 

299 # must always keep exp unsmoothed for the centroiding via qfm 

300 try: 

301 exp = self.butler.get(self.dataProductToPlot, dataId) 

302 except Exception: 

303 # oh no, that should never happen, but it does! Let's just skip 

304 logger.warning(f"Skipped {dataId}, because {self.dataProductToPlot} retrieval failed!") 

305 return 

306 toDisplay = exp 

307 if self.smoothImages: 

308 toDisplay = exp.clone() 

309 toDisplay = self._smoothExp(toDisplay, 2) 

310 

311 try: 

312 self.disp.mtv(toDisplay.image, title=self._titleFromExp(exp, dataId)) 

313 self.disp.scale("asinh", "zscale") 

314 except RuntimeError: # all-nan images slip through and don't display 

315 self.disp.scale("linear", 0, 1) 

316 self.disp.mtv(toDisplay.image, title=self._titleFromExp(exp, dataId)) 

317 self.disp.scale("asinh", "zscale") # set back for next image 

318 pass 

319 

320 if self.plotObjectCentroids: 

321 try: 

322 pixCoord = self.getStarPixCoord(exp) 

323 if pixCoord: 

324 self.disp.dot("x", *pixCoord, ctype="C1", size=50) 

325 self.disp.dot("o", *pixCoord, ctype="C1", size=50) 

326 else: 

327 self.disp.dot("x", 2000, 2000, ctype="red", size=2000) 

328 except Exception: 

329 logger.warning(f"Failed to find OBJECT location for {dataId}") 

330 

331 deltaH = -0.05 

332 deltaV = -0.05 

333 plt.subplots_adjust(right=1 + deltaH, left=0 - deltaH, top=1 + deltaV, bottom=0 - deltaV) 

334 self.fig.savefig(saveFilename) 

335 logger.info(f"Saved png for {dataId} to {saveFilename}") 

336 

337 del toDisplay 

338 del exp 

339 gc.collect() 

340 

341 def pngsToMp4(self, indir: str, outfile: str, framerate: float, verbose: bool = False) -> None: 

342 """Create the movie with ffmpeg, from files.""" 

343 # NOTE: the order of ffmpeg arguments *REALLY MATTERS*. 

344 # Reorder them at your own peril! 

345 pathPattern = f'"{os.path.join(indir, "*.png")}"' 

346 if verbose: 

347 ffmpeg_verbose = "info" 

348 else: 

349 ffmpeg_verbose = "error" 

350 cmd = [ 

351 "ffmpeg", 

352 "-v", 

353 ffmpeg_verbose, 

354 "-f", 

355 "image2", 

356 "-y", 

357 "-pattern_type glob", 

358 "-framerate", 

359 f"{framerate}", 

360 "-i", 

361 pathPattern, 

362 "-vcodec", 

363 "libx264", 

364 "-b:v", 

365 "20000k", 

366 "-profile:v", 

367 "main", 

368 "-pix_fmt", 

369 "yuv420p", 

370 "-threads", 

371 "10", 

372 "-r", 

373 f"{framerate}", 

374 os.path.join(outfile), 

375 ] 

376 

377 subprocess.check_call(r" ".join(cmd), shell=True) 

378 

379 def tidyUp(self, tempDir: str) -> None: 

380 shutil.rmtree(tempDir) 

381 return 

382 

383 def _smoothExp(self, exp: afwImage.Exposure, smoothing: float, kernelSize: int = 7) -> afwImage.Exposure: 

384 """Use for DISPLAY ONLY! 

385 

386 Return a smoothed copy of the exposure 

387 with the original mask plane in place.""" 

388 psf = measAlg.DoubleGaussianPsf(kernelSize, kernelSize, smoothing / (2 * math.sqrt(2 * math.log(2)))) 

389 newExp = exp.clone() 

390 originalMask = exp.mask 

391 

392 kernel = psf.getKernel() 

393 afwMath.convolve(newExp.maskedImage, newExp.maskedImage, kernel, afwMath.ConvolutionControl()) 

394 newExp.mask = originalMask 

395 return newExp 

396 

397 

398def animateDay( 

399 butler: dafButler.Butler, dayObs: int, outputPath: str, dataProductToPlot: str = "quickLookExp" 

400) -> str | None: 

401 outputFilename = f"{dayObs}.mp4" 

402 

403 onSkyIds = getLatissOnSkyDataIds(butler, startDate=dayObs, endDate=dayObs) 

404 logger.info(f"Found {len(onSkyIds)} on sky ids for {dayObs}") 

405 

406 onSkyIds = [updateDataIdOrDataCord(dataId, detector=0) for dataId in onSkyIds] 

407 

408 animator = Animator( 

409 butler, 

410 onSkyIds, 

411 outputPath, 

412 outputFilename, 

413 dataProductToPlot=dataProductToPlot, 

414 remakePngs=False, 

415 debug=False, 

416 clobberVideoAndGif=True, 

417 plotObjectCentroids=True, 

418 useQfmForCentroids=True, 

419 ) 

420 filename = animator.run() 

421 return filename 

422 

423 

424if __name__ == "__main__": 424 ↛ 425line 424 didn't jump to line 425, because the condition on line 424 was never true

425 outputPath = "/home/mfl/animatorOutput/main/" 

426 butler = makeDefaultLatissButler() 

427 

428 day = 20211104 

429 animateDay(butler, day, outputPath)