feat(simulator): make tqdm count all time in simulation

This commit is contained in:
Anthony Berg 2025-07-03 14:54:03 +02:00
parent cc9937c752
commit 3fdd4ab62b
2 changed files with 44 additions and 51 deletions

View File

@ -21,15 +21,15 @@ along with this program. If not, see <http://www.gnu.org/licenses/>.
"""
import os
import numpy as np
import time
import subprocess
import logging
import json
import numpy as np
from tqdm.auto import tqdm
from GPUSimulators.common.data_dumper import DataDumper
from GPUSimulators.common.progress_printer import ProgressPrinter
from GPUSimulators.common.timer import Timer
@ -160,43 +160,36 @@ def run_simulation(simulator, simulator_args, outfile, save_times, save_var_name
profiling_data_sim_runner["end"]["t_sim_init"] = time.time()
# Start simulation loop
progress_printer = ProgressPrinter(save_times[-1], print_every=10)
for k in range(len(save_times)):
# Get target time and step size there
t_step = t_steps[k]
t_end = save_times[k]
# Sanity check simulator
try:
sim.check()
except AssertionError as e:
logger.error(f"Error after {sim.sim_steps()} steps (t={sim.sim_time()}: {str(e)}")
return outdata.filename
with tqdm(total=save_times[-1], desc="Simulation progress", unit="sim s") as pbar:
# Start simulation loop
for k, t_step in enumerate(t_steps):
t_end = k
profiling_data_sim_runner["start"]["t_full_step"] += time.time()
# Sanity check simulator
try:
sim.check()
except AssertionError as e:
logger.error(f"Error after {sim.sim_steps()} steps (t={sim.sim_time()}: {str(e)}")
return outdata.filename
# Simulate
if t_step > 0.0:
sim.simulate(t_step, dt)
profiling_data_sim_runner["start"]["t_full_step"] += time.time()
profiling_data_sim_runner["end"]["t_full_step"] += time.time()
# Simulate
if t_step > 0.0:
sim.simulate(t_step, dt, pbar=pbar)
profiling_data_sim_runner["start"]["t_nc_write"] += time.time()
profiling_data_sim_runner["end"]["t_full_step"] += time.time()
#Download
save_vars = sim.download(download_vars)
#Save to file
for i, var_name in enumerate(save_var_names):
ncvars[var_name][k, :] = save_vars[i]
profiling_data_sim_runner["start"]["t_nc_write"] += time.time()
profiling_data_sim_runner["end"]["t_nc_write"] += time.time()
#Download
save_vars = sim.download(download_vars)
#Write progress to screen
print_string = progress_printer.get_print_string(t_end)
if print_string:
logger.debug(print_string)
#Save to file
for i, var_name in enumerate(save_var_names):
ncvars[var_name][k, :] = save_vars[i]
profiling_data_sim_runner["end"]["t_nc_write"] += time.time()
logger.debug(f"Simulated to t={t_end} in "
+ f"{sim.sim_steps()} timesteps (average dt={sim.sim_time() / sim.sim_steps()})")

View File

@ -2,7 +2,6 @@ import logging
import math
import numpy as np
from tqdm.auto import tqdm
from . import boundary
from GPUSimulators.gpu import KernelContext
@ -84,7 +83,7 @@ class BaseSimulator(object):
def __str__(self):
return f"{self.__class__.__name__} [{self.nx}x{self.ny}]"
def simulate(self, t, dt=None, tolerance=None):
def simulate(self, t, dt=None, tolerance=None, pbar=None):
"""
Function which simulates t_end seconds using the step function
Requires that the step() function is implemented in the subclasses
@ -93,6 +92,7 @@ class BaseSimulator(object):
t: How long the simulation should run for.
dt: Time steps.
tolerance: How small should the time steps be before considering it an infinite loop.
pbar: A tqdm progress bar to update time.
"""
t_start = self.sim_time()
@ -106,27 +106,27 @@ class BaseSimulator(object):
if tolerance is None:
tolerance = 0.000000001
with tqdm(total=t, desc="Running Simulator", leave=False) as pbar:
while self.sim_time() < t_end:
# Prevent an infinite loop from occurring from tiny numbers
if abs(t_end - self.sim_time()) < tolerance:
break
while self.sim_time() < t_end:
# Prevent an infinite loop from occurring from tiny numbers
if abs(t_end - self.sim_time()) < tolerance:
break
if update_dt and (self.sim_steps() % 100 == 0):
self.dt = self.compute_dt() * self.cfl_scale
if update_dt and (self.sim_steps() % 100 == 0):
self.dt = self.compute_dt() * self.cfl_scale
# Compute timestep for "this" iteration (i.e., shorten last timestep)
current_dt = np.float32(min(self.dt, t_end - self.sim_time()))
# Compute timestep for "this" iteration (i.e., shorten last timestep)
current_dt = np.float32(min(self.dt, t_end - self.sim_time()))
# Stop if end reached (should not happen)
if current_dt <= 0.0:
self.logger.warning(f"Timestep size {self.sim_steps()} is less than or equal to zero!")
break
# Stop if end reached (should not happen)
if current_dt <= 0.0:
self.logger.warning(f"Timestep size {self.sim_steps()} is less than or equal to zero!")
break
# Step forward in time
self.step(current_dt)
# Step forward in time
self.step(current_dt)
# Update the progress bar
# Update the progress bar
if pbar is not None:
pbar.update(float(current_dt))
def step(self, dt: int):