Remove some 'refresh' and make sure progress goes to 100%

Author: tomicapretto

pymc/sampling/forward.py

@@ -44,6 +44,7 @@
 )
 from pytensor.tensor.sharedvar import SharedVariable
 from rich.console import Console
+from rich.progress import BarColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn
 from rich.theme import Theme
 
 import pymc as pm
@@ -828,11 +829,21 @@ def sample_posterior_predictive(
     # All model variables have a name, but mypy does not know this
     _log.info(f"Sampling: {list(sorted(volatile_basic_rvs, key=lambda var: var.name))}")  # type: ignore
     ppc_trace_t = _DefaultTrace(samples)
+
+    progress = CustomProgress(
+        "[progress.description]{task.description}",
+        BarColumn(),
+        "[progress.percentage]{task.percentage:>3.0f}%",
+        TimeRemainingColumn(),
+        TextColumn("/"),
+        TimeElapsedColumn(),
+        console=Console(theme=progressbar_theme),
+        disable=not progressbar,
+    )
+
     try:
-        with CustomProgress(
-            console=Console(theme=progressbar_theme), disable=not progressbar
-        ) as progress:
-            task = progress.add_task("Sampling ...", total=samples)
+        with progress:
+            task = progress.add_task("Sampling ...", completed=0, total=samples)
             for idx in np.arange(samples):
                 if nchain > 1:
                     # the trace object will either be a MultiTrace (and have _straces)...
@@ -854,6 +865,7 @@ def sample_posterior_predictive(
                     ppc_trace_t.insert(k.name, v, idx)
 
                 progress.advance(task)
+            progress.update(task, refresh=True, completed=samples)
 
     except KeyboardInterrupt:
         pass

pymc/sampling/mcmc.py

@@ -36,6 +36,7 @@
 from arviz.data.base import make_attrs
 from pytensor.graph.basic import Variable
 from rich.console import Console
+from rich.progress import BarColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn
 from rich.theme import Theme
 from threadpoolctl import threadpool_limits
 from typing_extensions import Protocol
@@ -1075,16 +1076,28 @@ def _sample(
     )
     _pbar_data = {"chain": chain, "divergences": 0}
     _desc = "Sampling chain {chain:d}, {divergences:,d} divergences"
-    with CustomProgress(
-        console=Console(theme=progressbar_theme), disable=not progressbar
-    ) as progress:
+
+    progress = CustomProgress(
+        "[progress.description]{task.description}",
+        BarColumn(),
+        "[progress.percentage]{task.percentage:>3.0f}%",
+        TimeRemainingColumn(),
+        TextColumn("/"),
+        TimeElapsedColumn(),
+        console=Console(theme=progressbar_theme),
+        disable=not progressbar,
+    )
+
+    with progress:
         try:
-            task = progress.add_task(_desc.format(**_pbar_data), total=draws)
+            task = progress.add_task(_desc.format(**_pbar_data), completed=0, total=draws)
             for it, diverging in enumerate(sampling_gen):
                 if it >= skip_first and diverging:
                     _pbar_data["divergences"] += 1
-                progress.update(task)
-            progress.update(task, refresh=True, advance=1, completed=True)
+                progress.update(task, description=_desc.format(**_pbar_data), completed=it)
+            progress.update(
+                task, description=_desc.format(**_pbar_data), completed=draws, refresh=True
+            )
         except KeyboardInterrupt:
             pass
 

pymc/sampling/parallel.py

@@ -473,7 +473,6 @@ def __iter__(self):
                 self._completed_draws += 1
                 if not tuning and stats and stats[0].get("diverging"):
                     self._divergences += 1
-
                 progress.update(
                     task,
                     completed=self._completed_draws,

pymc/sampling/population.py

@@ -102,7 +102,6 @@ def _sample_population(
 
     with CustomProgress(disable=not progressbar) as progress:
         task = progress.add_task("[red]Sampling...", total=draws)
-
         for _ in sampling:
             progress.update(task)