Merge pull request #90 from aloctavodia/autocorrplot

autocorrplot: share x-y labels

arviz/plots/autocorrplot.py

@@ -1,6 +1,6 @@
 import matplotlib.pyplot as plt
 
-from .plot_utils import get_axis, _scale_text
+from .plot_utils import _scale_text
 from ..utils import get_varnames, trace_to_dataframe
 
 
@@ -41,13 +41,13 @@ def autocorrplot(trace, varnames=None, max_lag=100, symmetric_plot=False, combin
     varnames = get_varnames(trace, varnames)
 
     if figsize is None:
-        figsize = (6, len(varnames) * 2)
+        figsize = (12, len(varnames) * 2)
 
-    textsize, linewidth, _ = _scale_text(figsize, textsize=textsize)
+    textsize, linewidth, _ = _scale_text(figsize, textsize, 1)
 
     nchains = trace.columns.value_counts()[0]
-    ax = get_axis(ax, len(varnames), nchains, squeeze=False, sharex=True, sharey=True,
-                  figsize=figsize)
+    fig, ax = plt.subplots(len(varnames), nchains, squeeze=False, sharex=True, sharey=True,
+                           figsize=figsize)
 
     max_lag = min(len(trace) - 1, max_lag)
 
@@ -59,12 +59,6 @@ def autocorrplot(trace, varnames=None, max_lag=100, symmetric_plot=False, combin
                 data = trace[varname].values[:, j]
             ax[i, j].acorr(data, detrend=plt.mlab.detrend_mean, maxlags=max_lag, lw=linewidth)
 
-            if j == 0:
-                ax[i, j].set_ylabel("correlation", fontsize=textsize)
-
-            if i == len(varnames) - 1:
-                ax[i, j].set_xlabel("lag", fontsize=textsize)
-
             if not symmetric_plot:
                 ax[i, j].set_xlim(0, max_lag)
 
@@ -73,4 +67,10 @@ def autocorrplot(trace, varnames=None, max_lag=100, symmetric_plot=False, combin
             else:
                 ax[i, j].set_title(varname, fontsize=textsize)
             ax[i, j].tick_params(labelsize=textsize)
+
+    fig.add_subplot(111, frameon=False)
+    plt.tick_params(labelcolor='none', top='off', bottom='off', left='off', right='off')
+    plt.grid(False)
+    plt.xlabel("Lag", fontsize=textsize)
+    plt.ylabel("Correlation", fontsize=textsize)
     return ax

arviz/plots/plot_utils.py

@@ -2,27 +2,6 @@
 import matplotlib.pyplot as plt
 
 
-def get_axis(ax, default_rows, default_columns, **default_kwargs):
-    """Verifies the provided axis is of the correct shape, and creates one if needed.
-
-    Args:
-        ax: matplotlib axis or None
-        default_rows: int, expected rows in axis
-        default_columns: int, expected columns in axis
-        **default_kwargs: keyword arguments to pass to plt.subplot
-
-    Returns:
-        axis, or raises an error
-    """
-
-    default_shape = (default_rows, default_columns)
-    if ax is None:
-        _, ax = plt.subplots(*default_shape, **default_kwargs)
-    elif ax.shape != default_shape:
-        raise ValueError('Subplots with shape %r required' % (default_shape,))
-    return ax
-
-
 def make_2d(ary):
     """Convert any array into a 2d numpy array.
 

arviz/plots/traceplot.py

@@ -2,7 +2,7 @@
 import matplotlib.pyplot as plt
 
 from .kdeplot import fast_kde
-from .plot_utils import get_axis, make_2d, get_bins, _scale_text
+from .plot_utils import make_2d, get_bins, _scale_text
 from ..utils import get_varnames, trace_to_dataframe
 
 
@@ -67,7 +67,7 @@ def traceplot(trace, varnames=None, figsize=None, textsize=None, lines=None, com
 
     textsize, linewidth, _ = _scale_text(figsize, textsize=textsize, scale_ratio=1)
 
-    ax = get_axis(ax, len(varnames), 2, squeeze=False, figsize=figsize)
+    _, ax = plt.subplots(len(varnames), 2, squeeze=False, figsize=figsize)
 
     for i, varname in enumerate(varnames):
         if priors is not None: