Ensure stanfit objects can be deepcopied and pickled

Author: WardBrian

cmdstanpy/stanfit/mcmc.py

@@ -43,6 +43,7 @@
     get_logger,
     scan_generated_quantities_csv,
 )
+
 from .metadata import InferenceMetadata
 from .runset import RunSet
 
@@ -122,6 +123,8 @@ def __repr__(self) -> str:
 
     def __getattr__(self, attr: str) -> np.ndarray:
         """Synonymous with ``fit.stan_variable(attr)"""
+        if attr.startswith("_"):
+            raise AttributeError(f"Unknown variable name {attr}")
         try:
             return self.stan_variable(attr)
         except ValueError as e:
@@ -833,6 +836,8 @@ def __repr__(self) -> str:
 
     def __getattr__(self, attr: str) -> np.ndarray:
         """Synonymous with ``fit.stan_variable(attr)"""
+        if attr.startswith("_"):
+            raise AttributeError(f"Unknown variable name {attr}")
         try:
             return self.stan_variable(attr)
         except ValueError as e:

cmdstanpy/stanfit/mle.py

@@ -52,6 +52,8 @@ def __repr__(self) -> str:
 
     def __getattr__(self, attr: str) -> Union[np.ndarray, float]:
         """Synonymous with ``fit.stan_variable(attr)"""
+        if attr.startswith("_"):
+            raise AttributeError(f"Unknown variable name {attr}")
         try:
             return self.stan_variable(attr)
         except ValueError as e:

cmdstanpy/stanfit/vb.py

@@ -43,6 +43,8 @@ def __repr__(self) -> str:
 
     def __getattr__(self, attr: str) -> Union[np.ndarray, float]:
         """Synonymous with ``fit.stan_variable(attr)"""
+        if attr.startswith("_"):
+            raise AttributeError(f"Unknown variable name {attr}")
         try:
             return self.stan_variable(attr)
         except ValueError as e:

test/test_optimize.py

@@ -627,6 +627,28 @@ def test_attrs(self):
         with self.assertRaisesRegex(AttributeError, 'Unknown variable name:'):
             dummy = fit.c
 
+    def test_pickle_ability(self):
+        """Ensure fit objects are pickle-able and copy-able"""
+
+        import pickle
+
+        csvfiles_path = os.path.join(
+            DATAFILES_PATH, 'optimize', 'rosenbrock_mle.csv'
+        )
+        fit = from_csv(path=csvfiles_path)
+        pickled = pickle.dumps(fit)
+        unpickled = pickle.loads(pickled)
+        self.assertSequenceEqual(
+            fit.stan_variables().keys(), unpickled.stan_variables().keys()
+        )
+
+        import copy
+
+        fit2 = copy.deepcopy(fit)
+        self.assertSequenceEqual(
+            fit.stan_variables().keys(), fit2.stan_variables().keys()
+        )
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_sample.py

@@ -1881,6 +1881,26 @@ def test_diagnostics(self):
         self.assertEqual(fit.max_treedepths, None)
         self.assertEqual(fit.divergences, None)
 
+    def test_pickle_ability(self):
+        """Ensure fit objects are pickle-able and copy-able"""
+
+        import pickle
+
+        csvfiles_path = os.path.join(DATAFILES_PATH, 'lotka-volterra.csv')
+        fit = from_csv(path=csvfiles_path)
+        pickled = pickle.dumps(fit)
+        unpickled = pickle.loads(pickled)
+        self.assertSequenceEqual(
+            fit.stan_variables().keys(), unpickled.stan_variables().keys()
+        )
+
+        import copy
+
+        fit2 = copy.deepcopy(fit)
+        self.assertSequenceEqual(
+            fit.stan_variables().keys(), fit2.stan_variables().keys()
+        )
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_variational.py

@@ -284,6 +284,26 @@ def test_attrs(self):
         with self.assertRaisesRegex(AttributeError, 'Unknown variable name:'):
             dummy = fit.c
 
+    def test_pickle_ability(self):
+        """Ensure fit objects are pickle-able and copy-able"""
+
+        import pickle
+
+        csvfiles_path = os.path.join(DATAFILES_PATH, 'variational')
+        fit = from_csv(path=csvfiles_path)
+        pickled = pickle.dumps(fit)
+        unpickled = pickle.loads(pickled)
+        self.assertSequenceEqual(
+            fit.stan_variables().keys(), unpickled.stan_variables().keys()
+        )
+
+        import copy
+
+        fit2 = copy.deepcopy(fit)
+        self.assertSequenceEqual(
+            fit.stan_variables().keys(), fit2.stan_variables().keys()
+        )
+
 
 if __name__ == '__main__':
     unittest.main()