Let plot_posterior_predictive_glm work with inferencedata too

RELEASE-NOTES.md

@@ -24,6 +24,7 @@
 - Add Bayesian Additive Regression Trees (BARTs) [#4183](https://github.com/pymc-devs/pymc3/pull/4183))
 - Added a new `MixtureSameFamily` distribution to handle mixtures of arbitrary dimensions in vectorized form (see [#4185](https://github.com/pymc-devs/pymc3/issues/4185)).
 - Added semantically meaningful `str` representations to PyMC3 objects for console, notebook, and GraphViz use (see [#4076](https://github.com/pymc-devs/pymc3/pull/4076), [#4065](https://github.com/pymc-devs/pymc3/pull/4065), [#4159](https://github.com/pymc-devs/pymc3/pull/4159), and [#4217](https://github.com/pymc-devs/pymc3/pull/4217))
+- `plot_posterior_predictive_glm` now works with `arviz.InferenceData` as well (see [#4234](https://github.com/pymc-devs/pymc3/pull/4234))
 
 
 

pymc3/plots/posteriorplot.py

@@ -18,6 +18,8 @@
     pass
 import numpy as np
 
+from pymc3.backends.base import MultiTrace
+
 
 def plot_posterior_predictive_glm(trace, eval=None, lm=None, samples=30, **kwargs):
     """Plot posterior predictive of a linear model.
@@ -47,10 +49,23 @@ def plot_posterior_predictive_glm(trace, eval=None, lm=None, samples=30, **kwarg
     if "c" not in kwargs and "color" not in kwargs:
         kwargs["c"] = "k"
 
+    plotting_fn = _plot_multitrace if isinstance(trace, MultiTrace) else _plot_inferencedata
+    plotting_fn(trace, eval, lm, samples, kwargs)
+    plt.title("Posterior predictive")
+
+
+def _plot_multitrace(trace, eval, lm, samples, kwargs):
     for rand_loc in np.random.randint(0, len(trace), samples):
         rand_sample = trace[rand_loc]
         plt.plot(eval, lm(eval, rand_sample), **kwargs)
         # Make sure to not plot label multiple times
         kwargs.pop("label", None)
 
-    plt.title("Posterior predictive")
+
+def _plot_inferencedata(trace, eval, lm, samples, kwargs):
+    trace_df = trace.posterior.to_dataframe()
+    for rand_loc in np.random.randint(0, len(trace_df), samples):
+        rand_sample = trace_df.iloc[rand_loc]
+        plt.plot(eval, lm(eval, rand_sample), **kwargs)
+        # Make sure to not plot label multiple times
+        kwargs.pop("label", None)

pymc3/tests/test_plots.py

@@ -0,0 +1,54 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+
+from arviz import from_pymc3
+
+import pymc3 as pm
+
+from pymc3.backends.ndarray import point_list_to_multitrace
+from pymc3.plots import plot_posterior_predictive_glm
+
+
+@pytest.mark.parametrize("inferencedata", [True, False])
+def test_plot_posterior_predictive_glm_defaults(inferencedata):
+    with pm.Model() as model:
+        pm.Normal("x")
+        pm.Normal("Intercept")
+    trace = point_list_to_multitrace([{"x": np.array([1]), "Intercept": np.array([1])}], model)
+    if inferencedata:
+        trace = from_pymc3(trace, model=model)
+    _, ax = plt.subplots()
+    plot_posterior_predictive_glm(trace, samples=1)
+    lines = ax.get_lines()
+    expected_xvalues = np.linspace(0, 1, 100)
+    expected_yvalues = np.linspace(1, 2, 100)
+    for line in lines:
+        x_axis, y_axis = line.get_data()
+        np.testing.assert_array_equal(x_axis, expected_xvalues)
+        np.testing.assert_array_equal(y_axis, expected_yvalues)
+        assert line.get_lw() == 0.2
+        assert line.get_c() == "k"
+
+
+@pytest.mark.parametrize("inferencedata", [True, False])
+def test_plot_posterior_predictive_glm_non_defaults(inferencedata):
+    with pm.Model() as model:
+        pm.Normal("x")
+        pm.Normal("Intercept")
+    trace = point_list_to_multitrace([{"x": np.array([1]), "Intercept": np.array([1])}], model)
+    if inferencedata:
+        trace = from_pymc3(trace, model=model)
+    _, ax = plt.subplots()
+    plot_posterior_predictive_glm(
+        trace, samples=1, lm=lambda x, _: x, eval=np.linspace(0, 1, 10), lw=0.3, c="b"
+    )
+    lines = ax.get_lines()
+    expected_xvalues = np.linspace(0, 1, 10)
+    expected_yvalues = np.linspace(0, 1, 10)
+    for line in lines:
+        x_axis, y_axis = line.get_data()
+        np.testing.assert_array_equal(x_axis, expected_xvalues)
+        np.testing.assert_array_equal(y_axis, expected_yvalues)
+        assert line.get_lw() == 0.3
+        assert line.get_c() == "b"