Added tests for split rules

Author: velochy

pymc_bart/bart.py

@@ -143,7 +143,7 @@ def __new__(
                 alpha=alpha,
                 beta=beta,
                 split_prior=split_prior,
-                split_rules=split_rules
+                split_rules=split_rules,
             ),
         )()
 

pymc_bart/split_rules.py

@@ -87,9 +87,7 @@ def get_split_value(available_splitting_values):
         if available_splitting_values.size > 1 and not np.all(
             available_splitting_values == available_splitting_values[0]
         ):
-            unique_values = np.unique(available_splitting_values)[
-                :-1
-            ]  # Remove last one so it always goes to left
+            unique_values = np.unique(available_splitting_values)
             while True:
                 sample = np.random.randint(0, 2, size=len(unique_values)).astype(bool)
                 if np.any(sample):

tests/test_split_rules.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+from pymc_bart.split_rules import ContinuousSplitRule, OneHotSplitRule, SubsetSplitRule
+import pytest
+
+
+@pytest.mark.parametrize(
+    argnames="Rule",
+    argvalues=[ContinuousSplitRule, OneHotSplitRule, SubsetSplitRule],
+    ids=["continuous", "one_hot", "subset"],
+)
+def test_split_rule(Rule):
+
+    # Should return None if only one available value to pick from
+    assert Rule.get_split_value(np.zeros(1)) is None
+
+    # get_split should return a value divide can use
+    available_values = np.arange(10).astype(float)
+    sv = Rule.get_split_value(available_values)
+    left = Rule.divide(available_values, sv)
+
+    # divide should return a boolean numpy array
+    # This de facto ensures it is a binary split
+    assert len(left) == len(available_values)
+    assert left.dtype == "bool"
+
+    # divide should be deterministic
+    left_repeated = Rule.divide(available_values, sv)
+    assert (left == left_repeated).all()
+
+    # Most elements should have a chance to go either direction
+    # NB! This is not 100% necessary, but is a good proxy
+    probs = np.array(
+        [
+            Rule.divide(available_values, Rule.get_split_value(available_values))
+            for _ in range(10000)
+        ]
+    ).mean(axis=0)
+
+    assert (probs > 0.01).sum() >= len(available_values) - 1
+    assert (probs < 0.99).sum() >= len(available_values) - 1