From 8da34a01f85c8903d45692d182dd581c6b4c5a63 Mon Sep 17 00:00:00 2001 From: Anthony Berg Date: Thu, 7 Aug 2025 23:59:57 +0200 Subject: [PATCH] feat: add argument to show progress bar --- GPUSimulators/common/common.py | 4 ++-- mpi_testing_hip.py | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/GPUSimulators/common/common.py b/GPUSimulators/common/common.py index e99bf80..71e166b 100644 --- a/GPUSimulators/common/common.py +++ b/GPUSimulators/common/common.py @@ -84,7 +84,7 @@ def to_json(in_dict, compressed=True): return json.dumps(out_dict) -def run_simulation(simulator, simulator_args, outfile, save_times, save_var_names=[], dt=None): +def run_simulation(simulator, simulator_args, outfile, save_times, save_var_names=[], dt=None, progress_bar=False): """ Runs a simulation, and store output in a netcdf file. Stores the times given in save_times, and saves all the variables in list save_var_names. Elements in @@ -181,7 +181,7 @@ def run_simulation(simulator, simulator_args, outfile, save_times, save_var_name profiling_data_sim_runner["end"]["t_sim_init"] = time.time() - with tqdm(total=save_times[-1], desc="Simulation progress", unit="sim s") as pbar: + with tqdm(total=save_times[-1], desc="Simulation progress", unit="sim s", disable=not progress_bar) as pbar: # Start simulation loop for save_step, t_step in enumerate(t_steps): t_end = save_step diff --git a/mpi_testing_hip.py b/mpi_testing_hip.py index 4a4ca35..8881699 100644 --- a/mpi_testing_hip.py +++ b/mpi_testing_hip.py @@ -45,6 +45,8 @@ parser.add_argument('-nx', type=int, default=128) parser.add_argument('-ny', type=int, default=128) parser.add_argument('--profile', action='store_true') # default: False parser.add_argument('--compile_opts', type=str, help="Compiler options for HIP code.") +parser.add_argument('--progress', type=bool, default=False, + help="Displays a progress bar for the progress of the simulation.") args = parser.parse_args() @@ -128,7 +130,6 @@ compile_opts = args.compile_opts if compile_opts is not None: arguments['compile_opts'] += compile_opts - if args.profile: t_init_end = time.time() t_init = t_init_end - t_init_start @@ -139,6 +140,7 @@ if args.profile: #### logger.info("Running simulation") + # Helper function to create MPI simulator @@ -149,7 +151,7 @@ def gen_sim(grid, **kwargs): outfile, sim_runner_profiling_data, sim_profiling_data = run_simulation( - gen_sim, arguments, outfile, save_times, save_var_names, dt) + gen_sim, arguments, outfile, save_times, save_var_names, dt, progress_bar=args.progress) # Move NetCDF4 file to a unique file, for the next run. if rank == 0: