Correctly (re)order dimensions for bfmi and plot_energy

CHANGELOG.md

@@ -12,6 +12,7 @@
   ([2096](https://github.com/arviz-devs/arviz/pull/2096) and [2105](https://github.com/arviz-devs/arviz/pull/2105))
 * Bokeh kde contour plots started to use `contourpy` package ([2104](https://github.com/arviz-devs/arviz/pull/2104))
 * Update default Bokeh markers for rcparams ([2104](https://github.com/arviz-devs/arviz/pull/2104))
+* Correctly (re)order dimensions for `bfmi` and `plot_energy` ([2126](https://github.com/arviz-devs/arviz/pull/2126))
 
 ### Deprecation
 * Removed `fill_last`, `contour` and `plot_kwargs` arguments from `plot_pair` function ([2085](https://github.com/arviz-devs/arviz/pull/2085))

arviz/plots/energyplot.py

@@ -97,7 +97,7 @@ def plot_energy(
         >>> az.plot_energy(data, kind='hist')
 
     """
-    energy = convert_to_dataset(data, group="sample_stats").energy.values
+    energy = convert_to_dataset(data, group="sample_stats").energy.transpose("chain", "draw").values
 
     if kind == "histogram":
         warnings.warn(

arviz/stats/diagnostics.py

@@ -65,7 +65,7 @@ def bfmi(data):
     dataset = convert_to_dataset(data, group="sample_stats")
     if not hasattr(dataset, "energy"):
         raise TypeError("Energy variable was not found.")
-    return _bfmi(dataset.energy)
+    return _bfmi(dataset.energy.transpose("chain", "draw"))
 
 
 def ess(

arviz/tests/base_tests/test_diagnostics.py

@@ -52,6 +52,13 @@ def test_bfmi_dataset_bad(self):
         with pytest.raises(TypeError):
             bfmi(data)
 
+    def test_bfmi_correctly_transposed(self):
+        data = load_arviz_data("centered_eight")
+        vals1 = bfmi(data)
+        data.sample_stats["energy"] = data.sample_stats["energy"].T
+        vals2 = bfmi(data)
+        assert_almost_equal(vals1, vals2)
+
     def test_deterministic(self):
         """
         Test algorithm against posterior (R) convergence functions.

arviz/tests/base_tests/test_plots_matplotlib.py

@@ -384,6 +384,14 @@ def test_plot_energy_bad(models):
         plot_energy(models.model_1, kind="bad_kind")
 
 
+def test_plot_energy_correctly_transposed():
+    idata = load_arviz_data("centered_eight")
+    idata.sample_stats["energy"] = idata.sample_stats.energy.T
+    ax = plot_energy(idata)
+    # legend has one entry for each KDE and 1 BFMI for each chain
+    assert len(ax.legend_.texts) == 2 + len(idata.sample_stats.chain)
+
+
 def test_plot_parallel_raises_valueerror(df_trace):  # pylint: disable=invalid-name
     with pytest.raises(ValueError):
         plot_parallel(df_trace)