From 20bfdf2f18eb5174f9c84e44192a9d8d8d501f85 Mon Sep 17 00:00:00 2001 From: Ying-Li Niu <64801511@qq.com> Date: Thu, 11 Jun 2026 13:31:33 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=96=B0=E5=A2=9E=20plot=5Fwave=5Fsave?= =?UTF-8?q?=5Fgif/plot=5Fwave=5Fsave=5Fmp4=20=E5=BC=80=E5=85=B3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 所有 input.txt 新增 plot_wave_save_gif / plot_wave_save_mp4 参数 - dynamics.py 将参数传入 plot_wave() - plot_wave.py 根据 save_gif/save_mp4 标志条件保存文件 - 默认均为 0（不输出文件，只显示动画窗口） --- dynamics.py | 6 ++- examples/case01/input/input.txt | 4 +- examples/case02/input/input.txt | 4 +- examples/case03/input/input.txt | 4 +- examples/case04/input/input.txt | 4 +- examples/case05/input/input.txt | 4 +- examples/case06/input/input.txt | 4 +- plot_wave.py | 69 +++++++++++++++++++-------------- 8 files changed, 63 insertions(+), 36 deletions(-) diff --git a/dynamics.py b/dynamics.py index 93b96fc..2f6777f 100644 --- a/dynamics.py +++ b/dynamics.py @@ -507,7 +507,11 @@ def run_case(config_path, runtime_base, input_dir="input", output_dir="output", try: import plot_wave as pw print("[run] 正在绘制波形与能量图…") - pw.plot_wave(str(output_dir_abs)) + pw.plot_wave( + str(output_dir_abs), + save_gif=int(config.get("plot_wave_save_gif", 0)), + save_mp4=int(config.get("plot_wave_save_mp4", 0)), + ) except Exception as e: print(f"[run] 绘制波形图失败: {e}") diff --git a/examples/case01/input/input.txt b/examples/case01/input/input.txt index ab27154..c94c425 100644 --- a/examples/case01/input/input.txt +++ b/examples/case01/input/input.txt @@ -8,7 +8,9 @@ step_simulate: 1 # 运行物理模拟 → output/trajectory.txt step_sample: 1 # 抽帧 → output/display.txt step_plot: 1 # 绘制轨迹/能量图 → output/trajectory_plots.png -step_plot_wave: 0 # 绘制波形能量动画 → output/wave_animation.gif +step_plot_wave: 0 # 绘制波形能量动画 +plot_wave_save_gif: 0 # 输出波形 GIF（需 step_plot_wave=1） +plot_wave_save_mp4: 0 # 输出波形 MP4（需 step_plot_wave=1） step_animation: 1 # 自动播放 VisPy 3D 动画窗口（需安装 vispy） force_calc: 0 # 强制重新计算：1=跳过缓存强算，0=自动使用已有输出 diff --git a/examples/case02/input/input.txt b/examples/case02/input/input.txt index 292d59a..67ffab8 100644 --- a/examples/case02/input/input.txt +++ b/examples/case02/input/input.txt @@ -8,7 +8,9 @@ step_simulate: 1 # 运行物理模拟 → output/trajectory.txt step_sample: 1 # 抽帧 → output/display.txt step_plot: 1 # 绘制轨迹/能量图 → output/trajectory_plots.png -step_plot_wave: 0 # 绘制波形能量动画 → output/wave_animation.gif +step_plot_wave: 0 # 绘制波形能量动画 +plot_wave_save_gif: 0 # 输出波形 GIF（需 step_plot_wave=1） +plot_wave_save_mp4: 0 # 输出波形 MP4（需 step_plot_wave=1） step_animation: 1 # 自动播放 VisPy 3D 动画窗口（需安装 vispy） force_calc: 0 # 强制重新计算：1=跳过缓存强算，0=自动使用已有输出 diff --git a/examples/case03/input/input.txt b/examples/case03/input/input.txt index 40b0a9d..6d70e56 100644 --- a/examples/case03/input/input.txt +++ b/examples/case03/input/input.txt @@ -8,7 +8,9 @@ step_simulate: 1 # 运行物理模拟 → output/trajectory.txt step_sample: 1 # 抽帧 → output/display.txt step_plot: 0 # 绘制轨迹/能量图 → output/trajectory_plots.png -step_plot_wave: 0 # 绘制波形能量动画 → output/wave_animation.gif +step_plot_wave: 0 # 绘制波形能量动画 +plot_wave_save_gif: 0 # 输出波形 GIF（需 step_plot_wave=1） +plot_wave_save_mp4: 0 # 输出波形 MP4（需 step_plot_wave=1） step_animation: 1 # 自动播放 VisPy 3D 动画窗口（需安装 vispy） force_calc: 0 # 强制重新计算：1=跳过缓存强算，0=自动使用已有输出 diff --git a/examples/case04/input/input.txt b/examples/case04/input/input.txt index 40b0a9d..6d70e56 100644 --- a/examples/case04/input/input.txt +++ b/examples/case04/input/input.txt @@ -8,7 +8,9 @@ step_simulate: 1 # 运行物理模拟 → output/trajectory.txt step_sample: 1 # 抽帧 → output/display.txt step_plot: 0 # 绘制轨迹/能量图 → output/trajectory_plots.png -step_plot_wave: 0 # 绘制波形能量动画 → output/wave_animation.gif +step_plot_wave: 0 # 绘制波形能量动画 +plot_wave_save_gif: 0 # 输出波形 GIF（需 step_plot_wave=1） +plot_wave_save_mp4: 0 # 输出波形 MP4（需 step_plot_wave=1） step_animation: 1 # 自动播放 VisPy 3D 动画窗口（需安装 vispy） force_calc: 0 # 强制重新计算：1=跳过缓存强算，0=自动使用已有输出 diff --git a/examples/case05/input/input.txt b/examples/case05/input/input.txt index 38afc3b..3e13850 100644 --- a/examples/case05/input/input.txt +++ b/examples/case05/input/input.txt @@ -9,7 +9,9 @@ step_simulate: 0 # 运行物理模拟 → output/trajectory.txt step_sample: 0 # 抽帧 → output/display.txt step_plot: 0 # 绘制轨迹/能量图 → output/trajectory_plots.png step_animation: 0 # 自动播放 VisPy 3D 动画窗口（需安装 vispy） -step_plot_wave: 1 # 绘制波形图 +step_plot_wave: 1 # 绘制波形能量动画 +plot_wave_save_gif: 0 # 输出波形 GIF（需 step_plot_wave=1） +plot_wave_save_mp4: 0 # 输出波形 MP4（需 step_plot_wave=1） force_calc: 0 # 强制重新计算：1=跳过缓存强算，0=自动使用已有输出 # ── 计算引擎 ────────────────────────────────── diff --git a/examples/case06/input/input.txt b/examples/case06/input/input.txt index c3dda4e..4f78b81 100644 --- a/examples/case06/input/input.txt +++ b/examples/case06/input/input.txt @@ -8,7 +8,9 @@ step_simulate: 1 # 运行物理模拟 → output/trajectory.txt step_sample: 1 # 抽帧 → output/display.txt step_plot: 0 # 绘制轨迹/能量图 → output/trajectory_plots.png -step_plot_wave: 0 # 绘制波形能量动画 → output/wave_animation.gif +step_plot_wave: 0 # 绘制波形能量动画 +plot_wave_save_gif: 0 # 输出波形 GIF（需 step_plot_wave=1） +plot_wave_save_mp4: 0 # 输出波形 MP4（需 step_plot_wave=1） step_animation: 1 # 自动播放 VisPy 3D 动画窗口（需安装 vispy） force_calc: 0 # 强制重新计算：1=跳过缓存强算，0=自动使用已有输出 diff --git a/plot_wave.py b/plot_wave.py index 5e19cc7..f7ee65f 100644 --- a/plot_wave.py +++ b/plot_wave.py @@ -83,8 +83,14 @@ def compute_energy(x, y, z, vx, vy, vz, masses, mass_arr, return ek_sys, us_sys, ug_sys, ugr_sys -def plot_wave(output_dir): - """主绘图函数：读取 display.txt 并生成波形+能量动画。""" +def plot_wave(output_dir, save_gif=False, save_mp4=False): + """主绘图函数：读取 display.txt 并生成波形+能量动画。 + + Args: + output_dir: 输出目录（含 display.txt） + save_gif: 是否保存 GIF + save_mp4: 是否保存 MP4 + """ data = load_disp_data(output_dir) n_frames = int(data["n_frames"]) @@ -234,36 +240,41 @@ def plot_wave(output_dir): ani = FuncAnimation(fig, update, frames=n_frames, interval=50, blit=True) - # ── 先输出 GIF（自动循环）── - gif_path = os.path.join(output_dir, "wave_animation.gif") - ani.save(gif_path, writer="pillow", fps=min(20, max(1, n_frames // 5))) - print(f"[plot_wave] GIF 已保存: {gif_path}") + # ── 输出 GIF ── + if save_gif: + gif_path = os.path.join(output_dir, "wave_animation.gif") + ani.save(gif_path, writer="pillow", fps=min(20, max(1, n_frames // 5))) + print(f"[plot_wave] GIF 已保存: {gif_path}") - # ── 再输出 MP4（需要 ffmpeg）── - try: - import matplotlib.animation as manim - import matplotlib.pyplot as _plt - - # 尝试通过 imageio_ffmpeg 定位 ffmpeg - ffmpeg_path = None + # ── 输出 MP4（需要 ffmpeg）── + gif_path = None + if save_gif: + gif_path = os.path.join(output_dir, "wave_animation.gif") + if save_mp4: try: - import imageio_ffmpeg - ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe() - except Exception: - pass - if ffmpeg_path and os.path.exists(ffmpeg_path): - _plt.rcParams['animation.ffmpeg_path'] = ffmpeg_path + import matplotlib.animation as manim + import matplotlib.pyplot as _plt - ffps = min(20, max(1, n_frames // 5)) - writer = manim.FFMpegWriter(fps=ffps, codec="libx264", - extra_args=["-pix_fmt", "yuv420p"]) - mp4_path = os.path.join(output_dir, "wave_animation.mp4") - ani.save(mp4_path, writer=writer) - print(f"[plot_wave] MP4 已保存: {mp4_path}") - except FileNotFoundError: - print("[plot_wave] 警告: 未找到 ffmpeg，跳过 MP4 输出") - except Exception as e: - print(f"[plot_wave] 警告: MP4 输出失败 ({e})，跳过") + # 尝试通过 imageio_ffmpeg 定位 ffmpeg + ffmpeg_path = None + try: + import imageio_ffmpeg + ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe() + except Exception: + pass + if ffmpeg_path and os.path.exists(ffmpeg_path): + _plt.rcParams['animation.ffmpeg_path'] = ffmpeg_path + + ffps = min(20, max(1, n_frames // 5)) + writer = manim.FFMpegWriter(fps=ffps, codec="libx264", + extra_args=["-pix_fmt", "yuv420p"]) + mp4_path = os.path.join(output_dir, "wave_animation.mp4") + ani.save(mp4_path, writer=writer) + print(f"[plot_wave] MP4 已保存: {mp4_path}") + except FileNotFoundError: + print("[plot_wave] 警告: 未找到 ffmpeg，跳过 MP4 输出") + except Exception as e: + print(f"[plot_wave] 警告: MP4 输出失败 ({e})，跳过") # ── 最后显示动画窗口 ── plt.show()