Use PyMC v6, Pytensor v3 and ArviZ 1.1 (#269)

* use v6, v3 and 1.1

* update python version

Author: aloctavodia

.github/workflows/test.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.11", "3.12", "3.13"]
+        python-version: ["3.12", "3.13", "3.14"]
         linker: [cvm, numba]
 
     name: Tests py${{ matrix.python-version }} ${{ matrix.linker }}

pymc_bart/utils.py

@@ -90,8 +90,8 @@ def plot_convergence(
 
     Parameters
     ----------
-    idata : InferenceData
-        InferenceData object containing the posterior samples.
+    idata : DataTree
+        DataTree object containing the posterior samples.
     var_name : Optional[str]
         Name of the BART variable to plot. Defaults to None.
     kind : str
@@ -683,8 +683,8 @@ def get_variable_inclusion(idata, X, model=None, bart_var_name=None, labels=None
 
     Parameters
     ----------
-    idata : InferenceData
-        InferenceData with a variable "variable_inclusion" in ``sample_stats`` group
+    idata : DataTree
+        DataTree with a variable "variable_inclusion" in ``sample_stats`` group
     X : npt.NDArray
         The covariate matrix.
     model : Optional[pm.Model]
@@ -745,8 +745,8 @@ def plot_variable_inclusion(idata, X, labels=None, figsize=None, plot_kwargs=Non
 
     Parameters
     ----------
-    idata : InferenceData
-        InferenceData containing a collection of BART_trees in sample_stats group
+    idata : DataTree
+        DataTree containing a collection of BART_trees in sample_stats group
     X : npt.NDArray
         The covariate matrix.
     labels : Optional[list[str]]
@@ -813,8 +813,8 @@ def compute_variable_importance(  # noqa: PLR0915 PLR0912
 
     Parameters
     ----------
-    idata : InferenceData
-        InferenceData containing a "variable_inclusion" variable in the sample_stats group.
+    idata : DataTree
+        DataTree containing a "variable_inclusion" variable in the sample_stats group.
     bartrv : BART Random Variable
         BART variable once the model that include it has been fitted.
     X : npt.NDArray

requirements.txt

@@ -1,5 +1,6 @@
-pymc>=5.24.0
-arviz-stats[xarray]>=0.6.0
+pytensor @ git+https://github.com/pymc-devs/pytensor.git@v3
+pymc @ git+https://github.com/pymc-devs/pymc.git@v6
+arviz-stats[xarray]>=1.1.0
 numba
 matplotlib
 numpy>=2.0

setup.py

@@ -29,9 +29,9 @@
     "Development Status :: 5 - Production/Stable",
     "Programming Language :: Python",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
     "License :: OSI Approved :: Apache Software License",
     "Intended Audience :: Science/Research",
     "Topic :: Scientific/Engineering",
@@ -76,6 +76,6 @@ def get_version():
         packages=find_packages(),
         include_package_data=True,
         classifiers=classifiers,
-        python_requires=">=3.11",
+        python_requires=">=3.12",
         install_requires=install_reqs,
     )