Assign each step method a unique id, and track ids via stats emitted by Step.step

Author: jessegrabowski

pymc/backends/ndarray.py

@@ -113,6 +113,10 @@ def record(self, point, sampler_stats=None) -> None:
         if sampler_stats is not None:
             for data, vars in zip(self._stats, sampler_stats):
                 for key, val in vars.items():
+                    # step_meta is a key used by the progress bars to track which draw came from which step instance. It
+                    # should never be stored as a sampler statistic.
+                    if key == "step_meta":
+                        continue
                     data[key][draw_idx] = val
         elif self._stats is not None:
             raise ValueError("Expected sampler_stats")

pymc/sampling/mcmc.py

@@ -15,6 +15,7 @@
 """Functions for MCMC sampling."""
 
 import contextlib
+import itertools
 import logging
 import pickle
 import sys
@@ -111,6 +112,7 @@ def instantiate_steppers(
     step_kwargs: dict[str, dict] | None = None,
     initial_point: PointType | None = None,
     compile_kwargs: dict | None = None,
+    step_id_generator: Iterator[int] | None = None,
 ) -> Step | list[Step]:
     """Instantiate steppers assigned to the model variables.
 
@@ -139,6 +141,9 @@ def instantiate_steppers(
     if step_kwargs is None:
         step_kwargs = {}
 
+    if step_id_generator is None:
+        step_id_generator = itertools.count()
+
     used_keys = set()
     if selected_steps:
         if initial_point is None:
@@ -154,6 +159,7 @@ def instantiate_steppers(
                     model=model,
                     initial_point=initial_point,
                     compile_kwargs=compile_kwargs,
+                    step_id_generator=step_id_generator,
                     **kwargs,
                 )
                 steps.append(step)
@@ -853,16 +859,19 @@ def joined_blas_limiter():
         initial_points = [ipfn(seed) for ipfn, seed in zip(ipfns, random_seed_list)]
 
         # Instantiate automatically selected steps
+        # Use a counter to generate a unique id for each stepper used in the model.
+        step_id_generator = itertools.count()
         step = instantiate_steppers(
             model,
             steps=provided_steps,
             selected_steps=selected_steps,
             step_kwargs=kwargs,
             initial_point=initial_points[0],
             compile_kwargs=compile_kwargs,
+            step_id_generator=step_id_generator,
         )
         if isinstance(step, list):
-            step = CompoundStep(step)
+            step = CompoundStep(step, step_id_generator=step_id_generator)
 
     if var_names is not None:
         trace_vars = [v for v in model.unobserved_RVs if v.name in var_names]

pymc/step_methods/arraystep.py

@@ -13,7 +13,7 @@
 #   limitations under the License.
 
 from abc import abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from typing import cast
 
 import numpy as np
@@ -43,14 +43,25 @@ class ArrayStep(BlockedStep):
         :py:func:`pymc.util.get_random_generator` for more information.
     """
 
-    def __init__(self, vars, fs, allvars=False, blocked=True, rng: RandomGenerator = None):
+    def __init__(
+        self,
+        vars,
+        fs,
+        allvars=False,
+        blocked=True,
+        rng: RandomGenerator = None,
+        step_id_generator: Iterator[int] | None = None,
+    ):
         self.vars = vars
         self.fs = fs
         self.allvars = allvars
         self.blocked = blocked
         self.rng = get_random_generator(rng)
+        self._step_id = next(step_id_generator) if step_id_generator else None
 
-    def step(self, point: PointType) -> tuple[PointType, StatsType]:
+    def step(
+        self, point: PointType, step_parent_id: int | None = None
+    ) -> tuple[PointType, StatsType]:
         partial_funcs_and_point: list[Callable | PointType] = [
             DictToArrayBijection.mapf(x, start_point=point) for x in self.fs
         ]
@@ -61,6 +72,9 @@ def step(self, point: PointType) -> tuple[PointType, StatsType]:
         apoint = DictToArrayBijection.map(var_dict)
         apoint_new, stats = self.astep(apoint, *partial_funcs_and_point)
 
+        for sts in stats:
+            sts["step_meta"] = {"step_id": self._step_id, "step_parent_id": step_parent_id}
+
         if not isinstance(apoint_new, RaveledVars):
             # We assume that the mapping has stayed the same
             apoint_new = RaveledVars(apoint_new, apoint.point_map_info)
@@ -84,7 +98,14 @@ class ArrayStepShared(BlockedStep):
     and unmapping overhead as well as moving fewer variables around.
     """
 
-    def __init__(self, vars, shared, blocked=True, rng: RandomGenerator = None):
+    def __init__(
+        self,
+        vars,
+        shared,
+        blocked=True,
+        rng: RandomGenerator = None,
+        step_id_generator: Iterator[int] | None = None,
+    ):
         """
         Create the ArrayStepShared object.
 
@@ -103,8 +124,11 @@ def __init__(self, vars, shared, blocked=True, rng: RandomGenerator = None):
         self.shared = {get_var_name(var): shared for var, shared in shared.items()}
         self.blocked = blocked
         self.rng = get_random_generator(rng)
+        self._step_id = next(step_id_generator) if step_id_generator else None
 
-    def step(self, point: PointType) -> tuple[PointType, StatsType]:
+    def step(
+        self, point: PointType, step_parent_id: int | None = None
+    ) -> tuple[PointType, StatsType]:
         full_point = None
         if self.shared:
             for name, shared_var in self.shared.items():
@@ -115,6 +139,9 @@ def step(self, point: PointType) -> tuple[PointType, StatsType]:
         q = DictToArrayBijection.map(point)
         apoint, stats = self.astep(q)
 
+        for sts in stats:
+            sts["step_meta"] = {"step_id": self._step_id, "step_parent_id": step_parent_id}
+
         if not isinstance(apoint, RaveledVars):
             # We assume that the mapping has stayed the same
             apoint = RaveledVars(apoint, q.point_map_info)

pymc/step_methods/compound.py

@@ -21,7 +21,7 @@
 import warnings
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
 from dataclasses import field
 from enum import IntEnum, unique
 from typing import Any
@@ -125,6 +125,8 @@ class BlockedStep(ABC, WithSamplingState):
 
     def __new__(cls, *args, **kwargs):
         blocked = kwargs.get("blocked")
+        step_id_generator = kwargs.pop("step_id_generator", None)
+
         if blocked is None:
             # Try to look up default value from class
             blocked = getattr(cls, "default_blocked", True)
@@ -168,16 +170,19 @@ def __new__(cls, *args, **kwargs):
                 # call __init__
                 _kwargs = kwargs.copy()
                 _kwargs["rng"] = rng
+                _kwargs["step_id_generator"] = step_id_generator
                 step.__init__([var], *args, **_kwargs)
                 # Hack for creating the class correctly when unpickling.
                 step.__newargs = ([var], *args), _kwargs
                 steps.append(step)
 
-            return CompoundStep(steps)
+            return CompoundStep(steps, step_id_generator=step_id_generator)
         else:
             step = super().__new__(cls)
             step.stats_dtypes = stats_dtypes
             step.stats_dtypes_shapes = stats_dtypes_shapes
+            step._step_id = next(step_id_generator) if step_id_generator else None
+
             # Hack for creating the class correctly when unpickling.
             step.__newargs = (vars, *args), kwargs
             return step
@@ -223,7 +228,7 @@ def _make_update_stats_function(self) -> Callable[[dict, dict, int], dict]:
 
         def update_stats(
             displayed_stats: dict[str, np.ndarray],
-            step_stats: dict[str, str | float | int | bool | None],
+            step_stats_dict: dict[int, dict[str, str | float | int | bool | None]],
             chain_idx: int,
         ) -> dict[str, np.ndarray]:
             """
@@ -235,7 +240,7 @@ def update_stats(
                 Dictionary of statistics displayed in the progress bar. The keys are the names of the statistics and
                 the values are the current values of the statistics, with one value per chain being sampled.
 
-            step_stats: dict
+            step_stats_dict: dict
                 Dictionary of statistics generated by the step sampler when taking the current step. The keys are the
                 names of the statistics and the values are the values of the statistics generated by the step sampler.
 
@@ -256,7 +261,9 @@ def __getnewargs_ex__(self):
         return self.__newargs
 
     @abstractmethod
-    def step(self, point: PointType) -> tuple[PointType, StatsType]:
+    def step(
+        self, point: PointType, step_parent_id: int | None = None
+    ) -> tuple[PointType, StatsType]:
         """Perform a single step of the sampler."""
 
     @staticmethod
@@ -315,7 +322,7 @@ class CompoundStep(WithSamplingState):
 
     _state_class = CompoundStepState
 
-    def __init__(self, methods):
+    def __init__(self, methods, step_id_generator: Iterator[int] | None = None):
         self.methods = list(methods)
         self.stats_dtypes = []
         for method in self.methods:
@@ -325,11 +332,12 @@ def __init__(self, methods):
             f"Compound[{', '.join(getattr(m, 'name', 'UNNAMED_STEP') for m in self.methods)}]"
         )
         self.tune = True
+        self._step_id = next(step_id_generator) if step_id_generator else None
 
-    def step(self, point) -> tuple[PointType, StatsType]:
+    def step(self, point, step_parent_id: int | None = None) -> tuple[PointType, StatsType]:
         stats = []
         for method in self.methods:
-            point, sts = method.step(point)
+            point, sts = method.step(point, step_parent_id=self._step_id)
             stats.extend(sts)
         # Model logp can only be the logp of the _last_ stats,
         # if there is one. Pop all others.
@@ -409,7 +417,7 @@ def _progressbar_config(
 
         return columns, stats
 
-    def _make_update_stats_function(self) -> Callable[[dict, list[dict], int], dict]:
+    def _make_update_stats_function(self) -> Callable[[dict, dict[int, dict], int], dict]:
         """
         Create an update function used by the progress bar to update statistics during sampling.
 
@@ -419,11 +427,13 @@ def _make_update_stats_function(self) -> Callable[[dict, list[dict], int], dict]
             Function that updates displayed statistics for the current chain, given statistics generated by the step
             during the most recent step.
         """
-        update_fns = [method._make_update_stats_function() for method in self.methods]
+        update_fns = {
+            method._step_id: method._make_update_stats_function() for method in self.methods
+        }
 
         def update_stats(
             displayed_stats: dict[str, np.ndarray],
-            step_stats: list[dict[str, str | float | int | bool | None]],
+            step_stats_dict: dict[int, dict[str, str | float | int | bool | None]],
             chain_idx: int,
         ) -> dict[str, np.ndarray]:
             """
@@ -435,7 +445,7 @@ def update_stats(
                 Dictionary of statistics displayed in the progress bar. The keys are the names of the statistics and
                 the values are the current values of the statistics, with one value per chain being sampled.
 
-            step_stats: list of dict
+            step_stats_dict: dict of dict
                 List of dictionaries containing statistics generated by **each** step sampler in the CompoundStep when
                  taking the current step. For each dictionary, the keys are names of statistics and the values are
                  the values of the statistics generated by the step sampler.
@@ -452,11 +462,9 @@ def update_stats(
             #  In this case, the current loop logic is just overriding each Metropolis steps' stats with those of the
             #  next step (so the user only ever sees the 3rd step's stats). We should have a better way to aggregate
             #  the stats from each step.
-            if not isinstance(step_stats, list):
-                step_stats = [step_stats]
 
-            for step_stat, update_fn in zip(step_stats, update_fns):
-                displayed_stats = update_fn(displayed_stats, step_stat, chain_idx)
+            for step_id, update_fn in update_fns.items():
+                displayed_stats = update_fn(displayed_stats, step_stats_dict, chain_idx)
 
             return displayed_stats
 

pymc/step_methods/hmc/base_hmc.py

@@ -18,6 +18,7 @@
 import time
 
 from abc import abstractmethod
+from collections.abc import Iterator
 from typing import Any, NamedTuple
 
 import numpy as np
@@ -99,6 +100,7 @@ def __init__(
         step_rand=None,
         rng=None,
         initial_point: PointType | None = None,
+        step_id_generator: Iterator[int] | None = None,
         **pytensor_kwargs,
     ):
         """Set up Hamiltonian samplers with common structures.
@@ -133,6 +135,7 @@ def __init__(
         **pytensor_kwargs: passed to PyTensor functions
         """
         self._model = modelcontext(model)
+        self._step_id = next(step_id_generator) if step_id_generator else None
 
         if vars is None:
             vars = self._model.continuous_value_vars

pymc/step_methods/hmc/nuts.py

@@ -278,7 +278,7 @@ def _make_update_stats_function(self) -> Callable[[dict, dict, int], dict]:
 
         def update_stats(
             displayed_stats: dict[str, np.ndarray],
-            step_stats: dict[str, str | float | int | bool | None],
+            step_stats_dict: dict[int, dict[str, str | float | int | bool | None]],
             chain_idx: int,
         ) -> dict[str, np.ndarray]:
             """
@@ -290,7 +290,7 @@ def update_stats(
                 Dictionary of statistics displayed in the progress bar. The keys are the names of the statistics and
                 the values are the current values of the statistics, with one value per chain being sampled.
 
-            step_stats: dict
+            step_stats_dict: dict
                 Dictionary of statistics generated by the step sampler when taking the current step. The keys are the
                 names of the statistics and the values are the values of the statistics generated by the step sampler.
 
@@ -302,8 +302,7 @@ def update_stats(
             dict
                 The updated statistics dictionary to be displayed in the progress bar.
             """
-            if isinstance(step_stats, list):
-                step_stats = step_stats[0]
+            step_stats = step_stats_dict[self._step_id]
 
             if not step_stats["tune"]:
                 displayed_stats["divergences"][chain_idx] += step_stats["diverging"]