Coverage for pySDC/projects/GPU/run_experiment.py: 0%
41 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-04-01 13:12 +0000
« prev ^ index » next coverage.py v7.8.0, created at 2025-04-01 13:12 +0000
1def parse_args():
2 import argparse
4 def cast_to_bool(me):
5 return False if me in ['False', '0', 0] else True
7 def str_to_procs(me):
8 procs = me.split('/')
9 assert len(procs) == 3
10 return [int(p) for p in procs]
12 parser = argparse.ArgumentParser()
13 parser.add_argument('--useGPU', type=cast_to_bool, help='Toggle for GPUs', default=False)
14 parser.add_argument(
15 '--mode',
16 type=str,
17 help='Mode for this script',
18 default=None,
19 choices=['run', 'plot', 'render', 'plot_series', 'video'],
20 )
21 parser.add_argument('--config', type=str, help='Configuration to load', default=None)
22 parser.add_argument('--restart_idx', type=int, help='Restart from file by index', default=0)
23 parser.add_argument('--procs', type=str_to_procs, help='Processes in steps/sweeper/space', default='1/1/1')
24 parser.add_argument('--res', type=int, help='Space resolution along first axis', default=-1)
25 parser.add_argument(
26 '--logger_level', type=int, help='Logger level on the first rank in space and in the sweeper', default=15
27 )
28 parser.add_argument('-o', type=str, help='output path', default='./')
30 return vars(parser.parse_args())
33def run_experiment(args, config, **kwargs):
34 import pickle
36 # from pySDC.implementations.controller_classes.controller_MPI import controller_MPI
37 from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
38 from pySDC.helpers.stats_helper import filter_stats
40 type(config).base_path = args['o']
41 description = config.get_description(
42 useGPU=args['useGPU'], MPIsweeper=args['procs'][1] > 1, res=args['res'], **kwargs
43 )
44 controller_params = config.get_controller_params(logger_level=args['logger_level'])
46 if args['useGPU']:
47 from pySDC.implementations.hooks.log_timings import GPUTimings
49 controller_params['hook_class'].append(GPUTimings)
51 # controller = controller_MPI(controller_params, description, config.comms[0])
52 assert (
53 config.comms[0].size == 1
54 ), 'Have not figured out how to do MPI controller with GPUs yet because I need NCCL for that!'
55 controller = controller_nonMPI(num_procs=1, controller_params=controller_params, description=description)
56 prob = controller.MS[0].levels[0].prob
58 u0, t0 = config.get_initial_condition(prob, restart_idx=args['restart_idx'])
60 config.prepare_caches(prob)
62 uend, stats = controller.run(u0=u0, t0=t0, Tend=config.Tend)
64 combined_stats = filter_stats(stats, comm=config.comm_world)
66 if config.comm_world.rank == config.comm_world.size - 1:
67 path = f'{config.base_path}/data/{config.get_path()}-stats-whole-run.pickle'
68 # path = f'/p/project1/ccstma/baumann7/pySDC/pySDC/projects/GPU/data/{config.get_path()}-stats-whole-run.pickle'
69 with open(path, 'wb') as file:
70 pickle.dump(combined_stats, file)
71 print(f'Stored stats in {path}', flush=True)
73 return uend
76def plot_experiment(args, config): # pragma: no cover
77 import gc
78 import matplotlib.pyplot as plt
80 type(config).base_path = args['o']
81 description = config.get_description()
83 P = description['problem_class'](**description['problem_params'])
85 comm = config.comm_world
87 for idx in range(args['restart_idx'], 9999, comm.size):
88 try:
89 fig = config.plot(P=P, idx=idx + comm.rank, n_procs_list=args['procs'])
90 except FileNotFoundError:
91 break
93 path = f'simulation_plots/{config.get_path(ranks=[0,0,0])}-{idx+comm.rank:06d}.png'
94 fig.savefig(path, dpi=300, bbox_inches='tight')
95 print(f'{comm.rank} Stored figure {path!r}', flush=True)
97 if args['mode'] == 'render':
98 plt.pause(1e-9)
100 plt.close(fig)
101 del fig
102 gc.collect()
105def plot_series(args, config): # pragma: no cover
106 import numpy as np
107 import gc
108 import matplotlib.pyplot as plt
109 from pySDC.helpers.plot_helper import figsize_by_journal, setup_mpl
111 setup_mpl()
113 fig, axs = plt.subplots(
114 3, 3, figsize=figsize_by_journal("TUHH_thesis", 1.0, 1.2), sharex=True, sharey=True, layout="compressed"
115 )
117 type(config).base_path = args['o']
118 description = config.get_description()
120 P = description['problem_class'](**description['problem_params'])
122 idxs = np.linspace(0, config.num_frames * 0.9, 9, dtype=int)
124 for idx, ax in zip(idxs, axs.flatten()):
125 try:
126 _fig = config.plot(P=P, idx=idx, n_procs_list=args['procs'], ax=ax)
127 except FileNotFoundError:
128 break
130 plt.close(_fig)
131 del _fig
132 gc.collect()
134 for ax in axs.flatten():
135 if ax != axs[-1, 0]:
136 ax.set_xlabel('')
137 ax.set_ylabel('')
139 path = f'plots/{config.get_path(ranks=[0,0,0])}-series.pdf'
140 fig.savefig(path, dpi=300)
141 print(f'Stored figure {path!r}', flush=True)
142 plt.show()
145def make_video(args, config): # pragma: no cover
146 comm = config.comm_world
147 if comm.rank > 0:
148 return None
150 import subprocess
152 path = f'simulation_plots/{config.get_path(ranks=[0,0,0])}-%06d.png'
153 path_target = f'videos/{args["config"]}.mp4'
155 cmd = f'ffmpeg -i {path} -pix_fmt yuv420p -r 9 -s 2048:1536 -y {path_target}'.split()
157 subprocess.run(cmd)
160if __name__ == '__main__':
161 from pySDC.projects.GPU.configs.base_config import get_config
163 args = parse_args()
165 config = get_config(args)
167 if args['mode'] == 'run':
168 run_experiment(args, config)
169 elif args['mode'] in ['plot', 'render']: # pragma: no cover
170 plot_experiment(args, config)
171 elif args['mode'] == 'plot_series': # pragma: no cover
172 plot_series(args, config)
173 elif args['mode'] == 'video': # pragma: no cover
174 make_video(args, config)