diff --git a/examples/LinearMHD/itpa_tae_benchmark/pproc_TAE_benchmark.py b/examples/LinearMHD/itpa_tae_benchmark/pproc_TAE_benchmark.py index 509af7ce2..8cfae9a4d 100644 --- a/examples/LinearMHD/itpa_tae_benchmark/pproc_TAE_benchmark.py +++ b/examples/LinearMHD/itpa_tae_benchmark/pproc_TAE_benchmark.py @@ -10,6 +10,12 @@ 10% fails with an ``AssertionError`` after the figures are saved (not with ``--test-run``). Figures are written into ``/results``. The energy and harmonics plots use plasma-plots (``pip install "struphy[pproc]"``). + +Unless ``--no-frames``, the frames of three videos of the radial velocity u_r over the run (at most ``--max-frames`` +saved time steps) are written into ``/results/frames/``: ``u_r_rz`` (the poloidal plane), ``u_r_3d_cut`` +(3D, the camera on the poloidal cross-section of the cutaway) and ``u_r_3d`` (3D, the whole torus with the +cutaway), see ``struphy.post_processing.frames``. Make a video e.g. with +``ffmpeg -framerate 25 -i frame_%04d.png -pix_fmt yuv420p u_r_rz.mp4`` in a frames folder. """ import argparse @@ -21,6 +27,7 @@ from matplotlib import pyplot as plt from struphy import Output +from struphy.post_processing.frames import write_videos from struphy.post_processing.plot_3d import plot_field_3d from struphy.post_processing.slices import plot_physical_slices @@ -109,7 +116,7 @@ def finish(figs, results_dir, save_figs): plt.close(fig) -def main(path_out=DEFAULT_OUTPUT, save_figs=True): +def main(path_out=DEFAULT_OUTPUT, save_figs=True, frames=True, max_frames=500): run = Output(path_out) snapshots = run.standard_snapshots() # VTK files (for the 3D views) of the standard snapshots only @@ -179,6 +186,24 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): else: print(f"No 3D view of {species}/{variable}: its basis {basis!r} is not in the VTK files.") + if frames: + # frames of videos of the radial velocity u_r (the first component in the normalized basis, radial for the + # circular flux surfaces of the HollowTorus): the poloidal plane and two 3D views, at most max_frames saved + # time steps, in /results/frames/u_r_rz, u_r_3d_cut and u_r_3d (frame_*.png) + videos = write_videos( + run, + "mhd/velocity", + os.path.join(results_dir, "frames"), + component=0, + basis="norm", + prefix="u_r", + label="radial velocity u_r", + max_frames=max_frames, + ) + for video, files in videos.items(): + if files: # (empty on the ranks other than 0) + print(f"{len(files)} frames of {video} in {os.path.join(results_dir, 'frames', video)}") + if not save_figs: plt.show() @@ -204,5 +229,14 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): help="Simulation output folder (default: the output folder of the run flags, sim_data beside this script)", ) pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_parser.add_argument("--no-frames", action="store_true", help="Do not write the frames of the videos.") + pproc_parser.add_argument( + "--max-frames", type=int, default=500, help="Largest number of frames per video (default: 500)." + ) pproc_args = pproc_parser.parse_args(other_flags) - main(pproc_args.path_out, save_figs=not pproc_args.show) + main( + pproc_args.path_out, + save_figs=not pproc_args.show, + frames=not pproc_args.no_frames, + max_frames=pproc_args.max_frames, + ) diff --git a/examples/LinearMHD/lowmode_tae/pproc_lowmode_tae.py b/examples/LinearMHD/lowmode_tae/pproc_lowmode_tae.py index f8d2a789e..b0704d72c 100644 --- a/examples/LinearMHD/lowmode_tae/pproc_lowmode_tae.py +++ b/examples/LinearMHD/lowmode_tae/pproc_lowmode_tae.py @@ -9,6 +9,12 @@ and saved; a deviation of more than 10% fails with an ``AssertionError`` after the figures are saved (not with ``--test-run``). Figures are written into ``/results``. The energy and harmonics plots use plasma-plots (``pip install "struphy[pproc]"``). + +Unless ``--no-frames``, the frames of three videos of the radial velocity u_r over the run (at most ``--max-frames`` +saved time steps) are written into ``/results/frames/``: ``u_r_rz`` (the poloidal plane), ``u_r_3d_cut`` +(3D, the camera on the poloidal cross-section of the cutaway) and ``u_r_3d`` (3D, the whole torus with the +cutaway), see ``struphy.post_processing.frames``. Make a video e.g. with +``ffmpeg -framerate 25 -i frame_%04d.png -pix_fmt yuv420p u_r_rz.mp4`` in a frames folder. """ import argparse @@ -20,6 +26,7 @@ from matplotlib import pyplot as plt from struphy import Output +from struphy.post_processing.frames import write_videos from struphy.post_processing.plot_3d import plot_field_3d from struphy.post_processing.slices import plot_physical_slices @@ -105,7 +112,7 @@ def finish(figs, results_dir, save_figs): plt.close(fig) -def main(path_out=DEFAULT_OUTPUT, save_figs=True): +def main(path_out=DEFAULT_OUTPUT, save_figs=True, frames=True, max_frames=500): run = Output(path_out) snapshots = run.standard_snapshots() # VTK files (for the 3D views) of the standard snapshots only @@ -171,6 +178,24 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): else: print(f"No 3D view of {species}/{variable}: its basis {basis!r} is not in the VTK files.") + if frames: + # frames of videos of the radial velocity u_r (the first component in the normalized basis, radial for the + # circular flux surfaces of the HollowTorus): the poloidal plane and two 3D views, at most max_frames saved + # time steps, in /results/frames/u_r_rz, u_r_3d_cut and u_r_3d (frame_*.png) + videos = write_videos( + run, + "mhd/velocity", + os.path.join(results_dir, "frames"), + component=0, + basis="norm", + prefix="u_r", + label="radial velocity u_r", + max_frames=max_frames, + ) + for video, files in videos.items(): + if files: # (empty on the ranks other than 0) + print(f"{len(files)} frames of {video} in {os.path.join(results_dir, 'frames', video)}") + if not save_figs: plt.show() @@ -192,5 +217,14 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): help="Simulation output folder (default: the output folder of the run flags, sim_data beside this script)", ) pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_parser.add_argument("--no-frames", action="store_true", help="Do not write the frames of the videos.") + pproc_parser.add_argument( + "--max-frames", type=int, default=500, help="Largest number of frames per video (default: 500)." + ) pproc_args = pproc_parser.parse_args(other_flags) - main(pproc_args.path_out, save_figs=not pproc_args.show) + main( + pproc_args.path_out, + save_figs=not pproc_args.show, + frames=not pproc_args.no_frames, + max_frames=pproc_args.max_frames, + ) diff --git a/examples/LinearMHDDriftkineticCC/itpa_tae_benchmark/pproc_itpa_tae.py b/examples/LinearMHDDriftkineticCC/itpa_tae_benchmark/pproc_itpa_tae.py index 3e2d920db..0845204a8 100644 --- a/examples/LinearMHDDriftkineticCC/itpa_tae_benchmark/pproc_itpa_tae.py +++ b/examples/LinearMHDDriftkineticCC/itpa_tae_benchmark/pproc_itpa_tae.py @@ -11,6 +11,12 @@ folder of the run with ``--filter none``), the MHD kinetic energy and its part outside the toroidal harmonic n = 6 of both runs are overlaid in ``filter_comparison.png``. Figures are written into ``/results``. The energy and harmonics plots use plasma-plots (``pip install "struphy[pproc]"``). + +Unless ``--no-frames``, the frames of three videos of the radial velocity u_r over the run (at most ``--max-frames`` +saved time steps) are written into ``/results/frames/``: ``u_r_rz`` (the poloidal plane), ``u_r_3d_cut`` +(3D, the camera on the poloidal cross-section of the cutaway) and ``u_r_3d`` (3D, the whole torus with the +cutaway), see ``struphy.post_processing.frames``. Make a video e.g. with +``ffmpeg -framerate 25 -i frame_%04d.png -pix_fmt yuv420p u_r_rz.mp4`` in a frames folder. """ import argparse @@ -22,6 +28,7 @@ from matplotlib import pyplot as plt from struphy import Output +from struphy.post_processing.frames import write_videos from struphy.post_processing.plot_3d import plot_field_3d from struphy.post_processing.slices import plot_physical_slices @@ -270,7 +277,7 @@ def finish(figs, results_dir, save_figs): plt.close(fig) -def main(path_out=None, save_figs=True, compare=None, fit_window=FIT_WINDOW): +def main(path_out=None, save_figs=True, compare=None, fit_window=FIT_WINDOW, frames=True, max_frames=500): path_out = DEFAULT_OUTPUT if path_out is None else path_out run = Output(path_out) snapshots = run.standard_snapshots() @@ -361,6 +368,24 @@ def main(path_out=None, save_figs=True, compare=None, fit_window=FIT_WINDOW): else: print(f"No 3D view of {species}/{variable}: its basis {basis!r} is not in the VTK files.") + if frames: + # frames of videos of the radial velocity u_r (the first component in the normalized basis, radial for the + # circular flux surfaces of the HollowTorus): the poloidal plane and two 3D views, at most max_frames saved + # time steps, in /results/frames/u_r_rz, u_r_3d_cut and u_r_3d (frame_*.png) + videos = write_videos( + run, + "mhd/velocity", + os.path.join(results_dir, "frames"), + component=0, + basis="norm", + prefix="u_r", + label="radial velocity u_r", + max_frames=max_frames, + ) + for video, files in videos.items(): + if files: # (empty on the ranks other than 0) + print(f"{len(files)} frames of {video} in {os.path.join(results_dir, 'frames', video)}") + if not save_figs: plt.show() @@ -403,10 +428,16 @@ def cli(): help="Time window of the frequency and growth rate fits (default: %(default)s).", ) pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_parser.add_argument("--no-frames", action="store_true", help="Do not write the frames of the videos.") + pproc_parser.add_argument( + "--max-frames", type=int, default=500, help="Largest number of frames per video (default: 500)." + ) pproc_args = pproc_parser.parse_args(other_flags) main( pproc_args.path_out, save_figs=not pproc_args.show, compare=pproc_args.compare, fit_window=tuple(pproc_args.fit_window), + frames=not pproc_args.no_frames, + max_frames=pproc_args.max_frames, ) diff --git a/examples/LinearMHDDriftkineticCC/lowmode_tae/pproc_lowmode_tae.py b/examples/LinearMHDDriftkineticCC/lowmode_tae/pproc_lowmode_tae.py index 628355ceb..280715519 100644 --- a/examples/LinearMHDDriftkineticCC/lowmode_tae/pproc_lowmode_tae.py +++ b/examples/LinearMHDDriftkineticCC/lowmode_tae/pproc_lowmode_tae.py @@ -10,6 +10,12 @@ With ``--save-results`` (passed by ``submit_lowmode_tae_cpu_vs_gpu.py`` after each profiling run, ``--setup profiling``) the scope-profiler HTML report and a summary of the cost per step are written into ``/results`` instead; the profiling packaging uploads them with the run. + +Unless ``--no-frames``, the frames of three videos of the radial velocity u_r over the run (at most ``--max-frames`` +saved time steps) are written into ``/results/frames/``: ``u_r_rz`` (the poloidal plane), ``u_r_3d_cut`` +(3D, the camera on the poloidal cross-section of the cutaway) and ``u_r_3d`` (3D, the whole torus with the +cutaway), see ``struphy.post_processing.frames``. Make a video e.g. with +``ffmpeg -framerate 25 -i frame_%04d.png -pix_fmt yuv420p u_r_rz.mp4`` in a frames folder. """ import argparse @@ -23,6 +29,7 @@ from matplotlib import pyplot as plt from struphy import Output +from struphy.post_processing.frames import write_videos from struphy.post_processing.plot_3d import plot_field_3d from struphy.post_processing.slices import plot_physical_slices @@ -192,7 +199,7 @@ def save_profiling_results(path_out): results.print_summary(include=["^prop: ", "^kernel: "], sort="total", title="propagators and kernels") -def main(path_out=DEFAULT_OUTPUT, save_figs=True): +def main(path_out=DEFAULT_OUTPUT, save_figs=True, frames=True, max_frames=500): run = Output(path_out) snapshots = run.standard_snapshots() # VTK files (for the 3D views) of the standard snapshots only @@ -266,6 +273,24 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): else: print(f"No 3D view of {species}/{variable}: its basis {basis!r} is not in the VTK files.") + if frames: + # frames of videos of the radial velocity u_r (the first component in the normalized basis, radial for the + # circular flux surfaces of the HollowTorus): the poloidal plane and two 3D views, at most max_frames saved + # time steps, in /results/frames/u_r_rz, u_r_3d_cut and u_r_3d (frame_*.png) + videos = write_videos( + run, + "mhd/velocity", + os.path.join(results_dir, "frames"), + component=0, + basis="norm", + prefix="u_r", + label="radial velocity u_r", + max_frames=max_frames, + ) + for video, files in videos.items(): + if files: # (empty on the ranks other than 0) + print(f"{len(files)} frames of {video} in {os.path.join(results_dir, 'frames', video)}") + if not save_figs: plt.show() @@ -288,6 +313,10 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): help="Simulation output folder (default: the output folder of the run flags, sim_data beside this script)", ) pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_parser.add_argument("--no-frames", action="store_true", help="Do not write the frames of the videos.") + pproc_parser.add_argument( + "--max-frames", type=int, default=500, help="Largest number of frames per video (default: 500)." + ) pproc_parser.add_argument( "--save-results", action="store_true", @@ -299,4 +328,9 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): if pproc_args.save_results: save_profiling_results(pproc_args.path_out) else: - main(pproc_args.path_out, save_figs=not pproc_args.show) + main( + pproc_args.path_out, + save_figs=not pproc_args.show, + frames=not pproc_args.no_frames, + max_frames=pproc_args.max_frames, + ) diff --git a/examples/PerturbativeMHDWithQDriftkineticCC/lowmode_tae/pproc_lowmode_tae.py b/examples/PerturbativeMHDWithQDriftkineticCC/lowmode_tae/pproc_lowmode_tae.py index 8ac85e7b4..e84ee7df4 100644 --- a/examples/PerturbativeMHDWithQDriftkineticCC/lowmode_tae/pproc_lowmode_tae.py +++ b/examples/PerturbativeMHDWithQDriftkineticCC/lowmode_tae/pproc_lowmode_tae.py @@ -10,6 +10,12 @@ With ``--save-results`` (after a profiling run, ``--setup profiling``, as in the example of LinearMHDDriftkineticCC) the scope-profiler HTML report and a summary of the cost per step are written into ``/results`` instead; the profiling packaging uploads them with the run. + +Unless ``--no-frames``, the frames of three videos of the radial velocity u_r over the run (at most ``--max-frames`` +saved time steps) are written into ``/results/frames/``: ``u_r_rz`` (the poloidal plane), ``u_r_3d_cut`` +(3D, the camera on the poloidal cross-section of the cutaway) and ``u_r_3d`` (3D, the whole torus with the +cutaway), see ``struphy.post_processing.frames``. Make a video e.g. with +``ffmpeg -framerate 25 -i frame_%04d.png -pix_fmt yuv420p u_r_rz.mp4`` in a frames folder. """ import argparse @@ -23,6 +29,7 @@ from matplotlib import pyplot as plt from struphy import Output +from struphy.post_processing.frames import write_videos from struphy.post_processing.plot_3d import plot_field_3d from struphy.post_processing.slices import plot_physical_slices @@ -190,7 +197,7 @@ def save_profiling_results(path_out): results.print_summary(include=["^prop: ", "^kernel: "], sort="total", title="propagators and kernels") -def main(path_out=DEFAULT_OUTPUT, save_figs=True): +def main(path_out=DEFAULT_OUTPUT, save_figs=True, frames=True, max_frames=500): run = Output(path_out) snapshots = run.standard_snapshots() # VTK files (for the 3D views) of the standard snapshots only @@ -264,6 +271,24 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): else: print(f"No 3D view of {species}/{variable}: its basis {basis!r} is not in the VTK files.") + if frames: + # frames of videos of the radial velocity u_r (the first component in the normalized basis, radial for the + # circular flux surfaces of the HollowTorus): the poloidal plane and two 3D views, at most max_frames saved + # time steps, in /results/frames/u_r_rz, u_r_3d_cut and u_r_3d (frame_*.png) + videos = write_videos( + run, + "mhd/velocity", + os.path.join(results_dir, "frames"), + component=0, + basis="norm", + prefix="u_r", + label="radial velocity u_r", + max_frames=max_frames, + ) + for video, files in videos.items(): + if files: # (empty on the ranks other than 0) + print(f"{len(files)} frames of {video} in {os.path.join(results_dir, 'frames', video)}") + if not save_figs: plt.show() @@ -286,6 +311,10 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): help="Simulation output folder (default: the output folder of the run flags, sim_data beside this script)", ) pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_parser.add_argument("--no-frames", action="store_true", help="Do not write the frames of the videos.") + pproc_parser.add_argument( + "--max-frames", type=int, default=500, help="Largest number of frames per video (default: 500)." + ) pproc_parser.add_argument( "--save-results", action="store_true", @@ -297,4 +326,9 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): if pproc_args.save_results: save_profiling_results(pproc_args.path_out) else: - main(pproc_args.path_out, save_figs=not pproc_args.show) + main( + pproc_args.path_out, + save_figs=not pproc_args.show, + frames=not pproc_args.no_frames, + max_frames=pproc_args.max_frames, + ) diff --git a/examples/ShearAlfven/itpa_tae_benchmark/pproc_ShearAlfven.py b/examples/ShearAlfven/itpa_tae_benchmark/pproc_ShearAlfven.py index ec9db1c61..7272c4afc 100644 --- a/examples/ShearAlfven/itpa_tae_benchmark/pproc_ShearAlfven.py +++ b/examples/ShearAlfven/itpa_tae_benchmark/pproc_ShearAlfven.py @@ -10,6 +10,12 @@ 10% fails with an ``AssertionError`` after the figures are saved (not with ``--test-run``). Figures are written into ``/results``. The energy and harmonics plots use plasma-plots (``pip install "struphy[pproc]"``). + +Unless ``--no-frames``, the frames of three videos of the radial velocity u_r over the run (at most ``--max-frames`` +saved time steps) are written into ``/results/frames/``: ``u_r_rz`` (the poloidal plane), ``u_r_3d_cut`` +(3D, the camera on the poloidal cross-section of the cutaway) and ``u_r_3d`` (3D, the whole torus with the +cutaway), see ``struphy.post_processing.frames``. Make a video e.g. with +``ffmpeg -framerate 25 -i frame_%04d.png -pix_fmt yuv420p u_r_rz.mp4`` in a frames folder. """ import argparse @@ -21,6 +27,7 @@ from matplotlib import pyplot as plt from struphy import Output +from struphy.post_processing.frames import write_videos from struphy.post_processing.plot_3d import plot_field_3d from struphy.post_processing.slices import plot_physical_slices @@ -107,7 +114,7 @@ def finish(figs, results_dir, save_figs): plt.close(fig) -def main(path_out=DEFAULT_OUTPUT, save_figs=True): +def main(path_out=DEFAULT_OUTPUT, save_figs=True, frames=True, max_frames=500): run = Output(path_out) snapshots = run.standard_snapshots() # VTK files (for the 3D views) of the standard snapshots only @@ -177,6 +184,24 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): else: print(f"No 3D view of {species}/{variable}: its basis {basis!r} is not in the VTK files.") + if frames: + # frames of videos of the radial velocity u_r (the first component in the normalized basis, radial for the + # circular flux surfaces of the HollowTorus): the poloidal plane and two 3D views, at most max_frames saved + # time steps, in /results/frames/u_r_rz, u_r_3d_cut and u_r_3d (frame_*.png) + videos = write_videos( + run, + "mhd/velocity", + os.path.join(results_dir, "frames"), + component=0, + basis="norm", + prefix="u_r", + label="radial velocity u_r", + max_frames=max_frames, + ) + for video, files in videos.items(): + if files: # (empty on the ranks other than 0) + print(f"{len(files)} frames of {video} in {os.path.join(results_dir, 'frames', video)}") + if not save_figs: plt.show() @@ -202,5 +227,14 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): help="Simulation output folder (default: the output folder of the run flags, sim_data beside this script)", ) pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_parser.add_argument("--no-frames", action="store_true", help="Do not write the frames of the videos.") + pproc_parser.add_argument( + "--max-frames", type=int, default=500, help="Largest number of frames per video (default: 500)." + ) pproc_args = pproc_parser.parse_args(other_flags) - main(pproc_args.path_out, save_figs=not pproc_args.show) + main( + pproc_args.path_out, + save_figs=not pproc_args.show, + frames=not pproc_args.no_frames, + max_frames=pproc_args.max_frames, + ) diff --git a/examples/ShearAlfven/lowmode_tae/pproc_ShearAlfven.py b/examples/ShearAlfven/lowmode_tae/pproc_ShearAlfven.py index f907510bb..748277b86 100644 --- a/examples/ShearAlfven/lowmode_tae/pproc_ShearAlfven.py +++ b/examples/ShearAlfven/lowmode_tae/pproc_ShearAlfven.py @@ -9,6 +9,12 @@ and saved; a deviation of more than 10% fails with an ``AssertionError`` after the figures are saved (not with ``--test-run``). Figures are written into ``/results``. The energy and harmonics plots use plasma-plots (``pip install "struphy[pproc]"``). + +Unless ``--no-frames``, the frames of three videos of the radial velocity u_r over the run (at most ``--max-frames`` +saved time steps) are written into ``/results/frames/``: ``u_r_rz`` (the poloidal plane), ``u_r_3d_cut`` +(3D, the camera on the poloidal cross-section of the cutaway) and ``u_r_3d`` (3D, the whole torus with the +cutaway), see ``struphy.post_processing.frames``. Make a video e.g. with +``ffmpeg -framerate 25 -i frame_%04d.png -pix_fmt yuv420p u_r_rz.mp4`` in a frames folder. """ import argparse @@ -20,6 +26,7 @@ from matplotlib import pyplot as plt from struphy import Output +from struphy.post_processing.frames import write_videos from struphy.post_processing.plot_3d import plot_field_3d from struphy.post_processing.slices import plot_physical_slices @@ -103,7 +110,7 @@ def finish(figs, results_dir, save_figs): plt.close(fig) -def main(path_out=DEFAULT_OUTPUT, save_figs=True): +def main(path_out=DEFAULT_OUTPUT, save_figs=True, frames=True, max_frames=500): run = Output(path_out) snapshots = run.standard_snapshots() # VTK files (for the 3D views) of the standard snapshots only @@ -169,6 +176,24 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): else: print(f"No 3D view of {species}/{variable}: its basis {basis!r} is not in the VTK files.") + if frames: + # frames of videos of the radial velocity u_r (the first component in the normalized basis, radial for the + # circular flux surfaces of the HollowTorus): the poloidal plane and two 3D views, at most max_frames saved + # time steps, in /results/frames/u_r_rz, u_r_3d_cut and u_r_3d (frame_*.png) + videos = write_videos( + run, + "mhd/velocity", + os.path.join(results_dir, "frames"), + component=0, + basis="norm", + prefix="u_r", + label="radial velocity u_r", + max_frames=max_frames, + ) + for video, files in videos.items(): + if files: # (empty on the ranks other than 0) + print(f"{len(files)} frames of {video} in {os.path.join(results_dir, 'frames', video)}") + if not save_figs: plt.show() @@ -190,5 +215,14 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): help="Simulation output folder (default: the output folder of the run flags, sim_data beside this script)", ) pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_parser.add_argument("--no-frames", action="store_true", help="Do not write the frames of the videos.") + pproc_parser.add_argument( + "--max-frames", type=int, default=500, help="Largest number of frames per video (default: 500)." + ) pproc_args = pproc_parser.parse_args(other_flags) - main(pproc_args.path_out, save_figs=not pproc_args.show) + main( + pproc_args.path_out, + save_figs=not pproc_args.show, + frames=not pproc_args.no_frames, + max_frames=pproc_args.max_frames, + ) diff --git a/examples/VariationalLinearMHD/lowmode_tae/pproc_lowmode_tae.py b/examples/VariationalLinearMHD/lowmode_tae/pproc_lowmode_tae.py index 5b8b0e03b..9340c677f 100644 --- a/examples/VariationalLinearMHD/lowmode_tae/pproc_lowmode_tae.py +++ b/examples/VariationalLinearMHD/lowmode_tae/pproc_lowmode_tae.py @@ -9,6 +9,12 @@ and saved; a deviation of more than 10% fails with an ``AssertionError`` after the figures are saved (not with ``--test-run``). Figures are written into ``/results``. The energy and harmonics plots use plasma-plots (``pip install "struphy[pproc]"``). + +Unless ``--no-frames``, the frames of three videos of the radial velocity u_r over the run (at most ``--max-frames`` +saved time steps) are written into ``/results/frames/``: ``u_r_rz`` (the poloidal plane), ``u_r_3d_cut`` +(3D, the camera on the poloidal cross-section of the cutaway) and ``u_r_3d`` (3D, the whole torus with the +cutaway), see ``struphy.post_processing.frames``. Make a video e.g. with +``ffmpeg -framerate 25 -i frame_%04d.png -pix_fmt yuv420p u_r_rz.mp4`` in a frames folder. """ import argparse @@ -20,6 +26,7 @@ from matplotlib import pyplot as plt from struphy import Output +from struphy.post_processing.frames import write_videos from struphy.post_processing.plot_3d import plot_field_3d from struphy.post_processing.slices import plot_physical_slices @@ -105,7 +112,7 @@ def finish(figs, results_dir, save_figs): plt.close(fig) -def main(path_out=DEFAULT_OUTPUT, save_figs=True): +def main(path_out=DEFAULT_OUTPUT, save_figs=True, frames=True, max_frames=500): run = Output(path_out) snapshots = run.standard_snapshots() # VTK files (for the 3D views) of the standard snapshots only @@ -171,6 +178,24 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): filename=filename, ) + if frames: + # frames of videos of the radial velocity u_r (the first component in the normalized basis, radial for the + # circular flux surfaces of the HollowTorus): the poloidal plane and two 3D views, at most max_frames saved + # time steps, in /results/frames/u_r_rz, u_r_3d_cut and u_r_3d (frame_*.png) + videos = write_videos( + run, + "mhd/velocity", + os.path.join(results_dir, "frames"), + component=0, + basis="norm", + prefix="u_r", + label="radial velocity u_r", + max_frames=max_frames, + ) + for video, files in videos.items(): + if files: # (empty on the ranks other than 0) + print(f"{len(files)} frames of {video} in {os.path.join(results_dir, 'frames', video)}") + if not save_figs: plt.show() @@ -192,5 +217,14 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): help="Simulation output folder (default: the output folder of the run flags, sim_data beside this script)", ) pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_parser.add_argument("--no-frames", action="store_true", help="Do not write the frames of the videos.") + pproc_parser.add_argument( + "--max-frames", type=int, default=500, help="Largest number of frames per video (default: 500)." + ) pproc_args = pproc_parser.parse_args(other_flags) - main(pproc_args.path_out, save_figs=not pproc_args.show) + main( + pproc_args.path_out, + save_figs=not pproc_args.show, + frames=not pproc_args.no_frames, + max_frames=pproc_args.max_frames, + ) diff --git a/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/pproc_lowmode_tae.py b/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/pproc_lowmode_tae.py index e8439a2b8..1690a1d34 100644 --- a/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/pproc_lowmode_tae.py +++ b/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/pproc_lowmode_tae.py @@ -9,6 +9,12 @@ and saved; a deviation of more than 10% fails with an ``AssertionError`` after the figures are saved (not with ``--test-run``). Figures are written into ``/results``. The energy and harmonics plots use plasma-plots (``pip install "struphy[pproc]"``). + +Unless ``--no-frames``, the frames of three videos of the radial velocity u_r over the run (at most ``--max-frames`` +saved time steps) are written into ``/results/frames/``: ``u_r_rz`` (the poloidal plane), ``u_r_3d_cut`` +(3D, the camera on the poloidal cross-section of the cutaway) and ``u_r_3d`` (3D, the whole torus with the +cutaway), see ``struphy.post_processing.frames``. Make a video e.g. with +``ffmpeg -framerate 25 -i frame_%04d.png -pix_fmt yuv420p u_r_rz.mp4`` in a frames folder. """ import argparse @@ -20,6 +26,7 @@ from matplotlib import pyplot as plt from struphy import Output +from struphy.post_processing.frames import write_videos from struphy.post_processing.plot_3d import plot_field_3d from struphy.post_processing.slices import plot_physical_slices @@ -105,7 +112,7 @@ def finish(figs, results_dir, save_figs): plt.close(fig) -def main(path_out=DEFAULT_OUTPUT, save_figs=True): +def main(path_out=DEFAULT_OUTPUT, save_figs=True, frames=True, max_frames=500): run = Output(path_out) snapshots = run.standard_snapshots() # VTK files (for the 3D views) of the standard snapshots only @@ -171,6 +178,24 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): filename=filename, ) + if frames: + # frames of videos of the radial velocity u_r (the first component in the normalized basis, radial for the + # circular flux surfaces of the HollowTorus): the poloidal plane and two 3D views, at most max_frames saved + # time steps, in /results/frames/u_r_rz, u_r_3d_cut and u_r_3d (frame_*.png) + videos = write_videos( + run, + "mhd/velocity", + os.path.join(results_dir, "frames"), + component=0, + basis="norm", + prefix="u_r", + label="radial velocity u_r", + max_frames=max_frames, + ) + for video, files in videos.items(): + if files: # (empty on the ranks other than 0) + print(f"{len(files)} frames of {video} in {os.path.join(results_dir, 'frames', video)}") + if not save_figs: plt.show() @@ -192,5 +217,14 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True): help="Simulation output folder (default: the output folder of the run flags, sim_data beside this script)", ) pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_parser.add_argument("--no-frames", action="store_true", help="Do not write the frames of the videos.") + pproc_parser.add_argument( + "--max-frames", type=int, default=500, help="Largest number of frames per video (default: 500)." + ) pproc_args = pproc_parser.parse_args(other_flags) - main(pproc_args.path_out, save_figs=not pproc_args.show) + main( + pproc_args.path_out, + save_figs=not pproc_args.show, + frames=not pproc_args.no_frames, + max_frames=pproc_args.max_frames, + ) diff --git a/src/struphy/agent_guides/skills/postprocess-simulation/SKILL.md b/src/struphy/agent_guides/skills/postprocess-simulation/SKILL.md index e10ea4980..97c5c024e 100644 --- a/src/struphy/agent_guides/skills/postprocess-simulation/SKILL.md +++ b/src/struphy/agent_guides/skills/postprocess-simulation/SKILL.md @@ -92,6 +92,32 @@ With `pip install "struphy[pproc]"` (package plasma-plots), creating an `Output` To compare runs: `Output.compare(out_a, out_b, "em_fields/phi")` returns an `xarray.Dataset`. +## Videos of fields + +`struphy.post_processing.frames` writes the frames of a video of a raw FEEC field (one PNG per saved time step, +at most `max_frames`, evenly spaced; the same color scale, view and image size in all frames), evaluated from the +saved spline coefficients (no `pproc` needed): + +```python +from struphy.post_processing.frames import frames_to_video, write_3d_frames, write_slice_frames, write_videos + +# radial component (normalized basis) in the poloidal plane of a torus or cylinder +write_slice_frames(out, "mhd/velocity", "frames/u_r_rz", component=0, label="u_r") +# 3D with a cutaway; view="cut" (camera on the cut face) or "domain" (whole domain) +write_3d_frames(out, "mhd/velocity", "frames/u_r_3d", component=0, view="domain", label="u_r") +frames_to_video("frames/u_r_rz") # frames/u_r_rz.mp4, needs ffmpeg +# or all standard videos at once: frames/u_r_rz (or u_r_xy, u_r_xz, u_r_yz for a slab), frames/u_r_3d_cut, frames/u_r_3d +write_videos(out, "mhd/velocity", "frames", component=0, prefix="u_r", label="u_r") +``` + +Options: `component=None` for the magnitude of a vector field, `basis` (`"norm"` default, `"1"`, `"2"`, `"v"`, +`"physical"`), `plane` (`"rz"` for tori and cylinders, `"xy"`, `"xz"`, `"yz"`; for other domains than Cuboid, +Colella, Orthogonal, cylinders and tori the logical mid-surfaces `"eta3"`, `"eta2"`; `"rz"` for other tori such +as Tokamak), `clim` (None: one +scale over all frames, `"frame"`, or `(vmin, vmax)`), `steps`, `max_frames` (default 500). The titles show the time +as `t = /`. The frames go to `frame_0000.png`, ... with `frames.txt` (step and time per frame); they are +written on rank 0 only. + ## Common pitfalls - Don't re-run or re-import the simulation to get results; open the output folder. diff --git a/src/struphy/post_processing/frames.py b/src/struphy/post_processing/frames.py new file mode 100644 index 000000000..80b807e2c --- /dev/null +++ b/src/struphy/post_processing/frames.py @@ -0,0 +1,749 @@ +"""Frames of videos of FEEC fields: one PNG image per saved time step, for a video made e.g. with ffmpeg (see +:func:`frames_to_video`). + +* :func:`write_slice_frames`: the field in a physical plane (e.g. the poloidal plane ``rz`` of a torus, see + :func:`struphy.post_processing.slices.physical_planes`), with matplotlib; +* :func:`write_3d_frames`: the field in the 3D domain with a cutaway (see + :func:`struphy.post_processing.plot_3d.plot_field_3d`), with PyVista; +* :func:`write_videos`: the standard videos of a field (its planes and both 3D views). + +The planes are those of :func:`struphy.post_processing.slices.physical_planes`: ``rz`` for tori and cylinders, +``xy``, ``xz``, ``yz`` for a slab (a Cuboid), Colella or Orthogonal, the logical mid-surfaces ``eta3``, ``eta2`` for +other domains; domains covering only part of the poloidal angle have none (:func:`write_videos` then writes the 3D +views only). The 3D cutaway works for any domain: it removes a wedge around the vertical axis through the center of +the domain (the axis of a torus). + +What is shown is the same for both: a scalar field, one ``component`` of a vector field in a ``basis`` (``"norm"`` +by default: the components along the normalized logical basis vectors, e.g. the radial and poloidal components in a +torus with orthogonal coordinates; ``"1"``, ``"2"``, ``"v"`` or ``"physical"`` for the Cartesian components), or the +magnitude of a vector field. The fields are evaluated from the saved spline coefficients (no post-processing of the +run is needed). All frames of a video share the colors (one color scale, fixed over the video by default), the view +and the image size, and are written as ``frame_0000.png``, ``frame_0001.png``, ... into one folder, at most +``max_frames`` of them (evenly spaced saved time steps, see :func:`frame_steps`), together with ``frames.txt``, which +lists the time step and the time of each frame. + +Example (a pproc script):: + + from struphy import Output + from struphy.post_processing.frames import frames_to_video, write_3d_frames, write_slice_frames + + run = Output("sim_data") + write_slice_frames(run, "mhd/velocity", "frames/u_r_rz", component=0, label="u_r") + write_3d_frames(run, "mhd/velocity", "frames/u_r_3d", component=0, label="u_r") + frames_to_video("frames/u_r_rz") # frames/u_r_rz.mp4, needs ffmpeg + # or all standard videos at once: frames/u_r_rz, frames/u_r_3d_cut, frames/u_r_3d + write_videos(run, "mhd/velocity", "frames", component=0, prefix="u_r", label="u_r") + +The frames are written by the process of rank 0 only (the other ranks of an MPI run return an empty list). +""" + +import shutil +import subprocess +import warnings +from pathlib import Path + +import cunumpy as xp +import h5py +import numpy as np +import pyvista as pv +from matplotlib import pyplot as plt +from matplotlib.cm import ScalarMappable +from matplotlib.colors import Normalize +from matplotlib.image import imsave +from maybempi import MPI +from scipy.spatial import cKDTree + +from struphy.post_processing.output import Output +from struphy.post_processing.plot_3d import ( + TORUS_DOMAINS, + add_field_mesh, + crop_box, + cut_center, + cut_wedge, + default_colormap, + merge_seams, + new_plotter, + off_grid_azimuth, + peak_azimuth, + repeat_around_axis, + set_camera, + wedge_distance, +) +from struphy.post_processing.slices import BASES, evaluate_plane, physical_planes, rotate + +VECTOR_FORMS = ("1", "2", "v") +# default resolution of write_3d_frames: at most this many points of the 3D grid (with the copies of a torus sector) +MAX_POINTS_3D = 1e6 +# the values of the first pass of write_3d_frames are kept for the second one up to this size in bytes +FIRST_PASS_CACHE = 200e6 + + +def saved_times(out: Output) -> np.ndarray: + """The times of all saved time steps of the run (from the raw output, no post-processing needed), in the time + units of ``out`` (see :meth:`Output.with_time_units`).""" + with h5py.File(out.path_out / "data" / "data_proc0.hdf5", "r") as file: + return np.asarray(file["time/value"][:], dtype=float).reshape(-1) * out.time_scale + + +def frame_steps(out: Output, max_frames=500) -> list[int]: + """Positions of at most ``max_frames`` saved time steps, evenly spaced over the run, with the first and (for + ``max_frames`` > 1) the last one: all saved time steps if there are not more than ``max_frames``.""" + if int(max_frames) < 1: + raise ValueError(f"max_frames must be at least 1, got {max_frames}.") + count = len(saved_times(out)) + if count <= max_frames: + return list(range(count)) + return sorted({int(n) for n in np.round(np.linspace(0, count - 1, max_frames))}) + + +def _steps_and_times(out: Output, steps, max_frames): + """The positions of the saved time steps of the frames (``steps``, else :func:`frame_steps`) and their times.""" + steps = frame_steps(out, max_frames) if steps is None else [int(n) for n in steps] + if not steps: + raise ValueError("No time steps for the frames.") + return steps, saved_times(out)[steps] + + +def _time_title(out: Output, t) -> str: + """``t = / [unit]``, for the orientation in a video.""" + unit = out.time_unit + return f"t = {t:.4g}/{saved_times(out)[-1]:.4g}" + (f" {unit}" if unit else "") + + +def _finite(values): + """``values`` (eta1 along the first axis) as floats, with the non-finite values at eta1 = 0 (the singular axis + of polar domains, e.g. of a Tokamak or a HollowTorus with a1 = 0) replaced by those of the next ring of points + and the remaining non-finite values by NaN (not shown, ignored by the color scale).""" + values = np.array(values, dtype=float) + if values.ndim > 0 and values.shape[0] > 1: + values[0] = np.where(np.isfinite(values[0]), values[0], values[1]) + values[~np.isfinite(values)] = np.nan + return values + + +def _evaluate_uncached(out: Output, n, evaluate): + """``evaluate()`` without keeping the spline functions of the saved time step ``n`` in the cache of ``out`` + (unless they were there before): the cache would hold all fields of up to ``max_frames`` time steps.""" + cache = getattr(out, "_spline_snapshots", None) + key = n % len(saved_times(out)) + cached = cache is not None and key in cache + try: + return evaluate() + finally: + if cache is not None and not cached: + cache.pop(key, None) + + +def field_label(out: Output, name, component=None, basis=None) -> str: + """Default label of what is shown of the field ``name``: e.g. ``velocity_1 (normalized basis)`` or + ``|velocity|``.""" + variable = name.split("/")[-1] + if out.field_form(name) not in VECTOR_FORMS: + return variable + if component is None: + return f"|{variable}|" + basis = basis or "norm" + return f"{variable}_{component + 1} ({BASES.get(basis, basis)})" + + +def _kind(out: Output, name, component, basis): + """``(kind, components)`` of :func:`struphy.post_processing.slices.evaluate_plane` for what is shown.""" + form = out.field_form(name) + if form not in VECTOR_FORMS: + return None, False + if component is None: + return form, False + return basis or "norm", True + + +def _limits(clim, values, vmin, vmax, cmap): + """Colormap and color limits of a frame with ``values``: the fixed ``clim`` ``(vmin, vmax)``, else from the range + of ``values`` (``clim="frame"``) or the range ``vmin``, ``vmax`` of all frames (``clim`` None), see + :func:`struphy.post_processing.plot_3d.default_colormap`.""" + if clim is not None and not isinstance(clim, str): + return cmap or default_colormap(*clim)[0], (float(clim[0]), float(clim[1])) + if clim == "frame": + vmin, vmax = _range(values) + default_cmap, (lo, hi) = default_colormap(vmin, vmax) + if lo == hi: + # a field that is zero (or constant) everywhere + delta = abs(lo) or 1.0 + lo, hi = lo - delta, hi + delta + return cmap or default_cmap, (lo, hi) + + +def _range(values): + """Smallest and largest value, ignoring NaN; (0, 0) if all values are NaN.""" + if np.all(np.isnan(values)): + return 0.0, 0.0 + return float(np.nanmin(values)), float(np.nanmax(values)) + + +def _check_clim(clim): + if clim is None or (isinstance(clim, str) and clim == "frame"): + return + if isinstance(clim, str) or len(clim) != 2 or not clim[0] < clim[1]: + raise ValueError(f'clim must be None, "frame" or (vmin, vmax) with vmin < vmax, got {clim!r}.') + + +def _even(n): + """``n`` rounded down to an even number (video encoders need even image sizes).""" + return n - n % 2 + + +def _figure(width, height, dpi, **kwargs): + """A figure of ``width`` x ``height`` pixels (half a pixel more, against the truncation of the size of the + canvas; see :func:`_save_figure`).""" + return plt.figure(figsize=((width + 0.5) / dpi, (height + 0.5) / dpi), dpi=dpi, facecolor="white", **kwargs) + + +def _save_figure(fig, path, width, height): + """Save ``fig`` (of :func:`_figure`) as an image of exactly ``width`` x ``height`` pixels and close it.""" + fig.canvas.draw() + image = np.asarray(fig.canvas.buffer_rgba())[:height, :width, :3] + plt.close(fig) + imsave(path, image) + + +def _prepare_folder(folder) -> Path: + """``folder``, created, without the frames of a previous call.""" + folder = Path(folder) + folder.mkdir(parents=True, exist_ok=True) + for old in folder.glob("frame_*.png"): + old.unlink() + return folder + + +def _write_index(folder, steps, times): + with open(folder / "frames.txt", "w") as file: + file.write("# frame saved_step time\n") + for k, (n, t) in enumerate(zip(steps, times)): + file.write(f"{k:5d} {n:10d} {t:.10g}\n") + + +def _frame_name(k, count): + return f"frame_{k:0{max(4, len(str(count - 1)))}d}.png" + + +def write_slice_frames( + out: Output, + name, + folder, + *, + component=None, + basis=None, + plane="rz", + phi=np.pi, + steps=None, + max_frames=500, + clim=None, + cmap=None, + label=None, + points_per_cell=4, + height=6.0, + dpi=150, +) -> list[Path]: + """Write the frames of a video of the FEEC field ``name`` in a physical plane. + + Parameters + ---------- + out : Output + Output of the run. + name : str + Raw FEEC field, ``"species/variable"`` (e.g. ``"mhd/velocity"``). + folder : str | Path + Folder of the frames (created; frames of a previous call are removed). + component : int | None + Component (0, 1 or 2) of a vector field in ``basis``; None for the magnitude of a vector field (and for + scalar fields). + basis : str | None + Basis of ``component``: ``"norm"`` (default), ``"1"``, ``"2"``, ``"v"`` or ``"physical"`` (Cartesian). + plane : str + The plane, one of :func:`struphy.post_processing.slices.physical_planes` (``"eta1"``, ``"eta2"``, ``"eta3"`` + for the domains without physical planes; else ``"xy"``, ``"xz"``, ``"yz"``, or + ``"rz"`` for cylinders and tori). + phi : float + Angle of the poloidal plane ``rz``. + steps : list[int] | None + Positions of the saved time steps; by default :func:`frame_steps` with ``max_frames``. + max_frames : int + Largest number of frames. + clim : None | "frame" | tuple[float, float] + Color limits: None for one scale over all frames (symmetric about zero when the field changes sign), + ``"frame"`` for a scale per frame, or fixed ``(vmin, vmax)``. + cmap : str | None + Matplotlib colormap (default: zero light gray, see + :func:`struphy.post_processing.plot_3d.default_colormap`). + label : str | None + Name of what is shown, in the titles (default: :func:`field_label`). + points_per_cell : int + Evaluation points per cell of the simulation grid. + height : float + Height of the frames in inches (the width follows from the shape of the plane). + dpi : int + Resolution of the frames. + + Returns + ------- + list[Path] + The frames, in order. + """ + _check_clim(clim) + if MPI.COMM_WORLD.Get_rank() != 0: + return [] + planes = physical_planes(out.domain, phi, all_planes=True) + if plane not in planes: + raise ValueError(f"No plane {plane!r} for {type(out.domain).__name__}; available: {list(planes)}.") + plane_def = planes[plane] + steps, times = _steps_and_times(out, steps, max_frames) + kind, components = _kind(out, name, component, basis) + label = label or field_label(out, name, component, basis) + n_pts = points_per_cell * max(out.grid.num_elements) + 1 + + values = [] + for n in steps: + num, _, coords = _evaluate_uncached( + out, n, lambda: evaluate_plane(out, out.domain, name, n, None, plane_def, n_pts, kind, components) + ) + values.append(_finite(num[component] if components else num)) + vmin = min(_range(v)[0] for v in values) + vmax = max(_range(v)[1] for v in values) + + axis = ("eta1", "eta2", "eta3").index(plane_def.cut) + h, v = coords[plane_def.hlabel], coords[plane_def.vlabel] + aspect = (np.ptp(v) or 1.0) / (np.ptp(h) or 1.0) + # the plane, plus room for the color bar and the labels; even numbers of pixels + width_px = _even(int(dpi * (height / max(aspect, 0.2) + 1.6))) + height_px = _even(int(dpi * height)) + folder = _prepare_folder(folder) + files = [] + for k, (data, t) in enumerate(zip(values, times)): + frame_cmap, (lo, hi) = _limits(clim, data, vmin, vmax, cmap) + fig = _figure(width_px, height_px, dpi, layout="constrained") + ax = fig.add_subplot() + for i in range(data.shape[axis]): + pcm = ax.pcolormesh( + np.take(h, i, axis=axis), + np.take(v, i, axis=axis), + np.take(data, i, axis=axis), + shading="gouraud", + cmap=frame_cmap, + vmin=lo, + vmax=hi, + ) + fig.colorbar(pcm, ax=ax) + ax.set_aspect("equal") + ax.set_xlabel(plane_def.hlabel) + ax.set_ylabel(plane_def.vlabel) + ax.set_title(f"{label}, {plane_def.hlabel}-{plane_def.vlabel} plane ({plane_def.where}), {_time_title(out, t)}") + files.append(folder / _frame_name(k, len(steps))) + _save_figure(fig, files[-1], width_px, height_px) + _write_index(folder, steps, times) + return files + + +def _grid_values(out: Output, name, n, etas, component, basis): + """The values shown of the field ``name`` at the saved time step ``n`` on the tensor-product grid ``etas``: + a dict of point arrays (shape of the grid; the three Cartesian components ``c1``, ``c2``, ``c3`` for the + ``"physical"`` basis, so that copies of a torus sector can rotate them, else ``values``) and the Cartesian + coordinates X, Y, Z of the grid.""" + form = out.field_form(name) + vector = form in VECTOR_FORMS + physical = vector and (component is None or basis == "physical") + representation = form if physical else (basis or "norm") if vector else None + data = _evaluate_uncached(out, n, lambda: out.evaluate(name, t=n, representation=representation, **etas)) + xyz = tuple(data.coords[c].values for c in ("X", "Y", "Z")) + if not vector: + return {"values": _finite(data.values)}, xyz + if not physical: + return {"values": _finite(data.isel(component=component).values)}, xyz + grids = [data.coords[e].values for e in ("eta1", "eta2", "eta3")] + cartesian = out.domain.push([data.isel(component=k).values for k in range(3)], *grids, kind=form) + cartesian = [_finite(xp.to_numpy(c)) for c in cartesian] + if component is None: + return {"values": np.linalg.norm(np.stack(cartesian), axis=0)}, xyz + return {f"c{k + 1}": c for k, c in enumerate(cartesian)}, xyz + + +def _shown(arrays, component): + """The shown values from the arrays of :func:`_grid_values`.""" + return arrays["values"] if "values" in arrays else arrays[f"c{component + 1}"] + + +def _render_3d(grid, values, center, azimuth, toroidal, cmap, clim, options, surface=None, surface_values=None): + """Image (RGB array) of ``grid`` (with merged seams) colored by ``values``, cut at ``azimuth``, without title and + color bar; ``options`` are ``view``, ``window_size``, ``context`` and ``opacity`` of :func:`write_3d_frames`. The + faces of the cut are opaque, the rest of the boundary of the domain has the ``opacity``; ``surface`` (a surface + inside the domain, with merged seams) colored by ``surface_values`` is cut the same way and shown opaque.""" + scalars = "values" + grid.point_data[scalars] = values + grid.set_active_scalars(scalars) + shown = cut_wedge(grid, center, azimuth) + plotter = new_plotter(False, options["window_size"]) + opacity = options["opacity"] + if opacity < 1.0: + # translucent boundary: the faces of the cut (wedge distance zero) opaque, the rest with the opacity + plotter.enable_depth_peeling(number_of_peels=8, occlusion_ratio=0.0) + boundary = shown.extract_surface(algorithm="dataset_surface") + diagonal = np.linalg.norm(np.subtract(grid.bounds[1::2], grid.bounds[::2])) + # the faces of the cut: all their points on the faces of the wedge (the clip carries the wedge distance) + on_plane = np.abs(boundary.point_data["wedge"]) < 1e-6 * diagonal + on_cut = np.zeros(boundary.n_cells, dtype=bool) + if np.any(on_plane): + faces = boundary.extract_points(on_plane, adjacent_cells=False) + on_cut[np.asarray(faces.cell_data["vtkOriginalCellIds"])] = True + if np.any(on_cut): + add_field_mesh(plotter, boundary.extract_cells(on_cut), scalars, cmap, clim) + if np.any(~on_cut): + add_field_mesh( + plotter, boundary.extract_cells(~on_cut), scalars, cmap, clim, opacity=opacity, culling="back" + ) + else: + add_field_mesh(plotter, shown, scalars, cmap, clim) + if surface is not None: + surface.point_data[scalars] = surface_values + surface.set_active_scalars(scalars) + add_field_mesh(plotter, cut_wedge(surface, center, azimuth), scalars, cmap, clim) + set_camera( + plotter, grid, shown, center, azimuth, toroidal, cutaway=options["view"] == "cut", context=options["context"] + ) + image = plotter.screenshot(return_img=True) + plotter.close() + plotter.deep_clean() + return image + + +def _compose(image, title, cmap, clim, label, show_colorbar, dpi=100): + """Figure of ``image`` (pixel for pixel) with ``title`` above it and, with ``show_colorbar``, the color bar of + ``cmap`` and ``clim`` on its right, and its width and height in pixels (even numbers).""" + height, width = image.shape[:2] + title_px = int(0.5 * dpi) + bar_px = int(1.4 * dpi) if show_colorbar else 0 + fig_w, fig_h = width + bar_px, height + title_px + fig_w, fig_h = fig_w + fig_w % 2, fig_h + fig_h % 2 + fig = _figure(fig_w, fig_h, dpi) + ax = fig.add_axes((0.0, 0.0, width / fig_w, height / fig_h)) + ax.imshow(image, interpolation="none") + ax.set_axis_off() + fig.text(0.5 * width / fig_w, 1.0 - 0.5 * title_px / fig_h, title, ha="center", va="center", fontsize=13) + if show_colorbar: + cax = fig.add_axes( + (width / fig_w + 0.1 * dpi / fig_w, 0.1 * height / fig_h, 0.2 * dpi / fig_w, 0.8 * height / fig_h) + ) + fig.colorbar(ScalarMappable(norm=Normalize(*clim), cmap=cmap), cax=cax, label=label) + return fig, fig_w, fig_h + + +def write_3d_frames( + out: Output, + name, + folder, + *, + component=None, + basis=None, + steps=None, + max_frames=500, + clim=None, + cmap=None, + label=None, + azimuth=None, + points_per_cell=None, + view="cut", + opacity=0.35, + surface="peak", + window_size=(1200, 1200), + margin=10, + show_scalar_bar=True, + context=1.2, +) -> list[Path]: + """Write the frames of a video of the FEEC field ``name`` in 3D, with a cutaway: a wedge of 90 degrees around the + vertical axis is removed (for a torus, one face of the cut is a poloidal cross-section), as in + :func:`struphy.post_processing.plot_3d.plot_field_3d`. The view is the same in all frames. + + Parameters + ---------- + out : Output + Output of the run. + name : str + Raw FEEC field, ``"species/variable"`` (e.g. ``"mhd/velocity"``). + folder : str | Path + Folder of the frames (created; frames of a previous call are removed). + component : int | None + Component (0, 1 or 2) of a vector field in ``basis``; None for the magnitude of a vector field (and for + scalar fields). + basis : str | None + Basis of ``component``: ``"norm"`` (default), ``"1"``, ``"2"``, ``"v"`` or ``"physical"`` (Cartesian). + steps : list[int] | None + Positions of the saved time steps; by default :func:`frame_steps` with ``max_frames``. + max_frames : int + Largest number of frames. + clim : None | "frame" | tuple[float, float] + Color limits: None for one scale over all frames (symmetric about zero when the field changes sign), + ``"frame"`` for a scale per frame, or fixed ``(vmin, vmax)``. + cmap : str | None + Matplotlib colormap (default: zero light gray, see + :func:`struphy.post_processing.plot_3d.default_colormap`). + label : str | None + Name of what is shown, in the titles (default: :func:`field_label`). + azimuth : float | None + Azimuth (angle about the vertical axis, in radians) of the face of the cut that the camera looks at; by + default that of the largest absolute value over all frames. It is moved to the middle between two planes of + the evaluation grid (see :func:`struphy.post_processing.plot_3d.off_grid_azimuth`). + points_per_cell : int | None + Evaluation points per cell of the simulation grid in each direction; by default as many as possible (at + most 4) with at most :data:`MAX_POINTS_3D` points in the 3D grid (of the full torus). + view : str + ``"cut"``: the camera frames the cut (for a torus its face at ``azimuth``, a poloidal cross-section) and some + of the domain around it, see ``context``; ``"domain"``: the whole domain. + opacity : float + Opacity of the boundary of the domain (without the faces of the cut, which are opaque), so that ``surface`` + is visible through it; 1 for an opaque boundary. + surface : str | float | None + A surface eta1 = const inside the domain, colored by the field and opaque, cut like the domain: ``"peak"`` + through the largest absolute value over all frames (e.g. the flux surface of the mode in a torus, showing its + toroidal structure), a value of eta1 in [0, 1], or None for none. The nearest eta1 of the evaluation grid is + used, at least one point away from the boundary of the domain (and from the axis of a polar domain). + window_size : tuple[int, int] + Size of the render window in pixels; the render is cropped to the domain (the same box for all frames), the + title and the color bar are added around it. + margin : int + Margin of the cropped frames in pixels. + show_scalar_bar : bool + Show the color bar (on the right of the image). + context : float + With ``view="cut"``, how much of the domain around the cut is framed, see + :func:`struphy.post_processing.plot_3d.set_camera`. + + Returns + ------- + list[Path] + The frames, in order. + """ + _check_clim(clim) + if not (surface is None or surface == "peak" or (isinstance(surface, (int, float)) and 0 <= surface <= 1)): + raise ValueError(f'surface must be None, "peak" or a value of eta1 in [0, 1], got {surface!r}.') + if not 0 < opacity <= 1: + raise ValueError(f"opacity must be in (0, 1], got {opacity!r}.") + if view not in ("cut", "domain"): + raise ValueError(f'view must be "cut" or "domain", got {view!r}.') + if MPI.COMM_WORLD.Get_rank() != 0: + return [] + steps, times = _steps_and_times(out, steps, max_frames) + label = label or field_label(out, name, component, basis) + domain = out.domain + toroidal = any(cls.__name__ in TORUS_DOMAINS for cls in type(domain).__mro__) + copies = int(domain.params.get("tor_period", 1)) if toroidal else 1 + cells = out.grid.num_elements + if points_per_cell is None: + points_per_cell = int(np.clip(np.floor(np.cbrt(MAX_POINTS_3D / (copies * np.prod(cells))) + 1e-9), 1, 4)) + etas = {e: np.linspace(0.0, 1.0, points_per_cell * n + 1) for e, n in zip(("eta1", "eta2", "eta3"), cells)} + + first_xyz = [] + + def frame_grid(n): + """The structured grid of the domain (the full torus) with the arrays of :func:`_grid_values`.""" + arrays, xyz = _grid_values(out, name, n, etas, component, basis) + if not first_xyz: + first_xyz.extend(xyz) + grid = pv.StructuredGrid(*xyz) + for key, value in arrays.items(): + grid.point_data[key] = np.asarray(value, dtype=float).reshape(-1, order="F") + grid.set_active_scalars(next(iter(arrays))) + cartesian = ("c1", "c2") if "c1" in arrays else None + grid = repeat_around_axis(grid, copies, cartesian) + return grid, np.asarray(_shown(grid.point_data, component)) + + # first pass: the range of the values and the point of the largest absolute value over all frames (the values + # are kept for the second pass if they are not too large) + vmin, vmax, peak, peak_point = np.inf, -np.inf, -1.0, None + first_pass = {} + for n in steps: + grid, values = frame_grid(n) + if len(steps) * values.nbytes <= FIRST_PASS_CACHE: + first_pass[n] = values + low, high = _range(values) + vmin, vmax = min(vmin, low), max(vmax, high) + if not np.all(np.isnan(values)): + i = int(np.nanargmax(np.abs(values))) + if abs(values[i]) > peak: + peak, peak_point = abs(values[i]), grid.points[i] + + def frame_values(n): + return first_pass[n] if n in first_pass else frame_grid(n)[1] + + center, diagonal = cut_center(grid, toroidal) + if azimuth is None: + azimuth = 0.0 if peak_point is None else peak_azimuth(peak_point[None], [1.0], center, diagonal) + # not through a plane of grid points, which would leave dark specks of degenerate cells + azimuth = off_grid_azimuth(grid.points, center, azimuth) + # coincident points merged once; the values of each frame are taken from the nearest point of the grid + merged = merge_seams(grid, diagonal) + tree = cKDTree(grid.points) + _, index = tree.query(merged.points) + options = dict(view=view, window_size=window_size, context=context, opacity=opacity) + + # the surface eta1 = const inside the domain (e.g. a flux surface of a torus), with the copies of a torus sector + n1 = len(etas["eta1"]) + i1 = None + if surface == "peak" and peak_point is not None: + # eta1 of the peak: the nearest point of the sector to the peak rotated back into the sector (copies) + sector_tree = cKDTree(np.stack([c.reshape(-1, order="F") for c in first_xyz], axis=1)) + candidates = [] + for k in range(copies): + angle = -2 * np.pi * k / copies + x, y = ( + peak_point[0] * np.cos(angle) - peak_point[1] * np.sin(angle), + peak_point[0] * np.sin(angle) + peak_point[1] * np.cos(angle), + ) + candidates.append(sector_tree.query([x, y, peak_point[2]])) + i1 = int(min(candidates)[1]) % n1 + elif surface is not None and surface != "peak": + i1 = int(np.argmin(np.abs(etas["eta1"] - float(surface)))) + if i1 is not None and n1 > 2: + # not on the boundary of the domain (or on the axis of a polar domain) + i1 = min(max(i1, 1), n1 - 2) + surface_grid = surface_index = None + if i1 is not None and n1 > 2: + sector = pv.StructuredGrid(*(c[i1 : i1 + 1] for c in first_xyz)) + surface_grid = merge_seams(repeat_around_axis(sector, copies), diagonal) + # its values: those of the nearest points of the grid of the domain (the same points) + _, surface_index = tree.query(surface_grid.points) + + def render(values, opaque=False): + frame_cmap, limits = _limits(clim, values, vmin, vmax, cmap) + image = _render_3d( + merged, + values[index], + center, + azimuth, + toroidal, + frame_cmap, + limits, + {**options, "opacity": 1.0} if opaque else options, + None if opaque else surface_grid, + None if opaque or surface_grid is None else values[surface_index], + ) + return image, frame_cmap, limits + + # the crop box, the same for all frames (the rendered domain does not change): that of the first frame + # (from an opaque render: the translucent boundary may be nearly white) + top, bottom, left, right = crop_box(render(frame_values(steps[0]), opaque=True)[0], margin) + + folder = _prepare_folder(folder) + files = [] + for k, (n, t) in enumerate(zip(steps, times)): + image, frame_cmap, limits = render(frame_values(n)) + title = f"{label}, {_time_title(out, t)}" + fig, width, height = _compose(image[top:bottom, left:right], title, frame_cmap, limits, label, show_scalar_bar) + files.append(folder / _frame_name(k, len(steps))) + _save_figure(fig, files[-1], width, height) + _write_index(folder, steps, times) + return files + + +def write_videos( + out: Output, + name, + folder, + *, + component=None, + basis=None, + label=None, + prefix=None, + planes=None, + views=("cut", "domain"), + max_frames=500, + steps=None, + clim=None, + cmap=None, + slice_options=None, + options_3d=None, +) -> dict[str, list[Path]]: + """Write the frames of the standard videos of the FEEC field ``name``: one per physical plane + (:func:`write_slice_frames`) and one per 3D view (:func:`write_3d_frames`), into the subfolders + ``_``, ``_3d_cut`` (view ``"cut"``) and ``_3d`` (view ``"domain"``) of ``folder``. + + Parameters + ---------- + out : Output + Output of the run. + name : str + Raw FEEC field, ``"species/variable"`` (e.g. ``"mhd/velocity"``). + folder : str | Path + Folder of the videos. + component, basis, label, max_frames, steps, clim, cmap : + As in :func:`write_slice_frames` and :func:`write_3d_frames`. + prefix : str | None + Start of the names of the subfolders (default: the name of the variable). + planes : list[str] | None + Physical planes (see :func:`struphy.post_processing.slices.physical_planes`); by default the poloidal plane + ``rz`` of tori and cylinders, else all planes of the domain. + views : tuple[str, ...] + 3D views, ``"cut"`` and/or ``"domain"``. + slice_options : dict | None + Further keyword arguments of :func:`write_slice_frames` (e.g. ``points_per_cell``, ``dpi``). + options_3d : dict | None + Further keyword arguments of :func:`write_3d_frames` (e.g. ``window_size``, ``azimuth``). + + Returns + ------- + dict[str, list[Path]] + The frames of each video, by the name of its subfolder (empty lists on the ranks other than 0). + """ + prefix = prefix or name.split("/")[-1] + if planes is None: + try: + available = physical_planes(out.domain, all_planes=True) + planes = ["rz"] if "rz" in available else list(physical_planes(out.domain)) + except ValueError as error: + # e.g. a domain covering only part of the poloidal angle: the 3D views only + warnings.warn(f"No videos in planes: {error}", stacklevel=2) + planes = [] + common = dict( + component=component, basis=basis, label=label, max_frames=max_frames, steps=steps, clim=clim, cmap=cmap + ) + videos = {} + for plane in planes: + sub = f"{prefix}_{plane}" + videos[sub] = write_slice_frames(out, name, Path(folder) / sub, plane=plane, **common, **(slice_options or {})) + for view in views: + sub = f"{prefix}_3d_cut" if view == "cut" else f"{prefix}_3d" + videos[sub] = write_3d_frames(out, name, Path(folder) / sub, view=view, **common, **(options_3d or {})) + return videos + + +def frames_to_video(folder, output=None, fps=25, crf=18) -> Path: + """Make an MP4 video (H.264) of the frames in ``folder`` with ffmpeg, which must be installed. + + Parameters + ---------- + folder : str | Path + Folder of the frames (``frame_0000.png``, ...). + output : str | Path | None + The video file (default: ``.mp4``). + fps : int + Frames per second. + crf : int + Quality (lower is better, 18 is visually lossless). + + Returns + ------- + Path + The video file. + """ + folder = Path(folder) + if shutil.which("ffmpeg") is None: + raise FileNotFoundError("ffmpeg not found; install it or make the video from the frames yourself.") + frames = sorted(folder.glob("frame_*.png")) + if not frames: + raise FileNotFoundError(f"No frames in {folder}.") + digits = len(frames[0].stem.split("_")[1]) + output = Path(output) if output is not None else folder.with_suffix(".mp4") + # even image sizes for yuv420p (the frames of this module have them already) + command = [ + "ffmpeg", "-y", "-loglevel", "error", "-framerate", str(fps), "-i", str(folder / f"frame_%0{digits}d.png"), + "-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2:color=white", "-c:v", "libx264", "-pix_fmt", "yuv420p", + "-crf", str(crf), str(output), + ] # fmt: skip + subprocess.run(command, check=True) + return output diff --git a/src/struphy/post_processing/plot_3d.py b/src/struphy/post_processing/plot_3d.py index 29474dd16..0b9bc4057 100644 --- a/src/struphy/post_processing/plot_3d.py +++ b/src/struphy/post_processing/plot_3d.py @@ -1,8 +1,9 @@ """3D plots of FEEC fields from the VTK files of a run (written by ``Output.pproc(create_vtk=True)``), with PyVista. :func:`plot_field_3d` renders one field at one saved time step in the physical domain: a cutaway view, where a wedge of 90 -degrees around the vertical axis is removed, one of its faces passing through the largest absolute value, so that -the field is visible inside (for a torus, the poloidal cross-section through the maximum). +degrees around the vertical axis is removed, one of its faces passing close to the largest absolute value (between two +planes of grid points, see :func:`off_grid_azimuth`), so that the field is visible inside (for a torus, the +poloidal cross-section near the maximum). The resolution of the plot is that of the VTK files, set by ``celldivide`` of ``Output.pproc``. Vector fields are shown by their magnitude; they need the Cartesian components of ``pproc(physical=True)``. @@ -160,14 +161,22 @@ def default_colormap(vmin, vmax, noise=1e-2): return NEGATIVE, (vmin, 0.0) -def crop_image(image, margin=10, background=240): - """``image`` (an RGB array) cropped to the pixels darker than ``background`` (nearly white pixels, e.g. at the - border of an anti-aliased render, count as background), plus ``margin``.""" +def crop_box(image, margin=10, background=240) -> tuple[int, int, int, int]: + """The rows ``top:bottom`` and columns ``left:right`` of ``image`` (an RGB array) that hold the pixels darker + than ``background`` (nearly white pixels, e.g. at the border of an anti-aliased render, count as background), + plus ``margin``; the whole image if it is empty.""" rows, cols = np.nonzero(np.any(image[..., :3] < background, axis=-1)) if rows.size == 0: - return image + return 0, image.shape[0], 0, image.shape[1] top, bottom = max(rows.min() - margin, 0), min(rows.max() + margin + 1, image.shape[0]) left, right = max(cols.min() - margin, 0), min(cols.max() + margin + 1, image.shape[1]) + return int(top), int(bottom), int(left), int(right) + + +def crop_image(image, margin=10, background=240): + """``image`` (an RGB array) cropped to the pixels darker than ``background``, plus ``margin`` (see + :func:`crop_box`).""" + top, bottom, left, right = crop_box(image, margin, background) return image[top:bottom, left:right] @@ -226,6 +235,116 @@ def is_torus(path_out, grid=None) -> bool: return any(cls.__name__ in TORUS_DOMAINS for cls in type(domain).__mro__) +def cut_center(grid, toroidal) -> tuple[np.ndarray, float]: + """The center of the vertical axis of the cutaway (the center of ``grid``, on the axis x = y = 0 for a torus) + and the length of the diagonal of the bounding box of ``grid``.""" + x0, x1, y0, y1, z0, z1 = grid.bounds + center = np.array([0.5 * (x0 + x1), 0.5 * (y0 + y1), 0.5 * (z0 + z1)]) + if toroidal: + center[:2] = 0.0 + return center, float(np.linalg.norm([x1 - x0, y1 - y0, z1 - z0])) + + +def peak_azimuth(points, values, center, diagonal) -> float: + """The azimuth around the vertical axis through ``center`` of the point of the largest absolute value (0 for a + point on that axis).""" + offset = np.asarray(points)[np.nanargmax(np.abs(values))] - center + return float(np.arctan2(offset[1], offset[0])) if np.hypot(offset[0], offset[1]) > 1e-6 * diagonal else 0.0 + + +def off_grid_azimuth(points, center, azimuth, decimals=9) -> float: + """``azimuth`` moved to the middle between the two nearest azimuths (around the vertical axis through ``center``) + of ``points``, e.g. between two planes eta3 = const of the grid of a torus; if the other face of the wedge of + :func:`cut_wedge` (at ``azimuth + pi/2``) then lies on such a plane, a quarter of the gap from its first point + instead, so that neither face does. A cut through a plane of grid points leaves degenerate cells, which are drawn + as dark specks. Unchanged if the points have fewer than two distinct azimuths.""" + rel = np.asarray(points)[:, :2] - np.asarray(center)[:2] + radius = np.hypot(rel[:, 0], rel[:, 1]) + rel = rel[radius > 1e-9 * (np.max(radius) if radius.size else 1.0)] + angles = np.unique(np.round(np.mod(np.arctan2(rel[:, 1], rel[:, 0]), 2 * np.pi), decimals)) + if len(angles) < 2: + return float(azimuth) + angles = np.concatenate([angles - 2 * np.pi, angles, angles + 2 * np.pi]) + target = np.mod(azimuth, 2 * np.pi) + i = int(np.searchsorted(angles, target)) + lo, hi = angles[i - 1], angles[i] + if np.isclose(hi, target, rtol=0.0, atol=10.0**-decimals): + # on a plane of points: the gap above it + lo, hi = hi, angles[i + 1] + moved = 0.5 * (lo + hi) + # the other face of the wedge on a plane of points: a quarter of the gap instead of a half + other = np.mod(moved + np.pi / 2, 2 * np.pi) + if np.min(np.abs(angles - other)) < 1e-6 * (hi - lo): + moved = lo + 0.25 * (hi - lo) + return float(azimuth + moved - target) + + +def merge_seams(grid, diagonal): + """``grid`` as an unstructured grid with coincident points merged (periodic seams, the axis of polar domains), + so that the seams are not drawn as faces.""" + return grid.cast_to_unstructured_grid().clean(tolerance=1e-9 * diagonal) + + +def new_plotter(show=False, window_size=(1600, 1600)) -> pv.Plotter: + """A plotter with a white background and anti-aliasing, off screen unless ``show``.""" + plotter = pv.Plotter(off_screen=not show, window_size=list(window_size)) + plotter.set_background("white") + plotter.enable_anti_aliasing("ssaa") + return plotter + + +def add_field_mesh(plotter, mesh, scalars, cmap, clim, show_scalar_bar=False, scalar_bar_title=None, **kwargs): + """Add ``mesh`` colored by its point array ``scalars`` with the lighting of the 3D plots (further keyword + arguments, e.g. ``opacity``, go to ``plotter.add_mesh``); returns the actor.""" + return plotter.add_mesh( + mesh, + scalars=scalars, + cmap=cmap, + clim=clim, + smooth_shading=True, + # sharp edges (e.g. of a slab or of the cutaway) shaded as edges, not smoothed over + split_sharp_edges=True, + # lighting that keeps the colors of the colormap on surfaces facing the light (zero stays light gray) + ambient=0.2, + diffuse=0.8, + show_scalar_bar=show_scalar_bar, + scalar_bar_args={"title": scalar_bar_title or scalars, "vertical": True, "color": "black"}, + **kwargs, + ) + + +def set_camera(plotter, grid, shown, center, azimuth, toroidal, cutaway=True, context=1.2): + """Look into the wedge of the cutaway (see :func:`cut_wedge`) from above, framing the removed wedge and some of + the domain around it (enlarged by ``context``), so that the cut fills the view and the rest of the domain gives + the 3D context. For a torus, the face of the cut at ``azimuth`` (a poloidal cross-section) is framed instead, + enlarged by twice ``context``. ``grid`` is the whole domain, ``shown`` the rendered part of it.""" + diagonal = np.linalg.norm(np.subtract(grid.bounds[1::2], grid.bounds[::2])) + # elevation of the camera above the horizontal plane: about 39 degrees, half of it for a torus + elevation = np.arctan(0.8) / (2 if toroidal else 1) + # azimuth of the camera: along the bisector of the wedge, for a torus more in front of the cut face at `azimuth` + # (whose normal points to azimuth + pi/2) + turn = 3 * np.pi / 8 if toroidal else np.pi / 4 + view = np.array([np.cos(azimuth + turn), np.sin(azimuth + turn), np.tan(elevation)]) + plotter.camera.focal_point = center + plotter.camera.position = center + 2.0 * diagonal * view / np.linalg.norm(view) + plotter.camera.up = (0.0, 0.0, 1.0) + plotter.reset_camera() + if not cutaway: + return + framed = grid.points[wedge_distance(grid.points, center, azimuth) > 0] + if toroidal: + rel = shown.points - center + start = np.cos(azimuth) * rel[:, 0] + np.sin(azimuth) * rel[:, 1] + left = -np.sin(azimuth) * rel[:, 0] + np.cos(azimuth) * rel[:, 1] + face = shown.points[(np.abs(left) < 1e-6 * diagonal) & (start > 0)] + if len(face): + framed, context = face, 2 * context + if len(framed): + lo, hi = framed.min(axis=0), framed.max(axis=0) + mid, half = 0.5 * (lo + hi), 0.5 * context * (hi - lo) + plotter.reset_camera(bounds=[v for k in range(3) for v in (mid[k] - half[k], mid[k] + half[k])]) + + def plot_field_3d( path_out, variable, @@ -273,7 +392,8 @@ def plot_field_3d( Open an interactive window (otherwise render off screen). cutaway : bool Remove a wedge of 90 degrees around the vertical axis through the center of the domain, starting at the - azimuth of the largest absolute value, to show the field inside (one face of the cut passes through it). + azimuth of the largest absolute value, to show the field inside (one face of the cut passes close to it, see + :func:`off_grid_azimuth`). title : str | None Title of the image (default: none). cmap : str | None @@ -306,76 +426,21 @@ def plot_field_3d( default_cmap, clim = default_colormap(*grid.get_data_range(scalars)) cmap = cmap or default_cmap - # the vertical axis of the cut: through the center of the domain, for a torus its axis x = y = 0 - x0, x1, y0, y1, z0, z1 = grid.bounds - center = np.array([0.5 * (x0 + x1), 0.5 * (y0 + y1), 0.5 * (z0 + z1)]) - diagonal = np.linalg.norm([x1 - x0, y1 - y0, z1 - z0]) toroidal = is_torus(path_out, grid) - if toroidal: - center[:2] = 0.0 - # the azimuth of the largest |value| around that axis - offset = grid.points[np.nanargmax(np.abs(grid.point_data[scalars]))] - center - azimuth = np.arctan2(offset[1], offset[0]) if np.hypot(offset[0], offset[1]) > 1e-6 * diagonal else 0.0 - - # merge coincident points (periodic seams, the axis of polar domains), so that the seams are not drawn as faces - grid = grid.cast_to_unstructured_grid().clean(tolerance=1e-9 * diagonal) + center, diagonal = cut_center(grid, toroidal) + azimuth = peak_azimuth(grid.points, grid.point_data[scalars], center, diagonal) + if cutaway: + azimuth = off_grid_azimuth(grid.points, center, azimuth) + grid = merge_seams(grid, diagonal) - plotter = pv.Plotter(off_screen=not show, window_size=list(window_size)) - plotter.set_background("white") - plotter.enable_anti_aliasing("ssaa") + plotter = new_plotter(show, window_size) shown = cut_wedge(grid, center, azimuth) if cutaway else grid - plotter.add_mesh( - shown, - scalars=scalars, - cmap=cmap, - clim=clim, - smooth_shading=True, - # lighting that keeps the colors of the colormap on surfaces facing the light (zero stays light gray) - ambient=0.2, - diffuse=0.8, - show_scalar_bar=show_scalar_bar, - scalar_bar_args={"title": scalars, "vertical": True, "color": "black"}, - ) + add_field_mesh(plotter, shown, scalars, cmap, clim, show_scalar_bar) if show_axes: plotter.add_axes() if title is not None: plotter.add_text(title, font_size=12, color="black") - # look into the wedge from above, framing the removed wedge (between the two faces of the cut) and some of the - # domain around it, so that the cut fills the view and the rest of the domain gives the 3D context - # elevation of the camera above the horizontal plane: about 39 degrees, half of it for a torus - elevation = np.arctan(0.8) / (2 if toroidal else 1) - # azimuth of the camera: along the bisector of the wedge, for a torus more in front of the cut face at `azimuth` - # (whose normal points to azimuth + pi/2) - turn = 3 * np.pi / 8 if toroidal else np.pi / 4 - view = np.array([np.cos(azimuth + turn), np.sin(azimuth + turn), np.tan(elevation)]) - plotter.camera.focal_point = center - plotter.camera.position = center + 2.0 * diagonal * view / np.linalg.norm(view) - plotter.camera.up = (0.0, 0.0, 1.0) - plotter.reset_camera() - if cutaway: - framed = grid.points[wedge_distance(grid.points, center, azimuth) > 0] - # a domain around the vertical axis (a torus): frame only the face of the cut through the largest |value| - # (a poloidal cross-section), enlarged by twice the context - rel = shown.points - center - if toroidal: - start = np.cos(azimuth) * rel[:, 0] + np.sin(azimuth) * rel[:, 1] - left = -np.sin(azimuth) * rel[:, 0] + np.cos(azimuth) * rel[:, 1] - face = shown.points[(np.abs(left) < 1e-6 * diagonal) & (start > 0)] - if len(face): - framed, context = face, 2 * context - if len(framed): - lo, hi = framed.min(axis=0), framed.max(axis=0) - mid, half = 0.5 * (lo + hi), 0.5 * context * (hi - lo) - plotter.reset_camera( - bounds=[ - mid[0] - half[0], - mid[0] + half[0], - mid[1] - half[1], - mid[1] + half[1], - mid[2] - half[2], - mid[2] + half[2], - ] - ) + set_camera(plotter, grid, shown, center, azimuth, toroidal, cutaway=cutaway, context=context) if filename is not None: image = plotter.screenshot(return_img=True) diff --git a/src/struphy/post_processing/slices.py b/src/struphy/post_processing/slices.py index 104504a83..dba17fa93 100644 --- a/src/struphy/post_processing/slices.py +++ b/src/struphy/post_processing/slices.py @@ -52,8 +52,9 @@ class Plane: def physical_planes(domain, phi=np.pi, all_planes=False) -> dict[str, Plane]: - """The physical planes through ``domain`` as logical surfaces (Cuboid, HollowCylinder, IGAPolarCylinder, - HollowTorus, IGAPolarTorus). + """The physical planes through ``domain`` as logical surfaces (Cuboid, Colella, Orthogonal, HollowCylinder, + IGAPolarCylinder, HollowTorus, IGAPolarTorus); for other domains the logical mid-surfaces of + :func:`logical_mid_planes` (for other tori the poloidal plane ``rz`` at eta3 = 0.5, whatever ``phi``). For the spline mappings (IGAPolar...) the poloidal angle is only approximately 2 pi eta2, so the planes cut at eta2 are approximate. Tori may cover part of the toroidal angle (``tor_period`` > 1, see the module docstring); @@ -65,7 +66,8 @@ def physical_planes(domain, phi=np.pi, all_planes=False) -> dict[str, Plane]: periods = {k: domain.params.get(k, 1) for k in ("poc", "pol_period")} if any(v != 1 for v in periods.values()): raise ValueError(f"Physical planes of {kind} only for full poloidal angles (periods 1), got {periods}.") - if kind == "Cuboid": + if kind in ("Cuboid", "Colella", "Orthogonal"): + # the mid-planes eta_i = 0.5 are flat for these mappings too (their distortions vanish there) return { "xy": Plane("x", "y", "z at the center", "eta3", (0.5,)), "xz": Plane("x", "z", "y at the center", "eta2", (0.5,)), @@ -96,7 +98,32 @@ def toroidal(where, *angles): "rz": Plane("R", "z", **toroidal(f"phi = {phi:.3g}", phi)), } return planes if all_planes else {name: planes[name] for name in ("xy", "rz")} - raise NotImplementedError(f"No physical planes defined for the domain {kind}.") + return logical_mid_planes(domain, all_planes=all_planes) + + +def logical_mid_planes(domain, n=11, all_planes=False) -> dict[str, Plane]: + """The logical surfaces eta_i = 0.5 of a domain without the physical planes of :func:`physical_planes`. + + For the other tori around the vertical axis (e.g. Tokamak, GVEC and DESC equilibria, see + :data:`struphy.post_processing.plot_3d.TORUS_DOMAINS`), the poloidal cross-section eta3 = 0.5 as ``"rz"``, shown in + the major radius R and z. Otherwise the surfaces eta3 = 0.5 and eta2 = 0.5, with ``all_planes`` also eta1 = 0.5 + (e.g. a curved shell around the axis of a cylinder), named ``"eta3"``, ``"eta2"``, ``"eta1"``, each shown in the two + Cartesian coordinates that vary most on it (sampled at ``n`` x ``n`` points); a curved surface is shown by its + projection onto these coordinates.""" + from struphy.post_processing.plot_3d import TORUS_DOMAINS + + if any(cls.__name__ in TORUS_DOMAINS for cls in type(domain).__mro__): + return {"rz": Plane("R", "z", "eta3 = 0.5", "eta3", (0.5,))} + planes = {} + s = np.linspace(0.0, 1.0, n) + for cut in ("eta3", "eta2", "eta1") if all_planes else ("eta3", "eta2"): + etas = [s, s, s] + etas[("eta1", "eta2", "eta3").index(cut)] = np.array([0.5]) + xyz = [np.asarray(xp.to_numpy(c)) for c in domain(*etas)] + spread = [np.ptp(c) for c in xyz] + h, v = sorted(np.argsort(spread)[-2:]) + planes[cut] = Plane("xyz"[h], "xyz"[v], f"{cut} = 0.5", cut, (0.5,)) + return planes def evaluate_plane(out: Output, domain, name, t, exact, plane: Plane, n_pts, kind=None, components=False): diff --git a/src/struphy/post_processing/tests/test_frames.py b/src/struphy/post_processing/tests/test_frames.py new file mode 100644 index 000000000..0ed19765f --- /dev/null +++ b/src/struphy/post_processing/tests/test_frames.py @@ -0,0 +1,237 @@ +"""Frames of videos of FEEC fields (struphy.post_processing.frames) from short runs on a torus and other domains.""" + +import shutil + +import numpy as np +import pytest +from matplotlib.image import imread +from maybempi import MPI + +from struphy import BaseUnits, DerhamOptions, EnvironmentOptions, Output, Simulation, Time, domains, grids +from struphy.initial import perturbations +from struphy.post_processing.frames import ( + _check_clim, + _finite, + _grid_values, + _limits, + frame_steps, + frames_to_video, + saved_times, + write_3d_frames, + write_slice_frames, + write_videos, +) + +N_SAVED = 4 # saved time steps of the run (the initial one and three time steps) + + +def short_torus_run(out_folder, polar=False, domain=None) -> Output: + """A Maxwell run of three time steps on a coarse torus (with ``polar``, an IGAPolarTorus with polar splines, which + has an axis; or on ``domain``), with an m = 1, n = 1 electric field (2-form).""" + from struphy.models import Maxwell + + if domain is None: + domain = ( + domains.IGAPolarTorus(num_elements=(4, 9), degree=(2, 2), a=1.0, R0=3.0, tor_period=1) + if polar + else domains.HollowTorus(a1=0.2, a2=1.0, R0=3.0, tor_period=1) + ) + + model = Maxwell(base_units=BaseUnits()) + model.em_fields.e_field.save_data = True + model.em_fields.b_field.save_data = True + sim = Simulation( + model=model, + params_path=__file__, + env=EnvironmentOptions(out_folders=out_folder, sim_folder="sim_frames"), + time_opts=Time(dt=0.05, Tend=(N_SAVED - 1) * 0.05), + domain=domain, + grid=grids.TensorProductGrid(num_elements=(4, 9 if polar else 8, 4)), + # polar splines: clamped in eta1 (free on the axis, Dirichlet at the boundary) + derham_opts=DerhamOptions(degree=(2, 2, 1), bcs=(("free", "dirichlet"), None, None), polar_splines=True) + if polar + else DerhamOptions(degree=(2, 2, 1)), + ) + model.em_fields.e_field.add_perturbation( + perturbations.TorusModesSin(ms=(1,), ns=(1,), amps=(1e-3,), comp=0, given_in_basis="2") + ) + sim.run() + return Output(sim.env.path_out) + + +def check_frames(files, count): + """``count`` frames, all of the same even size, listed in frames.txt.""" + assert len(files) == count + shapes = {imread(f).shape for f in files} + assert len(shapes) == 1 + height, width = shapes.pop()[:2] + assert height % 2 == 0 and width % 2 == 0 + lines = (files[0].parent / "frames.txt").read_text().splitlines() + assert len(lines) == count + 1 + + +def test_frames(mpi_out_folder): + run = short_torus_run(mpi_out_folder) + rank = MPI.COMM_WORLD.Get_rank() + folder = run.path_out / "frames" + + # at most max_frames evenly spaced saved steps, with the first and the last one + assert frame_steps(run) == list(range(N_SAVED)) + assert frame_steps(run, max_frames=3) == [0, 2, 3] + assert len(saved_times(run)) == N_SAVED + with pytest.raises(ValueError): + frame_steps(run, max_frames=0) + + # the radial component (normalized basis) in the poloidal plane, one scale over all frames + files = write_slice_frames(run, "em_fields/e_field", folder / "e_r_rz", component=0, points_per_cell=2) + if rank == 0: + check_frames(files, N_SAVED) + else: + assert files == [] + + # the magnitude of the field in the plane z = 0 at three of the saved steps, a scale per frame + files = write_slice_frames(run, "em_fields/e_field", folder / "e_xy", plane="xy", max_frames=3, clim="frame") + if rank == 0: + check_frames(files, 3) + + # 3D, both views, the Cartesian component with fixed color limits + for view in ("cut", "domain"): + files = write_3d_frames( + run, + "em_fields/e_field", + folder / f"e_3d_{view}", + component=0, + view=view, + points_per_cell=2, + window_size=(400, 400), + ) + if rank == 0: + check_frames(files, N_SAVED) + files = write_3d_frames( + run, + "em_fields/e_field", + folder / "e_x_3d", + component=0, + basis="physical", + clim=(-1e-3, 1e-3), + steps=[0, N_SAVED - 1], + points_per_cell=1, + window_size=(300, 300), + ) + if rank == 0: + check_frames(files, 2) + + # an opaque boundary without the inner surface, and the surface at a given eta1 + for options in ({"opacity": 1.0, "surface": None}, {"surface": 0.5}): + files = write_3d_frames( + run, + "em_fields/e_field", + folder / f"e_3d_{options['surface']}", + component=0, + steps=[0, N_SAVED - 1], + points_per_cell=1, + window_size=(300, 300), + **options, + ) + if rank == 0: + check_frames(files, 2) + + if rank == 0 and shutil.which("ffmpeg") is not None: + video = frames_to_video(folder / "e_r_rz", fps=2) + assert video.exists() and video.stat().st_size > 0 + + with pytest.raises(ValueError): + write_slice_frames(run, "em_fields/e_field", folder / "bad", clim="global") + with pytest.raises(ValueError): + write_3d_frames(run, "em_fields/e_field", folder / "bad", view="side") + for bad in ({"surface": 1.5}, {"surface": "top"}, {"opacity": 0.0}): + with pytest.raises(ValueError): + write_3d_frames(run, "em_fields/e_field", folder / "bad", **bad) + if rank == 0: + with pytest.raises(ValueError): + write_slice_frames(run, "em_fields/e_field", folder / "bad", plane="ab") + MPI.COMM_WORLD.Barrier() + + +@pytest.mark.mpi_skip +def test_off_grid_azimuth(): + from struphy.post_processing.plot_3d import off_grid_azimuth + + # points on 8 vertical half planes around the z-axis (azimuths k pi/4), at two radii + angles = np.arange(8) * np.pi / 4 + points = np.array([[r * np.cos(a), r * np.sin(a), z] for a in angles for r in (1.0, 2.0) for z in (0.0, 1.0)]) + center = np.zeros(3) + # on a plane of points: moved to the middle of the next gap + assert np.isclose(off_grid_azimuth(points, center, np.pi / 2), np.pi / 2 + np.pi / 8) + # within a gap: its middle, on the branch of the given azimuth + assert np.isclose(off_grid_azimuth(points, center, 0.1), np.pi / 8) + assert np.isclose(off_grid_azimuth(points, center, -0.1), -np.pi / 8) + assert np.isclose(off_grid_azimuth(points, center, 2 * np.pi + 0.1), 2 * np.pi + np.pi / 8) + # 6 planes (every pi/3): the middle of the gap puts the other face (+ pi/2) on a plane, a quarter does not + sixths = np.arange(6) * np.pi / 3 + points6 = np.array([[np.cos(a), np.sin(a), 0.0] for a in sixths]) + assert np.isclose(off_grid_azimuth(points6, center, 0.0), np.pi / 12) + # points on the axis only: unchanged + assert off_grid_azimuth(np.zeros((3, 3)), center, 0.3) == 0.3 + + +def test_frames_polar_axis(mpi_out_folder): + """On a torus with a polar axis, the singular values of the push-forward on the axis are replaced by those of the + next ring, so that the color scale stays finite.""" + run = short_torus_run(mpi_out_folder, polar=True) + etas = {"eta1": np.linspace(0, 1, 5), "eta2": np.linspace(0, 1, 9), "eta3": np.linspace(0, 1, 5)} + for component, basis in ((0, "norm"), (None, None), (0, "physical")): + arrays, _ = _grid_values(run, "em_fields/e_field", N_SAVED - 1, etas, component, basis) + assert all(np.all(np.isfinite(values)) for values in arrays.values()) + files = write_slice_frames(run, "em_fields/e_field", run.path_out / "frames" / "rz", component=0, max_frames=2) + files += write_3d_frames( + run, + "em_fields/e_field", + run.path_out / "frames" / "3d", + max_frames=2, + points_per_cell=1, + window_size=(300, 300), + ) + if MPI.COMM_WORLD.Get_rank() == 0: + assert len(files) == 4 + MPI.COMM_WORLD.Barrier() + + +@pytest.mark.mpi_skip +def test_frame_helpers(): + # the axis (first index) takes the values of the next ring, other non-finite values become NaN + values = _finite([[np.inf, 1.0], [2.0, np.nan], [3.0, 4.0]]) + assert values[0, 0] == 2.0 and values[0, 1] == 1.0 and np.isnan(values[1, 1]) + # color limits: symmetric about zero for a field changing sign, never empty + assert _limits(None, None, -1.0, 2.0, None)[1] == (-2.0, 2.0) + assert _limits("frame", np.array([0.5, 1.0, np.nan]), -9.0, 9.0, None)[1] == (0.0, 1.0) + assert _limits(None, None, 0.0, 0.0, None)[1] == (-1.0, 1.0) + assert _limits(np.array([-3, 1]), None, 0.0, 1.0, "viridis") == ("viridis", (-3.0, 1.0)) + _check_clim(np.array([0.0, 1.0])) + for clim in ("global", (1.0, 0.0), np.array([1.0, 0.0]), (0.0, 1.0, 2.0)): + with pytest.raises(ValueError): + _check_clim(clim) + + +@pytest.mark.parametrize( + "domain", + [domains.Cuboid(r1=2.0, r2=1.0, r3=3.0), domains.Colella(Lx=2.0, Ly=1.0, alpha=0.1, Lz=3.0)], + ids=["slab", "colella"], +) +def test_videos_on_other_domains(mpi_out_folder, domain): + """All standard videos (the planes xy, xz, yz and both 3D views) of a slab and of the Colella mapping.""" + run = short_torus_run(mpi_out_folder, domain=domain) + videos = write_videos( + run, + "em_fields/e_field", + run.path_out / "frames", + component=0, + prefix="e1", + max_frames=2, + options_3d={"window_size": (300, 300)}, + ) + if MPI.COMM_WORLD.Get_rank() == 0: + assert set(videos) == {"e1_xy", "e1_xz", "e1_yz", "e1_3d_cut", "e1_3d"} + for files in videos.values(): + check_frames(files, 2) + MPI.COMM_WORLD.Barrier() diff --git a/src/struphy/post_processing/tests/test_slices.py b/src/struphy/post_processing/tests/test_slices.py index 6d73cddfd..17ef94bd8 100644 --- a/src/struphy/post_processing/tests/test_slices.py +++ b/src/struphy/post_processing/tests/test_slices.py @@ -15,6 +15,8 @@ DOMAINS = [ domains.Cuboid(r1=2.0, l2=-1.0, r2=1.0, r3=4.0), + domains.Colella(Lx=2.0, Ly=3.0, alpha=0.1, Lz=4.0), + domains.Orthogonal(Lx=2.0, Ly=3.0, alpha=0.1, Lz=4.0), domains.HollowCylinder(a1=0.1, a2=1.0, Lz=4.0), domains.IGAPolarCylinder(num_elements=(8, 9), degree=(3, 3), a=1.0, Lz=4.0), domains.HollowTorus(a1=0.1, a2=1.0, R0=4.0, tor_period=1), @@ -68,13 +70,34 @@ def test_tori_have_two_planes_by_default(): assert set(physical_planes(torus)) == {"xy", "rz"} assert set(physical_planes(torus, all_planes=True)) == {"xy", "xz", "yz", "rz"} # other domains keep all their planes - assert set(physical_planes(DOMAINS[1])) == {"xy", "xz", "yz", "rz"} + assert set(physical_planes(DOMAINS[3])) == {"xy", "xz", "yz", "rz"} @pytest.mark.mpi_skip -def test_unknown_domain_raises(): - with pytest.raises(NotImplementedError): - physical_planes(domains.Colella()) +def test_other_domains_have_logical_mid_planes(): + # an elliptic cylinder: the logical mid-surfaces, the surface eta3 = 0.5 is the plane z = const + domain = domains.PoweredEllipticCylinder() + planes = physical_planes(domain) + assert set(planes) == {"eta2", "eta3"} + assert set(physical_planes(domain, all_planes=True)) == {"eta1", "eta2", "eta3"} + assert (planes["eta3"].hlabel, planes["eta3"].vlabel) == ("x", "y") + assert all(plane.values == (0.5,) and plane.cut == key for key, plane in planes.items()) + + +@pytest.mark.mpi_skip +def test_other_tori_have_a_poloidal_plane(): + """The fallback for tori (e.g. Tokamak) is the poloidal plane eta3 = 0.5 in R and z, here checked on a torus with + tor_period = 4 (where a Cartesian projection would squeeze it by cos 45 degrees).""" + from struphy.post_processing.slices import logical_mid_planes + + domain = domains.HollowTorus(a1=0.1, a2=1.0, R0=4.0, tor_period=4) + (plane,) = logical_mid_planes(domain).values() + assert (plane.hlabel, plane.vlabel, plane.cut, plane.values) == ("R", "z", "eta3", (0.5,)) + xyz = surface(domain, plane, 0) + # the surface is a half plane through the vertical axis: R spans the minor diameter + radius = np.hypot(xyz["x"], xyz["y"]) + assert np.isclose(np.ptp(radius), 2.0) + assert np.isclose(np.max(np.hypot(radius - 4.0, xyz["z"])), 1.0) @pytest.mark.mpi_skip