Add serialization of samples.

Author: tillahoffmann

cmdstanpy/stanfit/mcmc.py

@@ -129,6 +129,10 @@ def __getattr__(self, attr: str) -> np.ndarray:
             # pylint: disable=raise-missing-from
             raise AttributeError(*e.args)
 
+    def __getstate__(self) -> dict:
+        self._assemble_draws()
+        return self.__dict__
+
     @property
     def chains(self) -> int:
         """Number of chains."""
@@ -259,8 +263,7 @@ def draws(
         CmdStanMCMC.draws_xr
         CmdStanGQ.draws
         """
-        if self._draws.shape == (0,):
-            self._assemble_draws()
+        self._assemble_draws()
 
         if inc_warmup and not self._save_warmup:
             get_logger().warning(
@@ -591,8 +594,7 @@ def draws_pd(
                 ' must run sampler with "save_warmup=True".'
             )
 
-        if self._draws.shape == (0,):
-            self._assemble_draws()
+        self._assemble_draws()
         cols = []
         if vars is not None:
             for var in dict.fromkeys(vars_list):
@@ -648,8 +650,7 @@ def draws_xr(
         else:
             vars_list = vars
 
-        if self._draws.shape == (0,):
-            self._assemble_draws()
+        self._assemble_draws()
 
         num_draws = self.num_draws_sampling
         meta = self._metadata.cmdstan_config
@@ -735,8 +736,7 @@ def stan_variable(
                 'Available variables are '
                 + ", ".join(self._metadata.stan_vars_dims)
             )
-        if self._draws.shape == (0,):
-            self._assemble_draws()
+        self._assemble_draws()
         draw1 = 0
         if not inc_warmup and self._save_warmup:
             draw1 = self.num_draws_warmup
@@ -783,8 +783,7 @@ def method_variables(self) -> Dict[str, np.ndarray]:
         containing per-draw diagnostic values.
         """
         result = {}
-        if self._draws.shape == (0,):
-            self._assemble_draws()
+        self._assemble_draws()
         for idxs in self.metadata.method_vars_cols.values():
             for idx in idxs:
                 result[self.column_names[idx]] = self._draws[:, :, idx]

test/test_sample.py

@@ -11,6 +11,7 @@
 import tempfile
 import unittest
 from multiprocessing import cpu_count
+import pickle
 from test import CustomTestCase
 from time import time
 
@@ -1928,6 +1929,30 @@ def test_json_edges(self):
         self.assertTrue(np.isnan(fit.stan_variable("nan_out")[0]))
         self.assertTrue(np.isinf(fit.stan_variable("inf_out")[0]))
 
+    def test_serialization(self, stanfile='bernoulli.stan'):
+        stan = os.path.join(DATAFILES_PATH, stanfile)
+        bern_model = CmdStanModel(stan_file=stan)
+
+        jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
+        bern_fit1 = bern_model.sample(
+            data=jdata,
+            chains=1,
+            iter_warmup=200,
+            iter_sampling=100,
+            show_progress=False,
+        )
+        # Dump the result (which assembles draws) and delete the source files.
+        dumped = pickle.dumps(bern_fit1)
+        for filename in bern_fit1.runset.csv_files:
+            os.unlink(filename)
+        # Load the serialized result and compare results.
+        bern_fit2: CmdStanMCMC = pickle.loads(dumped)
+        variables1 = bern_fit1.stan_variables()
+        variables2 = bern_fit2.stan_variables()
+        self.assertEqual(set(variables1), set(variables2))
+        for key, value1 in variables1.items():
+            np.testing.assert_array_equal(value1, variables2[key])
+
 
 if __name__ == '__main__':
     unittest.main()