[CmdStan 2.31] Support generate_quantities after non-sampling runs (#634)

* Checkpointing

* Basic refactor, passing mypy

* Testing new functionality and tweaking

* Maintain column headers in pd

* Fix docstring

Author: WardBrian

cmdstanpy/model.py

@@ -16,7 +16,17 @@
 from io import StringIO
 from multiprocessing import cpu_count
 from pathlib import Path
-from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    TypeVar,
+    Union,
+)
 
 import pandas as pd
 from tqdm.auto import tqdm
@@ -58,6 +68,7 @@
 from . import progress as progbar
 
 OptionalPath = Union[str, os.PathLike, None]
+Fit = TypeVar('Fit', CmdStanMCMC, CmdStanMLE, CmdStanVB)
 
 
 class CmdStanModel:
@@ -1202,22 +1213,26 @@ def sample(
     def generate_quantities(
         self,
         data: Union[Mapping[str, Any], str, os.PathLike, None] = None,
-        mcmc_sample: Union[CmdStanMCMC, List[str], None] = None,
+        previous_fit: Union[Fit, List[str], None] = None,
         seed: Optional[int] = None,
         gq_output_dir: OptionalPath = None,
         sig_figs: Optional[int] = None,
         show_console: bool = False,
         refresh: Optional[int] = None,
         time_fmt: str = "%Y%m%d%H%M%S",
         timeout: Optional[float] = None,
-    ) -> CmdStanGQ:
+        *,
+        mcmc_sample: Union[CmdStanMCMC, List[str], None] = None,
+    ) -> CmdStanGQ[Fit]:
         """
         Run CmdStan's generate_quantities method which runs the generated
         quantities block of a model given an existing sample.
 
-        This function takes a :class:`CmdStanMCMC` object and the dataset used
-        to generate that sample and calls to the CmdStan ``generate_quantities``
-        method to generate additional quantities of interest.
+        This function takes one of the Stan fit objects
+        :class:`CmdStanMCMC`, :class:`CmdStanMLE`, or :class:`CmdStanVB`
+        and the data required for the model and calls to the CmdStan
+        ``generate_quantities`` method to generate additional quantities of
+        interest.
 
         The :class:`CmdStanGQ` object records the command, the return code,
         and the paths to the generate method output CSV and console files.
@@ -1236,9 +1251,10 @@ def generate_quantities(
             either as a dictionary with entries matching the data variables,
             or as the path of a data file in JSON or Rdump format.
 
-        :param mcmc_sample: Can be either a :class:`CmdStanMCMC` object returned
-            by the :meth:`sample` method or a list of stan-csv files generated
-            by fitting the model to the data using any Stan interface.
+        :param previous_fit: Can be either a :class:`CmdStanMCMC`,
+            :class:`CmdStanMLE`, or :class:`CmdStanVB` or a list of
+            stan-csv files generated by fitting the model to the data
+            using any Stan interface.
 
         :param seed: The seed for random number generator. Must be an integer
             between 0 and 2^32 - 1. If unspecified,
@@ -1272,39 +1288,64 @@ def generate_quantities(
 
         :return: CmdStanGQ object
         """
-        if isinstance(mcmc_sample, CmdStanMCMC):
-            mcmc_fit = mcmc_sample
-            sample_csv_files = mcmc_sample.runset.csv_files
-        elif isinstance(mcmc_sample, list):
-            if len(mcmc_sample) < 1:
+        if mcmc_sample is not None:
+            if previous_fit:
+                raise ValueError(
+                    "Cannot supply both 'previous_fit' and "
+                    "deprecated argument 'mcmc_sample'"
+                )
+            get_logger().warning(
+                "Argument name `mcmc_sample` is deprecated, please "
+                "rename to `previous_fit`."
+            )
+
+            previous_fit = mcmc_sample  # type: ignore
+
+        if isinstance(previous_fit, (CmdStanMCMC, CmdStanMLE, CmdStanVB)):
+            fit_object = previous_fit
+            fit_csv_files = previous_fit.runset.csv_files
+        elif isinstance(previous_fit, list):
+            if len(previous_fit) < 1:
                 raise ValueError(
                     'Expecting list of Stan CSV files, found empty list'
                 )
             try:
-                sample_csv_files = mcmc_sample
-                sample_fit = from_csv(sample_csv_files)
-                mcmc_fit = sample_fit  # type: ignore
+                fit_csv_files = previous_fit
+                fit_object = from_csv(fit_csv_files)  # type: ignore
             except ValueError as e:
                 raise ValueError(
                     'Invalid sample from Stan CSV files, error:\n\t{}\n\t'
                     ' while processing files\n\t{}'.format(
-                        repr(e), '\n\t'.join(mcmc_sample)
+                        repr(e), '\n\t'.join(previous_fit)
                     )
                 ) from e
         else:
             raise ValueError(
-                'MCMC sample must be either CmdStanMCMC object'
-                ' or list of paths to sample Stan CSV files.'
-            )
-        chains = mcmc_fit.chains
-        chain_ids = mcmc_fit.chain_ids
-        if mcmc_fit.metadata.cmdstan_config['save_warmup']:
-            get_logger().warning(
-                'Sample contains saved warmup draws which will be used '
-                'to generate additional quantities of interest.'
+                'Previous fit must be either CmdStanPy fit object'
+                ' or list of paths to Stan CSV files.'
             )
+        if isinstance(fit_object, CmdStanMCMC):
+            chains = fit_object.chains
+            chain_ids = fit_object.chain_ids
+            if fit_object._save_warmup:
+                get_logger().warning(
+                    'Sample contains saved warmup draws which will be used '
+                    'to generate additional quantities of interest.'
+                )
+        elif isinstance(fit_object, CmdStanMLE):
+            chains = 1
+            chain_ids = [1]
+            if fit_object._save_iterations:
+                get_logger().warning(
+                    'MLE contains saved iterations which will be used '
+                    'to generate additional quantities of interest.'
+                )
+        else:  # isinstance(fit_object, CmdStanVB)
+            chains = 1
+            chain_ids = [1]
+
         generate_quantities_args = GenerateQuantitiesArgs(
-            csv_files=sample_csv_files
+            csv_files=fit_csv_files
         )
         generate_quantities_args.validate(chains)
         with MaybeDictToFilePath(data, None) as (_data, _inits):
@@ -1345,7 +1386,7 @@ def generate_quantities(
                     + ' output is unclear!'
                 )
                 raise RuntimeError(msg)
-            quantities = CmdStanGQ(runset=runset, mcmc_sample=mcmc_fit)
+            quantities = CmdStanGQ(runset=runset, previous_fit=fit_object)
         return quantities
 
     def variational(