followed sugestions made in #2861

Author: agustinaarroyuelo

RELEASE-NOTES.md

@@ -9,7 +9,7 @@
   `p=tt.nnet.sigmoid(logit_p)`.
 - Add `random` keyword to `pm.DensityDist` thus enabling users to pass custom random method
   which in turn makes sampling from a `DensityDist` possible. 
-- Add new 'scatterplot' function, for plotting scatter or hexbin matrices of sampled parameters.
+- Add new 'pairplot' function, for plotting scatter or hexbin matrices of sampled parameters.
   Optionally it can plot divergences. 
 
 ### Fixes

docs/source/notebooks/Diagnosing_biased_Inference_with_Divergences.ipynb

@@ -6,4 +6,4 @@
 from .traceplot import traceplot
 from .energyplot import energyplot
 from .densityplot import densityplot
-from .scatterplot import scatterplot
+from .pairplot import pairplot

pymc3/plots/__init__.py

@@ -167,4 +167,5 @@ def get_trace_dict(tr, varnames):
                     traces['_'.join([v, str(i)])] = vi
             else:
                 traces[v] = vals
-        return traces
+        return traces
+        

pymc3/plots/artists.py

@@ -0,0 +1,162 @@
+import warnings
+import matplotlib.gridspec as gridspec
+
+try:
+    import matplotlib.pyplot as plt
+except ImportError:  # mpl is optional
+    pass
+from ..util import get_default_varnames, is_transformed_name, get_untransformed_name
+from .artists import get_trace_dict, scale_text
+
+
+def pairplot(trace, varnames=None, figsize=None, text_size=None, 
+                gs=None, ax=None, hexbin=False, plot_transformed=False, 
+                divergences=False, kwargs_divergence=None, 
+                sub_varnames=None, **kwargs):
+
+    """
+    Plot a scatter or hexbin matrix of the sampled parameters.
+
+    Parameters
+    ----------
+
+    trace : result of MCMC run
+    varnames : list of variable names
+        Variables to be plotted, if None all variable are plotted
+    figsize : figure size tuple
+        If None, size is (8 + numvars, 8 + numvars)
+    text_size: int
+        Text size for labels
+    gs : Grid spec
+        Matplotlib Grid spec.
+    ax: axes
+        Matplotlib axes
+    hexbin : Boolean
+        If True draws an hexbin plot
+    plot_transformed : bool
+        Flag for plotting automatically transformed variables in addition to
+        original variables (defaults to False). Applies when varnames = None.
+        When a list of varnames is passed, transformed variables can be passed
+        using their names.
+    divergences : Boolean
+        If True divergences will be plotted in a diferent color
+    kwargs_divergence : dicts, optional
+        Aditional keywords passed to ax.scatter for divergences
+    sub_varnames : list, optional        
+        Aditional varnames passed for plotting subsets of multidimensional 
+        variables
+    Returns
+    -------
+
+    ax : matplotlib axes
+
+    """    
+    
+    if varnames is None:
+        if plot_transformed:
+
+            remove = [get_untransformed_name(var) for var in trace.varnames 
+                      if is_transformed_name(var)]
+            varnames_copy = trace.varnames
+
+            try:
+                [varnames_copy.remove(i) for i in remove]
+                varnames = varnames_copy
+            except:
+                varnames = trace.varnames
+                    
+            trace_dict = get_trace_dict(trace, 
+                get_default_varnames(varnames, True))
+
+        else:
+            trace_dict = get_trace_dict(trace, 
+                get_default_varnames(trace.varnames, False))
+
+        if sub_varnames is None:
+            varnames = list(trace_dict.keys())
+
+        else:
+            trace_dict = get_trace_dict(trace, 
+                get_default_varnames(trace.varnames, True))
+            varnames = sub_varnames
+
+    else:
+        trace_dict = get_trace_dict(trace, varnames)
+        varnames = list(trace_dict.keys())       
+
+    if text_size is None:
+        text_size = scale_text(figsize, text_size=text_size)
+
+    if kwargs_divergence is None:
+        kwargs_divergence = {}
+    
+    numvars = len(varnames)
+
+    if figsize is None:
+        figsize = (8+numvars, 8+numvars)    
+
+    if numvars < 2:
+        raise Exception('Number of variables to be plotted must be 2 or greater.')
+
+    if numvars == 2 and ax is not None:
+        if hexbin:
+            ax.hexbin(trace_dict[varnames[0]], 
+                      trace_dict[varnames[1]], mincnt=1, **kwargs)
+        else:
+            ax.scatter(trace_dict[varnames[0]], 
+                       trace_dict[varnames[1]], **kwargs)
+        
+        if divergences:
+            try:
+                divergent = trace['diverging']
+            except KeyError:
+                warnings.warn('No divergences were found.')
+
+            diverge = (divergent == 1)
+            ax.scatter(trace_dict[varnames[0]][diverge], 
+                       trace_dict[varnames[1]][diverge], **kwargs_divergence)
+
+    if gs is None and ax is None:
+
+        gs = gridspec.GridSpec(numvars -1, numvars-1)
+
+        for i in range(0, numvars-1):
+            var1 = trace_dict[varnames[i]]
+
+            for j in range(i, numvars-1):
+                var2 = trace_dict[varnames[j+1]]
+
+                ax = plt.subplot(gs[j, i])
+
+                if hexbin:
+                    ax.hexbin(var1, var2, mincnt=1, **kwargs)
+                else:
+                    ax.scatter(var1, var2, **kwargs)
+                
+                if divergences:
+                    try:
+                        divergent = trace['diverging']
+                    except KeyError:
+                        warnings.warn('No divergences were found.')
+                        return ax
+
+                    diverge = (divergent == 1)
+                    ax.scatter(trace_dict[varnames[0]][diverge], 
+                                     trace_dict[varnames[1]][diverge], 
+                                     **kwargs_divergence)
+
+                if j+1 != numvars-1:
+                    ax.set_xticks([])
+                else:
+                    ax.set_xlabel('{}'.format(varnames[i]), 
+                                            fontsize=text_size)
+                if i != 0:
+                    ax.set_yticks([])
+                else:
+                    ax.set_ylabel('{}'.format(varnames[j+1]), fontsize=text_size)
+                
+                ax.tick_params(labelsize=text_size)
+                ax.tick_params(labelsize=text_size)
+           
+    plt.tight_layout()
+    return ax

pymc3/plots/pairplot.py

@@ -4,7 +4,7 @@
     pass
 import numpy as np
 # plotting utilities can all be in this namespace
-from ..util import get_default_varnames  # pylint: disable=unused-import
+from ..util import get_default_varnames # pylint: disable=unused-import
 
 
 def identity_transform(x):

pymc3/plots/scatterplot.py

@@ -6,7 +6,7 @@
 from .checks import close_to
 
 from .models import multidimensional_model, simple_categorical
-from ..plots import traceplot, forestplot, autocorrplot, plot_posterior, energyplot, densityplot, scatterplot
+from ..plots import traceplot, forestplot, autocorrplot, plot_posterior, energyplot, densityplot, pairplot
 from ..plots.utils import make_2d
 from ..step_methods import Slice, Metropolis
 from ..sampling import sample
@@ -120,14 +120,14 @@ def test_plots_transformed():
     assert plot_posterior(trace).numCols == 1
     assert plot_posterior(trace, plot_transformed=True).shape == (2, )
 
-def test_scatterplot():
+def test_pairplot():
     with pm.Model() as model:
         a = pm.Normal('a', shape=2)
         c = pm.HalfNormal('c', shape=2)
         b = pm.Normal('b', a, c, shape=2)
         d = pm.Normal('d', 100, 1)
         trace = pm.sample(1000)
 
-    scatterplot(trace)
-    scatterplot(trace, hexbin=True, plot_transformed=True)
-    scatterplot(trace, sub_varnames=['a_0', 'c_0', 'b_1'])
+    pairplot(trace)
+    pairplot(trace, hexbin=True, plot_transformed=True)
+    pairplot(trace, sub_varnames=['a_0', 'c_0', 'b_1'])