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

1# This file is part of summit_extras.

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 gc

23import logging

24import math

25import os

26import shutil

27import subprocess

28import uuid

29from typing import Any

31import matplotlib.pyplot as plt

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 getDayObs,

42 getExpRecordFromDataId,

43 getLatissOnSkyDataIds,

44 getSeqNum,

45 makeDefaultLatissButler,

46 updateDataIdOrDataCord,

47)

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

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

51setupLogging()

54class Animator:

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

56 for the data product specified."""

58 def __init__(

59 self,

60 butler: dafButler.Butler,

61 dataIdList: list[dict],

62 outputPath: str,

63 outputFilename: str,

64 *,

65 remakePngs: bool = False,

66 clobberVideoAndGif: bool = False,

67 keepIntermediateGif: bool = False,

68 smoothImages: bool = True,

69 plotObjectCentroids: bool = True,

70 useQfmForCentroids: bool = False,

71 dataProductToPlot: str = "calexp",

72 debug: bool = False,

73 ):

74 self.butler = butler

75 self.dataIdList = dataIdList

76 self.outputPath = outputPath

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

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

79 self.outputFilename += ".mp4"

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

82 self.remakePngs = remakePngs

83 self.clobberVideoAndGif = clobberVideoAndGif

84 self.keepIntermediateGif = keepIntermediateGif

85 self.smoothImages = smoothImages

86 self.plotObjectCentroids = plotObjectCentroids

87 self.useQfmForCentroids = useQfmForCentroids

88 self.dataProductToPlot = dataProductToPlot

89 self.debug = debug

91 # zfilled at the start as animation is alphabetical

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

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

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

96 qfmTaskConfig = QuickFrameMeasurementTaskConfig()

97 self.qfmTask = QuickFrameMeasurementTask(config=qfmTaskConfig)

99 afwDisplay.setDefaultBackend("matplotlib")

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

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

102 self.disp.setImageColormap("gray")

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

104

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

106

107 self.preRun() # sets the above list

108

109 @staticmethod

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

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

112

113 Parameters

114 ----------

115 dataId : `dict`

116 The data id.

117

118 Returns

119 -------

120 strId : `str`

121 The data id as a string.

122 """

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

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

125

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

127 dIdStr = str(dataId)

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

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

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

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

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

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

134 return dIdStr

135

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

137 """Convert dataId to filename.

138

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

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

141 if includeNumber:

142 assert imNum is not None

143

144 dIdStr = self._strDataId(dataId)

145

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

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

148 return os.path.join(filename)

149 else:

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

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

152

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

154 if isinstance(obj, str):

155 return os.path.exists(obj)

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

157

158 def preRun(self) -> None:

159 # check the paths work

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

161 os.makedirs(self.pngPath)

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

163

164 if self.exists(self.outputFilename):

165 if self.clobberVideoAndGif:

166 os.remove(self.outputFilename)

167 else:

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

169

170 # make list of found & missing files

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

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

173 if self.debug:

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

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

176

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

178 missingData = [

179 d for d in dIdsWithoutPngs if not self.butler.exists(self.dataProductToPlot, d, detector=0)

180 ]

181

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

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

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

185

186 if missingData:

187 for dId in missingData:

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

189 logger.warning(msg)

190 self.dataIdList.remove(dId)

191 logger.info(

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

193 " did not have the corresponding dataset existing"

194 )

195

196 if self.remakePngs:

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

198 else:

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

200

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

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

203 msg += " because remakePngs=True"

204 logger.info(msg)

205

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

207 # make the missing pngs

208 if self.pngsToMakeDataIds:

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

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

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

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

213

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

215 if not self.dataIdList:

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

217 return None

218

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

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

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

222 assert self.exists(filename)

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

224 os.makedirs(tempDir)

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

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

227 srcFile = self.dataIdToFilename(dId)

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

229 shutil.copy(srcFile, destFile)

230 pngFileList.append(destFile)

231

232 # # create gif in temp dir

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

234 # self.pngsToGif(pngFileList, outputGifFilename)

235

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

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

238 # outputMp4Filename = self.outputFilename

239 # self.gifToMp4(outputGifFilename, outputMp4Filename)

240

241 # self.tidyUp(tempDir)

242 # logger.info('Finished!')

243

244 # create gif in temp dir

245

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

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

248 self.tidyUp(tempDir)

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

250 return self.outputFilename

251

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

253 expRecord = getExpRecordFromDataId(self.butler, dataId)

254 obj = expRecord.target_name

255 expTime = expRecord.exposure_time

256 filterCompound = expRecord.physical_filter

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

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

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

260 if not airmass:

261 airmass = -1

262 dayObs = dayObsIntToString(getDayObs(dataId))

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

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

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

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

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

268 return title

269

270 def getStarPixCoord(

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

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

273 target = exp.visitInfo.object

274

275 if self.useQfmForCentroids:

276 try:

277 result = self.qfmTask.run(exp)

278 pixCoord = result.brightestObjCentroid

279 expId = exp.info.id

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

281 except Exception:

282 return None

283 else:

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

285 return pixCoord

286

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

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

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

290

291 if self.debug:

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

293

294 self.fig.clear()

295

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

297 try:

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

299 except Exception:

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

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

302 return

303 toDisplay = exp

304 if self.smoothImages:

305 toDisplay = exp.clone()

306 toDisplay = self._smoothExp(toDisplay, 2)

307

308 try:

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

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

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

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

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

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

315 pass

316

317 if self.plotObjectCentroids:

318 try:

319 pixCoord = self.getStarPixCoord(exp)

320 if pixCoord:

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

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

323 else:

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

325 except Exception:

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

327

328 deltaH = -0.05

329 deltaV = -0.05

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

331 self.fig.savefig(saveFilename)

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

333

334 del toDisplay

335 del exp

336 gc.collect()

337

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

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

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

341 # Reorder them at your own peril!

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

343 if verbose:

344 ffmpeg_verbose = "info"

345 else:

346 ffmpeg_verbose = "error"

347 cmd = [

348 "ffmpeg",

349 "-v",

350 ffmpeg_verbose,

351 "-f",

352 "image2",

353 "-y",

354 "-pattern_type glob",

355 "-framerate",

356 f"{framerate}",

357 "-i",

358 pathPattern,

359 "-vcodec",

360 "libx264",

361 "-b:v",

362 "20000k",

363 "-profile:v",

364 "main",

365 "-pix_fmt",

366 "yuv420p",

367 "-threads",

368 "10",

369 "-r",

370 f"{framerate}",

371 os.path.join(outfile),

372 ]

373

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

375

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

377 shutil.rmtree(tempDir)

378 return

379

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

381 """Use for DISPLAY ONLY!

382

383 Return a smoothed copy of the exposure

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

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

386 newExp = exp.clone()

387 originalMask = exp.mask

388

389 kernel = psf.getKernel()

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

391 newExp.mask = originalMask

392 return newExp

393

394

395def animateDay(

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

397) -> str | None:

398 outputFilename = f"{dayObs}.mp4"

399

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

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

402

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

404

405 animator = Animator(

406 butler,

407 onSkyIds,

408 outputPath,

409 outputFilename,

410 dataProductToPlot=dataProductToPlot,

411 remakePngs=False,

412 debug=False,

413 clobberVideoAndGif=True,

414 plotObjectCentroids=True,

415 useQfmForCentroids=True,

416 )

417 filename = animator.run()

418 return filename

419

420

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

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

423 butler = makeDefaultLatissButler()

424

425 day = 20211104

426 animateDay(butler, day, outputPath)