Avoid Deepcopy on Tree and ParticleTree (#47)

* Implement a constructor of Tree with all the atributes(tree_structure, idx_leaf_nodes and output);
Create the class method new_tree with the previous behaviour of the constructor;
Change the implementation of copy on Tree, avoiding deepcopying all and just deepcopying the output;
Improve method trim(), previously we copy all and then we delete the atributes that we do not need, now we just create the tree copying the atributes that we care;
Insted of creating an empty list and append each particle we create an comprehensive list;
We implement the method copy on ParticleTree, instead of deepcopying;

* Changes Code Review;

Author: fjloyola

pymc_bart/pgbart.py

@@ -14,7 +14,6 @@
 
 import logging
 
-from copy import deepcopy
 from numba import njit
 
 import numpy as np
@@ -107,7 +106,7 @@ def __init__(
             config.floatX
         )
         self.sum_trees_noi = self.sum_trees - (init_mean / self.m)
-        self.a_tree = Tree(
+        self.a_tree = Tree.new_tree(
             leaf_node_value=init_mean / self.m,
             idx_data_points=np.arange(self.num_observations, dtype="int32"),
             num_observations=self.num_observations,
@@ -136,9 +135,7 @@ def __init__(
 
         shared = make_shared_replacements(initial_values, vars, model)
         self.likelihood_logp = logp(initial_values, [model.datalogp], vars, shared)
-        self.all_particles = []
-        for _ in range(self.m):
-            self.all_particles.append(ParticleTree(self.a_tree))
+        self.all_particles = list(ParticleTree(self.a_tree) for _ in range(self.m))
         self.all_trees = np.array([p.tree for p in self.all_particles])
         super().__init__(vars, shared)
 
@@ -239,7 +236,7 @@ def resample(self, particles, normalized_weights):
         new_particles = []
         for idx in new_indices:
             if idx in seen:
-                new_particles.append(deepcopy(particles[idx]))
+                new_particles.append(particles[idx].copy())
             else:
                 new_particles.append(particles[idx])
                 seen.append(idx)
@@ -274,7 +271,7 @@ def systematic(self, normalized_weights):
     def init_particles(self, tree_id: int) -> np.ndarray:
         """Initialize particles."""
         p0 = self.all_particles[tree_id]
-        p1 = deepcopy(p0)
+        p1 = p0.copy()
         p1.sample_leafs(
             self.sum_trees,
             self.m,
@@ -328,12 +325,22 @@ class ParticleTree:
     __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.tree = tree.copy()
         self.expansion_nodes = [0]
         self.log_weight = 0
         self.old_likelihood_logp = 0
         self.kfactor = 0.75
 
+    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,
+        )
+        return p
+
     def sample_tree(
         self,
         ssv,
@@ -500,7 +507,6 @@ def grow_tree(
 
 
 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]
@@ -509,7 +515,6 @@ def get_new_idx_data_points(split_value, idx_data_points, selected_predictor, X)
 
 
 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[

pymc_bart/tree.py

@@ -14,7 +14,6 @@
 
 import math
 
-from copy import deepcopy
 from functools import lru_cache
 
 from pytensor import config
@@ -42,9 +41,9 @@ class Tree:
 
     Parameters
     ----------
-    idx_data_points : array of integers
-    num_observations : integer
-    shape : int
+    tree_structure : Dictionary of nodes
+    idx_leaf_nodes :  List with the index of the leaf nodes of the tree.
+    output : Array of shape number of observations, shape
     """
 
     __slots__ = (
@@ -53,12 +52,20 @@ class Tree:
         "output",
     )
 
-    def __init__(self, leaf_node_value, idx_data_points, num_observations, shape):
-        self.tree_structure = {
-            0: Node.new_leaf_node(0, value=leaf_node_value, idx_data_points=idx_data_points)
-        }
-        self.idx_leaf_nodes = [0]
-        self.output = np.zeros((num_observations, shape)).astype(config.floatX).squeeze()
+    def __init__(self, tree_structure, idx_leaf_nodes, output):
+        self.tree_structure = tree_structure
+        self.idx_leaf_nodes = idx_leaf_nodes
+        self.output = output
+
+    @classmethod
+    def new_tree(cls, leaf_node_value, idx_data_points, num_observations, shape):
+        return cls(
+            tree_structure={
+                0: Node.new_leaf_node(0, value=leaf_node_value, idx_data_points=idx_data_points)
+            },
+            idx_leaf_nodes=[0],
+            output=np.zeros((num_observations, shape)).astype(config.floatX).squeeze(),
+        )
 
     def __getitem__(self, index):
         return self.get_node(index)
@@ -67,7 +74,11 @@ def __setitem__(self, index, node):
         self.set_node(index, node)
 
     def copy(self):
-        return deepcopy(self)
+        tree = {
+            k: Node(v.index, v.value, v.idx_data_points, v.idx_split_variable)
+            for k, v in self.tree_structure.items()
+        }
+        return Tree(tree, self.idx_leaf_nodes.copy(), self.output.copy())
 
     def get_node(self, index) -> "Node":
         return self.tree_structure[index]
@@ -82,14 +93,11 @@ def delete_leaf_node(self, index):
         del self.tree_structure[index]
 
     def trim(self):
-        a_tree = self.copy()
-        del a_tree.output
-        del a_tree.idx_leaf_nodes
-        for k in a_tree.tree_structure.keys():
-            current_node = a_tree[k]
-            if current_node.is_leaf_node():
-                del current_node.idx_data_points
-        return a_tree
+        tree = {
+            k: Node(v.index, v.value, None, v.idx_split_variable)
+            for k, v in self.tree_structure.items()
+        }
+        return Tree(tree, None, None)
 
     def get_split_variables(self):
         return [