Merge pull request pymc-devs#902 from jonsedar/issue_899_fix

Issue pymc-devs#899 fix

Author: twiecki

pymc3/examples/GHME 2013.ipynb

@@ -723,7 +723,7 @@
      "cell_type": "code",
      "collapsed": false,
      "input": [
-      "autocorrplot(trace, vars = [coeff_sd,sd ])"
+      "autocorrplot(trace, varnames = [coeff_sd,sd ])"
      ],
      "language": "python",
      "metadata": {},
@@ -739,4 +739,4 @@
    "metadata": {}
   }
  ]
-}
+}

pymc3/examples/GHME_2013.py

@@ -112,7 +112,7 @@ def run(n=3000):
 
     # <codecell>
 
-    autocorrplot(trace, vars = [coeff_sd,sd ])
+    autocorrplot(trace, varnames = [coeff_sd,sd ])
 
 if __name__ == '__main__':
     run()

pymc3/examples/gaussian_mixture_model.ipynb

@@ -242,7 +242,7 @@
    ],
    "source": [
     "# I prefer autocorrelation plots for serious confirmation of MCMC convergence\n",
-    "pm.autocorrplot(tr[5000::5], ['sd'])"
+    "pm.autocorrplot(tr[5000::5], varnames=['sd'])"
    ]
   },
   {

pymc3/examples/survival_analysis.ipynb

@@ -580,7 +580,7 @@
     }
    ],
    "source": [
-    "pm.autocorrplot(trace, vars=['beta']);"
+    "pm.autocorrplot(trace, varnames=['beta']);"
    ]
   },
   {

pymc3/plots.py

@@ -7,7 +7,7 @@
 
 
 def traceplot(trace, vars=None, figsize=None,
-              lines=None, combined=False, grid=True, 
+              lines=None, combined=False, grid=True,
               alpha=0.35, ax=None):
     """Plot samples histograms and values
 
@@ -149,82 +149,94 @@ def kde2plot(x, y, grid=200, ax=None):
     return ax
 
 
-def autocorrplot(trace, vars=None, max_lag=100, burn=0, ax=None,
-                 symmetric_plot=False):
+def autocorrplot(trace, varnames=None, max_lag=100, burn=0,
+                 symmetric_plot=False, ax=None, figsize=None):
     """Bar plot of the autocorrelation function for a trace
+
     Parameters
     ----------
     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.
         Vector-value stochastics are handled automatically.
     max_lag : int
         Maximum lag to calculate autocorrelation. Defaults to 100.
     burn : int
-        Number of samples to discard from the beginning of the trace. 
+        Number of samples to discard from the beginning of the trace.
         Defaults to 0.
-    ax : axes
-        Matplotlib axes. Defaults to None.
     symmetric_plot : boolean
         Plot from either [0, +lag] or [-lag, lag]. Defaults to False, [-, +lag].
-        
+    ax : axes
+        Matplotlib axes. Defaults to None.
+    figsize : figure size tuple
+        If None, size is (12, num of variables * 2) inches.
+        Note this is not used if ax is supplied.
+
     Returns
     -------
     ax : matplotlib axes
     """
-    
+
     import matplotlib.pyplot as plt
-        
-    def _handle_array_varnames(val):
-        if trace[0][val].__class__ is np.ndarray:
-            k = trace[val].shape[1]
-            for i in xrange(k):
-                yield val + '_{0}'.format(i)
+
+    def _handle_array_varnames(varname):
+        if trace[0][varname].__class__ is np.ndarray:
+            k = trace[varname].shape[1]
+            for i in range(k):
+                yield varname + '_{0}'.format(i)
         else:
-            yield val
-            
-    if vars is None:
-        vars = [item  for sub in [[i for i in _handle_array_varnames(var)]
-                    for var in trace.varnames] for item in sub]
+            yield varname
+
+    if varnames is None:
+        varnames = trace.varnames
     else:
-        vars = [str(var) for var in vars]
-        vars = [item  for sub in [[i for i in _handle_array_varnames(var)]
-                    for var in vars] for item in sub]
+        varnames = [str(v) for v in varnames]
 
-    chains = trace.nchains
+    varnames = [item for sub in [[i for i in _handle_array_varnames(v)]
+                    for v in varnames] for item in sub]
 
-    fig, ax = plt.subplots(len(vars), chains, squeeze=False,
-                    sharex=True, sharey=True)
+    nchains = trace.nchains
+
+    if figsize is None:
+        figsize = (12, len(varnames)*2)
+
+    if ax is None:
+        fig, ax = plt.subplots(len(varnames), nchains, squeeze=False,
+                                sharex=True, sharey=True, figsize=figsize)
+    elif ax.shape != (len(varnames), nchains):
+        raise ValueError('autocorrplot requires {}*{} subplots'.format(
+                            len(varnames), nchains))
+        return None
 
     max_lag = min(len(trace) - 1, max_lag)
 
-    for i, v in enumerate(vars):
-        for j in range(chains):
+    for i, v in enumerate(varnames):
+        for j in range(nchains):
             try:
                 d = np.squeeze(trace.get_values(v, chains=[j], burn=burn,
                                             combine=False))
             except KeyError:
                 k = int(v.split('_')[-1])
                 v_use = '_'.join(v.split('_')[:-1])
-                d = np.squeeze(trace.get_values(v_use, chains=[j], burn=burn,
-                                            combine=False)[:, k])
+                d = np.squeeze(trace.get_values(v_use, chains=[j],
+                                        burn=burn, combine=False)[:, k])
 
             ax[i, j].acorr(d, detrend=plt.mlab.detrend_mean, maxlags=max_lag)
 
             if not j:
                 ax[i, j].set_ylabel("correlation")
-            if i == len(vars) - 1:
+            if i == len(varnames) - 1:
                 ax[i, j].set_xlabel("lag")
-            
+
             ax[i, j].set_title(v)
-            
+
             if not symmetric_plot:
                 ax[i, j].set_xlim(0, max_lag)
-                
-            if chains > 1:
+
+            if nchains > 1:
                 ax[i, j].set_title("chain {0}".format(j+1))
-    
-    return (fig, ax)
+
+    return ax
 
 
 def var_str(name, shape):
@@ -294,10 +306,10 @@ def forestplot(trace_obj, vars=None, alpha=0.05, quartiles=True, rhat=True,
 
         vline (optional): numeric
             Location of vertical reference line (defaults to 0).
-            
+
         gs : GridSpec
             Matplotlib GridSpec object. Defaults to None.
-            
+
         Returns
         -------
 

pymc3/tests/test_plots.py

@@ -56,7 +56,7 @@ def test_multichain_plots():
 
     forestplot(ptrace, vars=['early_mean', 'late_mean'])
 
-    autocorrplot(ptrace, vars=['switchpoint'])
+    autocorrplot(ptrace, varnames=['switchpoint'])
 
 def test_make_2d():