From 448138a18d3bd534643406666fe088fd9057f5c1 Mon Sep 17 00:00:00 2001 From: hayk Date: Mon, 15 Jul 2024 19:06:38 -0400 Subject: [PATCH] mp changed + kwargs ffmpeg added --- nt2/__init__.py | 2 +- nt2/export.py | 66 +++++++++++++++++++++++++++---------------------- nt2/read.py | 36 ++++++++++++++++++--------- 3 files changed, 63 insertions(+), 41 deletions(-) diff --git a/nt2/__init__.py b/nt2/__init__.py index 3d26edf..df12433 100644 --- a/nt2/__init__.py +++ b/nt2/__init__.py @@ -1 +1 @@ -__version__ = "0.4.1" +__version__ = "0.4.2" diff --git a/nt2/export.py b/nt2/export.py index 6d164e8..5dcd850 100644 --- a/nt2/export.py +++ b/nt2/export.py @@ -15,8 +15,8 @@ def makeMovie(**ffmpeg_kwargs): tool with the given arguments. Examples -------- - >>> makeMovie(ffmpeg="/path/to/ffmpeg", framerate="30", start="0", input="step_", number=3, - extension="png", compression="1", overwrite=True, output="anim.mp4") + >>> makeMovie(ffmpeg="/path/to/ffmpeg", framerate=30, start=0, input="step_", number=3, + extension="png", compression=1, overwrite=True, output="anim.mp4") """ import subprocess @@ -24,20 +24,20 @@ def makeMovie(**ffmpeg_kwargs): ffmpeg_kwargs.get("ffmpeg", "ffmpeg"), "-nostdin", "-framerate", - ffmpeg_kwargs.get("framerate", "30"), + str(ffmpeg_kwargs.get("framerate", 30)), "-start_number", - ffmpeg_kwargs.get("start", "0"), + str(ffmpeg_kwargs.get("start", 0)), "-i", ffmpeg_kwargs.get("input", "step_") + f"%0{ffmpeg_kwargs.get('number', 3)}d.{ffmpeg_kwargs.get('extension', 'png')}", "-c:v", "libx264", "-crf", - ffmpeg_kwargs.get("compression", "1"), + str(ffmpeg_kwargs.get("compression", 1)), "-filter_complex", "[0:v]format=yuv420p,pad=ceil(iw/2)*2:ceil(ih/2)*2", "-y" if ffmpeg_kwargs.get("overwrite", False) else None, - ffmpeg_kwargs.get("output", "anim.mp4"), + ffmpeg_kwargs.get("output", "movie.mp4"), ] command = [str(c) for c in command if c is not None] print("Command:\n", " ".join(command)) @@ -51,6 +51,23 @@ def makeMovie(**ffmpeg_kwargs): return False +class PlotWorker: + def __init__(self, plot, fpath, data=None): + self.plot = plot + self.fpath = fpath + self.data = data + + def __call__(self, ti): + import matplotlib.pyplot as plt + + if self.data is None: + self.plot(ti) + else: + self.plot(ti, self.data) + plt.savefig(f"{self.fpath}/{ti:05d}.png") + plt.close() + + def makeFrames(plot, steps, fpath, data=None, num_cpus=None): """ Create plot frames from a set of timesteps of the same dataset. @@ -84,35 +101,26 @@ def makeFrames(plot, steps, fpath, data=None, num_cpus=None): >>> makeFrames(plot_func, range(100), 'output/', num_cpus=16) """ - from tqdm import tqdm + import tqdm import multiprocessing as mp - import matplotlib.pyplot as plt import os - global plotAndSave - - def plotAndSave(ti, t, fpath): - try: - if data is None: - plot(t) - else: - plot(t, data) - plt.savefig(f"{fpath}/{ti:05d}.png") - plt.close() - return True - except Exception as e: - print(f"Error: {e}") - return False + # if fpath doesn't exist, create it + if not os.path.exists(fpath): + os.makedirs(fpath) if num_cpus is None: num_cpus = mp.cpu_count() pool = mp.Pool(num_cpus) - # if fpath doesn't exist, create it - if not os.path.exists(fpath): - os.makedirs(fpath) - - tasks = [[ti, t, fpath] for ti, t in enumerate(steps)] - results = [pool.apply_async(plotAndSave, t) for t in tasks] - return [result.get() for result in tqdm(results)] + try: + for _ in tqdm.tqdm( + pool.imap_unordered(PlotWorker(plot, fpath, data), steps), + total=len(steps), + ): + ... + return True + except Exception as e: + print("Error:", e) + return False diff --git a/nt2/read.py b/nt2/read.py index dd9a427..e44e657 100644 --- a/nt2/read.py +++ b/nt2/read.py @@ -939,7 +939,7 @@ def plotGrid(self, ax, **kwargs): def makeMovie(self, plot, makeframes=True, **kwargs): """ - Makes a movie from a plot function + Makes a movie from a plot function. Additional keyword arguments are passed to `makeMovie`. Parameters ---------- @@ -949,20 +949,34 @@ def makeMovie(self, plot, makeframes=True, **kwargs): Whether to make the frames, or just proceed to making the movie. Default is True. num_cpus : int, optional The number of CPUs to use for making the frames. Default is None. - **kwargs : - Additional keyword arguments passed to `ffmpeg`. + ffmpeg : str, optional + Path to the ffmpeg executable. Default is "ffmpeg". + framerate : int | str, optional + The framerate of the movie. Default is 30. + start : int | str, optional + The start frame. Default is 0. + input : str, optional + The input pattern. Default is "step_". + number : int, optional + Zero-padding for the input pattern. Default is 3. + extension : str, optional + The extension of the input pattern. Default is "png". + compression : int | str, optional + The compression level. Default is 1. + overwrite : bool, optional + Whether to overwrite the output file. Default is False. + output : str, optional + The output file name. Default is "movie.mp4". """ import numpy as np if makeframes: - makemovie = all( - exp.makeFrames( - plot, - np.arange(len(self.t)), - f"{self.attrs['simulation.name']}/frames", - data=self, - num_cpus=kwargs.pop("num_cpus", None), - ) + makemovie = exp.makeFrames( + plot, + np.arange(len(self.t)), + f"{self.attrs['simulation.name']}/frames", + data=self, + num_cpus=kwargs.pop("num_cpus", None), ) else: makemovie = True