Make sampler stats accessible after loading (#3773)

* Make sampler stats accessible after loading

* Add check to verify the fix.

Co-authored-by: rpgoldman <[email protected]>

Author: ColCarroll

pymc3/backends/ndarray.py

@@ -101,16 +101,21 @@ def to_metadata(ndarray):
         """Extract ndarray metadata into json-serializable content"""
         if ndarray._stats is None:
             stats = ndarray._stats
+            sampler_vars = None
         else:
             stats = []
+            sampler_vars = []
             for stat in ndarray._stats:
                 stats.append({key: value.tolist() for key, value in stat.items()})
+                sampler_vars.append({key: str(value.dtype) for key, value in stat.items()})
+
 
         metadata = {
             'draw_idx': ndarray.draw_idx,
             'draws': ndarray.draws,
             '_stats': stats,
             'chain': ndarray.chain,
+            'sampler_vars': sampler_vars
         }
         return metadata
 
@@ -145,6 +150,9 @@ def load(self, model: Model) -> 'NDArray':
 
         metadata['_stats'] = [{k: np.array(v) for k, v in stat.items()} for stat in metadata['_stats']]
 
+        sampler_vars = metadata.pop('sampler_vars')
+        new_trace._set_sampler_vars(sampler_vars)
+
         for key, value in metadata.items():
             setattr(new_trace, key, value)
         new_trace.samples = dict(np.load(self.samples_path))

pymc3/tests/test_ndarray_backend.py

@@ -6,17 +6,9 @@
 import pytest
 
 
-STATS1 = [{
-    'a': np.float64,
-    'b': np.bool
-}]
+STATS1 = [{"a": np.float64, "b": np.bool}]
 
-STATS2 = [{
-    'a': np.float64
-}, {
-    'a': np.float64,
-    'b': np.int64,
-}]
+STATS2 = [{"a": np.float64}, {"a": np.float64, "b": np.int64,}]
 
 
 class TestNDArray0dSampling(bf.SamplingTestCase):
@@ -152,7 +144,7 @@ class TestMultiTrace_add_remove_values(bf.ModelBackendSampledTestCase):
     def test_add_values(self):
         mtrace = self.mtrace
         orig_varnames = list(mtrace.varnames)
-        name = 'new_var'
+        name = "new_var"
         vals = mtrace[orig_varnames[0]]
         mtrace.add_values({name: vals})
         assert len(orig_varnames) == len(mtrace.varnames) - 1
@@ -164,7 +156,6 @@ def test_add_values(self):
 
 
 class TestSqueezeCat:
-
     def setup_method(self):
         self.x = np.arange(10)
         self.y = np.arange(10, 20)
@@ -194,13 +185,14 @@ def test_combine_true_squeeze_true(self):
         result = base._squeeze_cat([self.x, self.y], True, True)
         npt.assert_equal(result, expected)
 
+
 class TestSaveLoad:
     @staticmethod
     def model():
         with pm.Model() as model:
-            x = pm.Normal('x', 0, 1)
-            y = pm.Normal('y', x, 1, observed=2)
-            z = pm.Normal('z', x + y, 1)
+            x = pm.Normal("x", 0, 1)
+            y = pm.Normal("y", x, 1, observed=2)
+            z = pm.Normal("z", x + y, 1)
         return model
 
     @classmethod
@@ -209,12 +201,12 @@ def setup_class(cls):
             cls.trace = pm.sample()
 
     def test_save_new_model(self, tmpdir_factory):
-        directory = str(tmpdir_factory.mktemp('data'))
+        directory = str(tmpdir_factory.mktemp("data"))
         save_dir = pm.save_trace(self.trace, directory, overwrite=True)
 
         assert save_dir == directory
         with pm.Model() as model:
-            w = pm.Normal('w', 0, 1)
+            w = pm.Normal("w", 0, 1)
             new_trace = pm.sample()
 
         with pytest.raises(OSError):
@@ -224,26 +216,32 @@ def test_save_new_model(self, tmpdir_factory):
         with model:
             new_trace_copy = pm.load_trace(directory)
 
-        assert (new_trace['w'] == new_trace_copy['w']).all()
+        assert (new_trace["w"] == new_trace_copy["w"]).all()
 
     def test_save_and_load(self, tmpdir_factory):
-        directory = str(tmpdir_factory.mktemp('data'))
+        directory = str(tmpdir_factory.mktemp("data"))
         save_dir = pm.save_trace(self.trace, directory, overwrite=True)
 
         assert save_dir == directory
 
         trace2 = pm.load_trace(directory, model=TestSaveLoad.model())
 
-        for var in ('x', 'z'):
+        for var in ("x", "z"):
             assert (self.trace[var] == trace2[var]).all()
 
+        assert self.trace.stat_names == trace2.stat_names
+        for stat in self.trace.stat_names:
+            assert all(self.trace[stat] == trace2[stat]), (
+                "Restored value of statistic %s does not match stored value" % stat
+            )
+
     def test_bad_load(self, tmpdir_factory):
-        directory = str(tmpdir_factory.mktemp('data'))
+        directory = str(tmpdir_factory.mktemp("data"))
         with pytest.raises(pm.TraceDirectoryError):
             pm.load_trace(directory, model=TestSaveLoad.model())
 
     def test_sample_posterior_predictive(self, tmpdir_factory):
-        directory = str(tmpdir_factory.mktemp('data'))
+        directory = str(tmpdir_factory.mktemp("data"))
         save_dir = pm.save_trace(self.trace, directory, overwrite=True)
 
         assert save_dir == directory