Refactor StatsBijection to expose information about object stats

Author: michaelosthege

pymc/step_methods/compound.py

@@ -191,6 +191,11 @@ def stop_tuning(self):
             self.tune = False
 
 
+def flat_statname(sampler_idx: int, sname: str) -> str:
+    """Get the flat-stats name for a samplers stat."""
+    return f"sampler_{sampler_idx}__{sname}"
+
+
 def get_stats_dtypes_shapes_from_steps(
     steps: Iterable[BlockedStep],
 ) -> Dict[str, Tuple[StatDtype, StatShape]]:
@@ -201,7 +206,7 @@ def get_stats_dtypes_shapes_from_steps(
     result = {}
     for s, step in enumerate(steps):
         for sname, (dtype, shape) in step.stats_dtypes_shapes.items():
-            result[f"sampler_{s}__{sname}"] = (dtype, shape)
+            result[flat_statname(s, sname)] = (dtype, shape)
     return result
 
 
@@ -262,10 +267,21 @@ class StatsBijection:
 
     def __init__(self, sampler_stats_dtypes: Sequence[Mapping[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)
-        ]
+        stat_groups = []
+        for s, names_dtypes in enumerate(sampler_stats_dtypes):
+            group = []
+            for statname, dtype in names_dtypes.items():
+                flatname = flat_statname(s, statname)
+                is_obj = np.dtype(dtype) == np.dtype(object)
+                group.append((flatname, statname, is_obj))
+            stat_groups.append(group)
+        self._stat_groups: List[List[Tuple[str, str, bool]]] = stat_groups
+        self.object_stats = {
+            fname: (s, sname)
+            for s, group in enumerate(self._stat_groups)
+            for fname, sname, is_obj in group
+            if is_obj
+        }
 
     @property
     def n_samplers(self) -> int:
@@ -275,9 +291,10 @@ def map(self, stats_list: Sequence[Mapping[str, Any]]) -> 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
+            for fname, sname, is_obj in self._stat_groups[s]:
+                if sname not in sts:
+                    continue
+                stats_dict[fname] = sts[sname]
         return stats_dict
 
     def rmap(self, stats_dict: Mapping[str, Any]) -> StatsType:
@@ -286,7 +303,11 @@ def rmap(self, stats_dict: Mapping[str, Any]) -> StatsType:
         The ``stats_dict`` can be a subset of all sampler stats.
         """
         stats_list = []
-        for namemap in self._stat_groups:
-            d = {statname: stats_dict[sname] for sname, statname in namemap if sname in stats_dict}
+        for group in self._stat_groups:
+            d = {}
+            for fname, sname, is_obj in group:
+                if fname not in stats_dict:
+                    continue
+                d[sname] = stats_dict[fname]
             stats_list.append(d)
         return stats_list

tests/step_methods/test_compound.py

@@ -164,20 +164,22 @@ def test_flatten_steps(self):
     def test_stats_bijection(self):
         step_stats_dtypes = [
             {"a": float, "b": int},
-            {"a": float, "c": int},
+            {"a": float, "c": Warning},
         ]
         bij = StatsBijection(step_stats_dtypes)
+        assert bij.object_stats == {"sampler_1__c": (1, "c")}
         assert bij.n_samplers == 2
+        w = Warning("hmm")
         stats_l = [
             dict(a=1.5, b=3),
-            dict(a=2.5, c=4),
+            dict(a=2.5, c=w),
         ]
         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
+        assert stats_d["sampler_1__c"] == w
         rev = bij.rmap(stats_d)
         assert isinstance(rev, list)
         assert len(rev) == len(stats_l)