From b656b99eaca655947942f0af5943828741fdc4a5 Mon Sep 17 00:00:00 2001 From: Daniel Zuegner Date: Tue, 10 Feb 2026 08:14:46 +0000 Subject: [PATCH 1/3] enable MD progress bar --- torch_sim/runners.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index 5dbd0f39..69811d4b 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -345,12 +345,15 @@ def integrate[T: SimState]( # noqa: C901 final_states: list[T] = [] og_filenames = trajectory_reporter.filenames if trajectory_reporter else None + # Extract pbar kwargs for both outer and inner progress bars + pbar_kwargs = pbar if isinstance(pbar, dict) else {} + tqdm_pbar = None if pbar and autobatcher: - pbar_kwargs = pbar if isinstance(pbar, dict) else {} - pbar_kwargs.setdefault("desc", "Integrate") - pbar_kwargs.setdefault("disable", None) - tqdm_pbar = tqdm(total=initial_state.n_systems, **pbar_kwargs) + outer_pbar_kwargs = pbar_kwargs.copy() + outer_pbar_kwargs.setdefault("desc", "Integrate") + outer_pbar_kwargs.setdefault("disable", None) + tqdm_pbar = tqdm(total=initial_state.n_systems, **outer_pbar_kwargs) # Handle both BinningAutoBatcher and list of tuples for state, system_indices in batch_iterator: @@ -373,7 +376,15 @@ def integrate[T: SimState]( # noqa: C901 _write_initial_state(trajectory_reporter, state, model) # run the simulation - for step in range(initial_step, initial_step + n_steps): + inner_pbar_kwargs = pbar_kwargs.copy() + inner_pbar_kwargs.setdefault("disable", pbar is False) + inner_pbar_kwargs.setdefault( + "desc", f"Batch systems {system_indices}" if system_indices else "Integrate" + ) + inner_pbar_kwargs.setdefault("leave", False) + for step in tqdm( + range(initial_step, initial_step + n_steps), **inner_pbar_kwargs + ): state = step_func( state=state, model=model, From 163a08aaf93418035ac896bfbee8ce75eaaba351 Mon Sep 17 00:00:00 2001 From: Daniel Zuegner Date: Tue, 10 Feb 2026 08:32:10 +0000 Subject: [PATCH 2/3] update comment --- torch_sim/runners.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index 69811d4b..b4332f3c 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -27,7 +27,6 @@ from torch_sim.typing import StateLike from torch_sim.units import UnitSystem - logger = logging.getLogger(__name__) @@ -298,8 +297,7 @@ def integrate[T: SimState]( # noqa: C901 tracking trajectory. If a dict, will be passed to the TrajectoryReporter constructor. autobatcher (BinningAutoBatcher | bool): Optional autobatcher to use - pbar (bool | dict[str, Any], optional): Show a progress bar. - Only works with an autobatcher in interactive shell. If a dict is given, + pbar (bool | dict[str, Any], optional): Show a progress bar. If a dict is given, it's passed to `tqdm` as kwargs. init_kwargs (dict[str, Any], optional): Additional keyword arguments for integrator init function. From 03585443531da5a9c5e39553d575e3e735c081a8 Mon Sep 17 00:00:00 2001 From: Daniel Zuegner Date: Tue, 10 Feb 2026 08:39:50 +0000 Subject: [PATCH 3/3] style --- torch_sim/runners.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index b4332f3c..8f277289 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -27,6 +27,7 @@ from torch_sim.typing import StateLike from torch_sim.units import UnitSystem + logger = logging.getLogger(__name__)