Improvements over Tree implementation (#40)

* Improvements over Tree implementation:
- Remove variable depth on the Node class to be a method;
- Remove variable leaf_node_value on the Tree class as is not need and was only writing and never read of it;

* Update pymc_bart/tree.py

Co-authored-by: Osvaldo A Martin <[email protected]>

Co-authored-by: Osvaldo A Martin <[email protected]>

Author: fjloyola

pymc_bart/pgbart.py

@@ -27,9 +27,8 @@
 from pymc.step_methods.arraystep import ArrayStepShared, Competence
 from pymc.pytensorf import inputvars, join_nonshared_inputs, make_shared_replacements
 
-
 from pymc_bart.bart import BARTRV
-from pymc_bart.tree import Tree, Node
+from pymc_bart.tree import Tree, Node, get_depth
 
 _log = logging.getLogger("pymc")
 
@@ -138,9 +137,7 @@ def __init__(
         self.likelihood_logp = logp(initial_values, [model.datalogp], vars, shared)
         self.all_particles = []
         for _ in range(self.m):
-            self.a_tree.leaf_node_value = init_mean / self.m
-            p = ParticleTree(self.a_tree)
-            self.all_particles.append(p)
+            self.all_particles.append(ParticleTree(self.a_tree))
         self.all_trees = np.array([p.tree for p in self.all_particles])
         super().__init__(vars, shared)
 
@@ -352,7 +349,7 @@ def sample_tree(
         if self.expansion_nodes:
             index_leaf_node = self.expansion_nodes.pop(0)
             # Probability that this node will remain a leaf node
-            prob_leaf = prior_prob_leaf_node[self.tree[index_leaf_node].depth]
+            prob_leaf = prior_prob_leaf_node[get_depth(index_leaf_node)]
 
             if prob_leaf < np.random.random():
                 index_selected_predictor = grow_tree(

pymc_bart/tree.py

@@ -15,6 +15,7 @@
 import math
 
 from copy import deepcopy
+from functools import lru_cache
 
 from pytensor import config
 import numpy as np
@@ -41,7 +42,6 @@ class Tree:
 
     Parameters
     ----------
-    leaf_node_value : int or float
     idx_data_points : array of integers
     num_observations : integer
     shape : int
@@ -51,7 +51,6 @@ class Tree:
         "tree_structure",
         "idx_leaf_nodes",
         "output",
-        "leaf_node_value",
     )
 
     def __init__(self, leaf_node_value, idx_data_points, num_observations, shape):
@@ -88,7 +87,6 @@ def trim(self):
         del a_tree.idx_leaf_nodes
         for k in a_tree.tree_structure.keys():
             current_node = a_tree[k]
-            del current_node.depth
             if current_node.is_leaf_node():
                 del current_node.idx_data_points
         return a_tree
@@ -174,11 +172,10 @@ def _traverse_leaf_values(self, leaf_values, node_index):
 
 
 class Node:
-    __slots__ = "index", "depth", "value", "idx_split_variable", "idx_data_points"
+    __slots__ = "index", "value", "idx_split_variable", "idx_data_points"
 
     def __init__(self, index: int, value=-1, idx_data_points=None, idx_split_variable=-1):
         self.index = index
-        self.depth = int(math.floor(math.log(index + 1, 2)))
         self.value = value
         self.idx_data_points = idx_data_points
         self.idx_split_variable = idx_split_variable
@@ -205,3 +202,8 @@ def is_split_node(self) -> bool:
 
     def is_leaf_node(self) -> bool:
         return not self.is_split_node()
+
+
+@lru_cache
+def get_depth(index: int) -> int:
+    return math.floor(math.log2(index + 1))

tests/test_tree.py

@@ -1,12 +1,12 @@
 import numpy as np
 
-from pymc_bart.tree import Node
+from pymc_bart.tree import Node, get_depth
 
 
 def test_split_node():
     split_node = Node.new_split_node(index=5, idx_split_variable=2, split_value=3.0)
     assert split_node.index == 5
-    assert split_node.depth == 2
+    assert get_depth(split_node.index) == 2
     assert split_node.value == 3.0
     assert split_node.idx_split_variable == 2
     assert split_node.idx_data_points is None
@@ -20,7 +20,7 @@ def test_split_node():
 def test_leaf_node():
     leaf_node = Node.new_leaf_node(index=5, value=3.14, idx_data_points=[1, 2, 3])
     assert leaf_node.index == 5
-    assert leaf_node.depth == 2
+    assert get_depth(leaf_node.index) == 2
     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])