Coverage for pySDC / projects / RayleighBenard / run_experiment.py: 53%

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', 'benchmark'], 

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', 'time_first'], 

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_nonMPI import controller_nonMPI 

46 from pySDC.helpers.stats_helper import filter_stats 

47 

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

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

50 

51 if args['mode'] == 'benchmark': 

52 config.prepare_for_benchmark() 

53 

54 description = config.get_description( 

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

56 ) 

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

58 

59 if args['mode'] == 'benchmark': 

60 config.prepare_description_for_benchmark(description, controller_params) 

61 

62 if args['useGPU']: 

63 from pySDC.implementations.hooks.log_timings import GPUTimings 

64 

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

66 

67 assert ( 

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

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

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

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

72 

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

74 

75 if args['mode'] == 'benchmark': 

76 config.prepare_caches_for_benchmark(prob, controller) 

77 

78 config.prepare_caches(prob) 

79 

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

81 

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

83 

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

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

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

87 pickle.dump(combined_stats, file) 

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

89 

90 return uend 

91 

92 

93if __name__ == '__main__': 

94 from pySDC.projects.RayleighBenard.RBC3D_configs import get_config 

95 

96 args = parse_args() 

97 

98 config = get_config(args) 

99 

100 if args['mode'] in ['run', 'benchmark']: 

101 run_experiment(args, config) 

102 else: 

103 raise NotImplementedError