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,42 +355,43 @@ 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
+        failing = False
+        all_step_stats = {}
 
-        if self.full_stats:
-            # TODO: Index by chain already?
-            chain_progress_stats = [
-                update_states_fn(step_stats)
-                for update_states_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 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:
-            all_step_stats = {}
+        chain_progress_stats = [
+            update_stats_fn(step_stats)
+            for update_stats_fn, step_stats in zip(self.update_stats_functions, stats, strict=True)
+        ]
+        for step_stats in chain_progress_stats:
+            for key, val in step_stats.items():
+                if key == "failing":
+                    failing |= val
+                    continue
+                if not self.full_stats:
+                    # Only care about the "failing" flag
+                    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
 
         self._progress.update(
             self.tasks[chain_idx],
             completed=draw,
             draws=draw,
             sampling_speed=speed,
             speed_unit=unit,
+            failing=failing,
             **all_step_stats,
         )
 
         if is_last:
             self._progress.update(
                 self.tasks[chain_idx],
                 draws=draw + 1 if not self.combined_progress else draw,
+                failing=failing,
                 **all_step_stats,
                 refresh=True,
             )
@@ -410,9 +412,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,15 @@ 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,
+            "divergences": self.divergences,
             "perf_counter_diff": perf_end - perf_start,
             "process_time_diff": process_end - process_start,
             "perf_counter_start": perf_start,
@@ -288,6 +293,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,32 @@ 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
+
+    @staticmethod
+    def _make_progressbar_update_functions():
+        def update_stats(stats):
+            return {
+                key: stats[key]
+                for key in (
+                    "divergences",
+                    "n_steps",
+                )
+            } | {
+                "failing": stats["divergences"] > 0,
+            }
+
+        return (update_stats,)

pymc/step_methods/hmc/nuts.py

@@ -115,6 +115,7 @@ class NUTS(BaseHMC):
         "step_size_bar": (np.float64, []),
         "tree_size": (np.float64, []),
         "diverging": (bool, []),
+        "divergences": (int, []),
         "energy_error": (np.float64, []),
         "energy": (np.float64, []),
         "max_energy_error": (np.float64, []),
@@ -250,7 +251,9 @@ def _progressbar_config(n_chains=1):
     @staticmethod
     def _make_progressbar_update_functions():
         def update_stats(stats):
-            return {key: stats[key] for key in ("diverging", "step_size", "tree_size")}
+            return {key: stats[key] for key in ("divergences", "step_size", "tree_size")} | {
+                "failing": stats["divergences"] > 0,
+            }
 
         return (update_stats,)
 

tests/step_methods/hmc/test_nuts.py

@@ -148,6 +148,7 @@ def test_sampler_stats(self):
         expected_stat_names = {
             "depth",
             "diverging",
+            "divergences",
             "energy",
             "energy_error",
             "model_logp",