Add Elapsed and Remaining columns

pymc-devs · Jan 10, 2025 · 667c78e · 667c78e
1 parent 2bf89ff
commit 667c78e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/pymc/sampling/parallel.py b/pymc/sampling/parallel.py
@@ -474,7 +474,7 @@ def __iter__(self):
                     self._desc.format(self),
                     completed=0,
                     draws=0,
-                    total=self._total_draws,
+                    total=self._total_draws - 1,
                     chain_idx=chain_idx,
                     sampling_speed=0,
                     speed_unit="draws/s",

diff --git a/pymc/util.py b/pymc/util.py
@@ -31,7 +31,14 @@
 from pytensor.graph.utils import ValidatingScratchpad
 from rich.box import SIMPLE_HEAD
 from rich.console import Console
-from rich.progress import BarColumn, Progress, Task, TextColumn
+from rich.progress import (
+    BarColumn,
+    Progress,
+    Task,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
 from rich.style import Style
 from rich.table import Column, Table
 from rich.theme import Theme
@@ -58,6 +65,8 @@ def __getattr__(name):
     {
         "bar.complete": "#1764f4",
         "bar.finished": "green",
+        "progress.remaining": "none",
+        "progress.elapsed": "none",
     }
 )
 
@@ -694,7 +703,9 @@ def create_progress_bar(step_columns, init_stat_dict, progressbar, progressbar_t
         TextColumn(
             "{task.fields[sampling_speed]:0.2f} {task.fields[speed_unit]}",
             table_column=Column("Sampling Speed", ratio=1),
-        )
+        ),
+        TimeElapsedColumn(table_column=Column("Elapsed", ratio=1)),
+        TimeRemainingColumn(table_column=Column("Remaining", ratio=1)),
     ]
 
     return CustomProgress(