We don't need to store index at each node. (#80)

* We don't need to store index at each node.

* Fast integer log2

---------

Co-authored-by: Si Yu How <siyuhow@siyu.colet10>

Author: howsiyu

pymc_bart/pgbart.py

@@ -25,7 +25,7 @@
 from pytensor.tensor.var import Variable
 
 from pymc_bart.bart import BARTRV
-from pymc_bart.tree import Node, Tree, get_depth
+from pymc_bart.tree import Node, Tree, get_idx_left_child, get_idx_right_child, get_depth
 
 
 class ParticleTree:
@@ -411,11 +411,10 @@ def grow_tree(
         available_splitting_values, split_value, idx_data_points
     )
     current_node_children = (
-        current_node.get_idx_left_child(),
-        current_node.get_idx_right_child(),
+        get_idx_left_child(index_leaf_node),
+        get_idx_right_child(index_leaf_node),
     )
 
-    new_nodes = np.empty(2, dtype=object)
     for idx in range(2):
         idx_data_point = new_idx_data_points[idx]
         node_value = draw_leaf_value(
@@ -426,17 +425,13 @@ def grow_tree(
         )
 
         new_node = Node.new_leaf_node(
-            index=current_node_children[idx],
             value=node_value,
             idx_data_points=idx_data_point,
         )
-        new_nodes[idx] = new_node
+        tree.set_node(current_node_children[idx], new_node)
 
     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]
+    return current_node_children
 
 
 @njit

pymc_bart/tree.py

@@ -12,7 +12,6 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-import math
 from functools import lru_cache
 from typing import Dict, Generator, List, Optional
 
@@ -26,41 +25,30 @@ class Node:
 
     Attributes
     ----------
-    index : int
     value : float
     idx_data_points : Optional[npt.NDArray[np.int_]]
     idx_split_variable : Optional[npt.NDArray[np.int_]]
     """
 
-    __slots__ = "index", "value", "idx_split_variable", "idx_data_points"
+    __slots__ = "value", "idx_split_variable", "idx_data_points"
 
     def __init__(
         self,
-        index: int,
         value: float = -1.0,
         idx_data_points: Optional[npt.NDArray[np.int_]] = None,
         idx_split_variable: int = -1,
     ) -> None:
-        self.index = index
         self.value = value
         self.idx_data_points = idx_data_points
         self.idx_split_variable = idx_split_variable
 
     @classmethod
-    def new_leaf_node(
-        cls, index: int, value: float, idx_data_points: Optional[npt.NDArray[np.int_]]
-    ) -> "Node":
-        return cls(index, value=value, idx_data_points=idx_data_points)
+    def new_leaf_node(cls, value: float, idx_data_points: Optional[npt.NDArray[np.int_]]) -> "Node":
+        return cls(value=value, idx_data_points=idx_data_points)
 
     @classmethod
-    def new_split_node(cls, index: int, split_value: float, idx_split_variable: int) -> "Node":
-        return cls(index=index, value=split_value, idx_split_variable=idx_split_variable)
-
-    def get_idx_left_child(self) -> int:
-        return self.index * 2 + 1
-
-    def get_idx_right_child(self) -> int:
-        return self.index * 2 + 2
+    def new_split_node(cls, split_value: float, idx_split_variable: int) -> "Node":
+        return cls(value=split_value, idx_split_variable=idx_split_variable)
 
     def is_split_node(self) -> bool:
         return self.idx_split_variable >= 0
@@ -69,9 +57,17 @@ def is_leaf_node(self) -> bool:
         return not self.is_split_node()
 
 
+def get_idx_left_child(index) -> int:
+    return index * 2 + 1
+
+
+def get_idx_right_child(index) -> int:
+    return index * 2 + 2
+
+
 @lru_cache
 def get_depth(index: int) -> int:
-    return math.floor(math.log2(index + 1))
+    return (index + 1).bit_length() - 1
 
 
 class Tree:
@@ -126,9 +122,7 @@ def new_tree(
     ) -> "Tree":
         return cls(
             tree_structure={
-                0: Node.new_leaf_node(
-                    index=0, value=leaf_node_value, idx_data_points=idx_data_points
-                )
+                0: Node.new_leaf_node(value=leaf_node_value, idx_data_points=idx_data_points)
             },
             idx_leaf_nodes=[0],
             output=np.zeros((num_observations, shape)).astype(config.floatX).squeeze(),
@@ -142,7 +136,7 @@ def __setitem__(self, index, node) -> None:
 
     def copy(self) -> "Tree":
         tree: Dict[int, Node] = {
-            k: Node(v.index, v.value, v.idx_data_points, v.idx_split_variable)
+            k: Node(v.value, v.idx_data_points, v.idx_split_variable)
             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
@@ -167,8 +161,7 @@ def grow_leaf_node(
 
     def trim(self) -> "Tree":
         tree: Dict[int, Node] = {
-            k: Node(v.index, v.value, None, v.idx_split_variable)
-            for k, v in self.tree_structure.items()
+            k: Node(v.value, None, v.idx_split_variable) for k, v in self.tree_structure.items()
         }
         return Tree(tree_structure=tree, idx_leaf_nodes=None, output=np.array([-1]))
 
@@ -241,9 +234,9 @@ def _traverse_tree(
             return np.mean(leaf_values, axis=0)
 
         if x[current_node.idx_split_variable] <= current_node.value:
-            next_node = current_node.get_idx_left_child()
+            next_node = get_idx_left_child(node_index)
         else:
-            next_node = current_node.get_idx_right_child()
+            next_node = get_idx_right_child(node_index)
         return self._traverse_tree(x=x, node_index=next_node, excluded=excluded)
 
     def _traverse_leaf_values(self, leaf_values: List[float], node_index: int) -> None:
@@ -259,5 +252,5 @@ def _traverse_leaf_values(self, leaf_values: List[float], node_index: int) -> No
         if node.is_leaf_node():
             leaf_values.append(node.value)
         else:
-            self._traverse_leaf_values(leaf_values, node.get_idx_left_child())
-            self._traverse_leaf_values(leaf_values, node.get_idx_right_child())
+            self._traverse_leaf_values(leaf_values, get_idx_left_child(node_index))
+            self._traverse_leaf_values(leaf_values, get_idx_right_child(node_index))

tests/test_tree.py

@@ -1,29 +1,29 @@
 import numpy as np
 
-from pymc_bart.tree import Node, get_depth
+from pymc_bart.tree import Node, get_idx_left_child, get_idx_right_child, 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 get_depth(split_node.index) == 2
+    index = 5
+    split_node = Node.new_split_node(idx_split_variable=2, split_value=3.0)
+    assert get_depth(index) == 2
     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_left_child() == 11
-    assert split_node.get_idx_right_child() == 12
+    assert get_idx_left_child(index) == 11
+    assert get_idx_right_child(index) == 12
     assert split_node.is_split_node() is True
     assert split_node.is_leaf_node() is False
 
 
 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 get_depth(leaf_node.index) == 2
+    index = 5
+    leaf_node = Node.new_leaf_node(value=3.14, idx_data_points=[1, 2, 3])
+    assert get_depth(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])
-    assert leaf_node.get_idx_left_child() == 11
-    assert leaf_node.get_idx_right_child() == 12
+    assert get_idx_left_child(index) == 11
+    assert get_idx_right_child(index) == 12
     assert leaf_node.is_split_node() is False
     assert leaf_node.is_leaf_node() is True