Repository: arviz-devs/arviz

Issue:
Energyplot needs update to work with xarray



Discussion:

Resolving patch:
diff --git a/arviz/plots/energyplot.py b/arviz/plots/energyplot.py
--- a/arviz/plots/energyplot.py
+++ b/arviz/plots/energyplot.py
@@ -3,20 +3,21 @@
 import matplotlib.pyplot as plt
 from .kdeplot import kdeplot
 from ..stats import bfmi as e_bfmi
-from ..utils import get_stats
+from .plot_utils import _scale_text
+from ..utils import convert_to_dataset
 
 
-def energyplot(trace, kind='kde', bfmi=True, figsize=None, legend=True, fill_alpha=(1, .75),
-               fill_color=('C0', 'C5'), bw=4.5, skip_first=0, fill_kwargs=None, ax=None,
-               **kwargs):
+def energyplot(data, kind='kde', bfmi=True, figsize=None, legend=True, fill_alpha=(1, .75),
+               fill_color=('C0', 'C5'), bw=4.5, textsize=None, fill_kwargs=None, plot_kwargs=None,
+               ax=None):
     """Plot energy transition distribution and marginal energy distribution in HMC algorithms.
 
     This may help to diagnose poor exploration by gradient-based algorithms like HMC or NUTS.
 
     Parameters
     ----------
-    trace : Pandas DataFrame or PyMC3 trace
-        Posterior samples
+    data : xarray dataset, or object that can be converted (must represent
+           `sample_stats` and have an `energy` variable)
     kind : str
         Type of plot to display (kde or histogram)
     bfmi : bool
@@ -35,10 +36,12 @@ def energyplot(trace, kind='kde', bfmi=True, figsize=None, legend=True, fill_alp
         Bandwidth scaling factor for the KDE. Should be larger than 0. The higher this number the
         smoother the KDE will be. Defaults to 4.5 which is essentially the same as the Scott's rule
         of thumb (the default rule used by SciPy). Only works if `kind='kde'`
-    skip_first : int
-        Number of first samples not shown in plots (burn-in).
+    textsize: int
+        Text size for labels
     fill_kwargs : dicts, optional
-        Additional keywords passed to `fill_between` (to control the shade)
+        Additional keywords passed to `arviz.kdeplot` (to control the shade)
+    plot_kwargs : dicts, optional
+        Additional keywords passed to `arviz.kdeplot` or `plt.hist` (if type='hist')
     ax : axes
         Matplotlib axes.
 
@@ -46,10 +49,10 @@ def energyplot(trace, kind='kde', bfmi=True, figsize=None, legend=True, fill_alp
     -------
     ax : matplotlib axes
     """
-    energy = get_stats(trace[skip_first:], 'energy')
+    energy = convert_to_dataset(data, group='sample_stats').energy.values
 
     if figsize is None:
-        figsize = (6, 6)
+        figsize = (8, 6)
 
     if ax is None:
         _, ax = plt.subplots(figsize=figsize)
@@ -57,6 +60,10 @@ def energyplot(trace, kind='kde', bfmi=True, figsize=None, legend=True, fill_alp
     if fill_kwargs is None:
         fill_kwargs = {}
 
+    if plot_kwargs is None:
+        plot_kwargs = {}
+    textsize, linewidth, _ = _scale_text(figsize, textsize=textsize, scale_ratio=1)
+
     series = zip(
         fill_alpha,
         fill_color,
@@ -68,24 +75,23 @@ def energyplot(trace, kind='kde', bfmi=True, figsize=None, legend=True, fill_alp
         for alpha, color, label, value in series:
             fill_kwargs['alpha'] = alpha
             fill_kwargs['color'] = color
-            plot_kwargs = {
-                'color': color,
-                'alpha': 0
-            }
-            plot_kwargs.update(kwargs)
-            kdeplot(value, bw=bw, label=label,
+            plot_kwargs.setdefault('color', color)
+            plot_kwargs.setdefault('alpha', 0)
+            plot_kwargs.setdefault('linewidth', linewidth)
+            kdeplot(value, bw=bw, label=label, textsize=textsize,
                     plot_kwargs=plot_kwargs, fill_kwargs=fill_kwargs, ax=ax)
 
     elif kind == 'hist':
         for alpha, color, label, value in series:
-            ax.hist(value, alpha=alpha, label=label, color=color, **kwargs)
+            ax.hist(value.flatten(), bins='auto', density=True, alpha=alpha,
+                    label=label, color=color, **plot_kwargs)
 
     else:
         raise ValueError('Plot type {} not recognized.'.format(kind))
 
     if bfmi:
-        for idx, val in enumerate(e_bfmi(trace)):
-            plt.plot([], label='chain {:>2} BFMI = {:.2f}'.format(idx, val), alpha=0)
+        for idx, val in enumerate(e_bfmi(energy)):
+            ax.plot([], label='chain {:>2} BFMI = {:.2f}'.format(idx, val), alpha=0)
 
     ax.set_xticks([])
     ax.set_yticks([])
diff --git a/arviz/stats/stats.py b/arviz/stats/stats.py
--- a/arviz/stats/stats.py
+++ b/arviz/stats/stats.py
@@ -7,13 +7,13 @@
 from scipy.stats import dirichlet, circmean, circstd
 from scipy.optimize import minimize
 
-from ..utils import get_stats, get_varnames, trace_to_dataframe, log_post_trace
+from ..utils import get_varnames, trace_to_dataframe, log_post_trace
 from .diagnostics import effective_n, gelman_rubin
 
 __all__ = ['bfmi', 'compare', 'hpd', 'loo', 'psislw', 'r2_score', 'summary', 'waic']
 
 
-def bfmi(trace):
+def bfmi(energy):
     R"""Calculate the estimated Bayesian fraction of missing information (BFMI).
 
     BFMI quantifies how well momentum resampling matches the marginal energy distribution. For more
@@ -24,8 +24,10 @@ def bfmi(trace):
 
     Parameters
     ----------
-    trace : Pandas DataFrame or PyMC3 trace
-        Result of an HMC/NUTS run, must contain energy information
+    energy : NumPy array
+        Should be extracted from a gradient based sampler, such as in Stan or PyMC3. Typically,
+        after converting a trace or fit to InferenceData, the energy will be in
+        `data.sample_stats.energy`.
 
     Returns
     -------
@@ -33,9 +35,8 @@ def bfmi(trace):
         The Bayesian fraction of missing information of the model and trace. One element per
         chain in the trace.
     """
-    energy = np.atleast_2d(get_stats(trace, 'energy', combined=False))
-
-    return np.square(np.diff(energy, axis=1)).mean(axis=1) / np.var(energy, axis=1)
+    energy_mat = np.atleast_2d(energy)
+    return np.square(np.diff(energy_mat, axis=1)).mean(axis=1) / np.var(energy_mat, axis=1)
 
 
 def compare(model_dict, ic='waic', method='stacking', b_samples=1000, alpha=1,


Test patch:
diff --git a/arviz/tests/test_plots.py b/arviz/tests/test_plots.py
--- a/arviz/tests/test_plots.py
+++ b/arviz/tests/test_plots.py
@@ -54,10 +54,10 @@ def test_forestplot(self):
             _, axes = forestplot(obj, kind='ridgeplot', r_hat=False, n_eff=False)
             assert axes.shape == (1,)
 
-    def test_energyplot(self):
-        with pytest.raises(AttributeError):
-            energyplot(self.df_trace)
-        assert energyplot(self.short_trace)
+    @pytest.mark.parametrize('kind', ['kde', 'hist'])
+    def test_energyplot(self, kind):
+        for obj in (self.short_trace, self.fit):
+            assert energyplot(obj, kind=kind)
 
     def test_parallelplot(self):
         with pytest.raises(ValueError):
diff --git a/arviz/tests/test_stats.py b/arviz/tests/test_stats.py
--- a/arviz/tests/test_stats.py
+++ b/arviz/tests/test_stats.py
@@ -11,8 +11,8 @@
 
 
 def test_bfmi():
-    trace = pd.DataFrame([1, 2, 3, 4], columns=['energy'])
-    assert_almost_equal(bfmi(trace), 0.8)
+    energy = np.array([1, 2, 3, 4])
+    assert_almost_equal(bfmi(energy), 0.8)
 
 
 def test_hpd():