add new scatterplot function

Author: agustinaarroyuelo

RELEASE-NOTES.md

@@ -9,6 +9,8 @@
   `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.
+  Optionally it can plot divergences. 
 
 ### Fixes
 

docs/source/notebooks/Diagnosing_biased_Inference_with_Divergences.ipynb

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

pymc3/plots/__init__.py

@@ -1,5 +1,6 @@
 import numpy as np
 from scipy.stats import mode
+from collections import OrderedDict
 
 from pymc3.stats import hpd
 from .kdeplot import fast_kde, kdeplot
@@ -144,3 +145,26 @@ def set_key_if_doesnt_exist(d, key, value):
         display_ref_val(ref_val)
     if rope is not None:
         display_rope(rope)
+
+def scale_text(figsize, text_size):
+        """Scale text to figsize."""
+
+        if text_size is None and figsize is not None:
+            if figsize[0] <= 11:
+                return 12
+            else:
+                return figsize[0]
+        else:
+            return text_size
+
+def get_trace_dict(tr, varnames):
+        traces = OrderedDict()
+        for v in varnames:
+            vals = tr.get_values(v, combine=True, squeeze=True)
+            if vals.ndim > 1:
+                vals_flat = vals.reshape(vals.shape[0], -1).T
+                for i, vi in enumerate(vals_flat):
+                    traces['_'.join([v, str(i)])] = vi
+            else:
+                traces[v] = vals
+        return traces

pymc3/plots/artists.py

@@ -6,7 +6,8 @@
     pass
 import numpy as np
 
-from .artists import plot_posterior_op
+from .artists import plot_posterior_op, get_trace_dict, scale_text
+
 from .utils import identity_transform, get_default_varnames
 
 
@@ -59,17 +60,6 @@ def plot_posterior(trace, varnames=None, transform=identity_transform, figsize=N
 
     """
 
-    def scale_text(figsize, text_size=text_size):
-        """Scale text to figsize."""
-
-        if text_size is None and figsize is not None:
-            if figsize[0] <= 11:
-                return 12
-            else:
-                return figsize[0]
-        else:
-            return text_size
-
     def create_axes_grid(figsize, traces):
         l_trace = len(traces)
         if l_trace == 1:
@@ -85,18 +75,6 @@ def create_axes_grid(figsize, traces):
                 ax = ax[:-1]
         return fig, ax
 
-    def get_trace_dict(tr, varnames):
-        traces = OrderedDict()
-        for v in varnames:
-            vals = tr.get_values(v, combine=True, squeeze=True)
-            if vals.ndim > 1:
-                vals_flat = vals.reshape(vals.shape[0], -1).T
-                for i, vi in enumerate(vals_flat):
-                    traces['_'.join([v, str(i)])] = vi
-            else:
-                traces[v] = vals
-        return traces
-
     if isinstance(trace, np.ndarray):
         if figsize is None:
             figsize = (6, 2)
@@ -106,7 +84,7 @@ def get_trace_dict(tr, varnames):
         plot_posterior_op(transform(trace), ax=ax, kde_plot=kde_plot,
                           point_estimate=point_estimate, round_to=round_to,
                           alpha_level=alpha_level, ref_val=ref_val, rope=rope,
-                          text_size=scale_text(figsize), **kwargs)
+                          text_size=scale_text(figsize, text_size), **kwargs)
     else:
         if varnames is None:
             varnames = get_default_varnames(trace.varnames, plot_transformed)
@@ -132,8 +110,8 @@ def get_trace_dict(tr, varnames):
             plot_posterior_op(tr_values, ax=a, kde_plot=kde_plot,
                               point_estimate=point_estimate, round_to=round_to,
                               alpha_level=alpha_level, ref_val=ref_val[idx],
-                              rope=rope[idx], text_size=scale_text(figsize), **kwargs)
-            a.set_title(v, fontsize=scale_text(figsize))
+                              rope=rope[idx], text_size=scale_text(figsize, text_size), **kwargs)
+            a.set_title(v, fontsize=scale_text(figsize, text_size))
 
         plt.tight_layout()
     return ax

pymc3/plots/posteriorplot.py

@@ -0,0 +1,136 @@
+import warnings
+
+try:
+    import matplotlib.pyplot as plt
+except ImportError:  # mpl is optional
+    pass
+import numpy as np
+from .utils import get_default_varnames
+from .artists import get_trace_dict, scale_text
+
+
+def scatterplot(trace, varnames=None, figsize=None, text_size=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
+    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/sub_varnames = None.
+        When a list of varnames/sub_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:
+            trace_dict = get_trace_dict(trace, get_default_varnames(trace.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:
+            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']
+                if np.any(divergent):
+                    ax.scatter(trace_dict[varnames[0]][divergent == 1], trace_dict[varnames[1]][divergent == 1], **kwargs_divergence)
+                else:
+                    print('No divergences were found.')
+            except KeyError:
+                warnings.warn('There is no divergence information in the passed trace.')
+                return ax
+
+    if ax is None:
+        _, ax = plt.subplots(nrows=numvars, ncols=numvars, figsize=figsize)
+
+        for i in range(numvars):
+            var1 = trace_dict[varnames[i]]            
+            for j in range(i, numvars):
+                var2 = trace_dict[varnames[j]]
+
+                if i == j:
+                    ax[i, j].axes.remove()
+
+                else:
+                    if hexbin:
+                        ax[j, i].hexbin(var1, var2, mincnt=1, **kwargs)
+                    else:
+                        ax[j, i].scatter(var1, var2, **kwargs)
+                    
+                    if divergences:
+                        try:
+                            divergent = trace['diverging']
+                            if np.any(divergent):
+                                ax[j, i].scatter(var1[divergent == 1], var2[divergent == 1], **kwargs_divergence)
+                            else:
+                                print('No divergences were found.')
+                        except KeyError:
+                            warnings.warn('There is no divergence information in the passed trace.')
+                            return ax
+                        
+                    ax[i, j].axes.remove()
+
+                    if j != numvars-1:
+                        ax[j, i].set_xticks([])
+                    if i != 0:
+                        ax[j, i].set_yticks([])
+
+                ax[numvars-1, j].set_xlabel('{}'.format(varnames[i]), fontsize=text_size)
+                ax[j, 0].set_ylabel('{}'.format(varnames[j]), fontsize=text_size)
+                ax[numvars-1, j].tick_params(labelsize=text_size)
+                ax[j, 0].tick_params(labelsize=text_size)
+           
+    plt.tight_layout()  
+    return ax      

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
+from ..plots import traceplot, forestplot, autocorrplot, plot_posterior, energyplot, densityplot, scatterplot
 from ..plots.utils import make_2d
 from ..step_methods import Slice, Metropolis
 from ..sampling import sample
@@ -66,6 +66,7 @@ def test_plots_multidimensional():
     forestplot(trace)
     densityplot(trace)
 
+
 @pytest.mark.xfail(condition=(theano.config.floatX == "float32"), reason="Fails on GPU due to cores=2")
 def test_multichain_plots():
     model = build_disaster_model()
@@ -118,3 +119,15 @@ def test_plots_transformed():
     assert autocorrplot(trace, plot_transformed=True).shape == (2, 2)
     assert plot_posterior(trace).numCols == 1
     assert plot_posterior(trace, plot_transformed=True).shape == (2, )
+
+def test_scatterplot():
+    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'])