show transformed variables when requested by the user (#1386)

* show transformed when requested by the user

* add tests plot_transformed and include_transformed

Author: aloctavodia

pymc3/plots.py

@@ -56,7 +56,10 @@ def traceplot(trace, varnames=None, transform=lambda x: x, figsize=None,
     """
 
     if varnames is None:
-        varnames = [name for name in trace.varnames if not name.endswith('_')]
+        if plot_transformed:
+            varnames = [name for name in trace.varnames]
+        else:
+            varnames = [name for name in trace.varnames if not name.endswith('_')]
 
     n = len(varnames)
 
@@ -217,7 +220,10 @@ def _handle_array_varnames(varname):
             yield varname
 
     if varnames is None:
-        varnames = [name for name in trace.varnames if not name.endswith('_')]
+        if plot_transformed:
+            varnames = [name for name in trace.varnames]
+        else:
+            varnames = [name for name in trace.varnames if not name.endswith('_')]
 
     varnames = [item for sub in [[i for i in _handle_array_varnames(v)]
                                  for v in varnames] for item in sub]
@@ -758,8 +764,10 @@ def create_axes_grid(figsize, varnames):
         plot_posterior_op(transform(trace), ax)
     else:
         if varnames is None:
-            varnames = [
-                name for name in trace.varnames if not name.endswith('_')]
+            if plot_transformed:
+                varnames = [name for name in trace.varnames]
+            else:
+                varnames = [name for name in trace.varnames if not name.endswith('_')]
 
         if ax is None:
             ax, fig = create_axes_grid(figsize, varnames)

pymc3/stats.py

@@ -485,7 +485,10 @@ def df_summary(trace, varnames=None, stat_funcs=None, extend=False, include_tran
     mu__1  0.067513 -0.159097 -0.045637  0.062912
     """
     if varnames is None:
-        varnames = [name for name in trace.varnames if not name.endswith('_')]
+        if include_transformed:
+            varnames = [name for name in trace.varnames]
+        else:
+            varnames = [name for name in trace.varnames if not name.endswith('_')]
 
     funcs = [lambda x: pd.Series(np.mean(x, 0), name='mean'),
              lambda x: pd.Series(np.std(x, 0), name='sd'),
@@ -550,7 +553,10 @@ def summary(trace, varnames=None, alpha=0.05, start=0, batches=100, roundto=3,
 
     """
     if varnames is None:
-        varnames = [name for name in trace.varnames if not name.endswith('_')]
+        if include_transformed:
+            varnames = [name for name in trace.varnames]
+        else:
+            varnames = [name for name in trace.varnames if not name.endswith('_')]
 
     stat_summ = _StatSummary(roundto, batches, alpha)
     pq_summ = _PosteriorQuantileSummary(roundto, alpha)

pymc3/tests/test_plots.py

@@ -2,10 +2,11 @@
 matplotlib.use('Agg', warn=False)
 
 import numpy as np
+import pymc3 as pm
 from .checks import close_to
 
 from .models import multidimensional_model
-from ..plots import traceplot, forestplot, autocorrplot, make_2d
+from ..plots import traceplot, forestplot, autocorrplot, plot_posterior, make_2d
 from ..step_methods import Slice, Metropolis
 from ..sampling import sample
 from ..tuning.scaling import find_hessian
@@ -64,3 +65,17 @@ def test_make_2d():
     assert res.shape == (n, 20)
     close_to(a[:, 0, 0], res[:, 0], 0)
     close_to(a[:, 3, 2], res[:, 2 * 4 + 3], 0)
+
+
+def test_plots_transformed():
+    with pm.Model() as model:
+        pm.Uniform('x', 0, 1)
+        step = pm.Metropolis()
+        trace = pm.sample(100, step=step)
+
+    assert traceplot(trace).shape == (1, 2)
+    assert traceplot(trace, plot_transformed=True).shape == (2, 2)
+    assert autocorrplot(trace).shape == (1, 1)
+    assert autocorrplot(trace, plot_transformed=True).shape == (2, 1)
+    assert plot_posterior(trace).shape == (1, )
+    assert plot_posterior(trace, plot_transformed=True).shape == (2, )

pymc3/tests/test_stats.py

@@ -366,3 +366,12 @@ def test_value_alignment(self):
                 else:
                     vidx = var
                 npt.assert_equal(val, ds.loc[vidx, 'mean'])
+
+    def test_row_names(self):
+        with Model() as model:
+            pm.Uniform('x', 0, 1)
+            step = Metropolis()
+            trace = pm.sample(100, step=step)
+        ds = df_summary(trace, batches=3, include_transformed=True)
+        npt.assert_equal(np.array(['x_interval_', 'x']),
+                         ds.index)