Use PyMC progress bar with nutpie

Author: ricardoV94

pymc/progress_bar/__init__.py

@@ -14,6 +14,7 @@
 
 from pymc.progress_bar.progress import (
     MCMCProgressBarManager,
+    NutpieProgressBarManager,
     ProgressBarManager,
     ProgressBarOptions,
     SMCProgressBarManager,

pymc/progress_bar/progress.py

@@ -329,6 +329,71 @@ def update(self, chain_idx: int, is_last: bool, draw: int, tuning: bool, stats)
         )
 
 
+class NutpieProgressBarManager(ProgressBarManager):
+    """Progress bar manager for nutpie NUTS sampling.
+
+    Bridges ``nutpie.sample``'s ``progress_callback`` (a callable that receives a
+    list of ``nutpie.ChainProgress`` objects) to PyMC's progress bar backends,
+    so nutpie draws through the same UI as the pymc sampler.
+    """
+
+    step_name: str = "Draw"
+
+    def __init__(
+        self,
+        chains: int,
+        draws: int,
+        tune: int,
+        progressbar: bool | ProgressBarOptions = True,
+        progressbar_theme: Theme | str | None = None,
+    ):
+        from pymc.step_methods.hmc import NUTS
+
+        super().__init__(
+            n_bars=chains,
+            progressbar=progressbar,
+            progressbar_theme=progressbar_theme,
+        )
+
+        progress_columns, progress_stats = NUTS._progressbar_config(chains)
+        progress_stats["draw"] = [0] * chains
+
+        self.total_draws = draws + tune
+        self._previous_finished = [0] * chains
+
+        self._backend = self._create_backend(
+            total=self.total_draws * chains if self.combined_progress else self.total_draws,
+            progress_columns=progress_columns,
+            progress_stats=progress_stats,
+        )
+
+    def update(self, chain_progresses) -> None:
+        """Consume a list of ``nutpie.ChainProgress`` objects and advance each bar."""
+        if not self._show_progress:
+            return
+
+        for chain_idx, cp in enumerate(chain_progresses):
+            delta = cp.finished_draws - self._previous_finished[chain_idx]
+            if delta <= 0:
+                continue
+            self._previous_finished[chain_idx] = cp.finished_draws
+            is_last = cp.finished_draws >= cp.total_draws
+            stats = {
+                "divergences": cp.divergences,
+                "step_size": cp.step_size,
+                "tree_size": cp.latest_num_steps,
+                "draw": cp.finished_draws,
+            }
+            task_id = 0 if self.combined_progress else chain_idx
+            self._backend.update(
+                task_id=task_id,
+                advance=delta,
+                failing=cp.divergences > 0,
+                stats=stats,
+                is_last=is_last,
+            )
+
+
 class SMCProgressBarManager(ProgressBarManager):
     """Progress bar manager for SMC sampling.
 

pymc/sampling/mcmc.py

@@ -14,6 +14,7 @@
 
 import contextlib
 import importlib.util
+import inspect
 import logging
 import multiprocessing
 import pickle
@@ -54,6 +55,7 @@
 from pymc.model import Model, modelcontext
 from pymc.progress_bar import (
     MCMCProgressBarManager,
+    NutpieProgressBarManager,
     ProgressBarOptions,
     default_progress_theme,
 )
@@ -333,7 +335,8 @@ def _sample_external_nuts(
     initvals: StartDict | Sequence[StartDict | None] | None,
     model: Model,
     var_names: Sequence[str] | None,
-    progressbar: bool,
+    progressbar: bool | ProgressBarOptions,
+    progressbar_theme: Theme | None,
     quiet: bool,
     idata_kwargs: dict | None,
     compute_convergence_checks: bool,
@@ -387,16 +390,38 @@ def _sample_external_nuts(
             var_names=var_names,
             **compile_kwargs,
         )
-        t_start = time.time()
-        idata = nutpie.sample(
-            compiled_model,
-            draws=draws,
-            tune=tune,
-            chains=chains,
-            seed=_get_seeds_per_chain(random_seed, 1)[0],
-            progress_bar=progressbar,
-            **nuts_kwargs,
+
+        nutpie_supports_progress_callback = (
+            "progress_callback" in inspect.signature(nutpie.sample).parameters
         )
+        if nutpie_supports_progress_callback:
+            pb_manager = NutpieProgressBarManager(
+                chains=chains,
+                draws=draws,
+                tune=tune,
+                progressbar=progressbar,
+                progressbar_theme=progressbar_theme,
+            )
+            nutpie_progress_kwargs = {
+                "progress_bar": False,
+                "progress_callback": pb_manager.update,
+            }
+            progress_cm: contextlib.AbstractContextManager = pb_manager
+        else:
+            nutpie_progress_kwargs = {"progress_bar": bool(progressbar)}
+            progress_cm = contextlib.nullcontext()
+
+        t_start = time.time()
+        with progress_cm:
+            idata = nutpie.sample(
+                compiled_model,
+                draws=draws,
+                tune=tune,
+                chains=chains,
+                seed=_get_seeds_per_chain(random_seed, 1)[0],
+                **nutpie_progress_kwargs,
+                **nuts_kwargs,
+            )
         t_sample = time.time() - t_start
         patch_nutpie_idata(
             idata,
@@ -438,7 +463,7 @@ def _sample_external_nuts(
             initvals=initvals,
             model=model,
             var_names=var_names,
-            progressbar=progressbar,
+            progressbar=bool(progressbar),
             quiet=quiet,
             nuts_sampler=sampler,
             nuts_kwargs=jax_nuts_kwargs,
@@ -877,7 +902,8 @@ def sample(
                 initvals=initvals,
                 model=model,
                 var_names=var_names,
-                progressbar=progress_bool,
+                progressbar=progressbar if nuts_sampler == "nutpie" else progress_bool,
+                progressbar_theme=progressbar_theme,
                 quiet=quiet,
                 idata_kwargs=idata_kwargs,
                 compute_convergence_checks=compute_convergence_checks,

tests/sampling/test_mcmc_external.py

@@ -12,12 +12,17 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+import unittest.mock as mock
+
+from types import SimpleNamespace
+
 import numpy as np
 import numpy.testing as npt
 import pytest
 import xarray as xr
 
 from pymc import Data, Deterministic, HalfNormal, Model, Normal, sample
+from pymc.progress_bar import NutpieProgressBarManager
 
 
 # temporarily skip nutpie
@@ -147,3 +152,69 @@ def test_sample_var_names(nuts_sampler):
         assert var in idata_2.posterior
 
         xr.testing.assert_allclose(idata_1.posterior[var], idata_2.posterior[var])
+
+
+def test_nutpie_progress_bar_manager_update():
+    pb = NutpieProgressBarManager(chains=2, draws=10, tune=10, progressbar=False)
+    pb._backend = mock.Mock()
+    pb._show_progress = True  # force the update path even without a real backend
+
+    cp0 = SimpleNamespace(
+        finished_draws=5,
+        total_draws=20,
+        divergences=0,
+        step_size=0.5,
+        latest_num_steps=3,
+    )
+    cp1 = SimpleNamespace(
+        finished_draws=4,
+        total_draws=20,
+        divergences=1,
+        step_size=0.4,
+        latest_num_steps=7,
+    )
+    pb.update([cp0, cp1])
+    assert pb._backend.update.call_count == 2
+    first_call = pb._backend.update.call_args_list[0].kwargs
+    assert first_call["task_id"] == 0
+    assert first_call["advance"] == 5
+    assert first_call["stats"]["divergences"] == 0
+    second_call = pb._backend.update.call_args_list[1].kwargs
+    assert second_call["task_id"] == 1
+    assert second_call["advance"] == 4
+    assert second_call["failing"] is True
+
+    # A second update only advances by the delta since the previous call.
+    cp0.finished_draws = 20
+    cp1.finished_draws = 20
+    pb._backend.update.reset_mock()
+    pb.update([cp0, cp1])
+    deltas = [c.kwargs["advance"] for c in pb._backend.update.call_args_list]
+    is_last_flags = [c.kwargs["is_last"] for c in pb._backend.update.call_args_list]
+    assert deltas == [15, 16]
+    assert is_last_flags == [True, True]
+
+
+def test_nutpie_end_to_end():
+    # Released nutpie 0.16.8 references `arviz.InferenceData` which arviz 1.0 removed,
+    # so `import nutpie` raises AttributeError on the current CI matrix. Skip until a
+    # nutpie release compatible with arviz 1.0 ships.
+    try:
+        import nutpie  # noqa: F401
+    except (ImportError, AttributeError):
+        pytest.skip("nutpie unavailable or incompatible with the installed arviz")
+    with Model() as m:
+        HalfNormal("sigma")
+        Normal("mu")
+        Normal("y", mu=0, sigma=1, observed=[1.0, 2.0, 3.0])
+        idata = sample(
+            nuts_sampler="nutpie",
+            tune=20,
+            draws=20,
+            chains=2,
+            progressbar=False,
+            random_seed=1411,
+        )
+    assert {"posterior", "sample_stats", "observed_data"} <= set(idata.children)
+    assert set(idata.posterior.data_vars) == {"mu", "sigma"}
+    assert idata.posterior.sizes == {"chain": 2, "draw": 20}