Merge pull request #3338 from ColCarroll/arviz_plot

Use arviz for plotting

Author: ericmjl

README.rst

@@ -78,6 +78,13 @@ Or via conda-forge:
 
     conda install -c conda-forge pymc3
 
+Plotting is done using `ArviZ <https://arviz-devs.github.io/arviz/>`__
+which may be installed separately, or along with PyMC3:
+
+::
+
+    pip install pymc3[plots]
+
 The current development branch of PyMC3 can be installed from GitHub, also using ``pip``:
 
 ::

docs/source/notebooks/BEST.ipynb

@@ -1,9 +1,74 @@
-from .autocorrplot import autocorrplot
-from .compareplot import compareplot
-from .forestplot import forestplot
-from .kdeplot import kdeplot
-from .posteriorplot import plot_posterior, plot_posterior_predictive_glm
-from .traceplot import traceplot
-from .energyplot import energyplot
-from .densityplot import densityplot
-from .pairplot import pairplot
+"""PyMC3 Plotting.
+
+Plots are delegated to the ArviZ library, a general purpose library for
+"exploratory analysis of Bayesian models." See https://arviz-devs.github.io/arviz/
+for details on plots.
+"""
+import functools
+import sys
+import warnings
+try:
+    import arviz as az
+except ImportError:  # arviz is optional, throw exception when used
+
+    class _ImportWarner:
+        __all__ = []
+
+        def __init__(self, attr):
+            self.attr = attr
+
+        def __call__(self, *args, **kwargs):
+            raise ImportError(
+                "ArviZ is not installed. In order to use `{0.attr}`:\npip install arviz".format(self)
+            )
+
+    class _ArviZ:
+        def __getattr__(self, attr):
+            return _ImportWarner(attr)
+
+
+    az = _ArviZ()
+
+def map_args(func):
+    swaps = [
+        ('varnames', 'var_names')
+    ]
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        for (old, new) in swaps:
+            if old in kwargs and new not in kwargs:
+                warnings.warn('Keyword argument `{old}` renamed to `{new}`, and will be removed in pymc3 3.8'.format(old=old, new=new))
+                kwargs[new] = kwargs.pop(old)
+            return func(*args, **kwargs)
+    return wrapped
+
+# pymc3 custom plots: override these names for custom behavior
+autocorrplot = map_args(az.plot_autocorr)
+compareplot = map_args(az.plot_compare)
+forestplot = map_args(az.plot_forest)
+kdeplot = map_args(az.plot_kde)
+plot_posterior = map_args(az.plot_posterior)
+traceplot = map_args(az.plot_trace)
+energyplot = map_args(az.plot_energy)
+densityplot = map_args(az.plot_density)
+pairplot = map_args(az.plot_pair)
+
+from .posteriorplot import plot_posterior_predictive_glm
+
+
+# Access to arviz plots: base plots provided by arviz
+for plot in az.plots.__all__:
+    setattr(sys.modules[__name__], plot, map_args(getattr(az.plots, plot)))
+
+__all__ = tuple(az.plots.__all__) + (
+    'autocorrplot',
+    'compareplot',
+    'forestplot',
+    'kdeplot',
+    'plot_posterior',
+    'traceplot',
+    'energyplot',
+    'densityplot',
+    'pairplot',
+    'plot_posterior_predictive_glm',
+)