Merge pull request #686 from stan-dev/feature/pathfinder

[Stan 2.33] Pathfinder

Author: WardBrian

.pylintrc

@@ -76,6 +76,7 @@ disable=consider-using-in,
     consider-using-dict-items,
     unspecified-encoding,
     consider-using-f-string,
+    duplicate-code,
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option

cmdstanpy/__init__.py

@@ -26,8 +26,10 @@ def _cleanup_tmpdir() -> None:
 from .model import CmdStanModel
 from .stanfit import (
     CmdStanGQ,
+    CmdStanLaplace,
     CmdStanMCMC,
     CmdStanMLE,
+    CmdStanPathfinder,
     CmdStanVB,
     InferenceMetadata,
     from_csv,
@@ -51,6 +53,8 @@ def _cleanup_tmpdir() -> None:
     'CmdStanMLE',
     'CmdStanGQ',
     'CmdStanVB',
+    'CmdStanLaplace',
+    'CmdStanPathfinder',
     'CmdStanModel',
     'InferenceMetadata',
     'from_csv',

cmdstanpy/cmdstan_args.py

@@ -44,6 +44,7 @@
     LaplaceArgs,
     Method,
     OptimizeArgs,
+    PathfinderArgs,
     SamplerArgs,
     VariationalArgs,
 )
@@ -53,6 +54,7 @@
     CmdStanLaplace,
     CmdStanMCMC,
     CmdStanMLE,
+    CmdStanPathfinder,
     CmdStanVB,
     RunSet,
     from_csv,
@@ -874,6 +876,7 @@ def sample(
             * dictionary - pairs parameter name : initial value.
             * string - pathname to a JSON or Rdump data file.
             * list of strings - per-chain pathname to data file.
+            * list of dictionaries - per-chain initial values.
 
         :param iter_warmup: Number of warmup iterations for each chain.
 
@@ -947,8 +950,7 @@ def sample(
             only those values in the generated quantities block that are
             produced by RNG functions may change.  This provides
             a way to use Stan programs to generate simulated data via the
-            generated quantities block.  This option must be used when the
-            parameters block is empty.  Default value is ``False``.
+            generated quantities block.  Default value is ``False``.
 
         :param output_dir: Name of the directory to which CmdStan output
             files are written. If unspecified, output files will be written
@@ -964,14 +966,14 @@ def sample(
             If ``True``, CSV outputs are written to an output file
             '<model_name>-<YYYYMMDDHHMM>-diagnostic-<chain_id>',
             e.g. 'bernoulli-201912081451-diagnostic-1.csv', see
-            https://mc-stan.org/docs/cmdstan-guide/stan-csv.html,
+            https://mc-stan.org/docs/cmdstan-guide/stan_csv.html,
             section "Diagnostic CSV output file" for details.
 
         :param save_profile: Whether or not to profile auto-diff operations in
             labelled blocks of code.  If ``True``, CSV outputs are written to
             file '<model_name>-<YYYYMMDDHHMM>-profile-<chain_id>'.
             Introduced in CmdStan-2.26, see
-            https://mc-stan.org/docs/cmdstan-guide/stan-csv.html,
+            https://mc-stan.org/docs/cmdstan-guide/stan_csv.html,
             section "Profiling CSV output file" for details.
 
         :param show_progress: If ``True``, display progress bar to track
@@ -1448,12 +1450,14 @@ def variational(
         adapt_iter: Optional[int] = None,
         tol_rel_obj: Optional[float] = None,
         eval_elbo: Optional[int] = None,
-        output_samples: Optional[int] = None,
+        draws: Optional[int] = None,
         require_converged: bool = True,
         show_console: bool = False,
         refresh: Optional[int] = None,
         time_fmt: str = "%Y%m%d%H%M%S",
         timeout: Optional[float] = None,
+        *,
+        output_samples: Optional[int] = None,
     ) -> CmdStanVB:
         """
         Run CmdStan's variational inference algorithm to approximate
@@ -1530,7 +1534,7 @@ def variational(
 
         :param eval_elbo: Number of iterations between ELBO evaluations.
 
-        :param output_samples: Number of approximate posterior output draws
+        :param draws: Number of approximate posterior output draws
             to save.
 
         :param require_converged: Whether or not to raise an error if Stan
@@ -1551,6 +1555,19 @@ def variational(
 
         :return: CmdStanVB object
         """
+        if output_samples is not None:
+            if draws is not None:
+                raise ValueError(
+                    "Cannot supply both 'draws' and deprecated argument "
+                    "'output_samples'"
+                )
+            get_logger().warning(
+                "Argument name `output_samples` is deprecated, please "
+                "rename to `draws`."
+            )
+
+            draws = output_samples
+
         variational_args = VariationalArgs(
             algorithm=algorithm,
             iter=iter,
@@ -1561,7 +1578,7 @@ def variational(
             adapt_iter=adapt_iter,
             tol_rel_obj=tol_rel_obj,
             eval_elbo=eval_elbo,
-            output_samples=output_samples,
+            output_samples=draws,
         )
 
         with temp_single_json(data) as _data, temp_inits(
@@ -1634,6 +1651,196 @@ def variational(
         vb = CmdStanVB(runset)
         return vb
 
+    def pathfinder(
+        self,
+        data: Union[Mapping[str, Any], str, os.PathLike, None] = None,
+        *,
+        init_alpha: Optional[float] = None,
+        tol_obj: Optional[float] = None,
+        tol_rel_obj: Optional[float] = None,
+        tol_grad: Optional[float] = None,
+        tol_rel_grad: Optional[float] = None,
+        tol_param: Optional[float] = None,
+        history_size: Optional[int] = None,
+        num_paths: Optional[int] = None,
+        max_lbfgs_iters: Optional[int] = None,
+        draws: Optional[int] = None,
+        num_single_draws: Optional[int] = None,
+        num_elbo_draws: Optional[int] = None,
+        # arguments standard to all methods
+        seed: Optional[int] = None,
+        inits: Union[Dict[str, float], float, str, os.PathLike, None] = None,
+        output_dir: OptionalPath = None,
+        sig_figs: Optional[int] = None,
+        save_profile: bool = False,
+        show_console: bool = False,
+        refresh: Optional[int] = None,
+        time_fmt: str = "%Y%m%d%H%M%S",
+        timeout: Optional[float] = None,
+    ) -> CmdStanPathfinder:
+        """
+        Run CmdStan's Pathfinder variational inference algorithm.
+
+        :param data: Values for all data variables in the model, specified
+            either as a dictionary with entries matching the data variables,
+            or as the path of a data file in JSON or Rdump format.
+
+        :param num_paths: Number of single-path Pathfinders to run.
+            Default is 4, when the number of paths is 1 then no importance
+            sampling is done.
+
+        :param draws: Number of approximate draws to return.
+
+        :param num_single_draws: Number of draws each single-pathfinder will
+            draw. By default, this is set to be equal to draws.
+            If ``num_paths`` is 1, only one of this and ``draws`` should be
+            used.
+
+        :param max_lbfgs_iters: Maximum number of L-BFGS iterations.
+
+        :param num_elbo_draws: Number of Monte Carlo draws to evaluate ELBO.
+
+        :param seed: The seed for random number generator. Must be an integer
+            between 0 and 2^32 - 1. If unspecified,
+            :class:`numpy.random.RandomState` is used to generate a seed.
+
+        :param inits: Specifies how the algorithm initializes parameter values.
+            Initialization is either uniform random on a range centered on 0,
+            exactly 0, or a dictionary or file of initial values for some or all
+            parameters in the model.  The default initialization behavior will
+            initialize all parameter values on range [-2, 2] on the
+            *unconstrained* support.  If the expected parameter values are
+            too far from this range, this option may improve adaptation.
+            The following value types are allowed:
+
+            * Single number n > 0 - initialization range is [-n, n].
+            * 0 - all parameters are initialized to 0.
+            * dictionary - pairs parameter name : initial value.
+            * string - pathname to a JSON or Rdump data file.
+            * list of strings - per-path pathname to data file.
+            * list of dictionaries - per-path initial values.
+
+        :param init_alpha: For internal L-BFGS: Line search step size for
+            first iteration
+
+        :param tol_obj: For internal L-BFGS: Convergence tolerance on changes
+            in objective function value
+
+        :param tol_rel_obj: For internal L-BFGS: Convergence tolerance on
+            relative changes in objective function value
+
+        :param tol_grad: For internal L-BFGS: Convergence tolerance on the
+            norm of the gradient
+
+        :param tol_rel_grad: For internal L-BFGS: Convergence tolerance on
+            the relative norm of the gradient
+
+        :param tol_param: For internal L-BFGS: Convergence tolerance on changes
+            in parameter value
+
+        :param history_size: For internal L-BFGS: Size of the history for LBFGS
+            Hessian approximation. The value should be less than the
+            dimensionality of the parameter space. 5-10 is usually sufficient
+
+        :param output_dir: Name of the directory to which CmdStan output
+            files are written. If unspecified, output files will be written
+            to a temporary directory which is deleted upon session exit.
+
+        :param sig_figs: Numerical precision used for output CSV and text files.
+            Must be an integer between 1 and 18.  If unspecified, the default
+            precision for the system file I/O is used; the usual value is 6.
+            Introduced in CmdStan-2.25.
+
+        :param save_profile: Whether or not to profile auto-diff operations in
+            labelled blocks of code.  If ``True``, CSV outputs are written to
+            file '<model_name>-<YYYYMMDDHHMM>-profile-<path_id>'.
+            Introduced in CmdStan-2.26, see
+            https://mc-stan.org/docs/cmdstan-guide/stan_csv.html,
+            section "Profiling CSV output file" for details.
+
+        :param show_console: If ``True``, stream CmdStan messages sent to stdout
+            and stderr to the console.  Default is ``False``.
+
+        :param refresh: Specify the number of iterations CmdStan will take
+            between progress messages. Default value is 100.
+
+        :param time_fmt: A format string passed to
+            :meth:`~datetime.datetime.strftime` to decide the file names for
+            output CSVs. Defaults to "%Y%m%d%H%M%S"
+
+        :param timeout: Duration at which Pathfinder times
+            out in seconds. Defaults to None.
+
+        :return: A :class:`CmdStanPathfinder` object
+
+        References
+        ----------
+
+        Zhang, L., Carpenter, B., Gelman, A., & Vehtari, A. (2022). Pathfinder:
+        Parallel quasi-Newton variational inference. Journal of Machine Learning
+        Research, 23(306), 1–49. Retrieved from
+        http://jmlr.org/papers/v23/21-0889.html
+        """
+        if cmdstan_version_before(2, 33, self.exe_info()):
+            raise ValueError(
+                "Method 'pathfinder' not available for CmdStan versions "
+                "before 2.33"
+            )
+
+        if num_single_draws is None:
+            num_single_draws = draws
+        elif num_paths == 1 and draws is not None and num_single_draws != draws:
+            raise ValueError(
+                "Cannot specify both 'draws' and 'num_single_draws'"
+                " when 'num_paths' is 1"
+            )
+
+        pathfinder_args = PathfinderArgs(
+            init_alpha=init_alpha,
+            tol_obj=tol_obj,
+            tol_rel_obj=tol_rel_obj,
+            tol_grad=tol_grad,
+            tol_rel_grad=tol_rel_grad,
+            tol_param=tol_param,
+            history_size=history_size,
+            num_psis_draws=draws,
+            num_paths=num_paths,
+            max_lbfgs_iters=max_lbfgs_iters,
+            num_draws=num_single_draws,
+            num_elbo_draws=num_elbo_draws,
+        )
+
+        with temp_single_json(data) as _data, temp_inits(inits) as _inits:
+            args = CmdStanArgs(
+                self._name,
+                self._exe_file,
+                chain_ids=None,
+                data=_data,
+                seed=seed,
+                inits=_inits,
+                output_dir=output_dir,
+                sig_figs=sig_figs,
+                save_profile=save_profile,
+                method_args=pathfinder_args,
+                refresh=refresh,
+            )
+            dummy_chain_id = 0
+            runset = RunSet(args=args, chains=1, time_fmt=time_fmt)
+            self._run_cmdstan(
+                runset,
+                dummy_chain_id,
+                show_console=show_console,
+                timeout=timeout,
+            )
+        runset.raise_for_timeouts()
+
+        if not runset._check_retcodes():
+            msg = "Error during Pathfinder! Command '{}' failed: {}".format(
+                ' '.join(runset.cmd(0)), runset.get_err_msgs()
+            )
+            raise RuntimeError(msg)
+        return CmdStanPathfinder(runset)
+
     def log_prob(
         self,
         params: Union[Dict[str, Any], str, os.PathLike],
@@ -1733,6 +1940,61 @@ def laplace_sample(
         timeout: Optional[float] = None,
         opt_args: Optional[Dict[str, Any]] = None,
     ) -> CmdStanLaplace:
+        """
+        Run a Laplace approximation around the posterior mode.
+
+        :param data: Values for all data variables in the model, specified
+            either as a dictionary with entries matching the data variables,
+            or as the path of a data file in JSON or Rdump format.
+
+        :param mode: The mode around which to place the approximation.
+            This can be:
+            * A :class:`CmdStanMLE` object
+            * A path to a CSV file containing the output of an optimization run.
+            * ``None`` - use default optimizer settings and/or any ``opt_args``.
+
+        :param draws: Number of approximate draws to return.
+            Defaults to 1000
+
+        :param jacobian: Whether or not to enable the Jacobian adjustment
+            for constrained parameters. Defaults to ``True``.
+            Note: This must match the argument used in the creation of
+            ``mode``, if supplied.
+
+        :param output_dir: Name of the directory to which CmdStan output
+            files are written. If unspecified, output files will be written
+            to a temporary directory which is deleted upon session exit.
+
+        :param sig_figs: Numerical precision used for output CSV and text files.
+            Must be an integer between 1 and 18.  If unspecified, the default
+            precision for the system file I/O is used; the usual value is 6.
+            Introduced in CmdStan-2.25.
+
+        :param save_profile: Whether or not to profile auto-diff operations in
+            labelled blocks of code.  If ``True``, CSV outputs are written to
+            file '<model_name>-<YYYYMMDDHHMM>-profile-<path_id>'.
+            Introduced in CmdStan-2.26, see
+            https://mc-stan.org/docs/cmdstan-guide/stan_csv.html,
+            section "Profiling CSV output file" for details.
+
+        :param show_console: If ``True``, stream CmdStan messages sent to stdout
+            and stderr to the console.  Default is ``False``.
+
+        :param refresh: Specify the number of iterations CmdStan will take
+            between progress messages. Default value is 100.
+
+        :param time_fmt: A format string passed to
+            :meth:`~datetime.datetime.strftime` to decide the file names for
+            output CSVs. Defaults to "%Y%m%d%H%M%S"
+
+        :param timeout: Duration at which Pathfinder times
+            out in seconds. Defaults to None.
+
+        :param opt_args: Dictionary of additional arguments
+            which will be passed to :meth:`~CmdStanModel.optimize`
+
+        :return: A :class:`CmdStanLaplace` object.
+        """
         if cmdstan_version_before(2, 32, self.exe_info()):
             raise ValueError(
                 "Method 'laplace_sample' not available for CmdStan versions "