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/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

@@ -55,6 +55,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, []),

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": (np.int64, []),
         "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_update_stats_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,)
 

pymc/util.py

@@ -657,22 +657,22 @@ class DivergenceBarColumn(BarColumn):
     def __init__(self, *args, diverging_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 = diverging_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:
@@ -794,6 +794,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()},
                 )
             ]
@@ -808,6 +809,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)
@@ -829,16 +831,22 @@ def update(self, chain_idx, is_last, draw, tuning, stats):
             self.divergences += 1
 
         if self.full_stats:
+            failing = False
+            all_step_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 == "failing":
+                        failing |= val
+                        continue
+
                     if key in all_step_stats:
                         continue
                         count = (
@@ -849,6 +857,7 @@ def update(self, chain_idx, is_last, draw, tuning, stats):
                         all_step_stats[key] = val
 
         else:
+            failing = False
             all_step_stats = {}
 
         # more_updates = (
@@ -863,6 +872,7 @@ def update(self, chain_idx, is_last, draw, tuning, stats):
             draws=draw,
             sampling_speed=speed,
             speed_unit=unit,
+            failing=failing,
             **all_step_stats,
         )