REF: Migrate to ArviZ 1.0 (InferenceData -> DataTree)

.gitignore

@@ -46,3 +46,9 @@ pytestdebug.log
 # Codespaces
 pythonenv*
 env/
+.venv/
+
+# Log files
+*.log
+log_errores*.txt
+test_results*.txt

notebooks/DFM_Example_(Coincident_Index).ipynb

@@ -8,21 +8,20 @@
    "outputs": [],
    "source": [
     "# Imports\n",
-    "import numpy as np\n",
-    "import pandas as pd\n",
-    "import matplotlib.pyplot as plt\n",
-    "from pandas_datareader.data import DataReader\n",
-    "from statsmodels.tsa.statespace.dynamic_factor import DynamicFactor\n",
     "import datetime\n",
-    "from statsmodels.tsa.stattools import adfuller\n",
-    "import statsmodels.api as sm\n",
+    "\n",
     "import arviz as az\n",
-    "import pytensor.tensor as pt\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
     "import pymc as pm\n",
     "import pytensor\n",
+    "import pytensor.tensor as pt\n",
+    "import statsmodels.api as sm\n",
+    "\n",
+    "from pandas_datareader.data import DataReader\n",
+    "from statsmodels.tsa.statespace.dynamic_factor import DynamicFactor\n",
     "\n",
-    "from pymc_extras.statespace.core.statespace import PyMCStateSpace\n",
-    "from pymc_extras.statespace.utils.constants import ALL_STATE_AUX_DIM, ALL_STATE_DIM\n",
     "from pymc_extras.statespace.models.DFM import BayesianDynamicFactor"
    ]
   },
@@ -248,7 +247,7 @@
     "            data.dropna(inplace=True)\n",
     "\n",
     "        print(series.center(110))\n",
-    "        print((\"=\" * 110))\n",
+    "        print(\"=\" * 110)\n",
     "        line = (\n",
     "            \"Specification\"\n",
     "            + \" \" * 15\n",
@@ -264,7 +263,7 @@
     "        )\n",
     "        line += \" \" * 10 + \"5%\" + \" \" * 8 + \"10%\"\n",
     "        print(line)\n",
-    "        print((\"-\" * 110))\n",
+    "        print(\"-\" * 110)\n",
     "        spec_fixed = False\n",
     "        for i, (name, reg) in enumerate(\n",
     "            zip([\"Constant and Trend\", \"Constant Only\", \"No Constant\"], [\"ct\", \"c\", \"n\"])\n",
@@ -279,7 +278,7 @@
     "            reg_tstat = pd.Series(regresult.resols.tvalues, index=names)\n",
     "            reg_pvals = pd.Series(regresult.resols.pvalues, index=names)\n",
     "\n",
-    "            line = f'{name:<21}{gamma:13.3f}{stat:15.3f}{p:13.3f}{n_lag:11}{crit[\"1%\"]:10.3f}{crit[\"5%\"]:12.3f}{crit[\"10%\"]:11.3f}'\n",
+    "            line = f\"{name:<21}{gamma:13.3f}{stat:15.3f}{p:13.3f}{n_lag:11}{crit['1%']:10.3f}{crit['5%']:12.3f}{crit['10%']:11.3f}\"\n",
     "            print(line)\n",
     "\n",
     "            for coef in reg_coefs.index:\n",
@@ -874,7 +873,7 @@
     "\n",
     "\n",
     "def print_model_ssm(mod, how=\"eval\"):\n",
-    "    nice_heading = f'{\"name\":<20}{\"__repr__\":<50}{\"shape\":<10}{\"value\":<20}'\n",
+    "    nice_heading = f\"{'name':<20}{'__repr__':<50}{'shape':<10}{'value':<20}\"\n",
     "    print(nice_heading)\n",
     "    print(\"=\" * len(nice_heading))\n",
     "    if how == \"eval\":\n",
@@ -930,19 +929,19 @@
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001B[3m                                     Model Requirements                                     \u001B[0m\n",
+       "\u001b[3m                                     Model Requirements                                     \u001b[0m\n",
        "                                                                                            \n",
-       " \u001B[1m \u001B[0m\u001B[1mVariable       \u001B[0m\u001B[1m \u001B[0m \u001B[1m \u001B[0m\u001B[1mShape   \u001B[0m\u001B[1m \u001B[0m \u001B[1m \u001B[0m\u001B[1mConstraints           \u001B[0m\u001B[1m \u001B[0m \u001B[1m \u001B[0m\u001B[1m                        Dimensions\u001B[0m\u001B[1m \u001B[0m \n",
+       " \u001b[1m \u001b[0m\u001b[1mVariable       \u001b[0m\u001b[1m \u001b[0m \u001b[1m \u001b[0m\u001b[1mShape   \u001b[0m\u001b[1m \u001b[0m \u001b[1m \u001b[0m\u001b[1mConstraints           \u001b[0m\u001b[1m \u001b[0m \u001b[1m \u001b[0m\u001b[1m                        Dimensions\u001b[0m\u001b[1m \u001b[0m \n",
        " ────────────────────────────────────────────────────────────────────────────────────────── \n",
-       "  x0                \u001B[1m(\u001B[0m\u001B[1;36m10\u001B[0m,\u001B[1m)\u001B[0m                                                       \u001B[1m(\u001B[0m\u001B[32m'state'\u001B[0m,\u001B[1m)\u001B[0m  \n",
-       "  P0                \u001B[1m(\u001B[0m\u001B[1;36m10\u001B[0m, \u001B[1;36m10\u001B[0m\u001B[1m)\u001B[0m   Positive Semi-definite               \u001B[1m(\u001B[0m\u001B[32m'state'\u001B[0m, \u001B[32m'state_aux'\u001B[0m\u001B[1m)\u001B[0m  \n",
-       "  factor_loadings   \u001B[1m(\u001B[0m\u001B[1;36m4\u001B[0m, \u001B[1;36m1\u001B[0m\u001B[1m)\u001B[0m                                    \u001B[1m(\u001B[0m\u001B[32m'observed_state'\u001B[0m, \u001B[32m'factor'\u001B[0m\u001B[1m)\u001B[0m  \n",
-       "  factor_ar         \u001B[1m(\u001B[0m\u001B[1;36m1\u001B[0m, \u001B[1;36m2\u001B[0m\u001B[1m)\u001B[0m                                            \u001B[1m(\u001B[0m\u001B[32m'factor'\u001B[0m, \u001B[32m'lag_ar'\u001B[0m\u001B[1m)\u001B[0m  \n",
-       "  error_ar          \u001B[1m(\u001B[0m\u001B[1;36m4\u001B[0m, \u001B[1;36m2\u001B[0m\u001B[1m)\u001B[0m                              \u001B[1m(\u001B[0m\u001B[32m'observed_state'\u001B[0m, \u001B[32m'error_lag_ar'\u001B[0m\u001B[1m)\u001B[0m  \n",
-       "  error_sigma       \u001B[1m(\u001B[0m\u001B[1;36m4\u001B[0m,\u001B[1m)\u001B[0m       Positive                                \u001B[1m(\u001B[0m\u001B[32m'observed_state'\u001B[0m,\u001B[1m)\u001B[0m  \n",
+       "  x0                \u001b[1m(\u001b[0m\u001b[1;36m10\u001b[0m,\u001b[1m)\u001b[0m                                                       \u001b[1m(\u001b[0m\u001b[32m'state'\u001b[0m,\u001b[1m)\u001b[0m  \n",
+       "  P0                \u001b[1m(\u001b[0m\u001b[1;36m10\u001b[0m, \u001b[1;36m10\u001b[0m\u001b[1m)\u001b[0m   Positive Semi-definite               \u001b[1m(\u001b[0m\u001b[32m'state'\u001b[0m, \u001b[32m'state_aux'\u001b[0m\u001b[1m)\u001b[0m  \n",
+       "  factor_loadings   \u001b[1m(\u001b[0m\u001b[1;36m4\u001b[0m, \u001b[1;36m1\u001b[0m\u001b[1m)\u001b[0m                                    \u001b[1m(\u001b[0m\u001b[32m'observed_state'\u001b[0m, \u001b[32m'factor'\u001b[0m\u001b[1m)\u001b[0m  \n",
+       "  factor_ar         \u001b[1m(\u001b[0m\u001b[1;36m1\u001b[0m, \u001b[1;36m2\u001b[0m\u001b[1m)\u001b[0m                                            \u001b[1m(\u001b[0m\u001b[32m'factor'\u001b[0m, \u001b[32m'lag_ar'\u001b[0m\u001b[1m)\u001b[0m  \n",
+       "  error_ar          \u001b[1m(\u001b[0m\u001b[1;36m4\u001b[0m, \u001b[1;36m2\u001b[0m\u001b[1m)\u001b[0m                              \u001b[1m(\u001b[0m\u001b[32m'observed_state'\u001b[0m, \u001b[32m'error_lag_ar'\u001b[0m\u001b[1m)\u001b[0m  \n",
+       "  error_sigma       \u001b[1m(\u001b[0m\u001b[1;36m4\u001b[0m,\u001b[1m)\u001b[0m       Positive                                \u001b[1m(\u001b[0m\u001b[32m'observed_state'\u001b[0m,\u001b[1m)\u001b[0m  \n",
        "                                                                                            \n",
-       "\u001B[2;3m  These parameters should be assigned priors inside a PyMC model block before calling the   \u001B[0m\n",
-       "\u001B[2;3m                               build_statespace_graph method.                               \u001B[0m\n"
+       "\u001b[2;3m  These parameters should be assigned priors inside a PyMC model block before calling the   \u001b[0m\n",
+       "\u001b[2;3m                               build_statespace_graph method.                               \u001b[0m\n"
       ]
      },
      "metadata": {},

notebooks/Making a Custom Statespace Model.ipynb

@@ -1127,12 +1127,12 @@
    "outputs": [],
    "source": [
     "from pymc_extras.statespace.core.properties import (\n",
+    "    Coord,\n",
     "    Parameter,\n",
-    "    State,\n",
     "    Shock,\n",
-    "    Coord,\n",
+    "    State,\n",
     ")\n",
-    "from pymc_extras.statespace.utils.constants import ALL_STATE_DIM, ALL_STATE_AUX_DIM, SHOCK_DIM"
+    "from pymc_extras.statespace.utils.constants import ALL_STATE_AUX_DIM, ALL_STATE_DIM, SHOCK_DIM"
    ]
   },
   {

notebooks/Structural Timeseries Modeling.ipynb

@@ -25,8 +25,8 @@
     "from patsy import dmatrix\n",
     "\n",
     "import pymc_extras.statespace as pmss\n",
+    "\n",
     "from pymc_extras.statespace import structural as st\n",
-    "from pymc_extras.statespace.utils.constants import SHORT_NAME_TO_LONG\n",
     "\n",
     "warnings.filterwarnings(\"ignore\", category=UserWarning, message=\"The RandomType SharedVariables\")\n",
     "\n",
@@ -62,10 +62,6 @@
    },
    "outputs": [],
    "source": [
-    "from pymc.pytensorf import inputvars\n",
-    "from pymc_extras.statespace.filters.distributions import LinearGaussianStateSpace\n",
-    "\n",
-    "\n",
     "def compile_component(component):\n",
     "    \"\"\"Helper function to call compile_statespace directly on Components\"\"\"\n",
     "    return pmss.compile_statespace(component.build(verbose=False))"
@@ -3174,6 +3170,7 @@
    "outputs": [],
    "source": [
     "import nutpie as ntp\n",
+    "\n",
     "from jax import nn\n",
     "\n",
     "# Uncomment to see arguments for with_transform_adapt\n",

pymc_extras/__init__.py

@@ -18,11 +18,7 @@
 from pymc_extras import gp, statespace, utils
 from pymc_extras.distributions import *
 from pymc_extras.inference import find_MAP, fit, fit_laplace, fit_pathfinder
-from pymc_extras.model.marginal.marginal_model import (
-    MarginalModel,
-    marginalize,
-    recover_marginals,
-)
+from pymc_extras.model.marginal.marginal_model import MarginalModel, marginalize, recover_marginals
 from pymc_extras.model.model_api import as_model
 
 _log = logging.getLogger("pmx")

pymc_extras/distributions/__init__.py

@@ -18,11 +18,7 @@
 """
 
 from pymc_extras.distributions.continuous import Chi, GenExtreme, Maxwell
-from pymc_extras.distributions.discrete import (
-    BetaNegativeBinomial,
-    GeneralizedPoisson,
-    Skellam,
-)
+from pymc_extras.distributions.discrete import BetaNegativeBinomial, GeneralizedPoisson, Skellam
 from pymc_extras.distributions.histogram_utils import histogram_approximation
 from pymc_extras.distributions.multivariate import R2D2M2CP
 from pymc_extras.distributions.timeseries import DiscreteMarkovChain

pymc_extras/distributions/multivariate/r2d2m2cp.py

@@ -429,7 +429,9 @@ def R2D2M2CP(
             dims=dims,
         )
         mask, psi = _psi(
-            positive_probs=positive_probs, positive_probs_std=positive_probs_std, dims=dims
+            positive_probs=positive_probs,
+            positive_probs_std=positive_probs_std,
+            dims=dims,
         )
 
     beta = _R2D2M2CP_beta(

pymc_extras/distributions/timeseries.py

@@ -12,11 +12,7 @@
     _support_point,
     support_point,
 )
-from pymc.distributions.shape_utils import (
-    _change_dist_size,
-    change_dist_size,
-    get_support_shape_1d,
-)
+from pymc.distributions.shape_utils import _change_dist_size, change_dist_size, get_support_shape_1d
 from pymc.logprob.abstract import _logprob
 from pymc.logprob.basic import logp
 from pymc.pytensorf import constant_fold, intX
@@ -136,7 +132,9 @@ def __new__(cls, *args, steps=None, n_lags=1, **kwargs):
     @classmethod
     def dist(cls, P=None, logit_P=None, steps=None, init_dist=None, n_lags=1, **kwargs):
         steps = get_support_shape_1d(
-            support_shape=steps, shape=kwargs.get("shape", None), support_shape_offset=n_lags
+            support_shape=steps,
+            shape=kwargs.get("shape", None),
+            support_shape_offset=n_lags,
         )
 
         if steps is None:

pymc_extras/inference/__init__.py

@@ -21,7 +21,7 @@
 __all__ = [
     "find_MAP",
     "fit",
+    "fit_dadvi",
     "fit_laplace",
     "fit_pathfinder",
-    "fit_dadvi",
 ]

pymc_extras/inference/advi/training.py

@@ -16,7 +16,10 @@ def __call__(self, draws: int, *params: np.ndarray) -> tuple[np.ndarray, ...]: .
 
 
 def compile_svi_training_fn(
-    model: Model, guide: AutoGuideModel, stick_the_landing: bool = True, **compile_kwargs
+    model: Model,
+    guide: AutoGuideModel,
+    stick_the_landing: bool = True,
+    **compile_kwargs,
 ) -> TrainingFn:
     draws = pt.scalar("draws", dtype=int)
     params = guide.params
@@ -33,7 +36,9 @@ def compile_svi_training_fn(
         compile_kwargs["trust_input"] = True
 
     f_loss_dloss = compile(
-        inputs=[draws, *params], outputs=[negative_elbo, *negative_elbo_grads], **compile_kwargs
+        inputs=[draws, *params],
+        outputs=[negative_elbo, *negative_elbo_grads],
+        **compile_kwargs,
     )
 
     return f_loss_dloss

pymc_extras/inference/dadvi/dadvi.py

@@ -1,16 +1,15 @@
-import arviz as az
 import numpy as np
 import pymc
 import pytensor
 import pytensor.tensor as pt
 
-from arviz import InferenceData
 from better_optimize import basinhopping, minimize
 from better_optimize.constants import minimize_method
 from pymc import DictToArrayBijection, Model, join_nonshared_inputs
 from pymc.blocking import RaveledVars
 from pymc.util import RandomSeed
 from pytensor.tensor.variable import TensorVariable
+from xarray import DataTree
 
 from pymc_extras.inference.laplace_approx.idata import (
     add_data_to_inference_data,
@@ -37,7 +36,7 @@ def fit_dadvi(
     random_seed: RandomSeed = None,
     progressbar: bool = True,
     **optimizer_kwargs,
-) -> az.InferenceData:
+) -> DataTree:
     """
     Does inference using Deterministic ADVI (Automatic Differentiation Variational Inference), DADVI for short.
 
@@ -197,9 +196,9 @@ def fit_dadvi(
         return_unconstrained=include_transformed,
         random_seed=random_seed,
     )
-    idata = InferenceData(posterior=posterior)
+    idata = DataTree.from_dict({"posterior": posterior})
     if include_transformed:
-        idata.add_groups(unconstrained_posterior=unconstrained_posterior)
+        idata["unconstrained_posterior"] = unconstrained_posterior
 
     var_name_to_model_var = {f"{var_name}_mu": var_name for var_name in initial_point_dict.keys()}
     var_name_to_model_var.update(

pymc_extras/inference/fit.py

@@ -11,10 +11,10 @@
 #   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.
-import arviz as az
+import xarray as xr
 
 
-def fit(method: str, **kwargs) -> az.InferenceData:
+def fit(method: str, **kwargs) -> xr.DataTree:
     """
     Fit a model with an inference algorithm.
     See :func:`fit_pathfinder` and :func:`fit_laplace` for more details.

pymc_extras/inference/laplace_approx/find_map.py

@@ -329,7 +329,11 @@ def find_MAP(
     idata = add_fit_to_inference_data(idata=idata, mu=raveled_optimized, H_inv=H_inv, model=model)
 
     idata = add_optimizer_result_to_inference_data(
-        idata=idata, result=optimizer_result, method=method, mu=raveled_optimized, model=model
+        idata=idata,
+        result=optimizer_result,
+        method=method,
+        mu=raveled_optimized,
+        model=model,
     )
 
     idata = add_data_to_inference_data(