Introduce StatsBijection to combine stats dicts

Author: michaelosthege

pymc/blocking.py

@@ -31,7 +31,8 @@
 
 T = TypeVar("T")
 PointType: TypeAlias = Dict[str, np.ndarray]
-StatsType: TypeAlias = List[Dict[str, Any]]
+StatsDict: TypeAlias = Dict[str, Any]
+StatsType: TypeAlias = List[StatsDict]
 
 # `point_map_info` is a tuple of tuples containing `(name, shape, dtype)` for
 # each of the raveled variables.

pymc/step_methods/compound.py

@@ -18,16 +18,15 @@
 @author: johnsalvatier
 """
 
-
 from abc import ABC, abstractmethod
 from enum import IntEnum, unique
-from typing import Dict, List, Tuple
+from typing import Dict, List, Sequence, Tuple, Union
 
 import numpy as np
 
 from pytensor.graph.basic import Variable
 
-from pymc.blocking import PointType, StatsType
+from pymc.blocking import PointType, StatsDict, StatsType
 from pymc.model import modelcontext
 
 __all__ = ("Competence", "CompoundStep")
@@ -165,3 +164,43 @@ def reset_tuning(self):
     @property
     def vars(self):
         return [var for method in self.methods for var in method.vars]
+
+
+def flatten_steps(step: Union[BlockedStep, CompoundStep]) -> List[BlockedStep]:
+    """Flatten a hierarchy of step methods to a list."""
+    if isinstance(step, BlockedStep):
+        return [step]
+    steps = []
+    if not isinstance(step, CompoundStep):
+        raise ValueError(f"Unexpected type of step method: {step}")
+    for sm in step.methods:
+        steps += flatten_steps(sm)
+    return steps
+
+
+class StatsBijection:
+    """Map between a `list` of stats to `dict` of stats."""
+
+    def __init__(self, sampler_stats_dtypes: Sequence[Dict[str, type]]) -> None:
+        # Keep a list of flat vs. original stat names
+        self._stat_groups: List[List[Tuple[str, str]]] = [
+            [(f"sampler_{s}__{statname}", statname) for statname, _ in names_dtypes.items()]
+            for s, names_dtypes in enumerate(sampler_stats_dtypes)
+        ]
+
+    def map(self, stats_list: StatsType) -> StatsDict:
+        """Combine stats dicts of multiple samplers into one dict."""
+        stats_dict = {}
+        for s, sts in enumerate(stats_list):
+            for statname, sval in sts.items():
+                sname = f"sampler_{s}__{statname}"
+                stats_dict[sname] = sval
+        return stats_dict
+
+    def rmap(self, stats_dict: StatsDict) -> StatsType:
+        """Split a global stats dict into a list of sampler-wise stats dicts."""
+        stats_list = []
+        for namemap in self._stat_groups:
+            d = {statname: stats_dict[sname] for sname, statname in namemap}
+            stats_list.append(d)
+        return stats_list

pymc/tests/step_methods/test_compound.py

@@ -25,6 +25,7 @@
     Metropolis,
     Slice,
 )
+from pymc.step_methods.compound import StatsBijection, flatten_steps
 from pymc.tests.helpers import StepMethodTester, fast_unstable_sampling_mode
 from pymc.tests.models import simple_2model_continuous
 
@@ -91,3 +92,41 @@ def test_compound_step(self):
                 step2 = NUTS([c2])
                 step = CompoundStep([step1, step2])
             assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set(step.vars)
+
+
+class TestStatsBijection:
+    def test_flatten_steps(self):
+        with pm.Model():
+            a = pm.Normal("a")
+            b = pm.Normal("b")
+            c = pm.Normal("c")
+            s1 = Metropolis([a])
+            s2 = Metropolis([b])
+            c1 = CompoundStep([s1, s2])
+            s3 = NUTS([c])
+            c2 = CompoundStep([c1, s3])
+        assert flatten_steps(s1) == [s1]
+        assert flatten_steps(c2) == [s1, s2, s3]
+        with pytest.raises(ValueError, match="Unexpected type"):
+            flatten_steps("not a step")
+
+    def test_stats_bijection(self):
+        step_stats_dtypes = [
+            {"a": float, "b": int},
+            {"a": float, "c": int},
+        ]
+        bij = StatsBijection(step_stats_dtypes)
+        stats_l = [
+            dict(a=1.5, b=3),
+            dict(a=2.5, c=4),
+        ]
+        stats_d = bij.map(stats_l)
+        assert isinstance(stats_d, dict)
+        assert stats_d["sampler_0__a"] == 1.5
+        assert stats_d["sampler_0__b"] == 3
+        assert stats_d["sampler_1__a"] == 2.5
+        assert stats_d["sampler_1__c"] == 4
+        rev = bij.rmap(stats_d)
+        assert isinstance(rev, list)
+        assert len(rev) == len(stats_l)
+        assert rev == stats_l