bart refactor (#43)

* bart refactor

* fix test

Author: aloctavodia

pymc_experimental/bart/pgbart.py

@@ -227,11 +227,12 @@ def init_particles(self, tree_id: int) -> np.ndarray:
             self.m,
             self.normal,
         )
+
         # The old tree and the one with new leafs do not grow so we update the weights only once
         self.update_weight(p0, old=True)
         self.update_weight(p1, old=True)
-
         particles = [p0, p1]
+
         for _ in self.indices:
             pt = ParticleTree(self.a_tree)
             if self.tune:
@@ -396,15 +397,9 @@ def grow_tree(
     index_selected_predictor = ssv.rvs()
     selected_predictor = available_predictors[index_selected_predictor]
     available_splitting_values = X[idx_data_points, selected_predictor]
-    if missing_data:
-        idx_data_points = idx_data_points[~np.isnan(available_splitting_values)]
-        available_splitting_values = available_splitting_values[
-            ~np.isnan(available_splitting_values)
-        ]
+    split_value = get_split_value(available_splitting_values, idx_data_points, missing_data)
 
-    if available_splitting_values.size > 0:
-        idx_selected_splitting_values = discrete_uniform_sampler(len(available_splitting_values))
-        split_value = available_splitting_values[idx_selected_splitting_values]
+    if split_value is not None:
 
         new_idx_data_points = get_new_idx_data_points(
             split_value, idx_data_points, selected_predictor, X
@@ -439,7 +434,7 @@ def grow_tree(
         )
 
         # update tree nodes and indexes
-        tree.delete_node(index_leaf_node)
+        tree.delete_leaf_node(index_leaf_node)
         tree.set_node(index_leaf_node, new_split_node)
         tree.set_node(new_nodes[0].index, new_nodes[0])
         tree.set_node(new_nodes[1].index, new_nodes[1])
@@ -456,6 +451,21 @@ def get_new_idx_data_points(split_value, idx_data_points, selected_predictor, X)
     return left_node_idx_data_points, right_node_idx_data_points
 
 
+def get_split_value(available_splitting_values, idx_data_points, missing_data):
+
+    if missing_data:
+        idx_data_points = idx_data_points[~np.isnan(available_splitting_values)]
+        available_splitting_values = available_splitting_values[
+            ~np.isnan(available_splitting_values)
+        ]
+
+    if available_splitting_values.size > 0:
+        idx_selected_splitting_values = discrete_uniform_sampler(len(available_splitting_values))
+        split_value = available_splitting_values[idx_selected_splitting_values]
+
+        return split_value
+
+
 def draw_leaf_value(Y_mu_pred, mean, m, normal, kf):
     """Draw Gaussian distributed leaf values."""
     if Y_mu_pred.size == 0:

pymc_experimental/bart/tree.py

@@ -68,10 +68,8 @@ def set_node(self, index, node):
         if isinstance(node, LeafNode):
             self.idx_leaf_nodes.append(index)
 
-    def delete_node(self, index):
-        current_node = self.get_node(index)
-        if isinstance(current_node, LeafNode):
-            self.idx_leaf_nodes.remove(index)
+    def delete_leaf_node(self, index):
+        self.idx_leaf_nodes.remove(index)
         del self.tree_structure[index]
 
     def trim(self):