Coverage for python/lsst/summit/extras/animation.py: 13%
229 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-11 05:39 -0700
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-11 05:39 -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/>.
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 datasetExists,
42 getDayObs,
43 getExpRecordFromDataId,
44 getLatissOnSkyDataIds,
45 getSeqNum,
46 makeDefaultLatissButler,
47 updateDataIdOrDataCord,
48)
49from lsst.summit.utils.utils import dayObsIntToString, 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
181 for d in dIdsWithoutPngs
182 if not datasetExists(self.butler, self.dataProductToPlot, d, detector=0)
183 ]
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")
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 )
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]
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)
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))
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
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)
235 # # create gif in temp dir
236 # outputGifFilename = os.path.join(tempDir, 'animation.gif')
237 # self.pngsToGif(pngFileList, outputGifFilename)
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)
244 # self.tidyUp(tempDir)
245 # logger.info('Finished!')
247 # create gif in temp dir
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
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
273 def getStarPixCoord(
274 self, exp: Any, doMotionCorrection: bool = True, useQfm: bool = False
275 ) -> tuple[float, float] | None:
276 target = exp.visitInfo.object
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
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?"
294 if self.debug:
295 logger.info(f"Creating {saveFilename}")
297 self.fig.clear()
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)
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
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}")
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}")
337 del toDisplay
338 del exp
339 gc.collect()
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 ]
377 subprocess.check_call(r" ".join(cmd), shell=True)
379 def tidyUp(self, tempDir: str) -> None:
380 shutil.rmtree(tempDir)
381 return
383 def _smoothExp(self, exp: afwImage.Exposure, smoothing: float, kernelSize: int = 7) -> afwImage.Exposure:
384 """Use for DISPLAY ONLY!
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
392 kernel = psf.getKernel()
393 afwMath.convolve(newExp.maskedImage, newExp.maskedImage, kernel, afwMath.ConvolutionControl())
394 newExp.mask = originalMask
395 return newExp
398def animateDay(
399 butler: dafButler.Butler, dayObs: int, outputPath: str, dataProductToPlot: str = "quickLookExp"
400) -> str | None:
401 outputFilename = f"{dayObs}.mp4"
403 onSkyIds = getLatissOnSkyDataIds(butler, startDate=dayObs, endDate=dayObs)
404 logger.info(f"Found {len(onSkyIds)} on sky ids for {dayObs}")
406 onSkyIds = [updateDataIdOrDataCord(dataId, detector=0) for dataId in onSkyIds]
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
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()
428 day = 20211104
429 animateDay(butler, day, outputPath)