Made existing tests work and added a new test for separate_trees

Author: velochy

pymc_bart/pgbart.py

@@ -175,7 +175,7 @@ def __init__(
             leaf_node_value=init_mean / self.m,
             idx_data_points=np.arange(self.num_observations, dtype="int32"),
             num_observations=self.num_observations,
-            shape=self.leaves_shape
+            shape=self.leaves_shape,
         )
 
         self.normal = NormalSampler(1, self.leaves_shape)
@@ -357,7 +357,8 @@ def update_weight(self, particle: ParticleTree, odim: int) -> None:
         """
 
         delta = (
-            np.identity(self.trees_shape)[odim][:, None, None] * particle.tree._predict()[None, :, :]
+            np.identity(self.trees_shape)[odim][:, None, None]
+            * particle.tree._predict()[None, :, :]
         )
 
         new_likelihood = self.likelihood_logp((self.sum_trees_noi + delta).flatten())

pymc_bart/utils.py

@@ -49,7 +49,7 @@ def _sample_posterior(
         X = X.eval()
 
     if size is None:
-        size_iter: Union[List, Tuple] = ()
+        size_iter: Union[List, Tuple] = (1,)
     elif isinstance(size, int):
         size_iter = [size]
     else:
@@ -69,7 +69,9 @@ def _sample_posterior(
     for ind, p in enumerate(pred):
         for oi, odim_trees in enumerate(stacked_trees[idx[ind]]):
             for tree in odim_trees:
-                p[oi] += np.vstack([tree.predict(x=x, excluded=excluded, shape=tree_shape) for x in X]).T
+                p[oi] += np.vstack(
+                    [tree.predict(x=x, excluded=excluded, shape=tree_shape) for x in X]
+                ).T
 
     # pred.reshape((*size_iter, shape, -1))
     return pred.transpose((0, 3, 1, 2)).reshape((*size_iter, -1, shape))

tests/test_bart.py

@@ -217,3 +217,23 @@ def test_bart_moment(size, expected):
     with pm.Model() as model:
         pmb.BART("x", X=X, Y=Y, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    argnames="separate_trees",
+    argvalues=[False, True],
+    ids=["shared-trees", "separate-trees"],
+)
+def test_categorical(separate_trees):
+
+    Y = np.array([0, 0, 0, 1, 1, 1, 2, 2, 2])
+    X = np.concatenate([Y[:, None], np.random.randint(0, 6, size=(9, 4))], axis=1)
+
+    with pm.Model() as model:
+        lo = pmb.BART("logodds", X, Y, m=2, shape=(3, 9), separate_trees=separate_trees)
+        y = pm.Categorical("y", p=pm.math.softmax(lo.T, axis=-1), observed=Y)
+        idata = pm.sample(random_seed=3415, tune=300, draws=300)
+        idata = pm.sample_posterior_predictive(idata, predictions=True, extend_inferencedata=True)
+
+    # Fit should be good enough so right category is selected over 50% of time
+    assert (idata.predictions.y.median(["chain", "draw"]) == Y).all()