Revive ArviZ aliases, but raise DeprecationWarning on old wrappers

RELEASE-NOTES.md

@@ -3,6 +3,8 @@
 ## PyMC3 vNext (4.0.0)
 ### Breaking Changes
 - ⚠ Theano-PyMC has been replaced with Aesara, so all external references to `theano`, `tt`, and `pymc3.theanof` need to be replaced with `aesara`, `aet`, and `pymc3.aesaraf` (see [4471](https://github.com/pymc-devs/pymc3/pull/4471)).
+- ArviZ `plots` and `stats` *wrappers* were removed. The functions are now just available by their original names (see [#4549](https://github.com/pymc-devs/pymc3/pull/4471) and `3.11.2` release notes).
+- ...
 
 ### New Features
 - The `CAR` distribution has been added to allow for use of conditional autoregressions which often are used in spatial and network models.

docs/source/api/plots.rst

@@ -7,8 +7,7 @@ Plots
 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.
+"exploratory analysis of Bayesian models".
 
-.. automodule:: pymc3.plots.posteriorplot
-   :members:
+Functions from the `arviz.plots` module are available through ``pymc3.<function>`` or ``pymc3.plots.<function>``,
+but for their API documentation please refer to the :ref:`ArviZ documentation <arviz:plot_api>`.

docs/source/api/stats.rst

@@ -1,8 +1,13 @@
 *****
 Stats
 *****
+
+.. currentmodule:: pymc3.stats
+
 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.
+"exploratory analysis of Bayesian models".
+
+Functions from the `arviz.stats` module are available through ``pymc3.<function>`` or ``pymc3.stats.<function>``,
+but for their API documentation please refer to the :ref:`ArviZ documentation <arviz:stats_api>`.

pymc3/__init__.py

@@ -62,6 +62,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.tuning import *

pymc3/plots/__init__.py

@@ -1,4 +1,4 @@
-#   Copyright 2020 The PyMC Developers
+#   Copyright 2021 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.
@@ -12,19 +12,61 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-"""PyMC3 Plotting.
+"""Alias for the `plots` submodule from ArviZ.
 
-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
 
+# Makes this module as identical to arviz.plots as possible
+for attr in az.plots.__all__:
+    obj = getattr(az.plots, attr)
+    if not attr.startswith("__"):
+        setattr(sys.modules[__name__], attr, obj)
+
+
+def alias_deprecation(func, alias: str):
+    original = func.__name__
+
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        raise DeprecationWarning(
+            f"The function `{alias}` from PyMC3 was an alias for `{original}` from ArviZ. "
+            "It was removed in PyMC3 4.0. "
+            f"Switch to `pymc3.{original}` or `arviz.{original}`."
+        )
+
+    return wrapped
+
+
+# Aliases of ArviZ functions
+autocorrplot = alias_deprecation(az.plot_autocorr, alias="autocorrplot")
+forestplot = alias_deprecation(az.plot_forest, alias="forestplot")
+kdeplot = alias_deprecation(az.plot_kde, alias="kdeplot")
+energyplot = alias_deprecation(az.plot_energy, alias="energyplot")
+densityplot = alias_deprecation(az.plot_density, alias="densityplot")
+pairplot = alias_deprecation(az.plot_pair, alias="pairplot")
+traceplot = alias_deprecation(az.plot_trace, alias="traceplot")
+compareplot = alias_deprecation(az.plot_compare, alias="compareplot")
+
+
 from pymc3.plots.posteriorplot import plot_posterior_predictive_glm
 
-__all__ = ["plot_posterior_predictive_glm"]
+__all__ = tuple(az.plots.__all__) + (
+    "autocorrplot",
+    "compareplot",
+    "forestplot",
+    "kdeplot",
+    "plot_posterior",
+    "traceplot",
+    "energyplot",
+    "densityplot",
+    "pairplot",
+    "plot_posterior_predictive_glm",
+)

pymc3/stats/__init__.py

@@ -0,0 +1,31 @@
+#   Copyright 2021 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.
+
+"""Alias for the `stats` submodule from ArviZ.
+
+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 sys
+
+import arviz as az
+
+for attr in az.stats.__all__:
+    obj = getattr(az.stats, attr)
+    if not attr.startswith("__"):
+        setattr(sys.modules[__name__], attr, obj)
+
+
+__all__ = tuple(az.stats.__all__)