Add progressbar config to HamiltonianMC

Author: jessegrabowski

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
@@ -203,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,)