Restore plots and stats wrappers

See pymc-devs#4528

Author: michaelosthege

Diff for: docs/source/api/plots.rst

@@ -8,7 +8,16 @@ Plots are delegated to the
 `ArviZ <https://arviz-devs.github.io/arviz/index.html>`_.
 library, a general purpose library for
 "exploratory analysis of Bayesian models."
-Refer to its documentation to use the plotting functions directly.
+For plots, ``pymc3.<function>`` are now aliases
+for ArviZ functions. Thus, the links below will redirect you to
+ArviZ docs:
 
-.. automodule:: pymc3.plots.posteriorplot
-   :members:
+- :func:`pymc3.traceplot <arviz:arviz.plot_trace>`
+- :func:`pymc3.plot_posterior <arviz:arviz.plot_posterior>`
+- :func:`pymc3.forestplot <arviz:arviz.plot_forest>`
+- :func:`pymc3.compareplot <arviz:arviz.plot_compare>`
+- :func:`pymc3.autocorrplot <arviz:arviz.plot_autocorr>`
+- :func:`pymc3.energyplot <arviz:arviz.plot_energy>`
+- :func:`pymc3.kdeplot <arviz:arviz.plot_kde>`
+- :func:`pymc3.densityplot <arviz:arviz.plot_density>`
+- :func:`pymc3.pairplot <arviz:arviz.plot_pair>`

Diff for: docs/source/api/stats.rst

@@ -5,4 +5,21 @@ Statistics and diagnostics are delegated to the
 `ArviZ <https://arviz-devs.github.io/arviz/index.html>`_.
 library, a general purpose library for
 "exploratory analysis of Bayesian models."
-Refer to its documentation to use the diagnostics functions directly.
+For statistics and diagnostics, ``pymc3.<function>`` are now aliases
+for ArviZ functions. Thus, the links below will redirect you to
+ArviZ docs:
+
+.. currentmodule:: pymc3.stats
+
+
+- :func:`pymc3.bfmi <arviz:arviz.bfmi>`
+- :func:`pymc3.compare <arviz:arviz.compare>`
+- :func:`pymc3.ess <arviz:arviz.ess>`
+- :data:`pymc3.geweke <arviz:arviz.geweke>`
+- :func:`pymc3.hpd <arviz:arviz.hpd>`
+- :func:`pymc3.loo <arviz:arviz.loo>`
+- :func:`pymc3.mcse <arviz:arviz.mcse>`
+- :func:`pymc3.r2_score <arviz:arviz.r2_score>`
+- :func:`pymc3.rhat <arviz:arviz.rhat>`
+- :func:`pymc3.summary <arviz:arviz.summary>`
+- :func:`pymc3.waic <arviz:arviz.waic>`

Diff for: pymc3/__init__.py

@@ -61,6 +61,7 @@ def __set_compiler_flags():
 from pymc3.plots import *
 from pymc3.sampling import *
 from pymc3.smc import *
+from pymc3.stats import *
 from pymc3.step_methods import *
 from pymc3.tests import test
 from pymc3.theanof import *

Diff for: pymc3/plots/__init__.py

@@ -14,17 +14,108 @@
 
 """PyMC3 Plotting.
 
-Plots are delegated to the `ArviZ <https://arviz-devs.github.io/arviz/>`_ library, a general purpose library for
-exploratory analysis of Bayesian models. For more details, see https://arviz-devs.github.io/arviz/.
-
-Only `plot_posterior_predictive_glm` is kept in the PyMC code base for now, but it will move to ArviZ once the latter adds features for regression plots.
+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
 
 import arviz as az
 
+
+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(
+                    f"Keyword argument `{old}` renamed to `{new}`, and will be removed in pymc3 3.8"
+                )
+                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)
+forestplot = map_args(az.plot_forest)
+kdeplot = map_args(az.plot_kde)
+plot_posterior = map_args(az.plot_posterior)
+energyplot = map_args(az.plot_energy)
+densityplot = map_args(az.plot_density)
+pairplot = map_args(az.plot_pair)
+
+# Use compact traceplot by default
+@map_args
+@functools.wraps(az.plot_trace)
+def traceplot(*args, **kwargs):
+    try:
+        kwargs.setdefault("compact", True)
+        return az.plot_trace(*args, **kwargs)
+    except TypeError:
+        kwargs.pop("compact")
+        return az.plot_trace(*args, **kwargs)
+
+
+# addition arg mapping for compare plot
+@functools.wraps(az.plot_compare)
+def compareplot(*args, **kwargs):
+    if "comp_df" in kwargs:
+        comp_df = kwargs["comp_df"].copy()
+    else:
+        args = list(args)
+        comp_df = args[0].copy()
+    if "WAIC" in comp_df.columns:
+        comp_df = comp_df.rename(
+            index=str,
+            columns={
+                "WAIC": "waic",
+                "pWAIC": "p_waic",
+                "dWAIC": "d_waic",
+                "SE": "se",
+                "dSE": "dse",
+                "var_warn": "warning",
+            },
+        )
+    elif "LOO" in comp_df.columns:
+        comp_df = comp_df.rename(
+            index=str,
+            columns={
+                "LOO": "loo",
+                "pLOO": "p_loo",
+                "dLOO": "d_loo",
+                "SE": "se",
+                "dSE": "dse",
+                "shape_warn": "warning",
+            },
+        )
+    if "comp_df" in kwargs:
+        kwargs["comp_df"] = comp_df
+    else:
+        args[0] = comp_df
+    return az.plot_compare(*args, **kwargs)
+
+
 from pymc3.plots.posteriorplot import plot_posterior_predictive_glm
 
-__all__ = ["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",
+)

Diff for: pymc3/stats/__init__.py

@@ -0,0 +1,69 @@
+#   Copyright 2020 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+"""Statistical utility functions for PyMC3
+
+Diagnostics and auxiliary statistical functions 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.
+"""
+import functools
+import warnings
+
+import arviz as az
+
+
+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.9".format(old=old, new=new)
+                )
+                kwargs[new] = kwargs.pop(old)
+            return func(*args, **kwargs)
+
+    return wrapped
+
+
+bfmi = map_args(az.bfmi)
+compare = map_args(az.compare)
+ess = map_args(az.ess)
+geweke = map_args(az.geweke)
+hpd = map_args(az.hpd)
+loo = map_args(az.loo)
+mcse = map_args(az.mcse)
+r2_score = map_args(az.r2_score)
+rhat = map_args(az.rhat)
+summary = map_args(az.summary)
+waic = map_args(az.waic)
+
+
+__all__ = [
+    "bfmi",
+    "compare",
+    "ess",
+    "geweke",
+    "hpd",
+    "loo",
+    "mcse",
+    "r2_score",
+    "rhat",
+    "summary",
+    "waic",
+]