add varnames argument for consintency (#1641)

* add varnames argument for consintency

* do not pass flat_names argument (or varnames)

Author: aloctavodia

pymc3/backends/text.py

@@ -192,12 +192,8 @@ def dump(name, trace, chains=None):
     if chains is None:
         chains = trace.chains
 
-    var_shapes = trace._straces[chains[0]].var_shapes
-    flat_names = {v: ttab.create_flat_names(v, shape)
-                  for v, shape in var_shapes.items()}
-
     for chain in chains:
         filename = os.path.join(name, 'chain-{}.csv'.format(chain))
         df = ttab.trace_to_dataframe(
-            trace, chains=chain, flat_names=flat_names,hide_transformed_vars=False)
+            trace, chains=chain, hide_transformed_vars=False)
         df.to_csv(filename, index=False)

pymc3/backends/tracetab.py

@@ -7,7 +7,7 @@
 __all__ = ['trace_to_dataframe']
 
 
-def trace_to_dataframe(trace, chains=None, flat_names=None, hide_transformed_vars=True):
+def trace_to_dataframe(trace, chains=None, varnames=None, hide_transformed_vars=True):
     """Convert trace to Pandas DataFrame.
 
     Parameters
@@ -16,21 +16,29 @@ def trace_to_dataframe(trace, chains=None, flat_names=None, hide_transformed_var
     chains : int or list of ints
         Chains to include. If None, all chains are used. A single
         chain value can also be given.
-    flat_names : dict or None
-        A dictionary that maps each variable name in `trace` to a list
+    varnames : list of variable names
+        Variables to be included in the DataFrame, if None all variable are 
+        included.
+    hide_transformed_vars: boolean
+        If true transformed variables will not be included in the resulting 
+        DataFrame.
     """
     var_shapes = trace._straces[0].var_shapes
-    if flat_names is None:
-        flat_names = {v: create_flat_names(v, shape)
-                      for v, shape in var_shapes.items()
-                      if not (hide_transformed_vars and v.endswith('_'))}
+    
+    if varnames is None:
+        varnames = var_shapes.keys()
 
+    flat_names = {v: create_flat_names(v, shape)
+                    for v, shape in var_shapes.items()
+                    if not (hide_transformed_vars and v.endswith('_'))}
+    
     var_dfs = []
-    for varname, shape in var_shapes.items():
-        if not hide_transformed_vars or not varname.endswith('_'):
-            vals = trace.get_values(varname, combine=True, chains=chains)
-            flat_vals = vals.reshape(vals.shape[0], -1)
-            var_dfs.append(pd.DataFrame(flat_vals, columns=flat_names[varname]))
+    for v, shape in var_shapes.items():
+        if v in varnames:
+            if not hide_transformed_vars or not v.endswith('_'):
+                vals = trace.get_values(v, combine=True, chains=chains)
+                flat_vals = vals.reshape(vals.shape[0], -1)
+                var_dfs.append(pd.DataFrame(flat_vals, columns=flat_names[v]))
     return pd.concat(var_dfs, axis=1)