Merge pull request #583 from stan-dev/fix/pickle-copy-stanfits

Ensure stanfit objects can be deepcopied and pickled

Author: WardBrian

cmdstanpy/stanfit/__init__.py

@@ -29,7 +29,8 @@
 
 
 def from_csv(
-    path: Union[str, List[str], None] = None, method: Optional[str] = None
+    path: Union[str, List[str], os.PathLike, None] = None,
+    method: Optional[str] = None,
 ) -> Union[CmdStanMCMC, CmdStanMLE, CmdStanVB, None]:
     """
     Instantiate a CmdStan object from a the Stan CSV files from a CmdStan run.
@@ -61,22 +62,22 @@ def from_csv(
     csvfiles = []
     if isinstance(path, list):
         csvfiles = path
-    elif isinstance(path, str):
-        if '*' in path:
-            splits = os.path.split(path)
-            if splits[0] is not None:
-                if not (os.path.exists(splits[0]) and os.path.isdir(splits[0])):
-                    raise ValueError(
-                        'Invalid path specification, {} '
-                        ' unknown directory: {}'.format(path, splits[0])
-                    )
-            csvfiles = glob.glob(path)
-        elif os.path.exists(path) and os.path.isdir(path):
+    elif isinstance(path, str) and '*' in path:
+        splits = os.path.split(path)
+        if splits[0] is not None:
+            if not (os.path.exists(splits[0]) and os.path.isdir(splits[0])):
+                raise ValueError(
+                    'Invalid path specification, {} '
+                    ' unknown directory: {}'.format(path, splits[0])
+                )
+        csvfiles = glob.glob(path)
+    elif isinstance(path, (str, os.PathLike)):
+        if os.path.exists(path) and os.path.isdir(path):
             for file in os.listdir(path):
-                if file.endswith(".csv"):
+                if os.path.splitext(file)[1] == ".csv":
                     csvfiles.append(os.path.join(path, file))
         elif os.path.exists(path):
-            csvfiles.append(path)
+            csvfiles.append(str(path))
         else:
             raise ValueError('Invalid path specification: {}'.format(path))
     else:
@@ -85,7 +86,7 @@ def from_csv(
     if len(csvfiles) == 0:
         raise ValueError('No CSV files found in directory {}'.format(path))
     for file in csvfiles:
-        if not (os.path.exists(file) and file.endswith('.csv')):
+        if not (os.path.exists(file) and os.path.splitext(file)[1] == ".csv"):
             raise ValueError(
                 'Bad CSV file path spec,'
                 ' includes non-csv file: {}'.format(file)

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_compliance.py

@@ -0,0 +1,67 @@
+"""Testing for things like pickleability, deep copying"""
+
+import copy
+import pathlib
+import pickle
+import unittest
+
+import cmdstanpy
+
+DATAFILES_PATH = pathlib.Path(__file__).parent.resolve() / 'data'
+
+
+class SampleCompliance(unittest.TestCase):
+    def test_sample_pickle_ability(self):
+        csvfiles_path = DATAFILES_PATH / 'lotka-volterra.csv'
+        fit = cmdstanpy.from_csv(path=csvfiles_path)
+        keys = fit.stan_variables().keys()
+        pickled = pickle.dumps(fit)
+        del fit
+        unpickled = pickle.loads(pickled)
+        self.assertSequenceEqual(keys, unpickled.stan_variables().keys())
+
+    def test_sample_copy_ability(self):
+        csvfiles_path = DATAFILES_PATH / 'lotka-volterra.csv'
+        fit = cmdstanpy.from_csv(path=csvfiles_path)
+        fit2 = copy.deepcopy(fit)
+        self.assertSequenceEqual(
+            fit.stan_variables().keys(), fit2.stan_variables().keys()
+        )
+
+
+class OptimizeCompliance(unittest.TestCase):
+    def test_optimize_pickle_ability(self):
+        csvfiles_path = DATAFILES_PATH / 'optimize' / 'rosenbrock_mle.csv'
+        fit = cmdstanpy.from_csv(path=csvfiles_path)
+        keys = fit.stan_variables().keys()
+        pickled = pickle.dumps(fit)
+        del fit
+        unpickled = pickle.loads(pickled)
+        self.assertSequenceEqual(keys, unpickled.stan_variables().keys())
+
+    def test_optimize_copy_ability(self):
+        csvfiles_path = DATAFILES_PATH / 'optimize' / 'rosenbrock_mle.csv'
+        fit = cmdstanpy.from_csv(path=csvfiles_path)
+        fit2 = copy.deepcopy(fit)
+        self.assertSequenceEqual(
+            fit.stan_variables().keys(), fit2.stan_variables().keys()
+        )
+
+
+class VariationalCompliance(unittest.TestCase):
+    def test_variational_pickle_ability(self):
+        csvfiles_path = DATAFILES_PATH / 'variational'
+        fit = cmdstanpy.from_csv(path=csvfiles_path)
+        keys = fit.stan_variables().keys()
+        pickled = pickle.dumps(fit)
+        del fit
+        unpickled = pickle.loads(pickled)
+        self.assertSequenceEqual(keys, unpickled.stan_variables().keys())
+
+    def test_variational_copy_ability(self):
+        csvfiles_path = DATAFILES_PATH / 'variational'
+        fit = cmdstanpy.from_csv(path=csvfiles_path)
+        fit2 = copy.deepcopy(fit)
+        self.assertSequenceEqual(
+            fit.stan_variables().keys(), fit2.stan_variables().keys()
+        )