Merge pull request #959 from pymc-devs/vars_varnames

Consistent use of vars and varnames

Author: Chris Fonnesbeck

pymc3/examples/GHME_2013.py

@@ -108,7 +108,7 @@ def run(n=3000):
 
     # <codecell>
 
-    traceplot(trace[100:], vars = [coeff_sd,sd ]);
+    traceplot(trace[100:], varnames = [coeff_sd,sd ]);
 
     # <codecell>
 

pymc3/examples/disaster_model_missing.py

@@ -60,7 +60,7 @@ def run(n=1000):
 
         tr = sample(n, tune=500, start=start, step=step)
 
-        summary(tr, vars=['disasters_missing'])
+        summary(tr, varnames=['disasters_missing'])
 
 if __name__ == '__main__':
     run()

pymc3/plots.py

@@ -6,7 +6,7 @@
 __all__ = ['traceplot', 'kdeplot', 'kde2plot', 'forestplot', 'autocorrplot']
 
 
-def traceplot(trace, vars=None, figsize=None,
+def traceplot(trace, varnames=None, figsize=None,
               lines=None, combined=False, grid=True,
               alpha=0.35, ax=None):
     """Plot samples histograms and values
@@ -15,7 +15,7 @@ def traceplot(trace, vars=None, figsize=None,
     ----------
 
     trace : result of MCMC run
-    vars : list of variable names
+    varnames : list of variable names
         Variables to be plotted, if None all variable are plotted
     figsize : figure size tuple
         If None, size is (12, num of variables * 2) inch
@@ -38,10 +38,10 @@ def traceplot(trace, vars=None, figsize=None,
 
     """
     import matplotlib.pyplot as plt
-    if vars is None:
-        vars = trace.varnames
+    if varnames is None:
+        varnames = trace.varnames
 
-    n = len(vars)
+    n = len(varnames)
 
     if figsize is None:
         figsize = (12, n*2)
@@ -52,7 +52,7 @@ def traceplot(trace, vars=None, figsize=None,
         print('traceplot requires n*2 subplots')
         return None
 
-    for i, v in enumerate(vars):
+    for i, v in enumerate(varnames):
         for d in trace.get_values(v, combine=combined, squeeze=False):
             d = np.squeeze(d)
             d = make_2d(d)
@@ -259,7 +259,7 @@ def var_str(name, shape):
     return names
 
 
-def forestplot(trace_obj, vars=None, alpha=0.05, quartiles=True, rhat=True,
+def forestplot(trace_obj, varnames=None, alpha=0.05, quartiles=True, rhat=True,
                main=None, xtitle=None, xrange=None, ylabels=None,
                chain_spacing=0.05, vline=0, gs=None):
     """ Forest plot (model summary plot)
@@ -271,7 +271,7 @@ def forestplot(trace_obj, vars=None, alpha=0.05, quartiles=True, rhat=True,
         trace_obj: NpTrace or MultiTrace object
             Trace(s) from an MCMC sample.
 
-        vars: list
+        varnames: list
             List of variables to plot (defaults to None, which results in all
             variables plotted).
 
@@ -339,14 +339,14 @@ def forestplot(trace_obj, vars=None, alpha=0.05, quartiles=True, rhat=True,
         from .diagnostics import gelman_rubin
 
         R = gelman_rubin(trace_obj)
-        if vars is not None:
-            R = {v: R[v] for v in vars}
+        if varnames is not None:
+            R = {v: R[v] for v in varnames}
     else:
         # Can't calculate Gelman-Rubin with a single trace
         rhat = False
 
-    if vars is None:
-        vars = trace_obj.varnames
+    if varnames is None:
+        varnames = trace_obj.varnames
 
     # Empty list for y-axis labels
     labels = []
@@ -370,7 +370,7 @@ def forestplot(trace_obj, vars=None, alpha=0.05, quartiles=True, rhat=True,
     for j, chain in enumerate(trace_obj.chains):
         # Counter for current variable
         var = 1
-        for varname in vars:
+        for varname in varnames:
             var_quantiles = trace_quantiles[chain][varname]
 
             quants = [var_quantiles[v] for v in qlist]
@@ -533,7 +533,7 @@ def forestplot(trace_obj, vars=None, alpha=0.05, quartiles=True, rhat=True,
         plt.yticks([-(l + 1) for l in range(len(labels))], "")
 
         i = 1
-        for varname in vars:
+        for varname in varnames:
 
             chain = trace_obj.chains[0]
             value = trace_obj.get_values(varname, chains=[chain])[0]

pymc3/stats.py

@@ -313,14 +313,14 @@ def quantiles(x, qlist=(2.5, 25, 50, 75, 97.5)):
         print("Too few elements for quantile calculation")
 
 
-def df_summary(trace, vars=None, stat_funcs=None, extend=False,
+def df_summary(trace, varnames=None, stat_funcs=None, extend=False,
                alpha=0.05, batches=100):
     """Create a data frame with summary statistics.
 
     Parameters
     ----------
     trace : MultiTrace instance
-    vars : list
+    varnames : list
         Names of variables to include in summary
     stat_funcs : None or list
         A list of functions used to calculate statistics. By default,
@@ -387,8 +387,8 @@ def df_summary(trace, vars=None, stat_funcs=None, extend=False,
     mu__0  0.066473  0.000312  0.105039  0.214242
     mu__1  0.067513 -0.159097 -0.045637  0.062912
     """
-    if vars is None:
-        vars = trace.varnames
+    if varnames is None:
+        varnames = trace.varnames
 
     funcs = [lambda x: pd.Series(np.mean(x, 0), name='mean'),
              lambda x: pd.Series(np.std(x, 0), name='sd'),
@@ -401,7 +401,7 @@ def df_summary(trace, vars=None, stat_funcs=None, extend=False,
         stat_funcs = funcs
 
     var_dfs = []
-    for var in vars:
+    for var in varnames:
         vals = trace.get_values(var, combine=True)
         flat_vals = vals.reshape(vals.shape[0], -1)
         var_df = pd.concat([f(flat_vals) for f in stat_funcs], axis=1)
@@ -416,7 +416,7 @@ def _hpd_df(x, alpha):
     return pd.DataFrame(hpd(x, alpha), columns=cnames)
 
 
-def summary(trace, vars=None, alpha=0.05, start=0, batches=100, roundto=3,
+def summary(trace, varnames=None, alpha=0.05, start=0, batches=100, roundto=3,
             to_file=None):
     """
     Generate a pretty-printed summary of the node.
@@ -425,7 +425,7 @@ def summary(trace, vars=None, alpha=0.05, start=0, batches=100, roundto=3,
     trace : Trace object
       Trace containing MCMC sample
 
-    vars : list of strings
+    varnames : list of strings
       List of variables to summarize. Defaults to None, which results
       in all variables summarized.
 
@@ -448,8 +448,8 @@ def summary(trace, vars=None, alpha=0.05, start=0, batches=100, roundto=3,
       File to write results to. If not given, print to stdout.
 
     """
-    if vars is None:
-        vars = trace.varnames
+    if varnames is None:
+        varnames = trace.varnames
 
     stat_summ = _StatSummary(roundto, batches, alpha)
     pq_summ = _PosteriorQuantileSummary(roundto, alpha)
@@ -459,7 +459,7 @@ def summary(trace, vars=None, alpha=0.05, start=0, batches=100, roundto=3,
     else:
         fh = open(to_file, mode='w')
 
-    for var in vars:
+    for var in varnames:
         # Extract sampled values
         sample = trace.get_values(var, burn=start, combine=True)
 

pymc3/tests/test_plots.py

@@ -54,7 +54,7 @@ def test_multichain_plots():
         start = {'early_mean': 2., 'late_mean': 3., 'switchpoint': 50}
         ptrace = sample(1000, [step1, step2], start, njobs=2)
 
-    forestplot(ptrace, vars=['early_mean', 'late_mean'])
+    forestplot(ptrace, varnames=['early_mean', 'late_mean'])
 
     autocorrplot(ptrace, varnames=['switchpoint'])