Merge pull request #359 from stan-dev/feature/enable_vars_in_misc._print_stanfit

Enable vars argument in misc._print_stanfit

Author: ariddell

pystan/__init__.py

@@ -7,7 +7,7 @@
 import logging
 
 from pystan.api import stanc, stan
-from pystan.misc import read_rdump, stan_rdump
+from pystan.misc import read_rdump, stan_rdump, stansummary
 from pystan.model import StanModel
 from pystan.lookup import lookup
 

pystan/misc.py

@@ -50,39 +50,74 @@
 logger = logging.getLogger('pystan')
 
 
-def _print_stanfit(fit, pars=None, probs=(0.025, 0.25, 0.5, 0.75, 0.975), digits_summary=2):
-        if fit.mode == 1:
-            return "Stan model '{}' is of mode 'test_grad';\n"\
-                   "sampling is not conducted.".format(fit.model_name)
-        elif fit.mode == 2:
-            return "Stan model '{}' does not contain samples.".format(fit.model_name)
-        if pars is None:
-            pars = fit.sim['pars_oi']
-            fnames = fit.sim['fnames_oi']
-        else:
-            # FIXME: does this case ever occur?
-            # need a way of getting fnames matching specified pars
-            raise NotImplementedError
-
-        n_kept = [s - w for s, w in zip(fit.sim['n_save'], fit.sim['warmup2'])]
-        header = "Inference for Stan model: {}.\n".format(fit.model_name)
-        header += "{} chains, each with iter={}; warmup={}; thin={}; \n"
-        header = header.format(fit.sim['chains'], fit.sim['iter'], fit.sim['warmup'],
-                               fit.sim['thin'], sum(n_kept))
-        header += "post-warmup draws per chain={}, total post-warmup draws={}.\n\n"
-        header = header.format(n_kept[0], sum(n_kept))
-        footer = "\n\nSamples were drawn using {} at {}.\n"\
-            "For each parameter, n_eff is a crude measure of effective sample size,\n"\
-            "and Rhat is the potential scale reduction factor on split chains (at \n"\
-            "convergence, Rhat=1)."
-        sampler = fit.sim['samples'][0]['args']['sampler_t']
-        date = fit.date.strftime('%c')  # %c is locale's representation
-        footer = footer.format(sampler, date)
-        s = _summary(fit, pars, probs)
-        body = _array_to_table(s['summary'], s['summary_rownames'],
-                               s['summary_colnames'], digits_summary)
-        return header + body + footer
+def stansummary(fit, pars=None, probs=(0.025, 0.25, 0.5, 0.75, 0.975), digits_summary=2):
+    """
+    Summary statistic table.
+    
+    Parameters
+    ----------
+    fit : StanFit4Model object
+    pars : str or sequence of str, optional
+        Parameter names. By default use all parameters
+    probs : sequence of float, optional
+        Quantiles. By default, (0.025, 0.25, 0.5, 0.75, 0.975)
+    digits_summary : int, optional
+        Number of significant digits. By default, 2
 
+    Returns
+    -------
+    summary : string
+        Table includes mean, se_mean, sd, probs_0, ..., probs_n, n_eff and Rhat.
+    
+    Examples
+    --------
+    >>> model_code = 'parameters {real y;} model {y ~ normal(0,1);}'
+    >>> m = StanModel(model_code=model_code, model_name="example_model")
+    >>> fit = m.sampling()
+    >>> print(stansummary(fit))
+    Inference for Stan model: example_model.
+    4 chains, each with iter=2000; warmup=1000; thin=1; 
+    post-warmup draws per chain=1000, total post-warmup draws=4000.
+    
+           mean se_mean     sd   2.5%    25%    50%    75%  97.5%  n_eff   Rhat
+    y      0.01    0.03    1.0  -2.01  -0.68   0.02   0.72   1.97   1330    1.0
+    lp__   -0.5    0.02   0.68  -2.44  -0.66  -0.24  -0.05-5.5e-4   1555    1.0
+    
+    Samples were drawn using NUTS at Thu Aug 17 00:52:25 2017.
+    For each parameter, n_eff is a crude measure of effective sample size,
+    and Rhat is the potential scale reduction factor on split chains (at 
+    convergence, Rhat=1).
+    """
+    if fit.mode == 1:
+        return "Stan model '{}' is of mode 'test_grad';\n"\
+               "sampling is not conducted.".format(fit.model_name)
+    elif fit.mode == 2:
+        return "Stan model '{}' does not contain samples.".format(fit.model_name)
+
+    n_kept = [s - w for s, w in zip(fit.sim['n_save'], fit.sim['warmup2'])]
+    header = "Inference for Stan model: {}.\n".format(fit.model_name)
+    header += "{} chains, each with iter={}; warmup={}; thin={}; \n"
+    header = header.format(fit.sim['chains'], fit.sim['iter'], fit.sim['warmup'],
+                           fit.sim['thin'], sum(n_kept))
+    header += "post-warmup draws per chain={}, total post-warmup draws={}.\n\n"
+    header = header.format(n_kept[0], sum(n_kept))
+    footer = "\n\nSamples were drawn using {} at {}.\n"\
+        "For each parameter, n_eff is a crude measure of effective sample size,\n"\
+        "and Rhat is the potential scale reduction factor on split chains (at \n"\
+        "convergence, Rhat=1)."
+    sampler = fit.sim['samples'][0]['args']['sampler_t']
+    date = fit.date.strftime('%c')  # %c is locale's representation
+    footer = footer.format(sampler, date)
+    s = _summary(fit, pars, probs)
+    body = _array_to_table(s['summary'], s['summary_rownames'],
+                           s['summary_colnames'], digits_summary)
+    return header + body + footer
+
+def _print_stanfit(fit, pars=None, probs=(0.025, 0.25, 0.5, 0.75, 0.975), digits_summary=2):
+    # warning added in PyStan 2.17.0
+    warnings.warn('Function `_print_stanfit` is deprecated and will be removed in a future version. '\
+                  'Use `stansummary` instead.', DeprecationWarning)
+    return stansummary(fit, pars=pars, probs=probs, digits_summary=digits_summary)
 
 def _array_to_table(arr, rownames, colnames, n_digits):
     """Print an array with row and column names

pystan/stanfit4model.pyx

@@ -614,10 +614,10 @@ cdef class StanFit4Model:
 
     def __unicode__(self):
         # for Python 2.x
-        return pystan.misc._print_stanfit(self)
+        return pystan.misc.stansummary(self)
 
     def __str__(self):
-        s = pystan.misc._print_stanfit(self)
+        s = pystan.misc.stansummary(self)
         return s.encode('utf-8') if PY2 else s
 
     def __repr__(self):
@@ -626,7 +626,46 @@ cdef class StanFit4Model:
     def __getitem__(self, key):
         extr = self.extract(pars=(key,))
         return extr[key]
+    
+    def stansummary(self, pars=None, probs=(0.025, 0.25, 0.5, 0.75, 0.975), digits_summary=2):
+        """
+        Summary statistic table.
 
+        Parameters
+        ----------
+        fit : StanFit4Model object
+        pars : str or sequence of str, optional
+            Parameter names. By default use all parameters
+        probs : sequence of float, optional
+            Quantiles. By default, (0.025, 0.25, 0.5, 0.75, 0.975)
+        digits_summary : int, optional
+            Number of significant digits. By default, 2
+        Returns
+        -------
+        summary : string
+            Table includes mean, se_mean, sd, probs_0, ..., probs_n, n_eff and Rhat.
+
+        Examples
+        --------
+        >>> model_code = 'parameters {real y;} model {y ~ normal(0,1);}'
+        >>> m = StanModel(model_code=model_code, model_name="example_model")
+        >>> fit = m.sampling()
+        >>> print(fit.stansummary())
+        Inference for Stan model: example_model.
+        4 chains, each with iter=2000; warmup=1000; thin=1; 
+        post-warmup draws per chain=1000, total post-warmup draws=4000.
+
+               mean se_mean     sd   2.5%    25%    50%    75%  97.5%  n_eff   Rhat
+        y      0.01    0.03    1.0  -2.01  -0.68   0.02   0.72   1.97   1330    1.0
+        lp__   -0.5    0.02   0.68  -2.44  -0.66  -0.24  -0.05-5.5e-4   1555    1.0
+
+        Samples were drawn using NUTS at Thu Aug 17 00:52:25 2017.
+        For each parameter, n_eff is a crude measure of effective sample size,
+        and Rhat is the potential scale reduction factor on split chains (at 
+        convergence, Rhat=1).
+        """
+        return pystan.misc.stansummary(fit=self, pars=pars, probs=probs, digits_summary=digits_summary)
+    
     def summary(self, pars=None, probs=None):
         return pystan.misc._summary(self, pars, probs)
 

pystan/tests/test_misc_args.py

@@ -9,7 +9,7 @@
 class TestArgs(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        model_code = 'parameters {real y;} model {y ~ normal(0,1);}'
+        model_code = 'parameters {real x;real y;real z;} model {x ~ normal(0,1);y ~ normal(0,1);z ~ normal(0,1);}'
         cls.model = pystan.StanModel(model_code=model_code)
 
     def test_control(self):
@@ -23,3 +23,20 @@ def test_control(self):
             model.sampling(control=control_invalid)
         with assertRaisesRegex(ValueError, '`metric` must be one of'):
             model.sampling(control={'metric': 'lorem-ipsum'})
+            
+    def test_print_summary(self):
+        model = self.model
+        fit = model.sampling(iter=100)
+        
+        summary_full = pystan.misc.stansummary(fit)
+        summary_one_par1 = pystan.misc.stansummary(fit, pars='z')
+        summary_one_par2 = pystan.misc.stansummary(fit, pars=['z'])
+        summary_pars = pystan.misc.stansummary(fit, pars=['x', 'y'])
+        
+        self.assertNotEqual(summary_full, summary_one_par1)
+        self.assertNotEqual(summary_full, summary_one_par2)
+        self.assertNotEqual(summary_full, summary_pars)
+        self.assertNotEqual(summary_one_par1, summary_pars)
+        self.assertNotEqual(summary_one_par2, summary_pars)
+        
+        self.assertEqual(summary_one_par1, summary_one_par2)