use slots (#17)

Author: aloctavodia

pymc_bart/pgbart.py

@@ -101,9 +101,10 @@ def __init__(
             aesara.config.floatX
         )
         self.sum_trees_noi = self.sum_trees - (self.init_mean / self.m)
-        self.a_tree = Tree.init_tree(
+        self.a_tree = Tree(
             leaf_node_value=self.init_mean / self.m,
             idx_data_points=np.arange(self.num_observations, dtype="int32"),
+            num_observations=self.num_observations,
             shape=self.shape,
         )
         self.normal = NormalSampler(mu_std, self.shape)
@@ -297,6 +298,8 @@ def competence(var, has_grad):
 class ParticleTree:
     """Particle tree."""
 
+    __slots__ = "tree", "expansion_nodes", "log_weight", "old_likelihood_logp", "kfactor"
+
     def __init__(self, tree):
         self.tree = tree.copy()  # keeps the tree that we care at the moment
         self.expansion_nodes = [0]

pymc_bart/tree.py

@@ -36,19 +36,31 @@ class Tree:
         of the tree itself.
     idx_leaf_nodes : list
         List with the index of the leaf nodes of the tree.
-    num_observations : int
-        Number of observations used to fit BART.
-    m : int
-        Number of trees
+    output: array
+        Array of shape number of observations, shape
 
     Parameters
     ----------
-    num_observations : int, optional
+    leaf_node_value : int or float
+    idx_data_points : array of integers
+    num_observations : integer
+    shape : int
     """
 
-    def __init__(self, num_observations=0, shape=1):
-        self.tree_structure = {}
-        self.idx_leaf_nodes = []
+    __slots__ = (
+        "tree_structure",
+        "idx_leaf_nodes",
+        "output",
+        "leaf_node_value",
+        "idx_data_points",
+        "shape",
+    )
+
+    def __init__(self, leaf_node_value, idx_data_points, num_observations, shape):
+        self.tree_structure = {
+            0: LeafNode(index=0, value=leaf_node_value, idx_data_points=idx_data_points)
+        }
+        self.idx_leaf_nodes = [0]
         self.output = np.zeros((num_observations, shape)).astype(aesara.config.floatX).squeeze()
 
     def __getitem__(self, index):
@@ -169,26 +181,10 @@ def _traverse_leaf_values(self, leaf_values, node_index):
         else:
             leaf_values.append(current_node.value)
 
-    @staticmethod
-    def init_tree(leaf_node_value, idx_data_points, shape):
-        """
-        Initialize tree.
-
-        Parameters
-        ----------
-        leaf_node_value
-        idx_data_points
-
-        Returns
-        -------
-        tree
-        """
-        new_tree = Tree(len(idx_data_points), shape)
-        new_tree[0] = LeafNode(index=0, value=leaf_node_value, idx_data_points=idx_data_points)
-        return new_tree
-
 
 class BaseNode:
+    __slots__ = "index", "depth"
+
     def __init__(self, index):
         self.index = index
         self.depth = int(math.floor(math.log(index + 1, 2)))
@@ -204,6 +200,8 @@ def get_idx_right_child(self):
 
 
 class SplitNode(BaseNode):
+    __slots__ = "index", "idx_split_variable", "split_value"
+
     def __init__(self, index, idx_split_variable, split_value):
         super().__init__(index)
 
@@ -212,6 +210,8 @@ def __init__(self, index, idx_split_variable, split_value):
 
 
 class LeafNode(BaseNode):
+    __slots__ = "index", "value", "idx_data_points"
+
     def __init__(self, index, value, idx_data_points):
         super().__init__(index)
         self.value = value