optimize - save_iterations, require_converged; unit tests

Author: mitzimorris

cmdstanpy/model.py

@@ -398,6 +398,7 @@ def optimize(
         history_size: Optional[int] = None,
         iter: Optional[int] = None,
         save_iterations: bool = False,
+        require_converged: bool = True,
         refresh: Optional[int] = None,
     ) -> CmdStanMLE:
         """
@@ -485,6 +486,9 @@ def optimize(
         :param save_iterations: When ``True``, save intermediate approximations
             to the output CSV file.  Default is ``False``.
 
+        :param require_converged: Whether or not to raise an error if Stan
+            reports that "The algorithm may not have converged".
+
         :param refresh: Specify the number of iterations cmdstan will take
             between progress messages. Default value is 100.
 
@@ -524,7 +528,10 @@ def optimize(
 
         if not runset._check_retcodes():
             msg = 'Error during optimization: {}'.format(runset.get_err_msgs())
-            get_logger().warn(msg)  # https://github.com/stan-dev/cmdstanr/issues/314
+            if 'Line search failed' in msg and not require_converged:
+                get_logger().warning(msg)
+            else:
+                raise RuntimeError(msg)
         mle = CmdStanMLE(runset)
         return mle
 
@@ -1119,7 +1126,7 @@ def variational(
         :param output_samples: Number of approximate posterior output draws
             to save.
 
-        :param require_converged: Whether or not to raise an error if stan
+        :param require_converged: Whether or not to raise an error if Stan
             reports that "The algorithm may not have converged".
 
         :param refresh: Specify the number of iterations cmdstan will take

cmdstanpy/stanfit.py

@@ -319,8 +319,8 @@ def get_err_msgs(self) -> str:
                             msgs.append(
                                 'chain_id {}:\n\t{}\n'.format(
                                     self._chain_ids[i], '\n\t'.join(errors)
-                                    )
                                 )
+                            )
             elif self._args.method == Method.OPTIMIZE:
                 msgs.append('console log output:\n')
                 with open(self._stdout_files[0], 'r') as fd:
@@ -800,7 +800,7 @@ def _assemble_draws(self) -> None:
                         line = fd.readline().strip()  # metric type
                         line = fd.readline().lstrip(' #\t')
                         num_unconstrained_params = len(line.split(','))
-                        if chain == 0:   # can't allocate w/o num params
+                        if chain == 0:  # can't allocate w/o num params
                             if self.metric_type == 'diag_e':
                                 self._metric = np.empty(
                                     (self.chains, num_unconstrained_params),
@@ -1235,6 +1235,13 @@ def __init__(self, runset: RunSet) -> None:
                 'found method {}'.format(runset.method)
             )
         self.runset = runset
+        # info from runset to be exposed
+        self.converged = runset._check_retcodes()
+        optimize_args = self.runset._args.method_args
+        assert isinstance(
+            optimize_args, OptimizeArgs
+        )  # make the typechecker happy
+        self.save_iterations = optimize_args.save_iterations
         self._set_mle_attrs(runset.csv_files[0])
 
     def __repr__(self) -> str:
@@ -1246,16 +1253,18 @@ def __repr__(self) -> str:
             '\n\t'.join(self.runset.csv_files),
             '\n\t'.join(self.runset.stdout_files),
         )
-        if not self.runset._check_retcodes():
+        if not self.converged:
             repr = '{}\n Warning: invalid estimate, '.format(repr)
             repr = '{} optimization failed to converge.'.format(repr)
         return repr
 
     def _set_mle_attrs(self, sample_csv_0: str) -> None:
-        meta = scan_optimize_csv(sample_csv_0)
+        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'] 
+        self._mle = meta['mle']
+        if self.save_iterations:
+            self._all_iters = meta['all_iters']
 
     @property
     def column_names(self) -> Tuple[str, ...]:
@@ -1276,36 +1285,90 @@ def metadata(self) -> InferenceMetadata:
 
     @property
     def optimized_params_np(self) -> np.ndarray:
-        """Returns optimized params as numpy array."""
-        if not self.runset._check_retcodes():
+        """
+        Returns all final estimates from the optimizer as a numpy.ndarray
+        which contains all optimizer outputs, i.e., the value for `lp__`
+        as well as all Stan program variables.
+        """
+        if not self.converged:
             get_logger().warning(
-                'invalid estimate, optimization failed to converge'
+                'Invalid estimate, optimization failed to converge.'
             )
-        # TODO: squeeze?
         return self._mle
 
+    @property
+    def optimized_iterations_np(self) -> np.ndarray:
+        """
+        Returns all saved iterations from the optimizer and final estimate
+        as a numpy.ndarray which contains all optimizer outputs, i.e.,
+        the value for `lp__` as well as all Stan program variables.
+
+        """
+        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.'
+            )
+            return None
+        if not self.converged:
+            get_logger().warning(
+                'Invalid estimate, optimization failed to converge.'
+            )
+        return self._all_iters
+
     @property
     def optimized_params_pd(self) -> pd.DataFrame:
-        """Returns optimized params as pandas DataFrame."""
+        """
+        Returns all final estimates from the optimizer as a pandas.DataFrame
+        which contains all optimizer outputs, i.e., the value for `lp__`
+        as well as all Stan program variables.
+        """
         if not self.runset._check_retcodes():
             get_logger().warning(
-                'invalid estimate, optimization failed to converge'
+                '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_iterations_pd(self) -> pd.DataFrame:
+        """
+        Returns all saved iterations from the optimizer and final estimate
+        as a pandas.DataFrame which contains all optimizer outputs, i.e.,
+        the value for `lp__` as well as all Stan program variables.
+
+        """
+        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.'
+            )
+            return None
+        if not self.converged:
+            get_logger().warning(
+                'Invalid estimate, optimization failed to converge.'
+            )
+        return pd.DataFrame(self._all_iters, columns=self.column_names)
 
     @property
     def optimized_params_dict(self) -> Dict[str, float]:
-        """Returns optimized params as Dict."""
+        """
+        Returns all estimates from the optimizer, including `lp__` as a
+        Python Dict.  Only returns estimate from final iteration.
+        """
         if not self.runset._check_retcodes():
             get_logger().warning(
-                'invalid estimate, optimization failed to converge'
+                'Invalid estimate, optimization failed to converge.'
             )
-        # TODO: return final estimate only
         return OrderedDict(zip(self.column_names, self._mle))
 
     def stan_variable(
-        self, var: Optional[str] = None,
-        check_convergence:bool = True, *, name: Optional[str] = None
+        self,
+        var: Optional[str] = None,
+        *,
+        warn: bool = True,
+        name: Optional[str] = None,
     ) -> np.ndarray:
         """
         Return a numpy.ndarray which contains the estimates for the
@@ -1314,11 +1377,6 @@ def stan_variable(
 
         :param var: variable name
 
-        :param check_convergence: Checks for failure to converge and
-            prints warning.failed to converge.  ``False`` will supress
-            check and warning, default is ``True``.
-
-
         See Also
         --------
         CmdStanMLE.stan_variables
@@ -1339,43 +1397,38 @@ 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 check_convergence and not self.runset._check_retcodes():
+        if warn and not self.runset._check_retcodes():
             get_logger().warning(
-                'invalid estimate, optimization failed to converge'
+                'Invalid estimate, optimization failed to converge.'
             )
 
         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, ...] = ()
         if len(col_idxs) > 0:
             shape = self._metadata.stan_vars_dims[var]
         return np.array(xs).reshape(shape)
 
-    def stan_variables(self, check_convergence:bool = True) -> Dict[str, np.ndarray]:
+    def stan_variables(self) -> Dict[str, np.ndarray]:
         """
         Return a dictionary mapping Stan program variables names
         to the corresponding numpy.ndarray containing the inferred values.
 
-        :param check_convergence: Checks for failure to converge and
-            prints warning.failed to converge.  ``False`` will supress
-            check and warning, default is ``True``.
-
         See Also
         --------
         CmdStanMLE.stan_variable
         CmdStanMCMC.stan_variables
         CmdStanVB.stan_variables
         CmdStanGQ.stan_variables
         """
-        if check_convergence and not self.runset._check_retcodes():
+        if not self.runset._check_retcodes():
             get_logger().warning(
-                'invalid estimate, optimization failed to converge'
+                'Invalid estimate, optimization failed to converge.'
             )
         result = {}
         for name in self._metadata.stan_vars_dims.keys():
-            result[name] = self.stan_variable(name, False)  # don't warn twice
+            result[name] = self.stan_variable(name, warn=False)
         return result
 
     def save_csvfiles(self, dir: Optional[str] = None) -> None:
@@ -2259,6 +2312,7 @@ def from_csv(
                 )
             optimize_args = OptimizeArgs(
                 algorithm=config_dict['algorithm'],
+                save_iterations=config_dict['save_iterations'],
             )
             cmdstan_args = CmdStanArgs(
                 model_name=config_dict['model'],

cmdstanpy/utils.py

@@ -589,10 +589,11 @@ def scan_optimize_csv(path: str, save_iters: bool = False) -> Dict[str, Any]:
         iters = 0
         for line in fd:
             iters += 1
+    if save_iters:
+        all_iters = np.empty(
+            (iters, len(dict['column_names'])), dtype=float, order='F'
+        )
     # 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()
@@ -605,8 +606,13 @@ def scan_optimize_csv(path: str, save_iters: bool = False) -> Dict[str, Any]:
                     )
                 )
             xs = line.split(',')
-            mle[i, :] = [float(x) for x in xs]
+            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)
     dict['mle'] = mle
+    if save_iters:
+        dict['all_iters'] = all_iters
     return dict
 
 

test/data/optimize/eight_schools.data.R

@@ -0,0 +1,4 @@
+J <- 8
+y <- c(28,  8, -3,  7, -1,  1, 18, 12)
+sigma <- c(15, 10, 16, 11,  9, 11, 10, 18)
+tau <- 25

test/data/optimize/eight_schools.stan

@@ -0,0 +1,14 @@
+data {
+  int<lower=0> J;          // number of schools
+  real y[J];               // estimated treatment effect (school j)
+  real<lower=0> sigma[J];  // std err of effect estimate (school j)
+}
+parameters {
+  real mu;
+  real theta[J];
+  real<lower=0> tau;
+}
+model {
+  theta ~ normal(mu, tau); 
+  y ~ normal(theta,sigma);
+}