improve code logic

Author: juanitorduz

pymc_bart/tree.py

@@ -29,7 +29,7 @@ class Node:
     index : int
     value : float
     idx_data_points : Optional[npt.NDArray[np.int_]]
-    idx_split_variable : Optional[npt.NDArray[np.int_]],
+    idx_split_variable : int
     linear_params: Optional[List[float]] = None
     """
 
@@ -190,7 +190,7 @@ def _predict(self) -> npt.NDArray[np.float_]:
         return output.T
 
     def predict(
-        self, x: npt.NDArray[np.float_], excluded: Optional[List[int]] = None
+        self, x: npt.NDArray[np.float_], m: int, excluded: Optional[List[int]] = None
     ) -> npt.NDArray[np.float_]:
         """
         Predict output of tree for an (un)observed point x.
@@ -199,6 +199,8 @@ def predict(
         ----------
         x : npt.NDArray[np.float_]
             Unobserved point
+        m : int
+            Number of trees
         excluded: Optional[List[int]]
             Indexes of the variables to exclude when computing predictions
 
@@ -209,12 +211,14 @@ def predict(
         """
         if excluded is None:
             excluded = []
-        return self._traverse_tree(x, 0, excluded)
+        return self._traverse_tree(x=x, m=m, node_index=0, split_variable=-1, excluded=excluded)
 
     def _traverse_tree(
         self,
         x: npt.NDArray[np.float_],
+        m: int,
         node_index: int,
+        split_variable: int = -1,
         excluded: Optional[List[int]] = None,
     ) -> npt.NDArray[np.float_]:
         """
@@ -224,8 +228,12 @@ def _traverse_tree(
         ----------
         x : npt.NDArray[np.float_]
             Unobserved point
+        m : int
+            Number of trees
         node_index : int
             Index of the node to start the traversal from
+        split_variable : int
+            Index of the variable used to split the node
         excluded: Optional[List[int]]
             Indexes of the variables to exclude when computing predictions
 
@@ -235,13 +243,15 @@ def _traverse_tree(
             Leaf node value or mean of leaf node values
         """
         current_node = self.get_node(node_index)
-
         if current_node.is_leaf_node():
-            if current_node.linear_params is not None:
+            if current_node.linear_params is None:
                 return np.array(current_node.value)
+
             x = x[split_variable].item()
             y_x = current_node.linear_params[0] + current_node.linear_params[1] * x
-            return y_x / m
+            return np.array(y_x / m)
+
+        split_variable = current_node.idx_split_variable
 
         if excluded is not None and current_node.idx_split_variable in excluded:
             leaf_values: List[float] = []
@@ -252,7 +262,9 @@ def _traverse_tree(
             next_node = current_node.get_idx_left_child()
         else:
             next_node = current_node.get_idx_right_child()
-        return self._traverse_tree(x=x, node_index=next_node, excluded=excluded)
+        return self._traverse_tree(
+            x=x, m=m, node_index=next_node, split_variable=split_variable, excluded=excluded
+        )
 
     def _traverse_leaf_values(self, leaf_values: List[float], node_index: int) -> None:
         """

pymc_bart/utils.py

@@ -21,6 +21,7 @@
 def _sample_posterior(
     all_trees: List[List[Tree]],
     X: TensorLike,
+    m: int,
     rng: np.random.Generator,
     size: Optional[Union[int, Tuple[int, ...]]] = None,
     excluded: Optional[npt.NDArray[np.int_]] = None,
@@ -35,6 +36,8 @@ def _sample_posterior(
     X : tensor-like
         A covariate matrix. Use the same used to fit BART for in-sample predictions or a new one for
         out-of-sample predictions.
+    m : int
+        Number of trees
     rng : NumPy RandomGenerator
     size : int or tuple
         Number of samples.
@@ -57,7 +60,7 @@ def _sample_posterior(
         flatten_size *= s
 
     idx = rng.integers(0, len(stacked_trees), size=flatten_size)
-    shape = stacked_trees[0][0].predict(X[0]).size
+    shape = stacked_trees[0][0].predict(x=X[0], m=m).size
 
     pred = np.zeros((flatten_size, X.shape[0], shape))
 
@@ -220,6 +223,8 @@ def plot_dependence(
     -------
     axes: matplotlib axes
     """
+    m: int = bartrv.owner.op.m
+
     if kind not in ["pdp", "ice"]:
         raise ValueError(f"kind={kind} is not suported. Available option are 'pdp' or 'ice'")
 
@@ -294,15 +299,15 @@ def plot_dependence(
                 new_X[:, indices_mi] = X[:, indices_mi]
                 new_X[:, i] = x_i
                 y_pred.append(
-                    np.mean(_sample_posterior(all_trees, X=new_X, rng=rng, size=samples), 1)
+                    np.mean(_sample_posterior(all_trees, X=new_X, m=m, rng=rng, size=samples), 1)
                 )
             new_x_target.append(new_x_i)
         else:
             for instance in instances_ary:
                 new_X = X[idx_s]
                 new_X[:, indices_mi] = X[:, indices_mi][instance]
                 y_pred.append(
-                    np.mean(_sample_posterior(all_trees, X=new_X, rng=rng, size=samples), 0)
+                    np.mean(_sample_posterior(all_trees, X=new_X, m=m, rng=rng, size=samples), 0)
                 )
             new_x_target.append(new_X[:, i])
         y_mins.append(np.min(y_pred))
@@ -445,6 +450,8 @@ def plot_variable_importance(
     """
     _, axes = plt.subplots(2, 1, figsize=figsize)
 
+    m: int = bartrv.owner.op.m
+
     if hasattr(X, "columns") and hasattr(X, "values"):
         labels = X.columns
         X = X.values
@@ -474,13 +481,13 @@ def plot_variable_importance(
 
     all_trees = bartrv.owner.op.all_trees
 
-    predicted_all = _sample_posterior(all_trees, X=X, rng=rng, size=samples, excluded=None)
+    predicted_all = _sample_posterior(all_trees, X=X, m=m, rng=rng, size=samples, excluded=None)
 
     ev_mean = np.zeros(len(var_imp))
     ev_hdi = np.zeros((len(var_imp), 2))
     for idx, subset in enumerate(subsets):
         predicted_subset = _sample_posterior(
-            all_trees=all_trees, X=X, rng=rng, size=samples, excluded=subset
+            all_trees=all_trees, X=X, m=m, rng=rng, size=samples, excluded=subset
         )
         pearson = np.zeros(samples)
         for j in range(samples):