Coverage for python / lsst / summit / extras / animation.py: 13%
230 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-25 09:04 +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 the list of dataIds in the order in which they are specified
57 for the data product specified."""
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/")
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
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"
97 qfmTaskConfig = QuickFrameMeasurementTaskConfig()
98 self.qfmTask = QuickFrameMeasurementTask(config=qfmTaskConfig)
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")
106 self.pngsToMakeDataIds: list[dict] = []
108 self.preRun() # sets the above list
110 @staticmethod
111 def _strDataId(dataId: dict) -> str:
112 """Make a dataId into a string suitable for use as a filename.
114 Parameters
115 ----------
116 dataId : `dict`
117 The data id.
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}"
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
137 def dataIdToFilename(self, dataId: dict, includeNumber: bool = False, imNum: int | None = None) -> str:
138 """Convert dataId to filename.
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
145 dIdStr = self._strDataId(dataId)
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)
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")
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}"
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")
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}")
178 # check the datasets exist for the pngs which need remaking
179 missingData = [
180 d for d in dIdsWithoutPngs if not self.butler.exists(self.dataProductToPlot, d, detector=0)
181 ]
183 logger.info(f"Of the provided {len(self.dataIdList)} dataIds:")
184 logger.info(f"{len(dIdsWithPngs)} existing pngs were found")
185 logger.info(f"{len(dIdsWithoutPngs)} do not yet exist")
187 if missingData:
188 for dId in missingData:
189 msg = f"Failed to find {self.dataProductToPlot} for {dId}"
190 logger.warning(msg)
191 self.dataIdList.remove(dId)
192 logger.info(
193 f"Of the {len(dIdsWithoutPngs)} dataIds without pngs, {len(missingData)}"
194 " did not have the corresponding dataset existing"
195 )
197 if self.remakePngs:
198 self.pngsToMakeDataIds = [d for d in self.dataIdList if d not in missingData]
199 else:
200 self.pngsToMakeDataIds = [d for d in dIdsWithoutPngs if d not in missingData]
202 msg = f"So {len(self.pngsToMakeDataIds)} will be made"
203 if self.remakePngs and len(dIdsWithPngs) > 0:
204 msg += " because remakePngs=True"
205 logger.info(msg)
207 def run(self) -> str | None:
208 # make the missing pngs
209 if self.pngsToMakeDataIds:
210 logger.info("Creating necessary pngs...")
211 for i, dataId in enumerate(self.pngsToMakeDataIds):
212 logger.info(f"Making png for file {i + 1} of {len(self.pngsToMakeDataIds)}")
213 self.makePng(dataId, self.dataIdToFilename(dataId))
215 # stage files in temp dir with numbers prepended to filenames
216 if not self.dataIdList:
217 logger.warning("No files to animate - nothing to do")
218 return None
220 logger.info("Copying files to ordered temp dir...")
221 pngFilesOriginal = [self.dataIdToFilename(d) for d in self.dataIdList]
222 for filename in pngFilesOriginal: # these must all now exist, but let's assert just in case
223 assert self.exists(filename)
224 tempDir = os.path.join(self.pngPath, f"{uuid.uuid1()}/"[0:8])
225 os.makedirs(tempDir)
226 pngFileList = [] # list of number-prepended files in the temp dir
227 for i, dId in enumerate(self.dataIdList):
228 srcFile = self.dataIdToFilename(dId)
229 destFile = os.path.join(tempDir, self.dataIdToFilename(dId, includeNumber=True, imNum=i))
230 shutil.copy(srcFile, destFile)
231 pngFileList.append(destFile)
233 # # create gif in temp dir
234 # outputGifFilename = os.path.join(tempDir, 'animation.gif')
235 # self.pngsToGif(pngFileList, outputGifFilename)
237 # # gif turn into mp4, optionally keep gif by moving up to output dir
238 # logger.info('Turning gif into mp4...')
239 # outputMp4Filename = self.outputFilename
240 # self.gifToMp4(outputGifFilename, outputMp4Filename)
242 # self.tidyUp(tempDir)
243 # logger.info('Finished!')
245 # create gif in temp dir
247 logger.info("Making mp4 of pngs...")
248 self.pngsToMp4(tempDir, self.outputFilename, 10, verbose=False)
249 self.tidyUp(tempDir)
250 logger.info(f"Finished! Output at {self.outputFilename}")
251 return self.outputFilename
253 def _titleFromExp(self, exp: afwImage.Exposure, dataId: dict) -> str:
254 expRecord = getExpRecordFromDataId(self.butler, dataId)
255 obj = expRecord.target_name
256 expTime = expRecord.exposure_time
257 filterCompound = expRecord.physical_filter
258 filt, grating = filterCompound.split("~")
259 rawMd = self.butler.get("raw.metadata", dataId)
260 airmass = airMassFromRawMetadata(rawMd) # XXX this could be improved a lot
261 if not airmass:
262 airmass = -1
263 dayObs = dayObsIntToString(getDayObs(dataId))
264 timestamp = expRecord.timespan.begin.to_datetime().strftime("%H:%M:%S") # no microseconds
265 ms = expRecord.timespan.begin.to_datetime().strftime("%f") # always 6 chars long, 000000 if zero
266 timestamp += f".{ms[0:2]}"
267 title = f"seqNum {getSeqNum(dataId)} - {dayObs} {timestamp}TAI - "
268 title += f"Object: {obj} expTime: {expTime}s Filter: {filt} Grating: {grating} Airmass: {airmass:.3f}"
269 return title
271 def getStarPixCoord(
272 self, exp: Any, doMotionCorrection: bool = True, useQfm: bool = False
273 ) -> tuple[float, float] | None:
274 target = exp.visitInfo.object
276 if self.useQfmForCentroids:
277 try:
278 result = self.qfmTask.run(exp)
279 pixCoord = result.brightestObjCentroid
280 expId = exp.info.id
281 logger.info(f"expId {expId} has centroid {pixCoord}")
282 except Exception:
283 return None
284 else:
285 pixCoord = getTargetCentroidFromWcs(exp, target, doMotionCorrection=doMotionCorrection)
286 return pixCoord
288 def makePng(self, dataId: dict, saveFilename: str) -> None:
289 if self.exists(saveFilename) and not self.remakePngs: # should not be possible due to prerun
290 assert False, f"Almost overwrote {saveFilename} - how is this possible?"
292 if self.debug:
293 logger.info(f"Creating {saveFilename}")
295 self.fig.clear()
297 # must always keep exp unsmoothed for the centroiding via qfm
298 try:
299 exp = self.butler.get(self.dataProductToPlot, dataId)
300 except Exception:
301 # oh no, that should never happen, but it does! Let's just skip
302 logger.warning(f"Skipped {dataId}, because {self.dataProductToPlot} retrieval failed!")
303 return
304 toDisplay = exp
305 if self.smoothImages:
306 toDisplay = exp.clone()
307 toDisplay = self._smoothExp(toDisplay, 2)
309 try:
310 self.disp.mtv(toDisplay.image, title=self._titleFromExp(exp, dataId))
311 self.disp.scale("asinh", "zscale")
312 except RuntimeError: # all-nan images slip through and don't display
313 self.disp.scale("linear", 0, 1)
314 self.disp.mtv(toDisplay.image, title=self._titleFromExp(exp, dataId))
315 self.disp.scale("asinh", "zscale") # set back for next image
316 pass
318 if self.plotObjectCentroids:
319 try:
320 pixCoord = self.getStarPixCoord(exp)
321 if pixCoord:
322 self.disp.dot("x", *pixCoord, ctype="C1", size=50)
323 self.disp.dot("o", *pixCoord, ctype="C1", size=50)
324 else:
325 self.disp.dot("x", 2000, 2000, ctype="red", size=2000)
326 except Exception:
327 logger.warning(f"Failed to find OBJECT location for {dataId}")
329 deltaH = -0.05
330 deltaV = -0.05
331 plt.subplots_adjust(right=1 + deltaH, left=0 - deltaH, top=1 + deltaV, bottom=0 - deltaV)
332 self.fig.savefig(saveFilename)
333 logger.info(f"Saved png for {dataId} to {saveFilename}")
335 del toDisplay
336 del exp
337 gc.collect()
339 def pngsToMp4(self, indir: str, outfile: str, framerate: float, verbose: bool = False) -> None:
340 """Create the movie with ffmpeg, from files."""
341 # NOTE: the order of ffmpeg arguments *REALLY MATTERS*.
342 # Reorder them at your own peril!
343 pathPattern = f'"{os.path.join(indir, "*.png")}"'
344 if verbose:
345 ffmpeg_verbose = "info"
346 else:
347 ffmpeg_verbose = "error"
348 cmd = [
349 "ffmpeg",
350 "-v",
351 ffmpeg_verbose,
352 "-f",
353 "image2",
354 "-y",
355 "-pattern_type glob",
356 "-framerate",
357 f"{framerate}",
358 "-i",
359 pathPattern,
360 "-vcodec",
361 "libx264",
362 "-b:v",
363 "20000k",
364 "-profile:v",
365 "main",
366 "-pix_fmt",
367 "yuv420p",
368 "-threads",
369 "10",
370 "-r",
371 f"{framerate}",
372 os.path.join(outfile),
373 ]
375 subprocess.check_call(r" ".join(cmd), shell=True)
377 def tidyUp(self, tempDir: str) -> None:
378 shutil.rmtree(tempDir)
379 return
381 def _smoothExp(self, exp: afwImage.Exposure, smoothing: float, kernelSize: int = 7) -> afwImage.Exposure:
382 """Use for DISPLAY ONLY!
384 Return a smoothed copy of the exposure
385 with the original mask plane in place."""
386 psf = measAlg.DoubleGaussianPsf(kernelSize, kernelSize, smoothing / (2 * math.sqrt(2 * math.log(2))))
387 newExp = exp.clone()
388 originalMask = exp.mask
390 kernel = psf.getKernel()
391 afwMath.convolve(newExp.maskedImage, newExp.maskedImage, kernel, afwMath.ConvolutionControl())
392 newExp.mask = originalMask
393 return newExp
396def animateDay(
397 butler: dafButler.Butler, dayObs: int, outputPath: str, dataProductToPlot: str = "quickLookExp"
398) -> str | None:
399 outputFilename = f"{dayObs}.mp4"
401 onSkyIds = getLatissOnSkyDataIds(butler, startDate=dayObs, endDate=dayObs)
402 logger.info(f"Found {len(onSkyIds)} on sky ids for {dayObs}")
404 onSkyIds = [updateDataIdOrDataCord(dataId, detector=0) for dataId in onSkyIds]
406 animator = Animator(
407 butler,
408 onSkyIds,
409 outputPath,
410 outputFilename,
411 dataProductToPlot=dataProductToPlot,
412 remakePngs=False,
413 debug=False,
414 clobberVideoAndGif=True,
415 plotObjectCentroids=True,
416 useQfmForCentroids=True,
417 )
418 filename = animator.run()
419 return filename
422if __name__ == "__main__": 422 ↛ 423line 422 didn't jump to line 423 because the condition on line 422 was never true
423 outputPath = "/home/mfl/animatorOutput/main/"
424 butler = makeDefaultLatissButler()
426 day = 20211104
427 animateDay(butler, day, outputPath)