Clean up marimo progress bar (#8184)

* Removed extra "s" from draws in progress bar

* Remove s from step_name

* Removed second Draws column

* Fix draw column: restore for Rich, deduplicate for marimo

* Fixed test failure

* Do not zero out final speed

* Remove rounding for sampling <1s; simplify rich progress label

* Remove guardrails on progbar stat columns; update column names

Author: fonnesbeck

pymc/progress_bar/marimo_progress.py

@@ -20,6 +20,8 @@
 
 def format_time(seconds: float) -> str:
     """Format elapsed time as mm:ss or hh:mm:ss."""
+    if seconds < 1:
+        return f"{seconds:.1f}s"
     minutes, secs = divmod(int(seconds), 60)
     hours, minutes = divmod(minutes, 60)
     if hours > 0:
@@ -201,24 +203,26 @@ def _render_html(self) -> str:
         """Generate HTML for all progress bars as a table with headers."""
         stat_keys = []
         if self.full_stats and self._task_state and self._task_state[0]["stats"]:
-            stat_keys = list(self._task_state[0]["stats"].keys())
+            stat_keys = [
+                k for k in self._task_state[0]["stats"].keys() if k != self.step_name.lower()
+            ]
 
         header_cells = ["Progress", self.step_name]
 
-        abbreviations = {
-            "divergences": "Div",
-            "diverging": "Div",
-            "step_size": "Step",
-            "tree_size": "Tree",
-            "tree_depth": "Depth",
-            "n_steps": "Steps",
-            "energy_error": "E-err",
-            "max_energy_error": "Max-E",
-            "mean_tree_accept": "Accept",
-            "scaling": "Scale",
+        column_names = {
+            "divergences": "Divergences",
+            "diverging": "Divergences",
+            "step_size": "Step size",
+            "tree_size": "Grad evals",
+            "tree_depth": "Tree depth",
+            "n_steps": "Grad evals",
+            "energy_error": "Energy error",
+            "max_energy_error": "Max energy error",
+            "mean_tree_accept": "Mean tree accept",
+            "scaling": "Scaling",
             "tune": "Tune",
         }
-        header_cells += [abbreviations.get(k, k[:6].capitalize()) for k in stat_keys]
+        header_cells += [column_names.get(k, k.replace("_", " ").capitalize()) for k in stat_keys]
 
         header_cells += ["Speed", "Elapsed"]
 

pymc/progress_bar/progress.py

@@ -214,7 +214,7 @@ class MCMCProgressBarManager(ProgressBarManager):
     Tracks progress via draw count with support for tuning phases.
     """
 
-    step_name: str = "Draws"
+    step_name: str = "Draw"
 
     def __init__(
         self,
@@ -249,7 +249,7 @@ def __init__(
         )
 
         progress_columns, progress_stats = step_method._progressbar_config(chains)
-        progress_stats["draws"] = [0] * chains
+        progress_stats["draw"] = [0] * chains
 
         self.progress_stats = progress_stats
         self.update_stats_functions = step_method._make_progressbar_update_functions()
@@ -318,7 +318,7 @@ def update(self, chain_idx: int, is_last: bool, draw: int, tuning: bool, stats)
             chain_idx = 0
 
         failing, all_step_stats = self._extract_stats(stats)
-        all_step_stats["draws"] = draw + 1 if not self.combined_progress else draw
+        all_step_stats["draw"] = draw + 1 if not self.combined_progress else draw
 
         self._backend.update(
             task_id=chain_idx,

pymc/progress_bar/rich_progress.py

@@ -199,7 +199,7 @@ def _create_progress_bar(
         columns += [
             TextColumn(
                 "{task.fields[sampling_speed]:0.2f} {task.fields[speed_unit]}",
-                table_column=Column("Sampling Speed", ratio=1),
+                table_column=Column("Speed", ratio=1),
             ),
             TimeElapsedColumn(table_column=Column("Elapsed", ratio=1)),
             TimeRemainingColumn(table_column=Column("Remaining", ratio=1)),
@@ -285,8 +285,6 @@ def update(
             self._progress.update(
                 rich_task_id,
                 completed=self.total if not self.combined else self.total * self.n_bars,
-                sampling_speed=0,
-                speed_unit="",
                 failing=failing,
                 refresh=True,
                 **stats,

tests/progress_bar/test_manager.py

@@ -63,7 +63,7 @@ def capturing_init(self, *args, **kwargs):
     for chain in range(chains):
         task = _get_task(manager, chain)
         assert task.completed == task.total == total
-        assert task.fields["draws"] == total
+        assert task.fields["draw"] == total
 
 
 def test_mcmc_combined_bar_ends_at_total():
@@ -92,7 +92,7 @@ def capturing_init(self, *args, **kwargs):
     total = (draws + tune) * chains
     task = _get_task(manager, 0)
     assert task.completed == task.total == total
-    assert task.fields["draws"] == total
+    assert task.fields["draw"] == total
 
 
 def test_mcmc_draws_stat_shows_completed_count():
@@ -104,13 +104,13 @@ def test_mcmc_draws_stat_shows_completed_count():
 
     with manager:
         manager.update(chain_idx=0, is_last=False, draw=0, tuning=True, stats=NUTS_DUMMY_STATS)
-        assert _get_task(manager).fields["draws"] == 1
+        assert _get_task(manager).fields["draw"] == 1
 
         for i in range(1, 15):
             manager.update(
                 chain_idx=0, is_last=i == 14, draw=i, tuning=i < 5, stats=NUTS_DUMMY_STATS
             )
-        assert _get_task(manager).fields["draws"] == 15
+        assert _get_task(manager).fields["draw"] == 15
 
 
 def test_smc_bar_starts_at_zero_ends_at_one(imh_kernel):