Show one progress bar per chain when sampling

pymc/sampling/mcmc.py

@@ -36,8 +36,6 @@
 from arviz import InferenceData, dict_to_dataset
 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
@@ -67,7 +65,7 @@
 from pymc.step_methods.arraystep import BlockedStep, PopulationArrayStepShared
 from pymc.step_methods.hmc import quadpotential
 from pymc.util import (
-    CustomProgress,
+    ProgressManager,
     RandomSeed,
     RandomState,
     _get_seeds_per_chain,
@@ -1138,34 +1136,44 @@ def _sample_many(
         Step function
     """
     initial_step_state = step.sampling_state
-    for i in range(chains):
-        step.sampling_state = initial_step_state
-        _sample(
-            draws=draws,
-            chain=i,
-            start=start[i],
-            step=step,
-            trace=traces[i],
-            rng=rngs[i],
-            callback=callback,
-            **kwargs,
-        )
+    progress_manager = ProgressManager(
+        step_method=step,
+        chains=chains,
+        draws=draws - kwargs.get("tune", 0),
+        tune=kwargs.get("tune", 0),
+        progressbar=kwargs.get("progressbar", True),
+        progressbar_theme=kwargs.get("progressbar_theme", default_progress_theme),
+    )
+
+    with progress_manager:
+        for i in range(chains):
+            step.sampling_state = initial_step_state
+            _sample(
+                draws=draws,
+                chain=i,
+                start=start[i],
+                step=step,
+                trace=traces[i],
+                rng=rngs[i],
+                callback=callback,
+                progress_manager=progress_manager,
+                **kwargs,
+            )
     return
 
 
 def _sample(
     *,
     chain: int,
-    progressbar: bool,
     rng: np.random.Generator,
     start: PointType,
     draws: int,
     step: Step,
     trace: IBaseTrace,
     tune: int,
     model: Model | None = None,
-    progressbar_theme: Theme | None = default_progress_theme,
     callback=None,
+    progress_manager: ProgressManager,
     **kwargs,
 ) -> None:
     """Sample one chain (singleprocess).
@@ -1176,27 +1184,23 @@ def _sample(
     ----------
     chain : int
         Number of the chain that the samples will belong to.
-    progressbar : bool
-        Whether or not to display a progress bar in the command line. The bar shows the percentage
-        of completion, the sampling speed in samples per second (SPS), and the estimated remaining
-        time until completion ("expected time of arrival"; ETA).
-    random_seed : single random seed
+    random_seed : Generator
+        Single random seed
     start : dict
         Starting point in parameter space (or partial point)
     draws : int
         The number of samples to draw
-    step : function
-        Step function
+    step : Step
+        Step class instance used to generate samples.
     trace
         A chain backend to record draws and stats.
     tune : int
         Number of iterations to tune.
-    model : Model (optional if in ``with`` context)
-    progressbar_theme : Theme
-        Optional custom theme for the progress bar.
+    model : Model, optional
+        PyMC model. If None, the model is taken from the current context.
+    progress_manager: ProgressManager
+        Helper class used to handle progress bar styling and updates
     """
-    skip_first = kwargs.get("skip_first", 0)
-
     sampling_gen = _iter_sample(
         draws=draws,
         step=step,
@@ -1208,32 +1212,19 @@ def _sample(
         rng=rng,
         callback=callback,
     )
-    _pbar_data = {"chain": chain, "divergences": 0}
-    _desc = "Sampling chain {chain:d}, {divergences:,d} divergences"
-
-    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:
+        for it, stats in enumerate(sampling_gen):
+            progress_manager.update(
+                chain_idx=chain, is_last=False, draw=it, stats=stats, tuning=it > tune
+            )
 
-    with progress:
-        try:
-            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, description=_desc.format(**_pbar_data), completed=it)
-            progress.update(
-                task, description=_desc.format(**_pbar_data), completed=draws, refresh=True
+        if not progress_manager.combined_progress or chain == progress_manager.chains - 1:
+            progress_manager.update(
+                chain_idx=chain, is_last=True, draw=it, stats=stats, tuning=False
             )
-        except KeyboardInterrupt:
-            pass
+
+    except KeyboardInterrupt:
+        pass
 
 
 def _iter_sample(
@@ -1293,22 +1284,25 @@ def _iter_sample(
                 step.iter_count = 0
             if i == tune:
                 step.stop_tuning()
+
             point, stats = step.step(point)
             trace.record(point, stats)
             log_warning_stats(stats)
-            diverging = i > tune and len(stats) > 0 and (stats[0].get("diverging") is True)
+
             if callback is not None:
                 callback(
                     trace=trace,
                     draw=Draw(chain, i == draws, i, i < tune, stats, point),
                 )
 
-            yield diverging
+            yield stats
+
     except (KeyboardInterrupt, BaseException):
         if isinstance(trace, ZarrChain):
             trace.record_sampling_state(step=step)
         trace.close()
         raise
+
     else:
         if isinstance(trace, ZarrChain):
             trace.record_sampling_state(step=step)

pymc/sampling/parallel.py

@@ -27,16 +27,14 @@
 import cloudpickle
 import numpy as np
 
-from rich.console import Console
-from rich.progress import BarColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn
 from rich.theme import Theme
 from threadpoolctl import threadpool_limits
 
 from pymc.backends.zarr import ZarrChain
 from pymc.blocking import DictToArrayBijection
 from pymc.exceptions import SamplingError
 from pymc.util import (
-    CustomProgress,
+    ProgressManager,
     RandomGeneratorState,
     default_progress_theme,
     get_state_from_generator,
@@ -485,23 +483,14 @@ def __init__(
         self._max_active = cores
 
         self._in_context = False
-
-        self._progress = CustomProgress(
-            "[progress.description]{task.description}",
-            BarColumn(),
-            "[progress.percentage]{task.percentage:>3.0f}%",
-            TimeRemainingColumn(),
-            TextColumn("/"),
-            TimeElapsedColumn(),
-            console=Console(theme=progressbar_theme),
-            disable=not progressbar,
+        self._progress = ProgressManager(
+            step_method=step_method,
+            chains=chains,
+            draws=draws,
+            tune=tune,
+            progressbar=progressbar,
+            progressbar_theme=progressbar_theme,
         )
-        self._show_progress = progressbar
-        self._divergences = 0
-        self._completed_draws = 0
-        self._total_draws = chains * (draws + tune)
-        self._desc = "Sampling {0._chains:d} chains, {0._divergences:,d} divergences"
-        self._chains = chains
 
     def _make_active(self):
         while self._inactive and len(self._active) < self._max_active:
@@ -516,32 +505,20 @@ def __iter__(self):
             raise ValueError("Use ParallelSampler as context manager.")
         self._make_active()
 
-        with self._progress as progress:
-            task = progress.add_task(
-                self._desc.format(self),
-                completed=self._completed_draws,
-                total=self._total_draws,
-            )
-
+        with self._progress:
             while self._active:
                 draw = ProcessAdapter.recv_draw(self._active)
                 proc, is_last, draw, tuning, stats = draw
-                self._completed_draws += 1
-                if not tuning and stats and stats[0].get("diverging"):
-                    self._divergences += 1
-                progress.update(
-                    task,
-                    completed=self._completed_draws,
-                    total=self._total_draws,
-                    description=self._desc.format(self),
+
+                self._progress.update(
+                    chain_idx=proc.chain, is_last=is_last, draw=draw, tuning=tuning, stats=stats
                 )
 
                 if is_last:
                     proc.join()
                     self._active.remove(proc)
                     self._finished.append(proc)
                     self._make_active()
-                    progress.update(task, description=self._desc.format(self), refresh=True)
 
                 # We could also yield proc.shared_point_view directly,
                 # and only call proc.write_next() after the yield returns.

pymc/step_methods/compound.py

@@ -181,6 +181,20 @@ def __new__(cls, *args, **kwargs):
             step.__newargs = (vars, *args), kwargs
             return step
 
+    @staticmethod
+    def _progressbar_config(n_chains=1):
+        columns = []
+        stats = {}
+
+        return columns, stats
+
+    @staticmethod
+    def _make_update_stats_function():
+        def update_stats(stats, step_stats, chain_idx):
+            return stats
+
+        return update_stats
+
     # Hack for creating the class correctly when unpickling.
     def __getnewargs_ex__(self):
         return self.__newargs
@@ -297,6 +311,38 @@ def set_rng(self, rng: RandomGenerator):
         for method, _rng in zip(self.methods, _rngs):
             method.set_rng(_rng)
 
+    def _progressbar_config(self, n_chains=1):
+        from functools import reduce
+
+        column_lists, stat_dict_list = zip(
+            *[method._progressbar_config(n_chains) for method in self.methods]
+        )
+        flat_list = reduce(lambda left_list, right_list: left_list + right_list, column_lists)
+
+        columns = []
+        headers = []
+
+        for col in flat_list:
+            name = col.get_table_column().header
+            if name not in headers:
+                headers.append(name)
+                columns.append(col)
+
+        stats = reduce(lambda left_dict, right_dict: left_dict | right_dict, stat_dict_list)
+
+        return columns, stats
+
+    def _make_update_stats_function(self):
+        update_fns = [method._make_update_stats_function() for method in self.methods]
+
+        def update_stats(stats, step_stats, chain_idx):
+            for step_stat, update_fn in zip(step_stats, update_fns):
+                stats = update_fn(stats, step_stat, chain_idx)
+
+            return stats
+
+        return update_stats
+
 
 def flatten_steps(step: BlockedStep | CompoundStep) -> list[BlockedStep]:
     """Flatten a hierarchy of step methods to a list."""

pymc/step_methods/hmc/nuts.py

@@ -20,6 +20,8 @@
 import numpy as np
 
 from pytensor import config
+from rich.progress import TextColumn
+from rich.table import Column
 
 from pymc.stats.convergence import SamplerWarning
 from pymc.step_methods.compound import Competence
@@ -229,6 +231,37 @@ def competence(var, has_grad):
             return Competence.PREFERRED
         return Competence.INCOMPATIBLE
 
+    @staticmethod
+    def _progressbar_config(n_chains=1):
+        columns = [
+            TextColumn("{task.fields[divergences]}", table_column=Column("Divergences", ratio=1)),
+            TextColumn("{task.fields[step_size]:0.2f}", table_column=Column("Step size", ratio=1)),
+            TextColumn("{task.fields[tree_size]}", table_column=Column("Grad evals", ratio=1)),
+        ]
+
+        stats = {
+            "divergences": [0] * n_chains,
+            "step_size": [0] * n_chains,
+            "tree_size": [0] * n_chains,
+        }
+
+        return columns, stats
+
+    @staticmethod
+    def _make_update_stats_function():
+        def update_stats(stats, step_stats, chain_idx):
+            if isinstance(step_stats, list):
+                step_stats = step_stats[0]
+
+            if not step_stats["tune"]:
+                stats["divergences"][chain_idx] += step_stats["diverging"]
+
+            stats["step_size"][chain_idx] = step_stats["step_size"]
+            stats["tree_size"][chain_idx] = step_stats["tree_size"]
+            return stats
+
+        return update_stats
+
 
 # A proposal for the next position
 Proposal = namedtuple("Proposal", "q, q_grad, energy, logp, index_in_trajectory")