Coverage for python / lsst / summit / extras / animation.py: 13%
232 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-06 09:16 +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/>.
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.dateTime import dayObsIntToString
49from lsst.summit.utils.utils import setupLogging
51logger = logging.getLogger("lsst.summit.extras.animation")
52setupLogging()
55class Animator:
56 """Animate a list of dataIds into an mp4 movie.
58 Iterates over the dataIds in the order supplied, renders each
59 exposure for the specified data product as a PNG (caching PNGs in
60 ``outputPath/pngs/``), stages them in a temporary numbered
61 directory, and assembles them into an mp4 with ``ffmpeg``.
63 Parameters
64 ----------
65 butler : `lsst.daf.butler.Butler`
66 Butler with access to the specified data product.
67 dataIdList : `list` [`dict`]
68 DataIds to animate, in display order.
69 outputPath : `str`
70 Directory to write the mp4 and the ``pngs/`` cache directory
71 to. Created if it does not exist.
72 outputFilename : `str`
73 Base name of the output mp4. The ``.mp4`` suffix is appended
74 if missing.
75 remakePngs : `bool`, optional
76 If `True`, regenerate PNGs even when cached copies exist.
77 clobberVideoAndGif : `bool`, optional
78 If `True`, overwrite any existing output mp4. Otherwise raise
79 if the file already exists.
80 keepIntermediateGif : `bool`, optional
81 Retained for backwards compatibility. Currently unused;
82 intermediate gif generation is disabled.
83 smoothImages : `bool`, optional
84 If `True`, smooth the displayed image with a small Gaussian
85 (for display only; centroiding still uses the raw exposure).
86 plotObjectCentroids : `bool`, optional
87 If `True`, overlay the target star centroid on each frame.
88 useQfmForCentroids : `bool`, optional
89 If `True`, use QuickFrameMeasurement for the overlay centroid
90 instead of the WCS-based target lookup.
91 dataProductToPlot : `str`, optional
92 The butler dataset type to retrieve for each dataId.
93 debug : `bool`, optional
94 If `True`, emit extra debug logging.
95 """
97 def __init__(
98 self,
99 butler: dafButler.Butler,
100 dataIdList: list[dict],
101 outputPath: str,
102 outputFilename: str,
103 *,
104 remakePngs: bool = False,
105 clobberVideoAndGif: bool = False,
106 keepIntermediateGif: bool = False,
107 smoothImages: bool = True,
108 plotObjectCentroids: bool = True,
109 useQfmForCentroids: bool = False,
110 dataProductToPlot: str = "calexp",
111 debug: bool = False,
112 ):
113 self.butler = butler
114 self.dataIdList = dataIdList
115 self.outputPath = outputPath
116 self.outputFilename = os.path.join(outputPath, outputFilename)
117 if not self.outputFilename.endswith(".mp4"):
118 self.outputFilename += ".mp4"
119 self.pngPath = os.path.join(outputPath, "pngs/")
121 self.remakePngs = remakePngs
122 self.clobberVideoAndGif = clobberVideoAndGif
123 self.keepIntermediateGif = keepIntermediateGif
124 self.smoothImages = smoothImages
125 self.plotObjectCentroids = plotObjectCentroids
126 self.useQfmForCentroids = useQfmForCentroids
127 self.dataProductToPlot = dataProductToPlot
128 self.debug = debug
130 # zfilled at the start as animation is alphabetical
131 # if you're doing more than 1e6 files you've got bigger problems
132 self.toAnimateTemplate = "%06d-%s-%s.png"
133 self.basicTemplate = "%s-%s.png"
135 qfmTaskConfig = QuickFrameMeasurementTaskConfig()
136 self.qfmTask = QuickFrameMeasurementTask(config=qfmTaskConfig)
138 afwDisplay.setDefaultBackend("matplotlib")
139 self.fig = plt.figure(figsize=(15, 15))
140 self.disp = afwDisplay.Display(self.fig)
141 self.disp.setImageColormap("gray")
142 self.disp.scale("asinh", "zscale")
144 self.pngsToMakeDataIds: list[dict] = []
146 self.preRun() # sets the above list
148 @staticmethod
149 def _strDataId(dataId: dict) -> str:
150 """Make a dataId into a string suitable for use as a filename.
152 Parameters
153 ----------
154 dataId : `dict`
155 The data id.
157 Returns
158 -------
159 strId : `str`
160 The data id as a string.
161 """
162 if (dayObs := getDayObs(dataId)) and (seqNum := getSeqNum(dataId)): # nicely ordered if easy
163 return f"{dayObsIntToString(dayObs)}-{seqNum:05d}"
165 # General case (and yeah, I should probably learn regex someday)
166 dIdStr = str(dataId)
167 dIdStr = dIdStr.replace(" ", "")
168 dIdStr = dIdStr.replace("{", "")
169 dIdStr = dIdStr.replace("}", "")
170 dIdStr = dIdStr.replace("'", "")
171 dIdStr = dIdStr.replace(":", "-")
172 dIdStr = dIdStr.replace(",", "-")
173 return dIdStr
175 def dataIdToFilename(self, dataId: dict, includeNumber: bool = False, imNum: int | None = None) -> str:
176 """Convert a dataId to a PNG filename.
178 Parameters
179 ----------
180 dataId : `dict`
181 The dataId to render.
182 includeNumber : `bool`, optional
183 If `True`, prepend a zero-padded frame number so that
184 alphabetical filename ordering matches the animation
185 order. In this mode the returned value is a bare filename
186 suitable for use inside the temporary staging directory.
187 Otherwise a full path inside ``self.pngPath`` is returned.
188 imNum : `int`, optional
189 Frame number used when ``includeNumber`` is `True`. Must
190 be provided in that case.
192 Returns
193 -------
194 filename : `str`
195 The PNG filename or full path, as described above.
196 """
197 if includeNumber:
198 assert imNum is not None
200 dIdStr = self._strDataId(dataId)
202 if includeNumber: # for use in temp dir, so not full path
203 filename = self.toAnimateTemplate % (imNum, dIdStr, self.dataProductToPlot)
204 return os.path.join(filename)
205 else:
206 filename = self.basicTemplate % (dIdStr, self.dataProductToPlot)
207 return os.path.join(self.pngPath, filename)
209 def exists(self, obj: Any) -> bool:
210 """Check whether a filesystem object exists.
212 Parameters
213 ----------
214 obj : `str`
215 The path to check.
217 Returns
218 -------
219 exists : `bool`
220 `True` if the path exists on disk.
222 Raises
223 ------
224 RuntimeError
225 Raised if ``obj`` is not a string; other existence checks
226 are not yet implemented.
227 """
228 if isinstance(obj, str):
229 return os.path.exists(obj)
230 raise RuntimeError("Other type checks not yet implemented")
232 def preRun(self) -> None:
233 """Prepare output directories and determine which PNGs to build.
235 Creates the PNG cache directory if missing, handles the output
236 mp4's pre-existence according to ``clobberVideoAndGif``, checks
237 the butler for each dataId's dataset, logs a summary, and
238 populates ``self.pngsToMakeDataIds`` with the dataIds whose
239 PNGs still need to be generated.
240 """
241 # check the paths work
242 if not os.path.exists(self.pngPath):
243 os.makedirs(self.pngPath)
244 assert os.path.exists(self.pngPath), f"Failed to create output dir: {self.pngPath}"
246 if self.exists(self.outputFilename):
247 if self.clobberVideoAndGif:
248 os.remove(self.outputFilename)
249 else:
250 raise RuntimeError(f"Output file {self.outputFilename} exists and clobber==False")
252 # make list of found & missing files
253 dIdsWithPngs = [d for d in self.dataIdList if self.exists(self.dataIdToFilename(d))]
254 dIdsWithoutPngs = [d for d in self.dataIdList if d not in dIdsWithPngs]
255 if self.debug:
256 logger.info(f"dIdsWithPngs = {dIdsWithPngs}")
257 logger.info(f"dIdsWithoutPngs = {dIdsWithoutPngs}")
259 # check the datasets exist for the pngs which need remaking
260 missingData = [
261 d for d in dIdsWithoutPngs if not self.butler.exists(self.dataProductToPlot, d, detector=0)
262 ]
264 logger.info(f"Of the provided {len(self.dataIdList)} dataIds:")
265 logger.info(f"{len(dIdsWithPngs)} existing pngs were found")
266 logger.info(f"{len(dIdsWithoutPngs)} do not yet exist")
268 if missingData:
269 for dId in missingData:
270 msg = f"Failed to find {self.dataProductToPlot} for {dId}"
271 logger.warning(msg)
272 self.dataIdList.remove(dId)
273 logger.info(
274 f"Of the {len(dIdsWithoutPngs)} dataIds without pngs, {len(missingData)}"
275 " did not have the corresponding dataset existing"
276 )
278 if self.remakePngs:
279 self.pngsToMakeDataIds = [d for d in self.dataIdList if d not in missingData]
280 else:
281 self.pngsToMakeDataIds = [d for d in dIdsWithoutPngs if d not in missingData]
283 msg = f"So {len(self.pngsToMakeDataIds)} will be made"
284 if self.remakePngs and len(dIdsWithPngs) > 0:
285 msg += " because remakePngs=True"
286 logger.info(msg)
288 def run(self) -> str | None:
289 """Render any missing PNGs and assemble the mp4 movie.
291 Returns
292 -------
293 outputFilename : `str` or `None`
294 Path to the produced mp4 file, or `None` if there were no
295 dataIds to animate.
296 """
297 # make the missing pngs
298 if self.pngsToMakeDataIds:
299 logger.info("Creating necessary pngs...")
300 for i, dataId in enumerate(self.pngsToMakeDataIds):
301 logger.info(f"Making png for file {i + 1} of {len(self.pngsToMakeDataIds)}")
302 self.makePng(dataId, self.dataIdToFilename(dataId))
304 # stage files in temp dir with numbers prepended to filenames
305 if not self.dataIdList:
306 logger.warning("No files to animate - nothing to do")
307 return None
309 logger.info("Copying files to ordered temp dir...")
310 pngFilesOriginal = [self.dataIdToFilename(d) for d in self.dataIdList]
311 for filename in pngFilesOriginal: # these must all now exist, but let's assert just in case
312 assert self.exists(filename)
313 tempDir = os.path.join(self.pngPath, uuid.uuid1().hex[:8])
314 os.makedirs(tempDir)
315 pngFileList = [] # list of number-prepended files in the temp dir
316 for i, dId in enumerate(self.dataIdList):
317 srcFile = self.dataIdToFilename(dId)
318 destFile = os.path.join(tempDir, self.dataIdToFilename(dId, includeNumber=True, imNum=i))
319 shutil.copy(srcFile, destFile)
320 pngFileList.append(destFile)
322 # # create gif in temp dir
323 # outputGifFilename = os.path.join(tempDir, 'animation.gif')
324 # self.pngsToGif(pngFileList, outputGifFilename)
326 # # gif turn into mp4, optionally keep gif by moving up to output dir
327 # logger.info('Turning gif into mp4...')
328 # outputMp4Filename = self.outputFilename
329 # self.gifToMp4(outputGifFilename, outputMp4Filename)
331 # self.tidyUp(tempDir)
332 # logger.info('Finished!')
334 # create gif in temp dir
336 logger.info("Making mp4 of pngs...")
337 self.pngsToMp4(tempDir, self.outputFilename, 10, verbose=False)
338 self.tidyUp(tempDir)
339 logger.info(f"Finished! Output at {self.outputFilename}")
340 return self.outputFilename
342 def _titleFromExp(self, exp: afwImage.Exposure, dataId: dict) -> str:
343 """Build the plot title line for a single exposure.
345 Parameters
346 ----------
347 exp : `lsst.afw.image.Exposure`
348 The exposure being plotted (unused except to mirror the
349 display path).
350 dataId : `dict`
351 The dataId for which to look up observing metadata.
353 Returns
354 -------
355 title : `str`
356 Human-readable title including seqNum, timestamp, target,
357 exposure time, filter, grating, and airmass.
358 """
359 expRecord = getExpRecordFromDataId(self.butler, dataId)
360 obj = expRecord.target_name
361 expTime = expRecord.exposure_time
362 filterCompound = expRecord.physical_filter
363 filt, grating = filterCompound.split("~")
364 rawMd = self.butler.get("raw.metadata", dataId)
365 airmass = airMassFromRawMetadata(rawMd) # XXX this could be improved a lot
366 if not airmass:
367 airmass = -1
368 dayObsInt = getDayObs(dataId)
369 assert dayObsInt is not None
370 dayObs = dayObsIntToString(dayObsInt)
371 timestamp = expRecord.timespan.begin.to_datetime().strftime("%H:%M:%S") # no microseconds
372 ms = expRecord.timespan.begin.to_datetime().strftime("%f") # always 6 chars long, 000000 if zero
373 timestamp += f".{ms[0:2]}"
374 title = f"seqNum {getSeqNum(dataId)} - {dayObs} {timestamp}TAI - "
375 title += f"Object: {obj} expTime: {expTime}s Filter: {filt} Grating: {grating} Airmass: {airmass:.3f}"
376 return title
378 def getStarPixCoord(
379 self, exp: Any, doMotionCorrection: bool = True, useQfm: bool = False
380 ) -> tuple[float, float] | None:
381 """Return the main star's pixel centroid for an exposure.
383 If ``self.useQfmForCentroids`` is set, this runs
384 QuickFrameMeasurement on the exposure; otherwise the target
385 position is derived from the WCS and the named target.
387 Parameters
388 ----------
389 exp : `lsst.afw.image.Exposure`
390 The exposure to measure.
391 doMotionCorrection : `bool`, optional
392 If `True`, apply proper-motion correction when looking up
393 the target in the catalog.
394 useQfm : `bool`, optional
395 Unused; retained for API stability. Centroid source is
396 controlled by ``self.useQfmForCentroids``.
398 Returns
399 -------
400 pixCoord : `tuple` [`float`, `float`] or `None`
401 The ``(x, y)`` pixel coordinate of the main star, or
402 `None` if the measurement failed.
403 """
404 target = exp.visitInfo.object
406 if self.useQfmForCentroids:
407 try:
408 result = self.qfmTask.run(exp)
409 pixCoord = result.brightestObjCentroid
410 expId = exp.info.id
411 logger.info(f"expId {expId} has centroid {pixCoord}")
412 except Exception:
413 return None
414 else:
415 pixCoord = getTargetCentroidFromWcs(exp, target, doMotionCorrection=doMotionCorrection)
416 return pixCoord
418 def makePng(self, dataId: dict, saveFilename: str) -> None:
419 """Render a single dataId to a PNG on disk.
421 Parameters
422 ----------
423 dataId : `dict`
424 The dataId to render.
425 saveFilename : `str`
426 Full path of the PNG to write.
427 """
428 if self.exists(saveFilename) and not self.remakePngs: # should not be possible due to prerun
429 raise RuntimeError(f"Almost overwrote {saveFilename} - how is this possible?")
431 if self.debug:
432 logger.info(f"Creating {saveFilename}")
434 self.fig.clear()
436 # must always keep exp unsmoothed for the centroiding via qfm
437 try:
438 exp = self.butler.get(self.dataProductToPlot, dataId)
439 except Exception:
440 # oh no, that should never happen, but it does! Let's just skip
441 logger.warning(f"Skipped {dataId}, because {self.dataProductToPlot} retrieval failed!")
442 return
443 toDisplay = exp
444 if self.smoothImages:
445 toDisplay = exp.clone()
446 toDisplay = self._smoothExp(toDisplay, 2)
448 try:
449 self.disp.mtv(toDisplay.image, title=self._titleFromExp(exp, dataId))
450 self.disp.scale("asinh", "zscale")
451 except RuntimeError: # all-nan images slip through and don't display
452 self.disp.scale("linear", 0, 1)
453 self.disp.mtv(toDisplay.image, title=self._titleFromExp(exp, dataId))
454 self.disp.scale("asinh", "zscale") # set back for next image
455 pass
457 if self.plotObjectCentroids:
458 try:
459 pixCoord = self.getStarPixCoord(exp)
460 if pixCoord:
461 self.disp.dot("x", *pixCoord, ctype="C1", size=50)
462 self.disp.dot("o", *pixCoord, ctype="C1", size=50)
463 else:
464 self.disp.dot("x", 2000, 2000, ctype="red", size=2000)
465 except Exception:
466 logger.warning(f"Failed to find OBJECT location for {dataId}")
468 deltaH = -0.05
469 deltaV = -0.05
470 plt.subplots_adjust(right=1 + deltaH, left=0 - deltaH, top=1 + deltaV, bottom=0 - deltaV)
471 self.fig.savefig(saveFilename)
472 logger.info(f"Saved png for {dataId} to {saveFilename}")
474 del toDisplay
475 del exp
476 gc.collect()
478 def pngsToMp4(self, indir: str, outfile: str, framerate: float, verbose: bool = False) -> None:
479 """Assemble all PNGs in ``indir`` into an mp4 file via ffmpeg.
481 Parameters
482 ----------
483 indir : `str`
484 Directory whose ``*.png`` files will be globbed and
485 animated in alphabetical order.
486 outfile : `str`
487 Path of the mp4 file to write.
488 framerate : `float`
489 Frames per second for the output movie.
490 verbose : `bool`, optional
491 If `True`, run ffmpeg at ``info`` verbosity; otherwise
492 only errors are printed.
493 """
494 # NOTE: the order of ffmpeg arguments *REALLY MATTERS*.
495 # Reorder them at your own peril!
496 pathPattern = f'"{os.path.join(indir, "*.png")}"'
497 if verbose:
498 ffmpeg_verbose = "info"
499 else:
500 ffmpeg_verbose = "error"
501 cmd = [
502 "ffmpeg",
503 "-v",
504 ffmpeg_verbose,
505 "-f",
506 "image2",
507 "-y",
508 "-pattern_type glob",
509 "-framerate",
510 f"{framerate}",
511 "-i",
512 pathPattern,
513 "-vcodec",
514 "libx264",
515 "-b:v",
516 "20000k",
517 "-profile:v",
518 "main",
519 "-pix_fmt",
520 "yuv420p",
521 "-threads",
522 "10",
523 "-r",
524 f"{framerate}",
525 os.path.join(outfile),
526 ]
528 subprocess.check_call(r" ".join(cmd), shell=True)
530 def tidyUp(self, tempDir: str) -> None:
531 """Remove the temporary staging directory.
533 Parameters
534 ----------
535 tempDir : `str`
536 Directory to remove, typically the staging directory
537 created inside ``self.pngPath``.
538 """
539 shutil.rmtree(tempDir)
540 return
542 def _smoothExp(self, exp: afwImage.Exposure, smoothing: float, kernelSize: int = 7) -> afwImage.Exposure:
543 """Return a smoothed copy of an exposure. For DISPLAY ONLY.
545 Convolves ``exp`` with a double-Gaussian kernel of the given
546 FWHM and kernel size, preserving the original mask plane on
547 the returned copy. Intended only for making on-screen images
548 legible; do not use for centroiding or measurement.
550 Parameters
551 ----------
552 exp : `lsst.afw.image.Exposure`
553 The exposure to smooth (not modified).
554 smoothing : `float`
555 FWHM of the smoothing kernel, in pixels.
556 kernelSize : `int`, optional
557 Size of the square convolution kernel, in pixels.
559 Returns
560 -------
561 newExp : `lsst.afw.image.Exposure`
562 A smoothed clone of ``exp`` with the original mask plane.
563 """
564 psf = measAlg.DoubleGaussianPsf(kernelSize, kernelSize, smoothing / (2 * math.sqrt(2 * math.log(2))))
565 newExp = exp.clone()
566 originalMask = exp.mask
568 kernel = psf.getKernel()
569 afwMath.convolve(newExp.maskedImage, newExp.maskedImage, kernel, afwMath.ConvolutionControl())
570 newExp.mask = originalMask
571 return newExp
574def animateDay(
575 butler: dafButler.Butler, dayObs: int, outputPath: str, dataProductToPlot: str = "quickLookExp"
576) -> str | None:
577 """Animate all LATISS on-sky exposures from a single dayObs.
579 Parameters
580 ----------
581 butler : `lsst.daf.butler.Butler`
582 Butler with access to the LATISS data.
583 dayObs : `int`
584 The dayObs to animate (``YYYYMMDD``).
585 outputPath : `str`
586 Directory in which the mp4 and cached PNGs will be written.
587 dataProductToPlot : `str`, optional
588 Butler dataset type to render. Defaults to ``quickLookExp``.
590 Returns
591 -------
592 filename : `str` or `None`
593 Path to the produced mp4, or `None` if no on-sky dataIds were
594 found for the day.
595 """
596 outputFilename = f"{dayObs}.mp4"
598 onSkyIds = getLatissOnSkyDataIds(butler, startDate=dayObs, endDate=dayObs) # type: ignore[arg-type]
599 logger.info(f"Found {len(onSkyIds)} on sky ids for {dayObs}")
601 onSkyIds = [updateDataIdOrDataCord(dataId, detector=0) for dataId in onSkyIds]
603 animator = Animator(
604 butler,
605 onSkyIds, # type: ignore[arg-type]
606 outputPath,
607 outputFilename,
608 dataProductToPlot=dataProductToPlot,
609 remakePngs=False,
610 debug=False,
611 clobberVideoAndGif=True,
612 plotObjectCentroids=True,
613 useQfmForCentroids=True,
614 )
615 filename = animator.run()
616 return filename
619if __name__ == "__main__": 619 ↛ 620line 619 didn't jump to line 620 because the condition on line 619 was never true
620 outputPath = "/home/mfl/animatorOutput/main/"
621 butler = makeDefaultLatissButler()
623 day = 20211104
624 animateDay(butler, day, outputPath)