diff --git a/kwave/kspaceFirstOrder.py b/kwave/kspaceFirstOrder.py index 3a4e9eac..2777c830 100644 --- a/kwave/kspaceFirstOrder.py +++ b/kwave/kspaceFirstOrder.py @@ -203,6 +203,7 @@ def kspaceFirstOrder( smooth_p0=False, pml_size=pml_size, pml_alpha=pml_alpha, + quiet=quiet, ).run() elif backend == "cpp": diff --git a/kwave/solvers/kspace_solver.py b/kwave/solvers/kspace_solver.py index b94f4bf2..4991ea91 100644 --- a/kwave/solvers/kspace_solver.py +++ b/kwave/solvers/kspace_solver.py @@ -7,6 +7,7 @@ from types import SimpleNamespace import numpy as np +from tqdm import tqdm try: import cupy as cp @@ -117,7 +118,19 @@ class Simulation: """ def __init__( - self, kgrid, medium, source, sensor, *, device="cpu", use_sg=True, use_kspace=True, smooth_p0=True, pml_size=None, pml_alpha=None + self, + kgrid, + medium, + source, + sensor, + *, + device="cpu", + use_sg=True, + use_kspace=True, + smooth_p0=True, + pml_size=None, + pml_alpha=None, + quiet=False, ): self.kgrid = kgrid self.medium = medium @@ -126,6 +139,7 @@ def __init__( self.use_sg = use_sg self.use_kspace = use_kspace self.smooth_p0 = smooth_p0 + self.quiet = quiet self._pml_size_override = pml_size self._pml_alpha_override = pml_alpha # kWaveGrid doesn't have pml_size_x attrs; warn if PML will silently be disabled @@ -632,8 +646,15 @@ def run(self): """Run simulation to completion. Returns results dict.""" if not self._is_setup: self.setup() - while self.t < self.Nt: - self.step() + remaining = self.Nt - self.t + if not self.quiet and remaining > 0: + with tqdm(total=remaining, desc="k-Wave", unit="step") as pbar: + while self.t < self.Nt: + self.step() + pbar.update(1) + else: + while self.t < self.Nt: + self.step() # Copy to CPU one-by-one, freeing GPU memory as we go result = {} for k in list(self.sensor_data): diff --git a/pyproject.toml b/pyproject.toml index f61d7313..1bd5d21f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,8 @@ dependencies = [ "matplotlib==3.10.7", "beartype==0.22.9", "jaxtyping==0.3.2", - "deprecated>=1.2.14" + "deprecated>=1.2.14", + "tqdm>=4.60" ] [project.urls] diff --git a/tests/test_native_solver.py b/tests/test_native_solver.py index 47a4dffe..30091149 100644 --- a/tests/test_native_solver.py +++ b/tests/test_native_solver.py @@ -240,6 +240,32 @@ def test_cartesian_1d(self, grid_1d): assert result["p"].shape == (3, 20) +class TestProgressBar: + def test_progress_bar_runs(self, grid_2d): + """quiet=False (default) shows tqdm progress bar without error.""" + result = kspaceFirstOrder( + grid_2d, + kWaveMedium(sound_speed=1500), + _p0_source((64, 64)), + kSensor(mask=np.ones((64, 64), dtype=bool)), + backend="python", + quiet=False, + ) + assert "p" in result + + def test_quiet_suppresses_bar(self, grid_2d): + """quiet=True suppresses progress bar.""" + result = kspaceFirstOrder( + grid_2d, + kWaveMedium(sound_speed=1500), + _p0_source((64, 64)), + kSensor(mask=np.ones((64, 64), dtype=bool)), + backend="python", + quiet=True, + ) + assert "p" in result + + class TestMatlabInterop: def test_simulate_from_dicts(self): from kwave.solvers.kspace_solver import simulate_from_dicts diff --git a/uv.lock b/uv.lock index 6491700c..0c28b447 100644 --- a/uv.lock +++ b/uv.lock @@ -770,6 +770,7 @@ dependencies = [ { name = "numpy" }, { name = "opencv-python" }, { name = "scipy" }, + { name = "tqdm" }, ] [package.optional-dependencies] @@ -824,6 +825,7 @@ requires-dist = [ { name = "sphinx-tabs", marker = "extra == 'docs'", specifier = "==3.4.7" }, { name = "sphinx-toolbox", marker = "extra == 'docs'", specifier = "==3.8.0" }, { name = "testfixtures", marker = "extra == 'test'", specifier = "==8.3.0" }, + { name = "tqdm", specifier = ">=4.60" }, ] [package.metadata.requires-dev]