conform to recent changes in pymc

Author: aloctavodia

pymc_bart/bart.py

@@ -37,8 +37,7 @@ class BARTRV(RandomVariable):
     """Base class for BART."""
 
     name: str = "BART"
-    ndim_supp = 1
-    ndims_params: List[int] = [2, 1, 0, 0, 0, 1]
+    signature = "(n,d),(n),(),(),(),(n)->(n)"
     dtype: str = "floatX"
     _print_name: Tuple[str, str] = ("BART", "\\operatorname{BART}")
     all_trees = List[List[List[Tree]]]
@@ -50,6 +49,8 @@ def _supp_shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=Non
     def rng_fn(  # pylint: disable=W0237
         cls, rng=None, X=None, Y=None, m=None, alpha=None, beta=None, split_prior=None, size=None
     ):
+        if not size:
+            size = None
         if not cls.all_trees:
             if size is not None:
                 return np.full((size[0], cls.Y.shape[0]), cls.Y.mean())

pymc_bart/pgbart.py

@@ -114,7 +114,10 @@ class PGBART(ArrayStepShared):
     name = "pgbart"
     default_blocked = False
     generates_stats = True
-    stats_dtypes = [{"variable_inclusion": object, "tune": bool}]
+    stats_dtypes_shapes: dict[str, tuple[type, list]] = {
+        "variable_inclusion": (object, []),
+        "tune": (bool, []),
+    }
 
     def __init__(  # noqa: PLR0915
         self,
@@ -227,7 +230,7 @@ def __init__(  # noqa: PLR0915
     def astep(self, _):
         variable_inclusion = np.zeros(self.num_variates, dtype="int")
 
-        upper = min(self.lower + self.batch[~self.tune], self.m)
+        upper = min(self.lower + self.batch[not self.tune], self.m)
         tree_ids = range(self.lower, upper)
         self.lower = upper if upper < self.m else 0
 

requirements.txt

@@ -1,4 +1,4 @@
-pymc<=5.16.2
+pymc>=5.16.2, <=5.18
 arviz>=0.18.0
 numba
 matplotlib

tests/test_bart.py

@@ -3,7 +3,7 @@
 import pytest
 from numpy.testing import assert_almost_equal, assert_array_equal
 from pymc.initial_point import make_initial_point_fn
-from pymc.logprob.basic import joint_logp
+from pymc.logprob.basic import transformed_conditional_logp
 
 import pymc_bart as pmb
 
@@ -12,7 +12,7 @@ def assert_moment_is_expected(model, expected, check_finite_logp=True):
     fn = make_initial_point_fn(
         model=model,
         return_transformed=False,
-        default_strategy="moment",
+        default_strategy="support_point",
     )
     moment = fn(0)["x"]
     expected = np.asarray(expected)
@@ -27,7 +27,7 @@ def assert_moment_is_expected(model, expected, check_finite_logp=True):
 
     if check_finite_logp:
         logp_moment = (
-            joint_logp(
+            transformed_conditional_logp(
                 (model["x"],),
                 rvs_to_values={model["x"]: pm.math.constant(moment)},
                 rvs_to_transforms={},
@@ -91,7 +91,7 @@ def test_shared_variable(response):
     Y = np.random.normal(0, 1, size=50)
 
     with pm.Model() as model:
-        data_X = pm.MutableData("data_X", X)
+        data_X = pm.Data("data_X", X)
         mu = pmb.BART("mu", data_X, Y, m=2, response=response)
         sigma = pm.HalfNormal("sigma", 1)
         y = pm.Normal("y", mu, sigma, observed=Y, shape=mu.shape)