Addressed most issues flagged in PR comments

Author: velochy

pymc_bart/__init__.py

@@ -15,7 +15,7 @@
 
 from pymc_bart.bart import BART
 from pymc_bart.pgbart import PGBART
-from pymc_bart.split_rules import *
+from pymc_bart.split_rules import ContinuousSplitRule, OneHotSplitRule, SubsetSplitRule
 from pymc_bart.utils import (
     plot_convergence,
     plot_pdp,

pymc_bart/bart.py

@@ -27,6 +27,7 @@
 
 from .tree import Tree
 from .utils import TensorLike, _sample_posterior
+from .split_rules import SplitRule
 
 __all__ = ["BART"]
 
@@ -109,7 +110,7 @@ def __new__(
         beta: float = 2.0,
         response: str = "constant",
         split_prior: Optional[List[float]] = None,
-        split_rules: Optional[List] = None,
+        split_rules: Optional[SplitRule] = None,
         **kwargs,
     ):
         manager = Manager()
@@ -138,7 +139,7 @@ def __new__(
                 alpha=alpha,
                 beta=beta,
                 split_prior=split_prior,
-                split_rules=split_rules
+                split_rules=split_rules,
             ),
         )()
 

pymc_bart/pgbart.py

@@ -28,6 +28,7 @@
 from pymc_bart.tree import Node, Tree, get_idx_left_child, get_idx_right_child, get_depth
 from pymc_bart.split_rules import ContinuousSplitRule
 
+
 class ParticleTree:
     """Particle tree."""
 
@@ -169,7 +170,7 @@ def __init__(
             idx_data_points=np.arange(self.num_observations, dtype="int32"),
             num_observations=self.num_observations,
             shape=self.shape,
-            split_rules = self.split_rules
+            split_rules=self.split_rules,
         )
 
         self.normal = NormalSampler(1, self.shape)
@@ -459,7 +460,7 @@ def grow_tree(
 
     to_left = split_rule.divide(available_splitting_values, split_value)
     new_idx_data_points = idx_data_points[to_left], idx_data_points[~to_left]
-    
+
     current_node_children = (
         get_idx_left_child(index_leaf_node),
         get_idx_right_child(index_leaf_node),
@@ -499,7 +500,6 @@ def filter_missing_values(available_splitting_values, idx_data_points, missing_d
     return idx_data_points, available_splitting_values
 
 
-
 def draw_leaf_value(
     y_mu_pred: npt.NDArray[np.float_],
     x_mu: npt.NDArray[np.float_],

pymc_bart/split_rules.py

@@ -14,8 +14,26 @@
 
 from numba import njit
 import numpy as np
+from abc import abstractmethod
 
-class ContinuousSplitRule:
+
+class SplitRule:
+    """
+    Abstract template class for a split rule
+    """
+
+    @staticmethod
+    @abstractmethod
+    def get_split_value(available_splitting_values):
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def divide(available_splitting_values, split_value):
+        pass
+
+
+class ContinuousSplitRule(SplitRule):
     """
     Standard continuous split rule: pick a pivot value and split depending on if variable is smaller or greater than the value picked.
     """
@@ -24,7 +42,9 @@ class ContinuousSplitRule:
     def get_split_value(available_splitting_values):
         split_value = None
         if available_splitting_values.size > 1:
-            idx_selected_splitting_values = int(np.random.random() * len(available_splitting_values))
+            idx_selected_splitting_values = int(
+                np.random.random() * len(available_splitting_values)
+            )
             split_value = available_splitting_values[idx_selected_splitting_values]
         return split_value
 
@@ -33,14 +53,19 @@ def get_split_value(available_splitting_values):
     def divide(available_splitting_values, split_value):
         return available_splitting_values <= split_value
 
-class OneHotSplitRule:
+
+class OneHotSplitRule(SplitRule):
     """Choose a single categorical value and branch on if the variable is that value or not"""
 
     @staticmethod
     def get_split_value(available_splitting_values):
         split_value = None
-        if available_splitting_values.size > 1 and not np.all(available_splitting_values==available_splitting_values[0]):
-            idx_selected_splitting_values = int(np.random.random() * len(available_splitting_values))
+        if available_splitting_values.size > 1 and not np.all(
+            available_splitting_values == available_splitting_values[0]
+        ):
+            idx_selected_splitting_values = int(
+                np.random.random() * len(available_splitting_values)
+            )
             split_value = available_splitting_values[idx_selected_splitting_values]
         return split_value
 
@@ -49,7 +74,8 @@ def get_split_value(available_splitting_values):
     def divide(available_splitting_values, split_value):
         return available_splitting_values == split_value
 
-class SubsetSplitRule:
+
+class SubsetSplitRule(SplitRule):
     """
     Choose a random subset of the categorical values and branch on if the value is within the chosen set.
     This is the approach taken in flexBART paper.
@@ -58,11 +84,16 @@ class SubsetSplitRule:
     @staticmethod
     def get_split_value(available_splitting_values):
         split_value = None
-        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
+        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
             while True:
                 sample = np.random.randint(0, 2, size=len(unique_values)).astype(bool)
-                if np.any(sample): break
+                if np.any(sample):
+                    break
             split_value = unique_values[sample]
         return split_value
 

pymc_bart/tree.py

@@ -114,12 +114,7 @@ class Tree:
     idx_leaf_nodes :  List with the index of the leaf nodes of the tree.
     """
 
-    __slots__ = (
-        "tree_structure",
-        "output",
-        "idx_leaf_nodes",
-        "split_rules"
-    )
+    __slots__ = ("tree_structure", "output", "idx_leaf_nodes", "split_rules")
 
     def __init__(
         self,
@@ -140,7 +135,7 @@ def new_tree(
         idx_data_points: Optional[npt.NDArray[np.int_]],
         num_observations: int,
         shape: int,
-        split_rules: List[object]
+        split_rules: List[object],
     ) -> "Tree":
         return cls(
             tree_structure={
@@ -152,7 +147,7 @@ def new_tree(
             },
             idx_leaf_nodes=[0],
             output=np.zeros((num_observations, shape)).astype(config.floatX).squeeze(),
-            split_rules=split_rules
+            split_rules=split_rules,
         )
 
     def __getitem__(self, index) -> Node:
@@ -173,7 +168,12 @@ def copy(self) -> "Tree":
             for k, v in self.tree_structure.items()
         }
         idx_leaf_nodes = self.idx_leaf_nodes.copy() if self.idx_leaf_nodes is not None else None
-        return Tree(tree_structure=tree, idx_leaf_nodes=idx_leaf_nodes, output=self.output, split_rules=self.split_rules)
+        return Tree(
+            tree_structure=tree,
+            idx_leaf_nodes=idx_leaf_nodes,
+            output=self.output,
+            split_rules=self.split_rules,
+        )
 
     def get_node(self, index: int) -> Node:
         return self.tree_structure[index]
@@ -207,7 +207,12 @@ def trim(self) -> "Tree":
             )
             for k, v in self.tree_structure.items()
         }
-        return Tree(tree_structure=tree, idx_leaf_nodes=None, output=np.array([-1]), split_rules=self.split_rules)
+        return Tree(
+            tree_structure=tree,
+            idx_leaf_nodes=None,
+            output=np.array([-1]),
+            split_rules=self.split_rules,
+        )
 
     def get_split_variables(self) -> Generator[int, None, None]:
         for node in self.tree_structure.values():
@@ -275,21 +280,25 @@ def _traverse_tree(
             Leaf node value or mean of leaf node values
         """
 
-        x_shape = (1,) if len(X.shape)==1 else X.shape[:-1]
+        x_shape = (1,) if len(X.shape) == 1 else X.shape[:-1]
 
-        stack = [(0, np.ones(x_shape)) ]  # (node_index, weight) initial state
-        p_d = np.zeros( shape + x_shape ) if isinstance(shape, tuple) else np.zeros( (shape,) + x_shape )
+        stack = [(0, np.ones(x_shape))]  # (node_index, weight) initial state
+        p_d = (
+            np.zeros(shape + x_shape) if isinstance(shape, tuple) else np.zeros((shape,) + x_shape)
+        )
         while stack:
             node_index, weights = stack.pop()
             node = self.get_node(node_index)
             if node.is_leaf_node():
                 params = node.linear_params
-                nd_dims = (...,)+(None,)*len(x_shape)
+                nd_dims = (...,) + (None,) * len(x_shape)
                 if params is None:
                     p_d += weights * node.value[nd_dims]
                 else:
                     # this produce nonsensical results
-                    p_d += weights * (params[0][nd_dims] + params[1][nd_dims] * X[...,node.idx_split_variable])
+                    p_d += weights * (
+                        params[0][nd_dims] + params[1][nd_dims] * X[..., node.idx_split_variable]
+                    )
                     # this produce reasonable result
                     # p_d += weight * node.value.mean()
             else:
@@ -301,15 +310,19 @@ def _traverse_tree(
                     stack.append((left_node_index, weights * prop_nvalue_left))
                     stack.append((right_node_index, weights * (1 - prop_nvalue_left)))
                 else:
-                    to_left = self.split_rules[node.idx_split_variable].divide(X[...,node.idx_split_variable],node.value).astype('float')
+                    to_left = (
+                        self.split_rules[node.idx_split_variable]
+                        .divide(X[..., node.idx_split_variable], node.value)
+                        .astype("float")
+                    )
                     stack.append((left_node_index, weights * to_left))
                     stack.append((right_node_index, weights * (1 - to_left)))
 
-        if len(X.shape)==1: p_d = p_d[...,0]
+        if len(X.shape) == 1:
+            p_d = p_d[..., 0]
 
         return p_d
 
-
     def _traverse_leaf_values(
         self, leaf_values: List[npt.NDArray[np.float_]], leaf_n_values: List[int], node_index: int
     ) -> None:
@@ -328,5 +341,3 @@ def _traverse_leaf_values(
         else:
             self._traverse_leaf_values(leaf_values, leaf_n_values, get_idx_left_child(node_index))
             self._traverse_leaf_values(leaf_values, leaf_n_values, get_idx_right_child(node_index))
-
-