From 58817f641f55c12e653b075e722d64ba79db6478 Mon Sep 17 00:00:00 2001 From: dutta-alankar Date: Mon, 4 May 2026 22:54:02 +0200 Subject: [PATCH 1/4] [ENH] replacement of `tqdm` progress bar with a modern `rich` progress bar --- binder/environment.yml | 2 +- pyproject.toml | 2 +- setup.py | 2 +- src/emcee/ensemble.py | 7 +--- src/emcee/pbar.py | 66 ++++++++++++++++++++++++------- src/emcee/tests/unit/test_pbar.py | 24 +++++------ 6 files changed, 66 insertions(+), 37 deletions(-) diff --git a/binder/environment.yml b/binder/environment.yml index 81495815..1c8e4d5f 100644 --- a/binder/environment.yml +++ b/binder/environment.yml @@ -8,7 +8,7 @@ dependencies: - h5py - matplotlib - corner - - tqdm + - rich - mpi4py - schwimmbad - pip diff --git a/pyproject.toml b/pyproject.toml index 6e314f05..bf9f26a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,5 +55,5 @@ exclude_lines = [ "logging.warning", "deprecation_warning", "deprecated", - "if tqdm is None" + "if Progress is None" ] diff --git a/setup.py b/setup.py index e2b1a91f..dc9a3dc7 100755 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ "wheel", ] EXTRA_REQUIRE = { - "extras": ["h5py", "scipy", "tqdm", "ipywidgets"], + "extras": ["h5py", "scipy", "rich", "ipywidgets"], "tests": ["pytest", "pytest-cov", "coverage[toml]"], } diff --git a/src/emcee/ensemble.py b/src/emcee/ensemble.py index c71f6ee5..0ceb0211 100644 --- a/src/emcee/ensemble.py +++ b/src/emcee/ensemble.py @@ -290,13 +290,10 @@ def sample( a file or if you don't need to analyze the samples after the fact (for burn-in for example) set ``store`` to ``False``. progress (Optional[bool or str]): If ``True``, a progress bar will - be shown as the sampler progresses. If a string, will select a - specific ``tqdm`` progress bar - most notable is - ``'notebook'``, which shows a progress bar suitable for - Jupyter notebooks. If ``False``, no progress bar will be + be shown as the sampler progresses. If ``False``, no progress bar will be shown. progress_kwargs (Optional[dict]): A ``dict`` of keyword arguments - to be passed to the tqdm call. + to be passed to the progress bar implementation. skip_initial_state_check (Optional[bool]): If ``True``, a check that the initial_state can fully explore the space will be skipped. (default: ``False``) diff --git a/src/emcee/pbar.py b/src/emcee/pbar.py index a0636349..7ec0df5e 100644 --- a/src/emcee/pbar.py +++ b/src/emcee/pbar.py @@ -1,6 +1,5 @@ # -*- coding: utf-8 -*- -import importlib import logging __all__ = ["get_progress_bar"] @@ -8,10 +7,10 @@ logger = logging.getLogger(__name__) try: - import tqdm - import tqdm.auto + from rich.console import Console + from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn except ImportError: - tqdm = None + Progress = None class _NoOpPBar(object): @@ -30,31 +29,68 @@ def update(self, count): pass +class _RichPBar(object): + """A wrapper that provides emcee's progress-bar interface over rich.""" + + def __init__(self, total, **kwargs): + self.total = total + self.description = kwargs.pop("desc", "Sampling") + leave = kwargs.pop("leave", True) + self.progress = None + self.task_id = None + + # leave=False means clearing the bar when complete. + self.transient = not leave + + # Preserve legacy behavior by writing to stderr by default. + self.console = kwargs.pop("console", Console(stderr=True)) + + if kwargs: + logger.warning( + "Ignoring unsupported progress bar kwargs for rich backend: %s", + ", ".join(sorted(kwargs.keys())), + ) + + def __enter__(self, *args, **kwargs): + self.progress = Progress( + TextColumn("{task.description}"), + BarColumn(), + TaskProgressColumn(), + console=self.console, + transient=self.transient, + ) + self.progress.__enter__() + self.task_id = self.progress.add_task(self.description, total=self.total) + return self + + def __exit__(self, *args, **kwargs): + self.progress.__exit__(*args, **kwargs) + + def update(self, count): + self.progress.update(self.task_id, advance=count) + + def get_progress_bar(display, total, **kwargs): """Get a progress bar interface with given properties - If the tqdm library is not installed, this will always return a "progress + If the rich library is not installed, this will always return a "progress bar" that does nothing. Args: - display (bool or str): Should the bar actually show the progress? Or a - string to indicate which tqdm bar (subomdule) to use. + display (bool or str): Should the bar actually show the progress? total (int): The total size of the progress bar. - kwargs (dict): Optional keyword arguments to be passed to the tqdm call. + kwargs (dict): Optional keyword arguments to be passed to the progress + bar implementation. """ if display: - if tqdm is None: + if Progress is None: logger.warning( - "You must install the tqdm library to use progress " + "You must install the rich library to use progress " "indicators with emcee" ) return _NoOpPBar() else: - if display is True: - return tqdm.auto.tqdm(total=total, **kwargs) - else: - tqdm_submodule = importlib.import_module(f"tqdm.{display}") - return tqdm_submodule.tqdm(total=total, **kwargs) + return _RichPBar(total=total, **kwargs) return _NoOpPBar() diff --git a/src/emcee/tests/unit/test_pbar.py b/src/emcee/tests/unit/test_pbar.py index 2ffa9bf0..541f3613 100644 --- a/src/emcee/tests/unit/test_pbar.py +++ b/src/emcee/tests/unit/test_pbar.py @@ -1,25 +1,21 @@ import pytest -from emcee.pbar import _NoOpPBar, get_progress_bar +from emcee.pbar import _NoOpPBar, _RichPBar, get_progress_bar try: - import tqdm + import rich except ImportError: - tqdm = None + rich = None def test_display_false(): assert isinstance(get_progress_bar(False, 100), _NoOpPBar) -@pytest.mark.skipif(tqdm is None, reason="tqdm not available") -def test_tqdm_modes(): - assert isinstance(get_progress_bar(True, 1000), tqdm.asyncio.tqdm_asyncio) - assert isinstance(get_progress_bar("std", 1000), tqdm.std.tqdm) - assert isinstance( - get_progress_bar("notebook", 1000), tqdm.notebook.tqdm_notebook - ) - assert isinstance( - get_progress_bar("auto", 1000), tqdm.asyncio.tqdm_asyncio - ) - assert isinstance(get_progress_bar("autonotebook", 1000), tqdm.std.tqdm) +@pytest.mark.skipif(rich is None, reason="rich not available") +def test_rich_modes(): + assert isinstance(get_progress_bar(True, 1000), _RichPBar) + assert isinstance(get_progress_bar("std", 1000), _RichPBar) + assert isinstance(get_progress_bar("notebook", 1000), _RichPBar) + assert isinstance(get_progress_bar("auto", 1000), _RichPBar) + assert isinstance(get_progress_bar("autonotebook", 1000), _RichPBar) From cf0e0ff56d9b81d6cb8bfea6bc5e4e49bb91b5e6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 4 May 2026 21:03:41 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/emcee/pbar.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/emcee/pbar.py b/src/emcee/pbar.py index 7ec0df5e..8555b32f 100644 --- a/src/emcee/pbar.py +++ b/src/emcee/pbar.py @@ -8,7 +8,12 @@ try: from rich.console import Console - from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn + from rich.progress import ( + BarColumn, + Progress, + TaskProgressColumn, + TextColumn, + ) except ImportError: Progress = None @@ -60,7 +65,9 @@ def __enter__(self, *args, **kwargs): transient=self.transient, ) self.progress.__enter__() - self.task_id = self.progress.add_task(self.description, total=self.total) + self.task_id = self.progress.add_task( + self.description, total=self.total + ) return self def __exit__(self, *args, **kwargs): From c12cc4e7d315eebf97586c7ef2fc3f87ba1dff6f Mon Sep 17 00:00:00 2001 From: dutta-alankar Date: Tue, 5 May 2026 18:04:36 +0200 Subject: [PATCH 3/4] [ENH] time estimate in rich progress --- src/emcee/pbar.py | 9 ++++++++- src/emcee/tests/unit/test_pbar.py | 13 +++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/emcee/pbar.py b/src/emcee/pbar.py index 7ec0df5e..56b8fa28 100644 --- a/src/emcee/pbar.py +++ b/src/emcee/pbar.py @@ -8,7 +8,13 @@ try: from rich.console import Console - from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn + from rich.progress import ( + BarColumn, + Progress, + TaskProgressColumn, + TextColumn, + TimeRemainingColumn, + ) except ImportError: Progress = None @@ -56,6 +62,7 @@ def __enter__(self, *args, **kwargs): TextColumn("{task.description}"), BarColumn(), TaskProgressColumn(), + TimeRemainingColumn(), console=self.console, transient=self.transient, ) diff --git a/src/emcee/tests/unit/test_pbar.py b/src/emcee/tests/unit/test_pbar.py index 541f3613..9cdd6c7b 100644 --- a/src/emcee/tests/unit/test_pbar.py +++ b/src/emcee/tests/unit/test_pbar.py @@ -4,8 +4,10 @@ try: import rich + from rich.progress import TimeRemainingColumn except ImportError: rich = None + TimeRemainingColumn = None def test_display_false(): @@ -19,3 +21,14 @@ def test_rich_modes(): assert isinstance(get_progress_bar("notebook", 1000), _RichPBar) assert isinstance(get_progress_bar("auto", 1000), _RichPBar) assert isinstance(get_progress_bar("autonotebook", 1000), _RichPBar) + + +@pytest.mark.skipif(rich is None, reason="rich not available") +def test_rich_progress_includes_eta(): + pbar = get_progress_bar(True, 1000) + + with pbar: + assert any( + isinstance(column, TimeRemainingColumn) + for column in pbar.progress.columns + ) From e62e6dc61d58373179e0215f0d0f3211056683fc Mon Sep 17 00:00:00 2001 From: dutta-alankar Date: Fri, 8 May 2026 08:14:06 +0200 Subject: [PATCH 4/4] [ENH] better progress bar representation --- src/emcee/pbar.py | 59 ++++++++++++++++++++++++++----- src/emcee/tests/unit/test_pbar.py | 18 +++++++--- 2 files changed, 63 insertions(+), 14 deletions(-) diff --git a/src/emcee/pbar.py b/src/emcee/pbar.py index 96f4e377..eaae2ff2 100644 --- a/src/emcee/pbar.py +++ b/src/emcee/pbar.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- import logging +from datetime import timedelta __all__ = ["get_progress_bar"] @@ -8,15 +9,51 @@ try: from rich.console import Console - from rich.progress import ( - BarColumn, - Progress, - TaskProgressColumn, - TextColumn, - TimeRemainingColumn, - ) + from rich.progress import BarColumn, Progress, ProgressColumn, TextColumn + from rich.text import Text + + _RICH_AVAILABLE = True except ImportError: Progress = None + _RICH_AVAILABLE = False + + +def _format_timer(seconds): + if seconds is None: + return "--:--" + + total_seconds = max(0, int(seconds)) + td = timedelta(seconds=total_seconds) + total_seconds = int(td.total_seconds()) + hours, remainder = divmod(total_seconds, 3600) + minutes, secs = divmod(remainder, 60) + + if hours > 0: + return f"{hours:d}:{minutes:02d}:{secs:02d}" + return f"{minutes:02d}:{secs:02d}" + + +if _RICH_AVAILABLE: + + class _ElapsedTimeColumn(ProgressColumn): + def render(self, task): + return Text( + f"elapsed {_format_timer(task.elapsed)}", + style="yellow", + ) + + class _RemainingTimeColumn(ProgressColumn): + def render(self, task): + if task.finished: + return Text("", style="blue") + return Text( + f"left {_format_timer(task.time_remaining)}", + style="blue", + ) + +else: + _ElapsedTimeColumn = object + _RemainingTimeColumn = object class _NoOpPBar(object): @@ -58,11 +95,13 @@ def __init__(self, total, **kwargs): ) def __enter__(self, *args, **kwargs): + assert Progress is not None self.progress = Progress( TextColumn("{task.description}"), BarColumn(), - TaskProgressColumn(), - TimeRemainingColumn(), + TextColumn("{task.completed:.0f}/{task.total:.0f}"), + _ElapsedTimeColumn(), + _RemainingTimeColumn(), console=self.console, transient=self.transient, ) @@ -73,9 +112,11 @@ def __enter__(self, *args, **kwargs): return self def __exit__(self, *args, **kwargs): + assert self.progress is not None self.progress.__exit__(*args, **kwargs) def update(self, count): + assert self.progress is not None self.progress.update(self.task_id, advance=count) diff --git a/src/emcee/tests/unit/test_pbar.py b/src/emcee/tests/unit/test_pbar.py index 9cdd6c7b..115f3931 100644 --- a/src/emcee/tests/unit/test_pbar.py +++ b/src/emcee/tests/unit/test_pbar.py @@ -1,13 +1,17 @@ import pytest -from emcee.pbar import _NoOpPBar, _RichPBar, get_progress_bar +from emcee.pbar import ( + _ElapsedTimeColumn, + _NoOpPBar, + _RemainingTimeColumn, + _RichPBar, + get_progress_bar, +) try: import rich - from rich.progress import TimeRemainingColumn except ImportError: rich = None - TimeRemainingColumn = None def test_display_false(): @@ -24,11 +28,15 @@ def test_rich_modes(): @pytest.mark.skipif(rich is None, reason="rich not available") -def test_rich_progress_includes_eta(): +def test_rich_progress_includes_elapsed_and_remaining(): pbar = get_progress_bar(True, 1000) with pbar: assert any( - isinstance(column, TimeRemainingColumn) + isinstance(column, _ElapsedTimeColumn) + for column in pbar.progress.columns + ) + assert any( + isinstance(column, _RemainingTimeColumn) for column in pbar.progress.columns )