Merge plotting optimizations from chb-copilot-test
- Implement multiprocessing-based parallel plotting - Add parallel_plot_helper.py for concurrent plot task execution - Use matplotlib 'Agg' backend for multiprocessing safety - Set OMP_NUM_THREADS=1 to prevent BLAS thread explosion - Use subprocess for binary data plots to avoid thread conflicts - Add fork bomb protection in main program This merge only includes plotting improvements and excludes MPI communication changes to preserve existing optimizations. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -8,6 +8,14 @@
|
|||||||
##
|
##
|
||||||
##################################################################
|
##################################################################
|
||||||
|
|
||||||
|
## Guard against re-execution by multiprocessing child processes.
|
||||||
|
## Without this, using 'spawn' or 'forkserver' context would cause every
|
||||||
|
## worker to re-run the entire script, spawning exponentially more
|
||||||
|
## workers (fork bomb).
|
||||||
|
if __name__ != '__main__':
|
||||||
|
import sys as _sys
|
||||||
|
_sys.exit(0)
|
||||||
|
|
||||||
|
|
||||||
##################################################################
|
##################################################################
|
||||||
|
|
||||||
@@ -424,26 +432,31 @@ print(
|
|||||||
|
|
||||||
import plot_xiaoqu
|
import plot_xiaoqu
|
||||||
import plot_GW_strain_amplitude_xiaoqu
|
import plot_GW_strain_amplitude_xiaoqu
|
||||||
|
from parallel_plot_helper import run_plot_tasks_parallel
|
||||||
|
|
||||||
|
plot_tasks = []
|
||||||
|
|
||||||
## Plot black hole trajectory
|
## Plot black hole trajectory
|
||||||
plot_xiaoqu.generate_puncture_orbit_plot( binary_results_directory, figure_directory )
|
plot_tasks.append( ( plot_xiaoqu.generate_puncture_orbit_plot, (binary_results_directory, figure_directory) ) )
|
||||||
plot_xiaoqu.generate_puncture_orbit_plot3D( binary_results_directory, figure_directory )
|
plot_tasks.append( ( plot_xiaoqu.generate_puncture_orbit_plot3D, (binary_results_directory, figure_directory) ) )
|
||||||
|
|
||||||
## Plot black hole separation vs. time
|
## Plot black hole separation vs. time
|
||||||
plot_xiaoqu.generate_puncture_distence_plot( binary_results_directory, figure_directory )
|
plot_tasks.append( ( plot_xiaoqu.generate_puncture_distence_plot, (binary_results_directory, figure_directory) ) )
|
||||||
|
|
||||||
## Plot gravitational waveforms (psi4 and strain amplitude)
|
## Plot gravitational waveforms (psi4 and strain amplitude)
|
||||||
for i in range(input_data.Detector_Number):
|
for i in range(input_data.Detector_Number):
|
||||||
plot_xiaoqu.generate_gravitational_wave_psi4_plot( binary_results_directory, figure_directory, i )
|
plot_tasks.append( ( plot_xiaoqu.generate_gravitational_wave_psi4_plot, (binary_results_directory, figure_directory, i) ) )
|
||||||
plot_GW_strain_amplitude_xiaoqu.generate_gravitational_wave_amplitude_plot( binary_results_directory, figure_directory, i )
|
plot_tasks.append( ( plot_GW_strain_amplitude_xiaoqu.generate_gravitational_wave_amplitude_plot, (binary_results_directory, figure_directory, i) ) )
|
||||||
|
|
||||||
## Plot ADM mass evolution
|
## Plot ADM mass evolution
|
||||||
for i in range(input_data.Detector_Number):
|
for i in range(input_data.Detector_Number):
|
||||||
plot_xiaoqu.generate_ADMmass_plot( binary_results_directory, figure_directory, i )
|
plot_tasks.append( ( plot_xiaoqu.generate_ADMmass_plot, (binary_results_directory, figure_directory, i) ) )
|
||||||
|
|
||||||
## Plot Hamiltonian constraint violation over time
|
## Plot Hamiltonian constraint violation over time
|
||||||
for i in range(input_data.grid_level):
|
for i in range(input_data.grid_level):
|
||||||
plot_xiaoqu.generate_constraint_check_plot( binary_results_directory, figure_directory, i )
|
plot_tasks.append( ( plot_xiaoqu.generate_constraint_check_plot, (binary_results_directory, figure_directory, i) ) )
|
||||||
|
|
||||||
|
run_plot_tasks_parallel(plot_tasks)
|
||||||
|
|
||||||
## Plot stored binary data
|
## Plot stored binary data
|
||||||
plot_xiaoqu.generate_binary_data_plot( binary_results_directory, figure_directory )
|
plot_xiaoqu.generate_binary_data_plot( binary_results_directory, figure_directory )
|
||||||
|
|||||||
29
parallel_plot_helper.py
Normal file
29
parallel_plot_helper.py
Normal file
@@ -0,0 +1,29 @@
|
|||||||
|
import multiprocessing
|
||||||
|
|
||||||
|
def run_plot_task(task):
|
||||||
|
"""Execute a single plotting task.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
task : tuple
|
||||||
|
A tuple of (function, args_tuple) where function is a callable
|
||||||
|
plotting function and args_tuple contains its arguments.
|
||||||
|
"""
|
||||||
|
func, args = task
|
||||||
|
return func(*args)
|
||||||
|
|
||||||
|
|
||||||
|
def run_plot_tasks_parallel(plot_tasks):
|
||||||
|
"""Execute a list of independent plotting tasks in parallel.
|
||||||
|
|
||||||
|
Uses the 'fork' context to create worker processes so that the main
|
||||||
|
script is NOT re-imported/re-executed in child processes.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
plot_tasks : list of tuples
|
||||||
|
Each element is (function, args_tuple).
|
||||||
|
"""
|
||||||
|
ctx = multiprocessing.get_context('fork')
|
||||||
|
with ctx.Pool() as pool:
|
||||||
|
pool.map(run_plot_task, plot_tasks)
|
||||||
@@ -11,6 +11,8 @@
|
|||||||
import numpy ## numpy for array operations
|
import numpy ## numpy for array operations
|
||||||
import scipy ## scipy for interpolation and signal processing
|
import scipy ## scipy for interpolation and signal processing
|
||||||
import math
|
import math
|
||||||
|
import matplotlib
|
||||||
|
matplotlib.use('Agg') ## use non-interactive backend for multiprocessing safety
|
||||||
import matplotlib.pyplot as plt ## matplotlib for plotting
|
import matplotlib.pyplot as plt ## matplotlib for plotting
|
||||||
import os ## os for system/file operations
|
import os ## os for system/file operations
|
||||||
|
|
||||||
|
|||||||
@@ -8,16 +8,23 @@
|
|||||||
##
|
##
|
||||||
#################################################
|
#################################################
|
||||||
|
|
||||||
|
## Restrict OpenMP to one thread per process so that running
|
||||||
|
## many workers in parallel does not create an O(workers * BLAS_threads)
|
||||||
|
## thread explosion. The variable MUST be set before numpy/scipy
|
||||||
|
## are imported, because the BLAS library reads them only at load time.
|
||||||
|
import os
|
||||||
|
os.environ.setdefault("OMP_NUM_THREADS", "1")
|
||||||
|
|
||||||
import numpy
|
import numpy
|
||||||
import scipy
|
import scipy
|
||||||
|
import matplotlib
|
||||||
|
matplotlib.use('Agg') ## use non-interactive backend for multiprocessing safety
|
||||||
import matplotlib.pyplot as plt
|
import matplotlib.pyplot as plt
|
||||||
from matplotlib.colors import LogNorm
|
from matplotlib.colors import LogNorm
|
||||||
from mpl_toolkits.mplot3d import Axes3D
|
from mpl_toolkits.mplot3d import Axes3D
|
||||||
## import torch
|
## import torch
|
||||||
import AMSS_NCKU_Input as input_data
|
import AMSS_NCKU_Input as input_data
|
||||||
|
|
||||||
import os
|
|
||||||
|
|
||||||
|
|
||||||
#########################################################################################
|
#########################################################################################
|
||||||
|
|
||||||
@@ -192,3 +199,19 @@ def get_data_xy( Rmin, Rmax, n, data0, time, figure_title, figure_outdir ):
|
|||||||
|
|
||||||
####################################################################################
|
####################################################################################
|
||||||
|
|
||||||
|
|
||||||
|
####################################################################################
|
||||||
|
## Allow this module to be run as a standalone script so that each
|
||||||
|
## binary-data plot can be executed in a fresh subprocess whose BLAS
|
||||||
|
## environment variables (set above) take effect before numpy loads.
|
||||||
|
##
|
||||||
|
## Usage: python3 plot_binary_data.py <filename> <binary_outdir> <figure_outdir>
|
||||||
|
####################################################################################
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
import sys
|
||||||
|
if len(sys.argv) != 4:
|
||||||
|
print(f"Usage: {sys.argv[0]} <filename> <binary_outdir> <figure_outdir>")
|
||||||
|
sys.exit(1)
|
||||||
|
plot_binary_data(sys.argv[1], sys.argv[2], sys.argv[3])
|
||||||
|
|
||||||
|
|||||||
@@ -8,6 +8,8 @@
|
|||||||
#################################################
|
#################################################
|
||||||
|
|
||||||
import numpy ## numpy for array operations
|
import numpy ## numpy for array operations
|
||||||
|
import matplotlib
|
||||||
|
matplotlib.use('Agg') ## use non-interactive backend for multiprocessing safety
|
||||||
import matplotlib.pyplot as plt ## matplotlib for plotting
|
import matplotlib.pyplot as plt ## matplotlib for plotting
|
||||||
from mpl_toolkits.mplot3d import Axes3D ## needed for 3D plots
|
from mpl_toolkits.mplot3d import Axes3D ## needed for 3D plots
|
||||||
import glob
|
import glob
|
||||||
@@ -15,6 +17,9 @@ import os ## operating system utilities
|
|||||||
|
|
||||||
import plot_binary_data
|
import plot_binary_data
|
||||||
import AMSS_NCKU_Input as input_data
|
import AMSS_NCKU_Input as input_data
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
import multiprocessing
|
||||||
|
|
||||||
# plt.rcParams['text.usetex'] = True ## enable LaTeX fonts in plots
|
# plt.rcParams['text.usetex'] = True ## enable LaTeX fonts in plots
|
||||||
|
|
||||||
@@ -50,10 +55,40 @@ def generate_binary_data_plot( binary_outdir, figure_outdir ):
|
|||||||
file_list.append(x)
|
file_list.append(x)
|
||||||
print(x)
|
print(x)
|
||||||
|
|
||||||
## Plot each file in the list
|
## Plot each file in parallel using subprocesses.
|
||||||
|
## Each subprocess is a fresh Python process where the BLAS thread-count
|
||||||
|
## environment variables (set at the top of plot_binary_data.py) take
|
||||||
|
## effect before numpy is imported. This avoids the thread explosion
|
||||||
|
## that occurs when multiprocessing.Pool with 'fork' context inherits
|
||||||
|
## already-initialized multi-threaded BLAS from the parent.
|
||||||
|
script = os.path.join( os.path.dirname(__file__), "plot_binary_data.py" )
|
||||||
|
max_workers = min( multiprocessing.cpu_count(), len(file_list) ) if file_list else 0
|
||||||
|
|
||||||
|
running = []
|
||||||
|
failed = []
|
||||||
for filename in file_list:
|
for filename in file_list:
|
||||||
print(filename)
|
print(filename)
|
||||||
plot_binary_data.plot_binary_data(filename, binary_outdir, figure_outdir)
|
proc = subprocess.Popen(
|
||||||
|
[sys.executable, script, filename, binary_outdir, figure_outdir],
|
||||||
|
)
|
||||||
|
running.append( (proc, filename) )
|
||||||
|
## Keep at most max_workers subprocesses active at a time
|
||||||
|
if len(running) >= max_workers:
|
||||||
|
p, fn = running.pop(0)
|
||||||
|
p.wait()
|
||||||
|
if p.returncode != 0:
|
||||||
|
failed.append(fn)
|
||||||
|
|
||||||
|
## Wait for all remaining subprocesses to finish
|
||||||
|
for p, fn in running:
|
||||||
|
p.wait()
|
||||||
|
if p.returncode != 0:
|
||||||
|
failed.append(fn)
|
||||||
|
|
||||||
|
if failed:
|
||||||
|
print( " WARNING: the following binary data plots failed:" )
|
||||||
|
for fn in failed:
|
||||||
|
print( " ", fn )
|
||||||
|
|
||||||
print( )
|
print( )
|
||||||
print( " Binary Data Plot Has been Finished " )
|
print( " Binary Data Plot Has been Finished " )
|
||||||
|
|||||||
Reference in New Issue
Block a user