Added script for running experiment to RBC project

Author: brownbaerchen

pySDC/projects/RayleighBenard/analysis_scripts/RBC3D_order.py

@@ -2,7 +2,7 @@
 import pickle
 import numpy as np
 from pySDC.helpers.fieldsIO import FieldsIO
-from pySDC.projects.GPU.configs.base_config import get_config
+from pySDC.projects.RayleighBenard.RBC3D_configs import get_config
 from pySDC.implementations.problem_classes.RayleighBenard3D import RayleighBenard3D
 from mpi4py import MPI
 import matplotlib.pyplot as plt
@@ -103,7 +103,7 @@ def compare_order(Ra):  # pragma: no cover
 
 
 def run(args, dt, Tend):
-    from pySDC.projects.GPU.run_experiment import run_experiment
+    from pySDC.projects.RayleighBenard.run_experiment import run_experiment
     from pySDC.core.errors import ConvergenceError
 
     args['mode'] = 'run'
@@ -134,7 +134,7 @@ def run(args, dt, Tend):
 
 
 if __name__ == '__main__':
-    from pySDC.projects.GPU.run_experiment import parse_args
+    from pySDC.projects.RayleighBenard.run_experiment import parse_args
 
     args = parse_args()
 

pySDC/projects/RayleighBenard/analysis_scripts/process_RBC3D_data.py

@@ -1,5 +1,5 @@
-from pySDC.projects.GPU.configs.base_config import get_config
-from pySDC.projects.GPU.run_experiment import parse_args
+from pySDC.projects.RayleighBenard.RBC3D_configs import get_config
+from pySDC.projects.RayleighBenard.run_experiment import parse_args
 from pySDC.helpers.fieldsIO import FieldsIO
 import matplotlib.pyplot as plt
 from tqdm import tqdm

pySDC/projects/RayleighBenard/run_experiment.py

@@ -0,0 +1,87 @@
+def parse_args():
+    import argparse
+
+    def cast_to_bool(me):
+        return False if me in ['False', '0', 0] else True
+
+    def str_to_procs(me):
+        procs = me.split('/')
+        assert len(procs) == 3
+        return [int(p) for p in procs]
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--useGPU', type=cast_to_bool, help='Toggle for GPUs', default=False)
+    parser.add_argument(
+        '--mode',
+        type=str,
+        help='Mode for this script',
+        default='run',
+        choices=['run'],
+    )
+    parser.add_argument('--config', type=str, help='Configuration to load', default=None)
+    parser.add_argument('--restart_idx', type=int, help='Restart from file by index', default=0)
+    parser.add_argument('--procs', type=str_to_procs, help='Processes in steps/sweeper/space', default='1/1/1')
+    parser.add_argument('--res', type=int, help='Space resolution along first axis', default=-1)
+    parser.add_argument('--dt', type=float, help='(Starting) Step size', default=-1)
+    parser.add_argument(
+        '--logger_level', type=int, help='Logger level on the first rank in space and in the sweeper', default=15
+    )
+    parser.add_argument('-o', type=str, help='output path', default='./')
+
+    return vars(parser.parse_args())
+
+
+def run_experiment(args, config, **kwargs):
+    import pickle
+    import os
+
+    from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
+    from pySDC.helpers.stats_helper import filter_stats
+
+    type(config).base_path = args['o']
+    os.makedirs(f'{args["o"]}/data', exist_ok=True)
+
+    description = config.get_description(
+        useGPU=args['useGPU'], MPIsweeper=args['procs'][1] > 1, res=args['res'], dt=args['dt'], **kwargs
+    )
+    controller_params = config.get_controller_params(logger_level=args['logger_level'])
+
+    if args['useGPU']:
+        from pySDC.implementations.hooks.log_timings import GPUTimings
+
+        controller_params['hook_class'].append(GPUTimings)
+
+    assert (
+        config.comms[0].size == 1
+    ), 'Have not figured out how to do MPI controller with GPUs yet because I need NCCL for that!'
+    controller = controller_nonMPI(num_procs=1, controller_params=controller_params, description=description)
+    prob = controller.MS[0].levels[0].prob
+
+    u0, t0 = config.get_initial_condition(prob, restart_idx=args['restart_idx'])
+
+    config.prepare_caches(prob)
+
+    uend, stats = controller.run(u0=u0, t0=t0, Tend=config.Tend)
+
+    combined_stats = filter_stats(stats, comm=config.comm_world)
+
+    if config.comm_world.rank == config.comm_world.size - 1:
+        path = f'{config.base_path}/data/{config.get_path()}-stats-whole-run.pickle'
+        with open(path, 'wb') as file:
+            pickle.dump(combined_stats, file)
+        print(f'Stored stats in {path}', flush=True)
+
+    return uend
+
+
+if __name__ == '__main__':
+    from pySDC.projects.RayleighBenard.RBC3D_configs import get_config
+
+    args = parse_args()
+
+    config = get_config(args)
+
+    if args['mode'] == 'run':
+        run_experiment(args, config)
+    else:
+        raise NotImplementedError

pySDC/projects/RayleighBenard/tests/test_RBC_3D_analysis.py

@@ -18,7 +18,7 @@ def get_args(path):
 
 
 def get_config(args):
-    from pySDC.projects.GPU.configs.base_config import get_config
+    from pySDC.projects.RayleighBenard.RBC3D_configs import get_config
 
     config = get_config(args)
     config.Tend = 1
@@ -27,7 +27,7 @@ def get_config(args):
 
 
 def generate_simulation_file(path, args=None):
-    from pySDC.projects.GPU.run_experiment import run_experiment
+    from pySDC.projects.RayleighBenard.run_experiment import run_experiment
 
     args = {**get_args(path), **args} if args is not None else get_args(path)
     config = get_config(args)
@@ -56,7 +56,7 @@ def tmp_processed_data(tmp_sim_data, tmp_path):
 
 
 def test_ic_interpolation(tmp_sim_data, tmp_path):
-    from pySDC.projects.GPU.run_experiment import run_experiment
+    from pySDC.projects.RayleighBenard.run_experiment import run_experiment
 
     args = get_args(tmp_path)