fix linear response bug (#101)

Author: aloctavodia

pymc_bart/pgbart.py

@@ -536,15 +536,14 @@ def draw_leaf_value(
         return np.zeros(shape), linear_params
 
     if y_mu_pred.size == 1:
-        mu_mean = np.full(shape, y_mu_pred.item() / m)
+        mu_mean = np.full(shape, y_mu_pred.item() / m) + norm
     else:
-        if response == "constant":
-            mu_mean = fast_mean(y_mu_pred) / m
-        if response == "linear":
-            mu_mean, linear_params = fast_linear_fit(x=x_mu, y=y_mu_pred, m=m)
+        if y_mu_pred.size < 3 or response == "constant":
+            mu_mean = fast_mean(y_mu_pred) / m + norm
+        else:
+            mu_mean, linear_params = fast_linear_fit(x=x_mu, y=y_mu_pred, m=m, norm=norm)
 
-    draw = mu_mean + norm
-    return draw, linear_params
+    return mu_mean, linear_params
 
 
 @njit
@@ -570,9 +569,10 @@ def fast_linear_fit(
     x: npt.NDArray[np.float_],
     y: npt.NDArray[np.float_],
     m: int,
+    norm: npt.NDArray[np.float_],
 ) -> Tuple[npt.NDArray[np.float_], List[npt.NDArray[np.float_]]]:
     n = len(x)
-    y = y / m
+    y = y / m + np.expand_dims(norm, axis=1)
 
     xbar = np.sum(x) / n
     ybar = np.sum(y, axis=1) / n

pymc_bart/tree.py

@@ -65,10 +65,6 @@ def new_leaf_node(
             linear_params=linear_params,
         )
 
-    @classmethod
-    def new_split_node(cls, split_value: npt.NDArray[np.float_], idx_split_variable: int) -> "Node":
-        return cls(value=split_value, idx_split_variable=idx_split_variable)
-
     def is_split_node(self) -> bool:
         return self.idx_split_variable >= 0
 
@@ -282,42 +278,42 @@ def _traverse_tree(
         """
 
         x_shape = (1,) if len(X.shape) == 1 else X.shape[:-1]
+        nd_dims = (...,) + (None,) * len(x_shape)
 
-        stack = [(0, np.ones(x_shape))]  # (node_index, weight) initial state
+        stack = [(0, np.ones(x_shape), 0)]  # (node_index, weight, idx_split_variable) initial state
         p_d = (
             np.zeros(shape + x_shape) if isinstance(shape, tuple) else np.zeros((shape,) + x_shape)
         )
         while stack:
-            node_index, weights = stack.pop()
+            node_index, weights, idx_split_variable = stack.pop()
             node = self.get_node(node_index)
             if node.is_leaf_node():
                 params = node.linear_params
-                nd_dims = (...,) + (None,) * len(x_shape)
                 if params is None:
                     p_d += weights * node.value[nd_dims]
                 else:
-                    # this produce nonsensical results
                     p_d += weights * (
-                        params[0][nd_dims] + params[1][nd_dims] * X[..., node.idx_split_variable]
+                        params[0][nd_dims] + params[1][nd_dims] * X[..., idx_split_variable]
                     )
-                    # this produce reasonable result
-                    # p_d += weight * node.value.mean()
             else:
                 left_node_index, right_node_index = get_idx_left_child(
                     node_index
                 ), get_idx_right_child(node_index)
+                idx_split_variable = node.idx_split_variable
                 if excluded is not None and node.idx_split_variable in excluded:
                     prop_nvalue_left = self.get_node(left_node_index).nvalue / node.nvalue
-                    stack.append((left_node_index, weights * prop_nvalue_left))
-                    stack.append((right_node_index, weights * (1 - prop_nvalue_left)))
+                    stack.append((left_node_index, weights * prop_nvalue_left, idx_split_variable))
+                    stack.append(
+                        (right_node_index, weights * (1 - prop_nvalue_left), idx_split_variable)
+                    )
                 else:
                     to_left = (
-                        self.split_rules[node.idx_split_variable]
-                        .divide(X[..., node.idx_split_variable], node.value)
+                        self.split_rules[idx_split_variable]
+                        .divide(X[..., idx_split_variable], node.value)
                         .astype("float")
                     )
-                    stack.append((left_node_index, weights * to_left))
-                    stack.append((right_node_index, weights * (1 - to_left)))
+                    stack.append((left_node_index, weights * to_left, idx_split_variable))
+                    stack.append((right_node_index, weights * (1 - to_left), idx_split_variable))
 
         if len(X.shape) == 1:
             p_d = p_d[..., 0]

tests/test_pgbart.py

@@ -54,7 +54,7 @@ def test_fast_mean():
     ids=["1d-id", "1d-const"],
 )
 def test_fast_linear_fit(x, y, a_expected, b_expected):
-    y_fit, linear_params = fast_linear_fit(x, y, m=1)
+    y_fit, linear_params = fast_linear_fit(x, y, m=1, norm=np.zeros(1))
     assert linear_params[0] == a_expected
     assert linear_params[1] == b_expected
     np.testing.assert_almost_equal(

tests/test_tree.py

@@ -5,7 +5,7 @@
 
 def test_split_node():
     index = 5
-    split_node = Node.new_split_node(idx_split_variable=2, split_value=3.0)
+    split_node = Node(idx_split_variable=2, value=3.0)
     assert get_depth(index) == 2
     assert split_node.value == 3.0
     assert split_node.idx_split_variable == 2