added arg 'inc_iterations' to stan_variable(s)

Author: mitzimorris

cmdstanpy/stanfit.py

@@ -1241,7 +1241,7 @@ def __init__(self, runset: RunSet) -> None:
         assert isinstance(
             optimize_args, OptimizeArgs
         )  # make the typechecker happy
-        self.save_iterations = optimize_args.save_iterations
+        self._save_iterations = optimize_args.save_iterations
         self._set_mle_attrs(runset.csv_files[0])
 
     def __repr__(self) -> str:
@@ -1259,11 +1259,11 @@ def __repr__(self) -> str:
         return repr
 
     def _set_mle_attrs(self, sample_csv_0: str) -> None:
-        meta = scan_optimize_csv(sample_csv_0, self.save_iterations)
+        meta = scan_optimize_csv(sample_csv_0, self._save_iterations)
         self._metadata = InferenceMetadata(meta)
         self._column_names: Tuple[str, ...] = meta['column_names']
         self._mle = meta['mle']
-        if self.save_iterations:
+        if self._save_iterations:
             self._all_iters = meta['all_iters']
 
     @property
@@ -1304,11 +1304,10 @@ def optimized_iterations_np(self) -> np.ndarray:
         the value for `lp__` as well as all Stan program variables.
 
         """
-        if not self.save_iterations:
+        if not self._save_iterations:
             get_logger().warning(
-                'Intermediate iterations not saved because optimizer argument '
-                '"save_iterations=True" not specified. You must rerun '
-                'the optimize method accordingly.'
+                'Intermediate iterations not saved to CSV output file. '
+                'Rerun the optimize method with "save_iterations=True".'
             )
             return None
         if not self.converged:
@@ -1338,11 +1337,10 @@ def optimized_iterations_pd(self) -> pd.DataFrame:
         the value for `lp__` as well as all Stan program variables.
 
         """
-        if not self.save_iterations:
+        if not self._save_iterations:
             get_logger().warning(
-                'Intermediate iterations not saved because optimizer argument '
-                '"save_iterations=True" not specified. You must rerun '
-                'the optimize method accordingly.'
+                'Intermediate iterations not saved to CSV output file. '
+                'Rerun the optimize method with "save_iterations=True".'
             )
             return None
         if not self.converged:
@@ -1367,6 +1365,7 @@ def stan_variable(
         self,
         var: Optional[str] = None,
         *,
+        inc_iterations: bool = False,
         warn: bool = True,
         name: Optional[str] = None,
     ) -> np.ndarray:
@@ -1377,6 +1376,11 @@ def stan_variable(
 
         :param var: variable name
 
+        :param inc_iterations: When ``True`` and the intermediate estimates
+            are included in the output, i.e., the optimizer was run with
+            ``save_iterations=True``, then intermediate estimates are included.
+            Default value is ``False``.
+
         See Also
         --------
         CmdStanMLE.stan_variables
@@ -1397,24 +1401,56 @@ def stan_variable(
             raise ValueError('no variable name specified.')
         if var not in self._metadata.stan_vars_dims:
             raise ValueError('unknown variable name: {}'.format(var))
+        if warn and inc_iterations and not self._save_iterations:
+            get_logger().warning(
+                'Intermediate iterations not saved to CSV output file. '
+                'Rerun the optimize method with "save_iterations=True".'
+            )
         if warn and not self.runset._check_retcodes():
             get_logger().warning(
                 'Invalid estimate, optimization failed to converge.'
             )
 
-        col_idxs = list(self._metadata.stan_vars_cols[var])
-        vals = list(self._mle)
-        xs = [vals[x] for x in col_idxs]
-        shape: Tuple[int, ...] = ()
+        col_idxs = self._metadata.stan_vars_cols[var]
+        if inc_iterations and self._save_iterations:
+            num_rows = self._all_iters.shape[0]
+        else:
+            num_rows = 1
+
+        # extract and reshape, container var
         if len(col_idxs) > 0:
-            shape = self._metadata.stan_vars_dims[var]
-        return np.array(xs).reshape(shape)
+            dims = (num_rows,) + self._metadata.stan_vars_dims[var]
+            # pylint: disable=redundant-keyword-arg
+            if num_rows > 1:
+                return self._all_iters[:, col_idxs].reshape(  # type: ignore
+                    dims, order='F'
+                )
+            else:
+                mle = np.expand_dims(self._mle, axis=0)  # hack for col indexing
+                return (
+                    mle[0, col_idxs]
+                    .reshape(dims, order='F')  # type: ignore
+                    .squeeze(axis=0)
+                )
 
-    def stan_variables(self) -> Dict[str, np.ndarray]:
+        # extract scalar var
+        if num_rows > 1:
+            return self._all_iters[:, col_idxs]
+        return mle[0, col_idxs]
+
+    def stan_variables(
+        self, inc_iterations: bool = False
+    ) -> Dict[str, np.ndarray]:
         """
         Return a dictionary mapping Stan program variables names
         to the corresponding numpy.ndarray containing the inferred values.
 
+        :param inc_iterations: When ``True`` and the intermediate estimates
+            are included in the output, i.e., the optimizer was run with
+            ``save_iterations=True``, then intermediate estimates are included.
+            Default value is ``False``.
+
+
         See Also
         --------
         CmdStanMLE.stan_variable
@@ -1428,7 +1464,9 @@ def stan_variables(self) -> Dict[str, np.ndarray]:
             )
         result = {}
         for name in self._metadata.stan_vars_dims.keys():
-            result[name] = self.stan_variable(name, warn=False)
+            result[name] = self.stan_variable(
+                name, inc_iterations=inc_iterations, warn=False
+            )
         return result
 
     def save_csvfiles(self, dir: Optional[str] = None) -> None:

cmdstanpy/utils.py

@@ -609,7 +609,7 @@ def scan_optimize_csv(path: str, save_iters: bool = False) -> Dict[str, Any]:
             if save_iters:
                 all_iters[i, :] = [float(x) for x in xs]
             if i == iters - 1:
-                mle = np.array([float(x) for x in xs], dtype=float)
+                mle = np.array([float(x) for x in xs])
     dict['mle'] = mle
     if save_iters:
         dict['all_iters'] = all_iters

test/data/optimize/rosenbrock_mle_iters.csv

@@ -0,0 +1,69 @@
+# stan_version_major = 2
+# stan_version_minor = 27
+# stan_version_patch = 0
+# model = rosenbrock_model
+# start_datetime = 2021-09-02 23:50:07 UTC
+# method = optimize
+#   optimize
+#     algorithm = lbfgs (Default)
+#       lbfgs
+#         init_alpha = 0.001 (Default)
+#         tol_obj = 9.9999999999999998e-13 (Default)
+#         tol_rel_obj = 10000 (Default)
+#         tol_grad = 1e-08 (Default)
+#         tol_rel_grad = 10000000 (Default)
+#         tol_param = 1e-08 (Default)
+#         history_size = 5 (Default)
+#     iter = 2000 (Default)
+#     save_iterations = 1
+# id = 0 (Default)
+# data
+#   file =  (Default)
+# init = 2 (Default)
+# random
+#   seed = 12345
+# output
+#   file = output.csv (Default)
+#   diagnostic_file =  (Default)
+#   refresh = 100 (Default)
+#   sig_figs = -1 (Default)
+#   profile_file = profile.csv (Default)
+# stanc_version = stanc3 v2.27.0
+# stancflags = 
+lp__,x,y
+-1732.6,1.98441,-0.2234
+-134.258,-1.32063,0.608857
+-17.9857,-1.03615,0.701595
+-5.53465,-0.93882,0.74813
+-3.59021,-0.888547,0.774151
+-3.53145,-0.878129,0.777499
+-3.52641,-0.875729,0.775875
+-3.5005,-0.86295,0.761978
+-3.43709,-0.832101,0.720763
+-2.19215,-0.451751,0.174997
+-2.1677,-0.456726,0.187233
+-1.7755,-0.228015,0.000272543
+-1.50767,-0.224409,0.041146
+-1.20549,-0.0725003,-0.0182459
+-1.0779,0.0346959,-0.0370183
+-0.86807,0.0898603,-0.0118539
+-0.698515,0.231494,0.0207393
+-0.481488,0.306423,0.0959909
+-0.385645,0.410491,0.148978
+-0.354102,0.436205,0.171239
+-0.214411,0.536994,0.288962
+-0.176819,0.612981,0.359304
+-0.0913081,0.710088,0.495705
+-0.0592561,0.757026,0.571606
+-0.0498401,0.824561,0.666095
+-0.0443989,0.826606,0.671305
+-0.0145855,0.887119,0.782687
+-0.00731916,0.920044,0.843437
+-0.00198537,0.955481,0.912759
+-0.00124069,0.97227,0.943138
+-0.000920772,0.975825,0.9504
+-6.69892e-05,0.992364,0.984492
+-5.87688e-06,0.998722,0.997239
+-8.49546e-07,0.99969,0.999466
+-3.31295e-08,0.999975,0.999969
+-5.1809e-11,1,1

test/test_optimize.py

@@ -113,9 +113,8 @@ def test_rosenbrock(self):
             (
                 'cmdstanpy',
                 'WARNING',
-                'Intermediate iterations not saved because optimizer argument '
-                '"save_iterations=True" not specified. You must rerun '
-                'the optimize method accordingly.',
+                'Intermediate iterations not saved to CSV output file. '
+                'Rerun the optimize method with "save_iterations=True".',
             )
         )
         with LogCapture() as log:
@@ -124,9 +123,8 @@ def test_rosenbrock(self):
             (
                 'cmdstanpy',
                 'WARNING',
-                'Intermediate iterations not saved because optimizer argument '
-                '"save_iterations=True" not specified. You must rerun '
-                'the optimize method accordingly.',
+                'Intermediate iterations not saved to CSV output file. '
+                'Rerun the optimize method with "save_iterations=True".',
             )
         )
 
@@ -142,6 +140,9 @@ def test_rosenbrock(self):
         self.assertAlmostEqual(mle.optimized_params_pd['x'][0], 1, places=3)
         self.assertAlmostEqual(mle.optimized_params_dict['x'], 1, places=3)
 
+        self.assertEqual(
+            mle.stan_variable('x', inc_iterations=True).shape, (36,)
+        )
         self.assertEqual(mle.optimized_iterations_np.shape, (36, 3))
         self.assertNotEqual(
             mle.optimized_iterations_np[0, 1],
@@ -166,7 +167,7 @@ def test_eight_schools(self):
         self.assertIn('method=optimize', mle.__repr__())
         self.assertFalse(mle.converged)
         with LogCapture() as log:
-            self.assertEqual(mle.optimized_params_pd.shape, (1,11))
+            self.assertEqual(mle.optimized_params_pd.shape, (1, 11))
         log.check_present(
             (
                 'cmdstanpy',
@@ -221,7 +222,6 @@ def test_variable_bern(self):
         )
 
     def test_variables_3d(self):
-        # construct fit using existing sampler output
         stan = os.path.join(DATAFILES_PATH, 'multidim_vars.stan')
         jdata = os.path.join(DATAFILES_PATH, 'logistic.data.R')
         multidim_model = CmdStanModel(stan_file=stan)
@@ -258,6 +258,31 @@ def test_variables_3d(self):
         self.assertTrue('frac_60' in vars)
         self.assertEqual(vars['frac_60'].shape, ())
 
+        multidim_mle_iters = multidim_model.optimize(
+            data=jdata,
+            seed=1239812093,
+            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,
+            save_iterations=True,
+        )
+        vars_iters = multidim_mle_iters.stan_variables(inc_iterations=True)
+        self.assertEqual(
+            len(vars_iters), len(multidim_mle_iters.metadata.stan_vars_dims)
+        )
+        self.assertTrue('y_rep' in vars_iters)
+        self.assertEqual(vars_iters['y_rep'].shape, (8, 5, 4, 3))
+        self.assertTrue('beta' in vars_iters)
+        self.assertEqual(vars_iters['beta'].shape, (8, 2))
+        self.assertTrue('frac_60' in vars_iters)
+        self.assertEqual(vars_iters['frac_60'].shape, (8,))
+
 
 class OptimizeTest(unittest.TestCase):
     def test_optimize_good(self):