avoid creating variables when it is posible; (#54)

change the get_split_variables() to be a generator;
instead of creating Particle and then put the value of the kfactor we refactor constructor to have the default value in 0.75 an the parameter;
grow_tree now returns the new leaf idx;
in grow_tree, change the if else sentence for early return, also, instead of creating new node split node and delete the leaf value there is a new method grow_leaf_node;
in the draw_leaf_value we simplify the expresions;
in _travearse_tree instead of having two returns we create a variable next_node;
adding @njit in the methods get_new_idx_data_points and draw_leaf_value;

Author: fjloyola

pymc_bart/pgbart.py

@@ -187,18 +187,18 @@ def astep(self, _):
 
             _, normalized_weights = self.normalize(particles)
             # Get the new tree and update
-            new_particle, new_tree = self.get_particle_tree(particles, normalized_weights)
-            self.all_particles[tree_id] = new_particle
+            self.all_particles[tree_id], new_tree = self.get_particle_tree(
+                particles, normalized_weights
+            )
             self.sum_trees = self.sum_trees_noi + new_tree._predict()
             self.all_trees[tree_id] = new_tree.trim()
-            used_variates = new_tree.get_split_variables()
 
             if self.tune:
                 self.ssv = SampleSplittingVariable(self.alpha_vec)
-                for index in used_variates:
+                for index in new_tree.get_split_variables():
                     self.alpha_vec[index] += 1
             else:
-                for index in used_variates:
+                for index in new_tree.get_split_variables():
                     variable_inclusion[index] += 1
 
         if not self.tune:
@@ -284,12 +284,9 @@ def init_particles(self, tree_id: int) -> np.ndarray:
         particles = [p0, p1]
 
         for _ in self.indices:
-            pt = ParticleTree(self.a_tree)
-            if self.tune:
-                pt.kfactor = self.uniform.random()
-            else:
-                pt.kfactor = p0.kfactor
-            particles.append(pt)
+            particles.append(
+                ParticleTree(self.a_tree, self.uniform.random() if self.tune else p0.kfactor)
+            )
 
         return np.array(particles)
 
@@ -305,10 +302,10 @@ def update_weight(self, particle, old=False):
         )
         if old:
             particle.log_weight = new_likelihood
-            particle.old_likelihood_logp = new_likelihood
         else:
             particle.log_weight += new_likelihood - particle.old_likelihood_logp
-            particle.old_likelihood_logp = new_likelihood
+
+        particle.old_likelihood_logp = new_likelihood
 
     @staticmethod
     def competence(var, has_grad):
@@ -324,21 +321,19 @@ class ParticleTree:
 
     __slots__ = "tree", "expansion_nodes", "log_weight", "old_likelihood_logp", "kfactor"
 
-    def __init__(self, tree):
+    def __init__(self, tree, kfactor=0.75):
         self.tree = tree.copy()
         self.expansion_nodes = [0]
         self.log_weight = 0
         self.old_likelihood_logp = 0
-        self.kfactor = 0.75
+        self.kfactor = kfactor
 
     def copy(self):
         p = ParticleTree(self.tree)
-        p.expansion_nodes, p.log_weight, p.old_likelihood_logp, p.kfactor = (
-            self.expansion_nodes.copy(),
-            self.log_weight,
-            self.old_likelihood_logp,
-            self.kfactor,
-        )
+        p.expansion_nodes = self.expansion_nodes.copy()
+        p.log_weight = self.log_weight
+        p.old_likelihood_logp = self.old_likelihood_logp
+        p.kfactor = self.kfactor
         return p
 
     def sample_tree(
@@ -360,7 +355,7 @@ def sample_tree(
             prob_leaf = prior_prob_leaf_node[get_depth(index_leaf_node)]
 
             if prob_leaf < np.random.random():
-                index_selected_predictor = grow_tree(
+                idx_new_nodes = grow_tree(
                     self.tree,
                     index_leaf_node,
                     ssv,
@@ -373,9 +368,8 @@ def sample_tree(
                     self.kfactor,
                     shape,
                 )
-                if index_selected_predictor is not None:
-                    new_indexes = self.tree.idx_leaf_nodes[-2:]
-                    self.expansion_nodes.extend(new_indexes)
+                if idx_new_nodes is not None:
+                    self.expansion_nodes.extend(idx_new_nodes)
                     tree_grew = True
 
         return tree_grew
@@ -389,8 +383,7 @@ def sample_leafs(self, sum_trees, m, normal, shape):
                 node_value = draw_leaf_value(
                     sum_trees[:, idx_data_points],
                     m,
-                    normal,
-                    self.kfactor,
+                    normal.random() * self.kfactor,
                     shape,
                 )
                 leaf.value = node_value
@@ -463,55 +456,43 @@ def grow_tree(
     split_value = get_split_value(available_splitting_values, idx_data_points, missing_data)
 
     if split_value is None:
-        index_selected_predictor = None
-    else:
-        new_idx_data_points = get_new_idx_data_points(
-            split_value, idx_data_points, selected_predictor, X
-        )
-        current_node_children = (
-            current_node.get_idx_left_child(),
-            current_node.get_idx_right_child(),
+        return None
+    new_idx_data_points = get_new_idx_data_points(
+        available_splitting_values, split_value, idx_data_points
+    )
+    current_node_children = (
+        current_node.get_idx_left_child(),
+        current_node.get_idx_right_child(),
+    )
+
+    new_nodes = []
+    for idx in range(2):
+        idx_data_point = new_idx_data_points[idx]
+        node_value = draw_leaf_value(
+            sum_trees[:, idx_data_point],
+            m,
+            normal.random() * kfactor,
+            shape,
         )
 
-        new_nodes = []
-        for idx in range(2):
-            idx_data_point = new_idx_data_points[idx]
-            node_value = draw_leaf_value(
-                sum_trees[:, idx_data_point],
-                m,
-                normal,
-                kfactor,
-                shape,
-            )
-
-            new_node = Node.new_leaf_node(
-                index=current_node_children[idx],
-                value=node_value,
-                idx_data_points=idx_data_point,
-            )
-            new_nodes.append(new_node)
-
-        new_split_node = Node.new_split_node(
-            index=index_leaf_node,
-            split_value=split_value,
-            idx_split_variable=selected_predictor,
+        new_node = Node.new_leaf_node(
+            index=current_node_children[idx],
+            value=node_value,
+            idx_data_points=idx_data_point,
         )
+        new_nodes.append(new_node)
 
-        # update tree nodes and indexes
-        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])
-
-    return index_selected_predictor
+    tree.grow_leaf_node(current_node, selected_predictor, split_value, index_leaf_node)
+    tree.set_node(new_nodes[0].index, new_nodes[0])
+    tree.set_node(new_nodes[1].index, new_nodes[1])
 
+    return [new_nodes[0].index, new_nodes[1].index]
 
-def get_new_idx_data_points(split_value, idx_data_points, selected_predictor, X):
-    left_idx = X[idx_data_points, selected_predictor] <= split_value
-    left_node_idx_data_points = idx_data_points[left_idx]
-    right_node_idx_data_points = idx_data_points[~left_idx]
 
-    return left_node_idx_data_points, right_node_idx_data_points
+@njit
+def get_new_idx_data_points(available_splitting_values, split_value, idx_data_points):
+    split_idx = available_splitting_values <= split_value
+    return idx_data_points[split_idx], idx_data_points[~split_idx]
 
 
 def get_split_value(available_splitting_values, idx_data_points, missing_data):
@@ -529,19 +510,18 @@ def get_split_value(available_splitting_values, idx_data_points, missing_data):
     return split_value
 
 
-def draw_leaf_value(y_mu_pred, m, normal, kfactor, shape):
+@njit
+def draw_leaf_value(y_mu_pred, m, norm, shape):
     """Draw Gaussian distributed leaf values."""
     if y_mu_pred.size == 0:
         return np.zeros(shape)
+
+    if y_mu_pred.size == 1:
+        mu_mean = np.full(shape, y_mu_pred.item() / m)
     else:
-        norm = normal.random() * kfactor
-        if y_mu_pred.size == 1:
-            mu_mean = np.full(shape, y_mu_pred.item() / m)
-        else:
-            mu_mean = fast_mean(y_mu_pred) / m
+        mu_mean = fast_mean(y_mu_pred) / m
 
-        draw = norm + mu_mean
-        return draw
+    return norm + mu_mean
 
 
 @njit

pymc_bart/tree.py

@@ -88,9 +88,11 @@ def set_node(self, index, node):
         if node.is_leaf_node():
             self.idx_leaf_nodes.append(index)
 
-    def delete_leaf_node(self, index):
-        self.idx_leaf_nodes.remove(index)
-        del self.tree_structure[index]
+    def grow_leaf_node(self, current_node, selected_predictor, split_value, index_leaf_node):
+        current_node.value = split_value
+        current_node.idx_split_variable = selected_predictor
+        current_node.idx_data_points = None
+        self.idx_leaf_nodes.remove(index_leaf_node)
 
     def trim(self):
         tree = {
@@ -100,9 +102,9 @@ def trim(self):
         return Tree(tree, None, None)
 
     def get_split_variables(self):
-        return [
-            node.idx_split_variable for node in self.tree_structure.values() if node.is_split_node()
-        ]
+        for node in self.tree_structure.values():
+            if node.is_split_node():
+                yield node.idx_split_variable
 
     def _predict(self):
         output = self.output
@@ -153,11 +155,10 @@ def _traverse_tree(self, x, node_index, excluded):
             return np.mean(leaf_values, 0)
 
         if x[current_node.idx_split_variable] <= current_node.value:
-            left_child = current_node.get_idx_left_child()
-            return self._traverse_tree(x, left_child, excluded)
+            next_node = current_node.get_idx_left_child()
         else:
-            right_child = current_node.get_idx_right_child()
-            return self._traverse_tree(x, right_child, excluded)
+            next_node = current_node.get_idx_right_child()
+        return self._traverse_tree(x, next_node, excluded)
 
     def _traverse_leaf_values(self, leaf_values, node_index):
         """
@@ -196,14 +197,11 @@ def new_leaf_node(cls, index: int, value, idx_data_points) -> "Node":
     def new_split_node(cls, index: int, split_value, idx_split_variable) -> "Node":
         return cls(index, value=split_value, idx_split_variable=idx_split_variable)
 
-    def get_idx_parent_node(self) -> int:
-        return (self.index - 1) // 2
-
     def get_idx_left_child(self) -> int:
         return self.index * 2 + 1
 
     def get_idx_right_child(self) -> int:
-        return self.get_idx_left_child() + 1
+        return self.index * 2 + 2
 
     def is_split_node(self) -> bool:
         return self.idx_split_variable >= 0

tests/test_tree.py

@@ -10,7 +10,6 @@ def test_split_node():
     assert split_node.value == 3.0
     assert split_node.idx_split_variable == 2
     assert split_node.idx_data_points is None
-    assert split_node.get_idx_parent_node() == 2
     assert split_node.get_idx_left_child() == 11
     assert split_node.get_idx_right_child() == 12
     assert split_node.is_split_node() is True
@@ -24,7 +23,6 @@ def test_leaf_node():
     assert leaf_node.value == 3.14
     assert leaf_node.idx_split_variable == -1
     assert np.array_equal(leaf_node.idx_data_points, [1, 2, 3])
-    assert leaf_node.get_idx_parent_node() == 2
     assert leaf_node.get_idx_left_child() == 11
     assert leaf_node.get_idx_right_child() == 12
     assert leaf_node.is_split_node() is False