descriptive variable name, add type hints

Author: aloctavodia

pymc/step_methods/pgbart.py

@@ -32,6 +32,59 @@
 _log = logging.getLogger("pymc")
 
 
+class ParticleTree:
+    """
+    Particle tree
+    """
+
+    def __init__(self, tree, log_weight, likelihood):
+        self.tree = tree.copy()  # keeps the tree that we care at the moment
+        self.expansion_nodes = [0]
+        self.log_weight = log_weight
+        self.old_likelihood_logp = likelihood
+        self.used_variates = []
+
+    def sample_tree_sequential(
+        self,
+        ssv,
+        available_predictors,
+        prior_prob_leaf_node,
+        X,
+        missing_data,
+        sum_trees_output,
+        mean,
+        m,
+        normal,
+        mu_std,
+    ):
+        tree_grew = False
+        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]
+
+            if prob_leaf < np.random.random():
+                tree_grew, index_selected_predictor = grow_tree(
+                    self.tree,
+                    index_leaf_node,
+                    ssv,
+                    available_predictors,
+                    X,
+                    missing_data,
+                    sum_trees_output,
+                    mean,
+                    m,
+                    normal,
+                    mu_std,
+                )
+                if tree_grew:
+                    new_indexes = self.tree.idx_leaf_nodes[-2:]
+                    self.expansion_nodes.extend(new_indexes)
+                    self.used_variates.append(index_selected_predictor)
+
+        return tree_grew
+
+
 class PGBART(ArrayStepShared):
     """
     Particle Gibss BART sampling step
@@ -152,7 +205,7 @@ def astep(self, q: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
             for t in range(self.max_stages):
                 # Sample each particle (try to grow each tree), except for the first one.
                 for p in particles[1:]:
-                    clp = p.sample_tree_sequential(
+                    tree_grew = p.sample_tree_sequential(
                         self.ssv,
                         self.available_predictors,
                         self.prior_prob_leaf_node,
@@ -164,7 +217,7 @@ def astep(self, q: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
                         self.normal,
                         self.mu_std,
                     )
-                    if clp:  # update weights only if p has changed from the previous iteration
+                    if tree_grew:
                         self.update_weight(p)
                 # Normalize weights
                 W_t, normalized_weights = self.normalize(particles)
@@ -222,7 +275,7 @@ def competence(var, has_grad):
             return Competence.IDEAL
         return Competence.INCOMPATIBLE
 
-    def normalize(self, particles):
+    def normalize(self, particles: List[ParticleTree]) -> Tuple[float, np.ndarray]:
         """
         Use logsumexp trick to get W_t and softmax to get normalized_weights
         """
@@ -238,7 +291,7 @@ def normalize(self, particles):
 
         return W_t, normalized_weights
 
-    def init_particles(self, tree_id):
+    def init_particles(self, tree_id: int) -> np.ndarray:
         """
         Initialize particles
         """
@@ -259,12 +312,12 @@ def init_particles(self, tree_id):
 
         return np.array(particles)
 
-    def update_weight(self, particle):
+    def update_weight(self, particle: List[ParticleTree]) -> None:
         """
         Update the weight of a particle
 
         Since the prior is used as the proposal,the weights are updated additively as the ratio of
-        the new and old log_likelihoods.
+        the new and old log-likelihoods.
         """
         new_likelihood = self.likelihood_logp(
             self.sum_trees_output_noi + particle.tree.predict_output()
@@ -273,59 +326,6 @@ def update_weight(self, particle):
         particle.old_likelihood_logp = new_likelihood
 
 
-class ParticleTree:
-    """
-    Particle tree
-    """
-
-    def __init__(self, tree, log_weight, likelihood):
-        self.tree = tree.copy()  # keeps the tree that we care at the moment
-        self.expansion_nodes = [0]
-        self.log_weight = log_weight
-        self.old_likelihood_logp = likelihood
-        self.used_variates = []
-
-    def sample_tree_sequential(
-        self,
-        ssv,
-        available_predictors,
-        prior_prob_leaf_node,
-        X,
-        missing_data,
-        sum_trees_output,
-        mean,
-        m,
-        normal,
-        mu_std,
-    ):
-        clp = False
-        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]
-
-            if prob_leaf < np.random.random():
-                clp, index_selected_predictor = grow_tree(
-                    self.tree,
-                    index_leaf_node,
-                    ssv,
-                    available_predictors,
-                    X,
-                    missing_data,
-                    sum_trees_output,
-                    mean,
-                    m,
-                    normal,
-                    mu_std,
-                )
-                if clp:
-                    new_indexes = self.tree.idx_leaf_nodes[-2:]
-                    self.expansion_nodes.extend(new_indexes)
-                    self.used_variates.append(index_selected_predictor)
-
-        return clp
-
-
 def preprocess_XY(X, Y):
     if isinstance(Y, (Series, DataFrame)):
         Y = Y.to_numpy()