Merge pull request #632 from tillahoffmann/getstate

Add serialization of samples.

Author: WardBrian

cmdstanpy/stanfit/gq.py

@@ -85,6 +85,10 @@ def __getattr__(self, attr: str) -> np.ndarray:
             # pylint: disable=raise-missing-from
             raise AttributeError(*e.args)
 
+    def __getstate__(self) -> dict:
+        self._assemble_generated_quantities()
+        return self.__dict__
+
     def _validate_csv_files(self) -> Dict[str, Any]:
         """
         Checks that Stan CSV output files for all chains are consistent
@@ -189,8 +193,7 @@ def draws(
         CmdStanGQ.draws_xr
         CmdStanMCMC.draws
         """
-        if self._draws.shape == (0,):
-            self._assemble_generated_quantities()
+        self._assemble_generated_quantities()
         if (
             inc_warmup
             and not self.mcmc_sample.metadata.cmdstan_config['save_warmup']
@@ -277,8 +280,7 @@ def draws_pd(
                 'Draws from warmup iterations not available,'
                 ' must run sampler with "save_warmup=True".'
             )
-        if self._draws.shape == (0,):
-            self._assemble_generated_quantities()
+        self._assemble_generated_quantities()
 
         gq_cols = []
         mcmc_vars = []
@@ -400,8 +402,7 @@ def draws_xr(
         for var in dup_vars:
             vars_list.remove(var)
 
-        if self._draws.shape == (0,):
-            self._assemble_generated_quantities()
+        self._assemble_generated_quantities()
 
         num_draws = self.mcmc_sample.num_draws_sampling
         sample_config = self.mcmc_sample.metadata.cmdstan_config
@@ -505,8 +506,7 @@ def stan_variable(
         if var not in gq_var_names:
             return self.mcmc_sample.stan_variable(var, inc_warmup=inc_warmup)
         else:  # is gq variable
-            if self._draws.shape == (0,):
-                self._assemble_generated_quantities()
+            self._assemble_generated_quantities()
             draw1 = 0
             if (
                 not inc_warmup
@@ -561,6 +561,8 @@ def stan_variables(self, inc_warmup: bool = False) -> Dict[str, np.ndarray]:
         return result
 
     def _assemble_generated_quantities(self) -> None:
+        if self._draws.shape != (0,):
+            return
         # use numpy loadtxt
         warmup = self.mcmc_sample.metadata.cmdstan_config['save_warmup']
         num_draws = self.mcmc_sample.draws(inc_warmup=warmup).shape[0]

cmdstanpy/stanfit/mcmc.py

@@ -129,6 +129,14 @@ def __getattr__(self, attr: str) -> np.ndarray:
             # pylint: disable=raise-missing-from
             raise AttributeError(*e.args)
 
+    def __getstate__(self) -> dict:
+        # This function returns the mapping of objects to serialize with pickle.
+        # See https://docs.python.org/3/library/pickle.html#object.__getstate__
+        # for details. We call _assemble_draws to ensure posterior samples have
+        # been loaded prior to serialization.
+        self._assemble_draws()
+        return self.__dict__
+
     @property
     def chains(self) -> int:
         """Number of chains."""
@@ -259,8 +267,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 +598,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 +654,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 +740,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 +787,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_generate_quantities.py

@@ -5,15 +5,19 @@
 import json
 import logging
 import os
+import pickle
+import shutil
 import unittest
 from test import CustomTestCase
 
 import numpy as np
 import pandas as pd
+import pytest
 from numpy.testing import assert_array_equal, assert_raises
 from testfixtures import LogCapture
 
 import cmdstanpy.stanfit
+from cmdstanpy.stanfit import CmdStanGQ
 from cmdstanpy.cmdstan_args import Method
 from cmdstanpy.model import CmdStanModel
 
@@ -485,6 +489,26 @@ def test_timeout(self):
                 timeout=0.1, mcmc_sample=fit, data={'loop': 1}
             )
 
+    @pytest.mark.order(before="test_no_xarray")
+    def test_serialization(self):
+        stan_bern = os.path.join(DATAFILES_PATH, 'bernoulli.stan')
+        model_bern = CmdStanModel(stan_file=stan_bern)
+        jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
+        fit_sampling = model_bern.sample(chains=1, iter_sampling=10, data=jdata)
+
+        stan = os.path.join(DATAFILES_PATH, 'named_output.stan')
+        model = CmdStanModel(stan_file=stan)
+        fit1 = model.generate_quantities(data=jdata, mcmc_sample=fit_sampling)
+
+        dumped = pickle.dumps(fit1)
+        shutil.rmtree(fit1.runset._output_dir)
+        fit2: CmdStanGQ = pickle.loads(dumped)
+        variables1 = fit1.stan_variables()
+        variables2 = 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()

test/test_optimize.py

@@ -4,6 +4,7 @@
 import io
 import json
 import os
+import pickle
 import shutil
 import unittest
 
@@ -649,6 +650,32 @@ def test_timeout(self):
         with self.assertRaises(TimeoutError):
             timeout_model.optimize(data={'loop': 1}, timeout=0.1)
 
+    def test_serialization(self):
+        stan = os.path.join(DATAFILES_PATH, 'bernoulli.stan')
+        model = CmdStanModel(stan_file=stan)
+        jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
+        jinit = os.path.join(DATAFILES_PATH, 'bernoulli.init.json')
+        mle1 = model.optimize(
+            data=jdata,
+            seed=1239812093,
+            inits=jinit,
+            algorithm='LBFGS',
+            init_alpha=0.001,
+            iter=100,
+            tol_obj=1e-12,
+            tol_rel_obj=1e4,
+            tol_grad=1e-8,
+            tol_rel_grad=1e7,
+            tol_param=1e-8,
+            history_size=5,
+        )
+        dumped = pickle.dumps(mle1)
+        shutil.rmtree(mle1.runset._output_dir)
+        mle2: CmdStanMLE = pickle.loads(dumped)
+        np.testing.assert_array_equal(
+            mle1.optimized_params_np, mle2.optimized_params_np
+        )
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_sample.py

@@ -11,10 +11,12 @@
 import tempfile
 import unittest
 from multiprocessing import cpu_count
+import pickle
 from test import CustomTestCase
 from time import time
 
 import numpy as np
+import pytest
 from testfixtures import LogCapture, StringComparison
 
 import cmdstanpy.stanfit
@@ -1928,6 +1930,33 @@ 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]))
 
+    @pytest.mark.order(before="test_no_xarray")
+    def test_serialization(self, stanfile='bernoulli.stan'):
+        # This test must before any test that uses the `without_import` context
+        # manager because the latter uses `reload` with side effects that affect
+        # the consistency of classes.
+        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)
+        shutil.rmtree(bern_fit1.runset._output_dir)
+        # 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()

test/test_variational.py

@@ -3,6 +3,7 @@
 import contextlib
 import io
 import os
+import pickle
 import shutil
 import unittest
 from math import fabs
@@ -297,6 +298,23 @@ def test_timeout(self):
         with self.assertRaises(TimeoutError):
             timeout_model.variational(timeout=0.1, data={'loop': 1})
 
+    def test_serialization(self):
+        stan = os.path.join(
+            DATAFILES_PATH, 'variational', 'eta_should_be_big.stan'
+        )
+        model = CmdStanModel(stan_file=stan)
+        variational1 = model.variational(algorithm='meanfield', seed=999999)
+        dumped = pickle.dumps(variational1)
+        shutil.rmtree(variational1.runset._output_dir)
+        variational2: CmdStanVB = pickle.loads(dumped)
+        np.testing.assert_array_equal(
+            variational1.variational_sample, variational2.variational_sample
+        )
+        self.assertEqual(
+            variational1.variational_params_dict,
+            variational2.variational_params_dict,
+        )
+
 
 if __name__ == '__main__':
     unittest.main()