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', type=str, help='Mode for this script', default=None, choices=['run', 'plot', 'render', 'video'] 

16 ) 

17 parser.add_argument('--config', type=str, help='Configuration to load', default=None) 

18 parser.add_argument('--restart_idx', type=int, help='Restart from file by index', default=0) 

19 parser.add_argument('--procs', type=str_to_procs, help='Processes in steps/sweeper/space', default='1/1/1') 

20 parser.add_argument('--res', type=int, help='Space resolution along first axis', default=-1) 

21 parser.add_argument( 

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

23 ) 

24 

25 return vars(parser.parse_args()) 

26 

27 

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

29 import pickle 

30 

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

32 from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

33 from pySDC.helpers.stats_helper import filter_stats 

34 

35 description = config.get_description( 

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

37 ) 

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

39 

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

41 assert ( 

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

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

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

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

46 

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

48 

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

50 

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

52 

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

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

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

56 pickle.dump(combined_stats, file) 

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

58 

59 return uend 

60 

61 

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

63 import gc 

64 import matplotlib.pyplot as plt 

65 

66 description = config.get_description() 

67 

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

69 

70 comm = config.comm_world 

71 

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

73 try: 

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

75 except FileNotFoundError: 

76 break 

77 

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

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

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

81 

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

83 plt.pause(1e-9) 

84 

85 plt.close(fig) 

86 del fig 

87 gc.collect() 

88 

89 

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

91 comm = config.comm_world 

92 if comm.rank > 0: 

93 return None 

94 

95 import subprocess 

96 

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

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

99 

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

101 

102 subprocess.run(cmd) 

103 

104 

105if __name__ == '__main__': 

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

107 

108 args = parse_args() 

109 

110 config = get_config(args) 

111 

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

113 run_experiment(args, config) 

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

115 plot_experiment(args, config) 

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

117 make_video(args, config)