Coverage for pySDC / projects / GPU / run_experiment.py: 0%

1def parse_args(): 

2 import argparse 

3 

4 def cast_to_bool(me): 

5 return False if me in ['False', '0', 0] else True 

6 

7 def str_to_procs(me): 

8 procs = me.split('/') 

9 assert len(procs) == 3 

10 return [int(p) for p in procs] 

11 

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('--dt', type=float, help='(Starting) Step size', default=-1) 

26 parser.add_argument( 

27 '--logger_level', type=int, help='Logger level on the first rank in space and in the sweeper', default=15 

28 ) 

29 parser.add_argument('-o', type=str, help='output path', default='./') 

30 parser.add_argument( 

31 '--distribution', 

32 type=str, 

33 help='distribute tasks', 

34 default='space_first', 

35 choices=['space_first', 'space_major', 'time_first', 'time_major'], 

36 ) 

37 

38 return vars(parser.parse_args()) 

39 

40 

41def run_experiment(args, config, **kwargs): 

42 import pickle 

43 import os 

44 

45 # from pySDC.implementations.controller_classes.controller_MPI import controller_MPI 

46 from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

47 from pySDC.helpers.stats_helper import filter_stats 

48 

49 type(config).base_path = args['o'] 

50 os.makedirs(f'{args["o"]}/data', exist_ok=True) 

51 

52 description = config.get_description( 

53 useGPU=args['useGPU'], MPIsweeper=args['procs'][1] > 1, res=args['res'], dt=args['dt'], **kwargs 

54 ) 

55 controller_params = config.get_controller_params(logger_level=args['logger_level']) 

56 

57 if args['useGPU']: 

58 from pySDC.implementations.hooks.log_timings import GPUTimings 

59 

60 controller_params['hook_class'].append(GPUTimings) 

61 

62 # controller = controller_MPI(controller_params, description, config.comms[0]) 

63 assert ( 

64 config.comms[0].size == 1 

65 ), 'Have not figured out how to do MPI controller with GPUs yet because I need NCCL for that!' 

66 controller = controller_nonMPI(num_procs=1, controller_params=controller_params, description=description) 

67 prob = controller.MS[0].levels[0].prob 

68 

69 u0, t0 = config.get_initial_condition(prob, restart_idx=args['restart_idx']) 

70 

71 config.prepare_caches(prob) 

72 

73 uend, stats = controller.run(u0=u0, t0=t0, Tend=config.Tend) 

74 

75 combined_stats = filter_stats(stats, comm=config.comm_world) 

76 

77 if config.comm_world.rank == config.comm_world.size - 1: 

78 path = f'{config.base_path}/data/{config.get_path()}-stats-whole-run.pickle' 

79 # path = f'/p/project1/ccstma/baumann7/pySDC/pySDC/projects/GPU/data/{config.get_path()}-stats-whole-run.pickle' 

80 with open(path, 'wb') as file: 

81 pickle.dump(combined_stats, file) 

82 print(f'Stored stats in {path}', flush=True) 

83 

84 return uend 

85 

86 

87def plot_experiment(args, config): # pragma: no cover 

88 import gc 

89 import matplotlib.pyplot as plt 

90 

91 type(config).base_path = args['o'] 

92 description = config.get_description() 

93 

94 P = description['problem_class'](**description['problem_params']) 

95 

96 comm = config.comm_world 

97 

98 for idx in range(args['restart_idx'], 9999, comm.size): 

99 try: 

100 fig = config.plot(P=P, idx=idx + comm.rank, n_procs_list=args['procs']) 

101 except FileNotFoundError: 

102 break 

103 

104 path = f'simulation_plots/{config.get_path(ranks=[0,0,0])}-{idx+comm.rank:06d}.png' 

105 fig.savefig(path, dpi=300, bbox_inches='tight') 

106 print(f'{comm.rank} Stored figure {path!r}', flush=True) 

107 

108 if args['mode'] == 'render': 

109 plt.pause(1e-9) 

110 

111 plt.close(fig) 

112 del fig 

113 gc.collect() 

114 

115 

116def plot_series(args, config): # pragma: no cover 

117 import numpy as np 

118 import gc 

119 import matplotlib.pyplot as plt 

120 from pySDC.helpers.plot_helper import figsize_by_journal, setup_mpl 

121 

122 setup_mpl() 

123 

124 fig, axs = plt.subplots( 

125 3, 3, figsize=figsize_by_journal("TUHH_thesis", 1.0, 1.2), sharex=True, sharey=True, layout="compressed" 

126 ) 

127 

128 type(config).base_path = args['o'] 

129 description = config.get_description() 

130 

131 P = description['problem_class'](**description['problem_params']) 

132 

133 idxs = np.linspace(0, config.num_frames * 0.9, 9, dtype=int) 

134 

135 for idx, ax in zip(idxs, axs.flatten(), strict=True): 

136 try: 

137 _fig = config.plot(P=P, idx=idx, n_procs_list=args['procs'], ax=ax) 

138 except FileNotFoundError: 

139 break 

140 

141 plt.close(_fig) 

142 del _fig 

143 gc.collect() 

144 

145 for ax in axs.flatten(): 

146 if ax != axs[-1, 0]: 

147 ax.set_xlabel('') 

148 ax.set_ylabel('') 

149 

150 path = f'plots/{config.get_path(ranks=[0,0,0])}-series.pdf' 

151 fig.savefig(path, dpi=300) 

152 print(f'Stored figure {path!r}', flush=True) 

153 plt.show() 

154 

155 

156def make_video(args, config): # pragma: no cover 

157 comm = config.comm_world 

158 if comm.rank > 0: 

159 return None 

160 

161 import subprocess 

162 

163 path = f'simulation_plots/{config.get_path(ranks=[0,0,0])}-%06d.png' 

164 path_target = f'videos/{args["config"]}.mp4' 

165 

166 cmd = f'ffmpeg -i {path} -pix_fmt yuv420p -r 9 -s 2048:1536 -y {path_target}'.split() 

167 

168 subprocess.run(cmd) 

169 print(f'Made video {path_target!r}') 

170 

171 

172if __name__ == '__main__': 

173 from pySDC.projects.GPU.configs.base_config import get_config 

174 

175 args = parse_args() 

176 

177 config = get_config(args) 

178 

179 if args['mode'] == 'run': 

180 run_experiment(args, config) 

181 elif args['mode'] in ['plot', 'render']: # pragma: no cover 

182 plot_experiment(args, config) 

183 elif args['mode'] == 'plot_series': # pragma: no cover 

184 plot_series(args, config) 

185 elif args['mode'] == 'video': # pragma: no cover 

186 make_video(args, config)