clean and refactor (#18)

Author: aloctavodia

pymc_bart/bart.py

@@ -34,7 +34,6 @@ class BARTRV(RandomVariable):
     ndims_params = [2, 1, 0, 0, 1]
     dtype = "floatX"
     _print_name = ("BART", "\\operatorname{BART}")
-    all_trees = None
 
     def _supp_shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=None):
         return (self.X.shape[0],)

pymc_bart/pgbart.py

@@ -123,6 +123,7 @@ def __init__(
             else:
                 self.batch = (batch, batch)
 
+        self.num_particles = num_particles
         self.log_num_particles = np.log(num_particles)
         self.indices = list(range(2, num_particles))
         self.len_indices = len(self.indices)
@@ -185,14 +186,10 @@ def astep(self, _):
 
             _, normalized_weights = self.normalize(particles)
             # Get the new tree and update
-            new_particle = np.random.choice(particles, p=normalized_weights)
-            new_tree = new_particle.tree
-
-            new_particle.log_weight = new_particle.old_likelihood_logp - self.log_num_particles
+            new_particle, new_tree = self.get_particle_tree(particles, normalized_weights)
             self.all_particles[tree_id] = new_particle
             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:
@@ -230,7 +227,7 @@ def resample(self, particles, normalized_weights):
 
         Ensure particles are copied only if needed.
         """
-        new_indices = systematic(normalized_weights)
+        new_indices = self.systematic(normalized_weights)
         seen = []
         new_particles = []
         for idx in new_indices:
@@ -244,6 +241,29 @@ def resample(self, particles, normalized_weights):
 
         return particles
 
+    def get_particle_tree(self, particles, normalized_weights):
+        """
+        Sample a new particle, new tree and update log_weight
+        """
+        new_index = self.systematic(normalized_weights)[
+            discrete_uniform_sampler(self.num_particles)
+        ]
+        new_particle = particles[new_index - 2]
+        new_particle.log_weight = new_particle.old_likelihood_logp - self.log_num_particles
+        return new_particle, new_particle.tree
+
+    def systematic(self, normalized_weights):
+        """
+        Systematic resampling.
+
+        Return indices in the range 2, ..., len(normalized_weights)+2
+
+        Note: adapted from https://github.com/nchopin/particles
+        """
+        lnw = len(normalized_weights)
+        single_uniform = (self.uniform.random() + np.arange(lnw)) / lnw
+        return inverse_cdf(single_uniform, normalized_weights) + 2
+
     def init_particles(self, tree_id: int) -> np.ndarray:
         """Initialize particles."""
         p0 = self.all_particles[tree_id]
@@ -584,19 +604,6 @@ def update(self):
         )
 
 
-def systematic(normalized_weights):
-    """
-    Systematic resampling.
-
-    Return indices in the range 2, ..., len(normalized_weights)+2
-
-    Note: adapted from https://github.com/nchopin/particles
-    """
-    lnw = len(normalized_weights)
-    single_uniform = (np.random.rand(1) + np.arange(lnw)) / lnw
-    return inverse_cdf(single_uniform, normalized_weights) + 2
-
-
 @njit
 def inverse_cdf(single_uniform, normalized_weights):
     """

pymc_bart/tree.py

@@ -52,8 +52,6 @@ class Tree:
         "idx_leaf_nodes",
         "output",
         "leaf_node_value",
-        "idx_data_points",
-        "shape",
     )
 
     def __init__(self, leaf_node_value, idx_data_points, num_observations, shape):
@@ -200,7 +198,7 @@ def get_idx_right_child(self):
 
 
 class SplitNode(BaseNode):
-    __slots__ = "index", "idx_split_variable", "split_value"
+    __slots__ = "idx_split_variable", "split_value"
 
     def __init__(self, index, idx_split_variable, split_value):
         super().__init__(index)
@@ -210,7 +208,7 @@ def __init__(self, index, idx_split_variable, split_value):
 
 
 class LeafNode(BaseNode):
-    __slots__ = "index", "value", "idx_data_points"
+    __slots__ = "value", "idx_data_points"
 
     def __init__(self, index, value, idx_data_points):
         super().__init__(index)