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

226 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-30 12:43 +0000

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 

29 

30import matplotlib.pyplot as plt 

31 

32import lsst.afw.display as afwDisplay 

33import lsst.afw.math as afwMath 

34import lsst.meas.algorithms as measAlg 

35from lsst.atmospec.utils import airMassFromRawMetadata, getTargetCentroidFromWcs 

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

37from lsst.summit.utils.butlerUtils import ( 

38 datasetExists, 

39 getDayObs, 

40 getExpRecordFromDataId, 

41 getLatissOnSkyDataIds, 

42 getSeqNum, 

43 makeDefaultLatissButler, 

44 updateDataIdOrDataCord, 

45) 

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

47 

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

49setupLogging() 

50 

51 

52class Animator: 

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

54 for the data product specified.""" 

55 

56 def __init__( 

57 self, 

58 butler, 

59 dataIdList, 

60 outputPath, 

61 outputFilename, 

62 *, 

63 remakePngs=False, 

64 clobberVideoAndGif=False, 

65 keepIntermediateGif=False, 

66 smoothImages=True, 

67 plotObjectCentroids=True, 

68 useQfmForCentroids=False, 

69 dataProductToPlot="calexp", 

70 debug=False, 

71 ): 

72 self.butler = butler 

73 self.dataIdList = dataIdList 

74 self.outputPath = outputPath 

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

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

77 self.outputFilename += ".mp4" 

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

79 

80 self.remakePngs = remakePngs 

81 self.clobberVideoAndGif = clobberVideoAndGif 

82 self.keepIntermediateGif = keepIntermediateGif 

83 self.smoothImages = smoothImages 

84 self.plotObjectCentroids = plotObjectCentroids 

85 self.useQfmForCentroids = useQfmForCentroids 

86 self.dataProductToPlot = dataProductToPlot 

87 self.debug = debug 

88 

89 # zfilled at the start as animation is alphabetical 

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

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

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

93 

94 qfmTaskConfig = QuickFrameMeasurementTaskConfig() 

95 self.qfmTask = QuickFrameMeasurementTask(config=qfmTaskConfig) 

96 

97 afwDisplay.setDefaultBackend("matplotlib") 

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

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

100 self.disp.setImageColormap("gray") 

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

102 

103 self.pngsToMakeDataIds = [] 

104 self.preRun() # sets the above list 

105 

106 @staticmethod 

107 def _strDataId(dataId): 

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

109 

110 Parameters 

111 ---------- 

112 dataId : `dict` 

113 The data id. 

114 

115 Returns 

116 ------- 

117 strId : `str` 

118 The data id as a string. 

119 """ 

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

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

122 

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

124 dIdStr = str(dataId) 

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

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

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

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

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

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

131 return dIdStr 

132 

133 def dataIdToFilename(self, dataId, includeNumber=False, imNum=None): 

134 """Convert dataId to filename. 

135 

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

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

138 if includeNumber: 

139 assert imNum is not None 

140 

141 dIdStr = self._strDataId(dataId) 

142 

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

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

145 return os.path.join(filename) 

146 else: 

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

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

149 

150 def exists(self, obj): 

151 if isinstance(obj, str): 

152 return os.path.exists(obj) 

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

154 

155 def preRun(self): 

156 # check the paths work 

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

158 os.makedirs(self.pngPath) 

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

160 

161 if self.exists(self.outputFilename): 

162 if self.clobberVideoAndGif: 

163 os.remove(self.outputFilename) 

164 else: 

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

166 

167 # make list of found & missing files 

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

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

170 if self.debug: 

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

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

173 

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

175 missingData = [ 

176 d 

177 for d in dIdsWithoutPngs 

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

179 ] 

180 

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

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

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

184 

185 if missingData: 

186 for dId in missingData: 

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

188 logger.warning(msg) 

189 self.dataIdList.remove(dId) 

190 logger.info( 

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

192 " did not have the corresponding dataset existing" 

193 ) 

194 

195 if self.remakePngs: 

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

197 else: 

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

199 

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

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

202 msg += " because remakePngs=True" 

203 logger.info(msg) 

204 

205 def run(self): 

206 # make the missing pngs 

207 if self.pngsToMakeDataIds: 

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

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

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

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

212 

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

214 if not self.dataIdList: 

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

216 return 

217 

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

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

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

221 assert self.exists(filename) 

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

223 os.makedirs(tempDir) 

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

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

226 srcFile = self.dataIdToFilename(dId) 

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

228 shutil.copy(srcFile, destFile) 

229 pngFileList.append(destFile) 

230 

231 # # create gif in temp dir 

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

233 # self.pngsToGif(pngFileList, outputGifFilename) 

234 

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

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

237 # outputMp4Filename = self.outputFilename 

238 # self.gifToMp4(outputGifFilename, outputMp4Filename) 

239 

240 # self.tidyUp(tempDir) 

241 # logger.info('Finished!') 

242 

243 # create gif in temp dir 

244 

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

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

247 self.tidyUp(tempDir) 

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

249 return self.outputFilename 

250 

251 def _titleFromExp(self, exp, dataId): 

252 expRecord = getExpRecordFromDataId(self.butler, dataId) 

253 obj = expRecord.target_name 

254 expTime = expRecord.exposure_time 

255 filterCompound = expRecord.physical_filter 

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

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

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

259 if not airmass: 

260 airmass = -1 

261 dayObs = dayObsIntToString(getDayObs(dataId)) 

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

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

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

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

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

267 return title 

268 

269 def getStarPixCoord(self, exp, doMotionCorrection=True, useQfm=False): 

270 target = exp.visitInfo.object 

271 

272 if self.useQfmForCentroids: 

273 try: 

274 result = self.qfmTask.run(exp) 

275 pixCoord = result.brightestObjCentroid 

276 expId = exp.info.id 

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

278 except Exception: 

279 return None 

280 else: 

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

282 return pixCoord 

283 

284 def makePng(self, dataId, saveFilename): 

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

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

287 

288 if self.debug: 

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

290 

291 self.fig.clear() 

292 

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

294 try: 

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

296 except Exception: 

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

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

299 return 

300 toDisplay = exp 

301 if self.smoothImages: 

302 toDisplay = exp.clone() 

303 toDisplay = self._smoothExp(toDisplay, 2) 

304 

305 try: 

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

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

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

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

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

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

312 pass 

313 

314 if self.plotObjectCentroids: 

315 try: 

316 pixCoord = self.getStarPixCoord(exp) 

317 if pixCoord: 

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

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

320 else: 

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

322 except Exception: 

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

324 

325 deltaH = -0.05 

326 deltaV = -0.05 

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

328 self.fig.savefig(saveFilename) 

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

330 

331 del toDisplay 

332 del exp 

333 gc.collect() 

334 

335 def pngsToMp4(self, indir, outfile, framerate, verbose=False): 

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

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

338 # Reorder them at your own peril! 

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

340 if verbose: 

341 ffmpeg_verbose = "info" 

342 else: 

343 ffmpeg_verbose = "error" 

344 cmd = [ 

345 "ffmpeg", 

346 "-v", 

347 ffmpeg_verbose, 

348 "-f", 

349 "image2", 

350 "-y", 

351 "-pattern_type glob", 

352 "-framerate", 

353 f"{framerate}", 

354 "-i", 

355 pathPattern, 

356 "-vcodec", 

357 "libx264", 

358 "-b:v", 

359 "20000k", 

360 "-profile:v", 

361 "main", 

362 "-pix_fmt", 

363 "yuv420p", 

364 "-threads", 

365 "10", 

366 "-r", 

367 f"{framerate}", 

368 os.path.join(outfile), 

369 ] 

370 

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

372 

373 def tidyUp(self, tempDir): 

374 shutil.rmtree(tempDir) 

375 return 

376 

377 def _smoothExp(self, exp, smoothing, kernelSize=7): 

378 """Use for DISPLAY ONLY! 

379 

380 Return a smoothed copy of the exposure 

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

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

383 newExp = exp.clone() 

384 originalMask = exp.mask 

385 

386 kernel = psf.getKernel() 

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

388 newExp.mask = originalMask 

389 return newExp 

390 

391 

392def animateDay(butler, dayObs, outputPath, dataProductToPlot="quickLookExp"): 

393 outputFilename = f"{dayObs}.mp4" 

394 

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

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

397 

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

399 

400 animator = Animator( 

401 butler, 

402 onSkyIds, 

403 outputPath, 

404 outputFilename, 

405 dataProductToPlot=dataProductToPlot, 

406 remakePngs=False, 

407 debug=False, 

408 clobberVideoAndGif=True, 

409 plotObjectCentroids=True, 

410 useQfmForCentroids=True, 

411 ) 

412 filename = animator.run() 

413 return filename 

414 

415 

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

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

418 butler = makeDefaultLatissButler() 

419 

420 day = 20211104 

421 animateDay(butler, day, outputPath)