checkpointing; needs TODOs plus unit tests

Author: mitzimorris

cmdstanpy/stanfit.py

@@ -1255,7 +1255,7 @@ def _set_mle_attrs(self, sample_csv_0: str) -> None:
         meta = scan_optimize_csv(sample_csv_0)
         self._metadata = InferenceMetadata(meta)
         self._column_names: Tuple[str, ...] = meta['column_names']
-        self._mle: List[float] = meta['mle']
+        self._mle = meta['mle'] 
 
     @property
     def column_names(self) -> Tuple[str, ...]:
@@ -1281,7 +1281,8 @@ def optimized_params_np(self) -> np.ndarray:
             get_logger().warning(
                 'invalid estimate, optimization failed to converge'
             )
-        return np.asarray(self._mle)
+        # TODO: squeeze?
+        return self._mle
 
     @property
     def optimized_params_pd(self) -> pd.DataFrame:
@@ -1290,7 +1291,7 @@ def optimized_params_pd(self) -> pd.DataFrame:
             get_logger().warning(
                 'invalid estimate, optimization failed to converge'
             )
-        return pd.DataFrame([self._mle], columns=self.column_names)
+        return pd.DataFrame(self._mle, columns=self.column_names)
 
     @property
     def optimized_params_dict(self) -> Dict[str, float]:
@@ -1299,6 +1300,7 @@ def optimized_params_dict(self) -> Dict[str, float]:
             get_logger().warning(
                 'invalid estimate, optimization failed to converge'
             )
+        # TODO: return final estimate only
         return OrderedDict(zip(self.column_names, self._mle))
 
     def stan_variable(
@@ -1343,6 +1345,7 @@ def stan_variable(
             )
 
         col_idxs = list(self._metadata.stan_vars_cols[var])
+        # TODO: return final estimate only
         vals = list(self._mle)
         xs = [vals[x] for x in col_idxs]
         shape: Tuple[int, ...] = ()

cmdstanpy/utils.py

@@ -590,14 +590,23 @@ def scan_optimize_csv(path: str, save_iters: bool = False) -> Dict[str, Any]:
         for line in fd:
             iters += 1
     # rescan to capture estimates
+    mle = np.empty(
+        (iters, len(dict['column_names'])), dtype=float, order='F'
+    )
     with open(path, 'r') as fd:
         for i in range(lineno):
             fd.readline()
-        line = fd.readline()
-        # allocate numpy ndarray for 
-        print(line)
-        xs = line.split(',')
-        dict['mle'] = [float(x) for x in xs]
+        for i in range(iters):
+            line = fd.readline().strip()
+            if len(line) < 1:
+                raise ValueError(
+                    'cannot parse CSV file {}, error at line {}'.format(
+                        path, lineno + i
+                    )
+                )
+            xs = line.split(',')
+            mle[i, :] = [float(x) for x in xs]
+    dict['mle'] = mle
     return dict