Abstract special behavior of NUTS divergences in ProgressBar

Every step sampler can now decide whether sampling is failing or not by setting "failing" in the returned update dict

Author: ricardoV94

pymc/progress_bar.py

@@ -168,28 +168,28 @@ def call_column(column, task):
         return table
 
 
-class DivergenceBarColumn(BarColumn):
-    """Rich colorbar that changes color when a chain has detected a divergence."""
+class RecolorOnFailureBarColumn(BarColumn):
+    """Rich colorbar that changes color when a chain has detected a failure."""
 
-    def __init__(self, *args, diverging_color="red", **kwargs):
+    def __init__(self, *args, failing_color="red", **kwargs):
         from matplotlib.colors import to_rgb
 
-        self.diverging_color = diverging_color
-        self.diverging_rgb = [int(x * 255) for x in to_rgb(self.diverging_color)]
+        self.failing_color = failing_color
+        self.failing_rgb = [int(x * 255) for x in to_rgb(self.failing_color)]
 
         super().__init__(*args, **kwargs)
 
-        self.non_diverging_style = self.complete_style
-        self.non_diverging_finished_style = self.finished_style
+        self.default_complete_style = self.complete_style
+        self.default_finished_style = self.finished_style
 
     def callbacks(self, task: "Task"):
-        divergences = task.fields.get("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))
+        if task.fields["failing"]:
+            self.complete_style = Style.parse("rgb({},{},{})".format(*self.failing_rgb))
+            self.finished_style = Style.parse("rgb({},{},{})".format(*self.failing_rgb))
         else:
-            self.complete_style = self.non_diverging_style
-            self.finished_style = self.non_diverging_finished_style
+            # Recovered from failing yay
+            self.complete_style = self.default_complete_style
+            self.finished_style = self.default_finished_style
 
 
 class ProgressBarManager:
@@ -284,7 +284,6 @@ def __init__(
         self.update_stats_functions = step_method._make_progressbar_update_functions()
 
         self._show_progress = show_progress
-        self.divergences = 0
         self.completed_draws = 0
         self.total_draws = draws + tune
         self.desc = "Sampling chain"
@@ -311,6 +310,7 @@ def _initialize_tasks(self):
                     chain_idx=0,
                     sampling_speed=0,
                     speed_unit="draws/s",
+                    failing=False,
                     **{stat: value[0] for stat, value in self.progress_stats.items()},
                 )
             ]
@@ -325,6 +325,7 @@ def _initialize_tasks(self):
                     chain_idx=chain_idx,
                     sampling_speed=0,
                     speed_unit="draws/s",
+                    failing=False,
                     **{stat: value[chain_idx] for stat, value in self.progress_stats.items()},
                 )
                 for chain_idx in range(self.chains)
@@ -354,27 +355,30 @@ def update(self, chain_idx, is_last, draw, tuning, stats):
         elapsed = self._progress.tasks[chain_idx].elapsed
         speed, unit = self.compute_draw_speed(elapsed, draw)
 
-        if not tuning and stats and stats[0].get("diverging"):
-            self.divergences += 1
-
         if self.full_stats:
-            # TODO: Index by chain already?
+            failing = False
+            all_step_stats = {}
+
             chain_progress_stats = [
-                update_states_fn(step_stats)
-                for update_states_fn, step_stats in zip(
+                update_stats_fn(step_stats)
+                for update_stats_fn, step_stats in zip(
                     self.update_stats_functions, stats, strict=True
                 )
             ]
-            all_step_stats = {}
             for step_stats in chain_progress_stats:
                 for key, val in step_stats.items():
+                    if key == "failing":
+                        failing |= val
+                        continue
+
                     if key in all_step_stats:
                         # TODO: Figure out how to integrate duplicate / non-scalar keys, ignoring them for now
                         continue
                     else:
                         all_step_stats[key] = val
 
         else:
+            failing = False
             all_step_stats = {}
 
         self._progress.update(
@@ -383,6 +387,7 @@ def update(self, chain_idx, is_last, draw, tuning, stats):
             draws=draw,
             sampling_speed=speed,
             speed_unit=unit,
+            failing=failing,
             **all_step_stats,
         )
 
@@ -410,9 +415,9 @@ def create_progress_bar(self, step_columns, progressbar, progressbar_theme):
         ]
 
         return CustomProgress(
-            DivergenceBarColumn(
+            RecolorOnFailureBarColumn(
                 table_column=Column("Progress", ratio=2),
-                diverging_color="tab:red",
+                failing_color="tab:red",
                 complete_style=Style.parse("rgb(31,119,180)"),  # tab:blue
                 finished_style=Style.parse("rgb(31,119,180)"),  # tab:blue
             ),

pymc/step_methods/hmc/base_hmc.py

@@ -184,6 +184,7 @@ def __init__(
 
         self._step_rand = step_rand
         self._num_divs_sample = 0
+        self.divergences = 0
 
     @abstractmethod
     def _hamiltonian_step(self, start, p0, step_size) -> HMCStepData:
@@ -266,11 +267,14 @@ def astep(self, q0: RaveledVars) -> tuple[RaveledVars, StatsType]:
                 divergence_info=info_store,
             )
 
+        diverging = bool(hmc_step.divergence_info)
+        if not self.tune:
+            self.divergences += diverging
         self.iter_count += 1
 
         stats: dict[str, Any] = {
             "tune": self.tune,
-            "diverging": bool(hmc_step.divergence_info),
+            "diverging": diverging,
             "perf_counter_diff": perf_end - perf_start,
             "process_time_diff": process_end - process_start,
             "perf_counter_start": perf_start,
@@ -288,6 +292,8 @@ def reset_tuning(self, start=None):
         self.reset(start=None)
 
     def reset(self, start=None):
+        self.iter_count = 0
+        self.divergences = 0
         self.tune = True
         self.potential.reset()
 

pymc/step_methods/hmc/hmc.py

@@ -19,6 +19,9 @@
 
 import numpy as np
 
+from rich.progress import TextColumn
+from rich.table import Column
+
 from pymc.stats.convergence import SamplerWarning
 from pymc.step_methods.compound import Competence
 from pymc.step_methods.hmc.base_hmc import BaseHMC, BaseHMCState, DivergenceInfo, HMCStepData
@@ -55,6 +58,7 @@ class HamiltonianMC(BaseHMC):
         "accept": (np.float64, []),
         "diverging": (bool, []),
         "energy_error": (np.float64, []),
+        "divergences": (np.int64, []),
         "energy": (np.float64, []),
         "path_length": (np.float64, []),
         "accepted": (bool, []),
@@ -202,3 +206,27 @@ def competence(var, has_grad):
         if var.dtype in discrete_types or not has_grad:
             return Competence.INCOMPATIBLE
         return Competence.COMPATIBLE
+
+    @staticmethod
+    def _progressbar_config(n_chains=1):
+        columns = [
+            TextColumn("{task.fields[divergences]}", table_column=Column("Divergences", ratio=1)),
+            TextColumn("{task.fields[n_steps]}", table_column=Column("Grad evals", ratio=1)),
+        ]
+
+        stats = {
+            "divergences": [0] * n_chains,
+            "n_steps": [0] * n_chains,
+        }
+
+        return columns, stats
+
+    def _make_progressbar_update_functions(self):
+        def update_stats(stats):
+            divergences = self.divergences
+            return {key: stats[key] for key in ("n_steps",)} | {
+                "failing": divergences > 0,
+                "divergences": divergences,
+            }
+
+        return (update_stats,)

pymc/step_methods/hmc/nuts.py

@@ -247,10 +247,13 @@ def _progressbar_config(n_chains=1):
 
         return columns, stats
 
-    @staticmethod
-    def _make_update_stats_functions():
+    def _make_update_stats_functions(self):
         def update_stats(stats):
-            return {key: stats[key] for key in ("diverging", "step_size", "tree_size")}
+            divergences = self.divergences
+            return {key: stats[key] for key in ("step_size", "tree_size")} | {
+                "failing": divergences > 0,
+                "divergences": divergences,
+            }
 
         return (update_stats,)