Cleanup root namespace

Author: ricardoV94

pymc/__init__.py

@@ -46,36 +46,19 @@ def __set_compiler_flags():
 
 __set_compiler_flags()
 
-from pymc import _version, gp, ode, sampling
-from pymc.backends import *
-from pymc.blocking import *
+from pymc import _version, gp, ode, plots, sampling, stats
 from pymc.data import *
 from pymc.distributions import *
-from pymc.exceptions import *
 from pymc.func_utils import find_constrained_prior
 from pymc.logprob import *
-from pymc.math import (
-    expand_packed_triangular,
-    invlogit,
-    invprobit,
-    logaddexp,
-    logit,
-    logsumexp,
-    probit,
-)
 from pymc.model.core import *
 from pymc.model.transform.conditioning import do, observe
 from pymc.model_graph import model_to_graphviz, model_to_networkx
-from pymc.plots import *
-from pymc.printing import *
 from pymc.pytensorf import *
 from pymc.sampling import *
 from pymc.smc import *
-from pymc.stats import *
 from pymc.step_methods import *
 from pymc.tuning import *
-from pymc.util import drop_warning_stat
 from pymc.variational import *
-from pymc.vartypes import *
 
 __version__ = _version.get_versions()["version"]

pymc/blocking.py

@@ -37,9 +37,6 @@
 
 from typing_extensions import TypeAlias
 
-__all__ = ["DictToArrayBijection"]
-
-
 T = TypeVar("T")
 PointType: TypeAlias = Dict[str, np.ndarray]
 StatsDict: TypeAlias = Dict[str, Any]

pymc/data.py

@@ -40,7 +40,6 @@
 
 __all__ = [
     "get_data",
-    "GeneratorAdapter",
     "Minibatch",
     "Data",
     "ConstantData",

pymc/exceptions.py

@@ -14,8 +14,6 @@
 
 __all__ = [
     "SamplingError",
-    "IncorrectArgumentsError",
-    "TraceDirectoryError",
     "ImputationWarning",
     "ShapeWarning",
     "ShapeError",
@@ -26,16 +24,6 @@ class SamplingError(RuntimeError):
     pass
 
 
-class IncorrectArgumentsError(ValueError):
-    pass
-
-
-class TraceDirectoryError(ValueError):
-    """Error from trying to load a trace from an incorrectly-structured directory,"""
-
-    pass
-
-
 class ImputationWarning(UserWarning):
     """Warning that there are missing values that will be imputed."""
 

pymc/plots/__init__.py

@@ -20,7 +20,6 @@
 """
 import functools
 import sys
-import warnings
 
 import arviz as az
 
@@ -29,40 +28,3 @@
     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 FutureWarning(
-            f"The function `{alias}` from PyMC was an alias for `{original}` from ArviZ. "
-            "It was removed in PyMC 4.0. "
-            f"Switch to `pymc.{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")
-
-
-__all__ = tuple(az.plots.__all__) + (
-    "autocorrplot",
-    "compareplot",
-    "forestplot",
-    "kdeplot",
-    "traceplot",
-    "energyplot",
-    "densityplot",
-    "pairplot",
-)

pymc/pytensorf.py

@@ -76,16 +76,9 @@
     "hessian",
     "hessian_diag",
     "inputvars",
-    "cont_inputs",
     "floatX",
     "intX",
-    "smartfloatX",
     "jacobian",
-    "CallableTensor",
-    "join_nonshared_inputs",
-    "make_shared_replacements",
-    "generator",
-    "convert_observed_data",
     "compile_pymc",
 ]
 

pymc/sampling/forward.py

@@ -53,9 +53,11 @@
 from pytensor.tensor.sharedvar import SharedVariable
 from typing_extensions import TypeAlias
 
-import pymc as pm
-
-from pymc.backends.arviz import _DefaultTrace
+from pymc.backends.arviz import (
+    _DefaultTrace,
+    predictions_to_inference_data,
+    to_inference_data,
+)
 from pymc.backends.base import MultiTrace
 from pymc.blocking import PointType
 from pymc.model import Model, modelcontext
@@ -438,7 +440,7 @@ def sample_prior_predictive(
     ikwargs: Dict[str, Any] = dict(model=model)
     if idata_kwargs:
         ikwargs.update(idata_kwargs)
-    return pm.to_inference_data(prior=prior, **ikwargs)
+    return to_inference_data(prior=prior, **ikwargs)
 
 
 def sample_posterior_predictive(
@@ -669,8 +671,8 @@ def sample_posterior_predictive(
         if extend_inferencedata:
             ikwargs.setdefault("idata_orig", idata)
             ikwargs.setdefault("inplace", True)
-        return pm.predictions_to_inference_data(ppc_trace, **ikwargs)
-    idata_pp = pm.to_inference_data(posterior_predictive=ppc_trace, **ikwargs)
+        return predictions_to_inference_data(ppc_trace, **ikwargs)
+    idata_pp = to_inference_data(posterior_predictive=ppc_trace, **ikwargs)
 
     if extend_inferencedata and idata is not None:
         idata.extend(idata_pp)

pymc/sampling/mcmc.py

@@ -52,6 +52,7 @@
     coords_and_dims_for_inferencedata,
     find_constants,
     find_observations,
+    to_inference_data,
 )
 from pymc.backends.base import IBaseTrace, MultiTrace, _choose_chains
 from pymc.blocking import DictToArrayBijection
@@ -892,7 +893,7 @@ def _sample_return(
     if compute_convergence_checks or return_inferencedata:
         ikwargs: Dict[str, Any] = dict(model=model, save_warmup=not discard_tuned_samples)
         ikwargs.update(idata_kwargs)
-        idata = pm.to_inference_data(mtrace, **ikwargs)
+        idata = to_inference_data(mtrace, **ikwargs)
 
         if compute_convergence_checks:
             warns = run_convergence_checks(idata, model)

pymc/stats/__init__.py

@@ -28,5 +28,3 @@
         setattr(sys.modules[__name__], attr, obj)
 
 from pymc.stats.log_likelihood import compute_log_likelihood
-
-__all__ = ("compute_log_likelihood",) + tuple(az.stats.__all__)

pymc/step_methods/metropolis.py

@@ -31,6 +31,7 @@
     compile_pymc,
     floatX,
     join_nonshared_inputs,
+    make_shared_replacements,
     replace_rng_nodes,
 )
 from pymc.step_methods.arraystep import (
@@ -804,7 +805,7 @@ def __init__(
 
         self.mode = mode
 
-        shared = pm.make_shared_replacements(initial_values, vars, model)
+        shared = make_shared_replacements(initial_values, vars, model)
         self.delta_logp = delta_logp(initial_values, model.logp(), vars, shared)
         super().__init__(vars, shared)
 
@@ -960,7 +961,7 @@ def __init__(
 
         self.mode = mode
 
-        shared = pm.make_shared_replacements(initial_values, vars, model)
+        shared = make_shared_replacements(initial_values, vars, model)
         self.delta_logp = delta_logp(initial_values, model.logp(), vars, shared)
         super().__init__(vars, shared)
 

pymc/tuning/__init__.py

@@ -14,3 +14,5 @@
 
 from pymc.tuning.scaling import find_hessian, guess_scaling, trace_cov
 from pymc.tuning.starting import find_MAP
+
+__all__ = ("find_MAP", "find_hessian")

pymc/variational/opvi.py

@@ -64,6 +64,7 @@
 
 import pymc as pm
 
+from pymc.backends import to_inference_data
 from pymc.backends.base import MultiTrace
 from pymc.backends.ndarray import NDArray
 from pymc.blocking import DictToArrayBijection
@@ -1578,7 +1579,7 @@ def sample(
         if not return_inferencedata:
             return multi_trace
         else:
-            return pm.to_inference_data(multi_trace, model=self.model, **kwargs)
+            return to_inference_data(multi_trace, model=self.model, **kwargs)
 
     @property
     def ndim(self):