Step samplers are responsible for setting up progress bars

Author: jessegrabowski

pymc/sampling/parallel.py

@@ -27,23 +27,20 @@
 import cloudpickle
 import numpy as np
 
-from rich.console import Console
-from rich.progress import TextColumn
-from rich.style import Style
-from rich.table import Column
 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,
-    DivergenceBarColumn,
-    RandomGeneratorState,
+    compute_draw_speed,
+    create_progress_bar,
     default_progress_theme,
+    RandomGeneratorState,
     get_state_from_generator,
     random_generator_from_state,
+
 )
 
 logger = logging.getLogger(__name__)
@@ -489,33 +486,21 @@ def __init__(
 
         self._in_context = False
 
-        self._progress = CustomProgress(
-            DivergenceBarColumn(
-                table_column=Column("Progress", ratio=2),
-                diverging_color="tab:red",
-                complete_style=Style.parse("rgb(31,119,180)"),  # tab:blue
-                finished_style=Style.parse("rgb(44,160,44)"),  # tab:green
-            ),
-            TextColumn("{task.fields[draws]:,d}", table_column=Column("Draws", ratio=1)),
-            TextColumn(
-                "{task.fields[divergences]:,d}", table_column=Column("Divergences", ratio=1)
-            ),
-            TextColumn("{task.fields[step_size]:0.2f}", table_column=Column("Step size", ratio=1)),
-            TextColumn("{task.fields[tree_depth]:,d}", table_column=Column("Tree depth", ratio=1)),
-            TextColumn(
-                "{task.fields[sampling_speed]:0.2f} {task.fields[speed_unit]}",
-                table_column=Column("Sampling Speed", ratio=1),
-            ),
-            console=Console(theme=progressbar_theme),
-            disable=not progressbar,
-            include_headers=True,
+        progress_columns, progress_stats = step_method._progressbar_config(chains)
+
+        self._progress = create_progress_bar(
+            progress_columns,
+            progress_stats,
+            progressbar=progressbar,
+            progressbar_theme=progressbar_theme,
         )
 
+        self.progress_stats = progress_stats
+        self.update_stats = step_method._make_update_stat_function()
+
         self._show_progress = progressbar
         self._divergences = 0
-        self._divergences_by_chain = [0] * chains
         self._completed_draws = 0
-        self._completed_draws_by_chain = [0] * chains
         self._total_draws = draws + tune
         self._desc = "Sampling chain"
         self._chains = chains
@@ -537,66 +522,48 @@ def __iter__(self):
             tasks = [
                 progress.add_task(
                     self._desc.format(self),
-                    completed=self._completed_draws,
+                    completed=0,
+                    draws=0,
                     total=self._total_draws,
                     chain_idx=chain_idx,
-                    draws=0,
-                    divergences=0,
-                    step_size=0.0,
-                    tree_depth=0,
                     sampling_speed=0,
                     speed_unit="draws/s",
+                    **{stat: value[chain_idx] for stat, value in self.progress_stats.items()},
                 )
                 for chain_idx in range(self._chains)
             ]
 
             while self._active:
                 draw = ProcessAdapter.recv_draw(self._active)
                 proc, is_last, draw, tuning, stats = draw
-                speed = 0
-                unit = "draws/s"
 
                 self._completed_draws += 1
-                self._completed_draws_by_chain[proc.chain] += 1
+
+                speed, unit = compute_draw_speed(progress._tasks[proc.chain].elapsed, draw)
 
                 if not tuning and stats and stats[0].get("diverging"):
                     self._divergences += 1
-                    self._divergences_by_chain[proc.chain] += 1
-
-                if self._show_progress:
-                    elapsed = max(progress._tasks[proc.chain].elapsed, 1e-4)
-                    speed = self._completed_draws_by_chain[proc.chain] / elapsed
 
-                    if speed > 1:
-                        unit = "draws/s"
-                    else:
-                        unit = "s/draws"
-                        speed = 1 / speed
+                self.progress_stats = self.update_stats(self.progress_stats, stats, proc.chain)
 
                 progress.update(
                     tasks[proc.chain],
-                    completed=self._completed_draws_by_chain[proc.chain],
+                    completed=draw,
                     draws=draw,
-                    divergences=self._divergences_by_chain[proc.chain],
-                    step_size=stats[0].get("step_size", 0),
-                    tree_depth=stats[0].get("tree_size", 0),
                     sampling_speed=speed,
                     speed_unit=unit,
+                    **{stat: value[proc.chain] for stat, value in self.progress_stats.items()},
                 )
 
                 if is_last:
-                    self._completed_draws_by_chain[proc.chain] += 1
-
                     proc.join()
                     self._active.remove(proc)
                     self._finished.append(proc)
                     self._make_active()
                     progress.update(
                         tasks[proc.chain],
                         draws=draw + 1,
-                        divergences=self._divergences_by_chain[proc.chain],
-                        step_size=stats[0].get("step_size", 0),
-                        tree_depth=stats[0].get("tree_size", 0),
+                        **{stat: value[proc.chain] for stat, value in self.progress_stats.items()},
                         refresh=True,
                     )
 

pymc/step_methods/compound.py

@@ -297,6 +297,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_stat_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,35 @@ 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_stat_function():
+        def update_stats(stats, step_stats, chain_idx):
+            if isinstance(step_stats, list):
+                step_stats = step_stats[0]
+
+            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")

pymc/step_methods/metropolis.py

@@ -24,6 +24,8 @@
 from pytensor import tensor as pt
 from pytensor.graph.fg import MissingInputError
 from pytensor.tensor.random.basic import BernoulliRV, CategoricalRV
+from rich.progress import TextColumn
+from rich.table import Column
 
 import pymc as pm
 
@@ -325,6 +327,38 @@ def astep(self, q0: RaveledVars) -> tuple[RaveledVars, StatsType]:
     def competence(var, has_grad):
         return Competence.COMPATIBLE
 
+    @staticmethod
+    def _progressbar_config(n_chains=1):
+        columns = [
+            TextColumn("{task.fields[tune]}", table_column=Column("Tuning", ratio=1)),
+            TextColumn("{task.fields[scaling]:0.2f}", table_column=Column("Scaling", ratio=1)),
+            TextColumn(
+                "{task.fields[accept_rate]:0.2f}", table_column=Column("Accept Rate", ratio=1)
+            ),
+        ]
+
+        stats = {
+            "tune": [True] * n_chains,
+            "scaling": [0] * n_chains,
+            "accept_rate": [0.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]
+
+            stats["tune"][chain_idx] = step_stats["tune"]
+            stats["accept_rate"][chain_idx] = step_stats["accept"]
+            stats["scaling"][chain_idx] = step_stats["scaling"]
+
+            return stats
+
+        return update_stats
+
 
 def tune(scale, acc_rate):
     """

pymc/step_methods/slicer.py

@@ -17,6 +17,9 @@
 
 import numpy as np
 
+from rich.progress import TextColumn
+from rich.table import Column
+
 from pymc.blocking import RaveledVars, StatsType
 from pymc.initial_point import PointType
 from pymc.model import modelcontext
@@ -195,3 +198,29 @@ def competence(var, has_grad):
                 return Competence.PREFERRED
             return Competence.COMPATIBLE
         return Competence.INCOMPATIBLE
+
+    @staticmethod
+    def _progressbar_config(n_chains=1):
+        columns = [
+            TextColumn("{task.fields[tune]}", table_column=Column("Tuning", ratio=1)),
+            TextColumn("{task.fields[nstep_out]}", table_column=Column("Steps out", ratio=1)),
+            TextColumn("{task.fields[nstep_in]}", table_column=Column("Steps in", ratio=1)),
+        ]
+
+        stats = {"tune": [True] * n_chains, "nstep_out": [0] * n_chains, "nstep_in": [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]
+
+            stats["tune"][chain_idx] = step_stats["tune"]
+            stats["nstep_out"][chain_idx] = step_stats["nstep_out"]
+            stats["nstep_in"][chain_idx] = step_stats["nstep_in"]
+
+            return stats
+
+        return update_stats

pymc/util.py

@@ -31,7 +31,8 @@
 from pytensor.compile import SharedVariable
 from pytensor.graph.utils import ValidatingScratchpad
 from rich.box import SIMPLE_HEAD
-from rich.progress import BarColumn, Progress, Task
+from rich.console import Console
+from rich.progress import BarColumn, Progress, Task, TextColumn
 from rich.style import Style
 from rich.table import Column, Table
 from rich.theme import Theme
@@ -684,14 +685,50 @@ def __init__(self, *args, diverging_color="red", **kwargs):
 
     def callbacks(self, task: "Task"):
         divergences = task.fields.get("divergences", 0)
-        if divergences > 0:
+        if isinstance(divergences, float | int) and divergences > 0:
             self.complete_style = Style.parse("rgb({},{},{})".format(*self.diverging_rgb))
             self.finished_style = Style.parse("rgb({},{},{})".format(*self.diverging_rgb))
         else:
             self.complete_style = self.non_diverging_style
             self.finished_style = self.non_diverging_finished_style
 
 
+def create_progress_bar(step_columns, init_stat_dict, progressbar, progressbar_theme):
+    columns = [TextColumn("{task.fields[draws]}", table_column=Column("Draws", ratio=1))]
+    columns += step_columns
+    columns += [
+        TextColumn(
+            "{task.fields[sampling_speed]:0.2f} {task.fields[speed_unit]}",
+            table_column=Column("Sampling Speed", ratio=1),
+        )
+    ]
+
+    return CustomProgress(
+        DivergenceBarColumn(
+            table_column=Column("Progress", ratio=2),
+            diverging_color="tab:red",
+            complete_style=Style.parse("rgb(31,119,180)"),  # tab:blue
+            finished_style=Style.parse("rgb(44,160,44)"),  # tab:green
+        ),
+        *columns,
+        console=Console(theme=progressbar_theme),
+        disable=not progressbar,
+        include_headers=True,
+    )
+
+
+def compute_draw_speed(elapsed, draws):
+    speed = draws / max(elapsed, 1e-6)
+
+    if speed > 1 or speed == 0:
+        unit = "draws/s"
+    else:
+        unit = "s/draws"
+        speed = 1 / speed
+
+    return speed, unit
+
+
 RandomGeneratorState = namedtuple("RandomGeneratorState", ["bit_generator_state", "seed_seq_state"])