Rename variables to comply with black and improve the test to include split_rules

Author: velochy

pymc_bart/utils.py

@@ -61,15 +61,15 @@ def _sample_posterior(
 
     idx = rng.integers(0, len(stacked_trees), size=flatten_size)
 
-    st_shape = len(stacked_trees[0])
-    tree_shape = shape // st_shape
+    trees_shape = len(stacked_trees[0])
+    leaves_shape = shape // trees_shape
 
-    pred = np.zeros((flatten_size, st_shape, tree_shape, X.shape[0]))
+    pred = np.zeros((flatten_size, trees_shape, leaves_shape, X.shape[0]))
 
     for ind, p in enumerate(pred):
         for odim, odim_trees in enumerate(stacked_trees[idx[ind]]):
             for tree in odim_trees:
-                p[odim] += tree.predict(x=X, excluded=excluded, shape=tree_shape)
+                p[odim] += tree.predict(x=X, excluded=excluded, shape=leaves_shape)
 
     # pred.reshape((*size_iter, shape, -1))
     return pred.transpose((0, 3, 1, 2)).reshape((*size_iter, -1, shape))

tests/test_bart.py

@@ -220,17 +220,24 @@ def test_bart_moment(size, expected):
 
 
 @pytest.mark.parametrize(
-    argnames="separate_trees",
-    argvalues=[False, True],
-    ids=["shared-trees", "separate-trees"],
+    argnames="separate_trees,split_rule",
+    argvalues=[
+        (False,pmb.ContinuousSplitRule),
+        (False,pmb.OneHotSplitRule),
+        (False,pmb.SubsetSplitRule),
+        (True,pmb.ContinuousSplitRule)
+    ],
+    ids=["continuous", "one-hot", "subset", "separate-trees"],
 )
-def test_categorical(separate_trees):
+def test_categorical_model(separate_trees,split_rule):
 
     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)
+        lo = pmb.BART("logodds", X, Y, m=2, shape=(3, 9),
+            split_rules=[split_rule]*5,
+            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)