feat: add argument to show progress bar

This commit is contained in:
Anthony Berg 2025-08-07 23:59:57 +02:00
parent 5ba88b81d6
commit 8da34a01f8
2 changed files with 6 additions and 4 deletions

View File

@ -84,7 +84,7 @@ def to_json(in_dict, compressed=True):
return json.dumps(out_dict) 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 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 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() 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 # Start simulation loop
for save_step, t_step in enumerate(t_steps): for save_step, t_step in enumerate(t_steps):
t_end = save_step t_end = save_step

View File

@ -45,6 +45,8 @@ parser.add_argument('-nx', type=int, default=128)
parser.add_argument('-ny', type=int, default=128) parser.add_argument('-ny', type=int, default=128)
parser.add_argument('--profile', action='store_true') # default: False 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('--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() args = parser.parse_args()
@ -128,7 +130,6 @@ compile_opts = args.compile_opts
if compile_opts is not None: if compile_opts is not None:
arguments['compile_opts'] += compile_opts arguments['compile_opts'] += compile_opts
if args.profile: if args.profile:
t_init_end = time.time() t_init_end = time.time()
t_init = t_init_end - t_init_start t_init = t_init_end - t_init_start
@ -139,6 +140,7 @@ if args.profile:
#### ####
logger.info("Running simulation") logger.info("Running simulation")
# Helper function to create MPI simulator # 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( 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. # Move NetCDF4 file to a unique file, for the next run.
if rank == 0: if rank == 0: