refactor so predict does not need m (#88)

Author: aloctavodia

pymc_bart/bart.py

@@ -56,7 +56,7 @@ def rng_fn(cls, rng=None, X=None, Y=None, m=None, alpha=None, split_prior=None,
                 shape = size[0]
             else:
                 shape = 1
-            return _sample_posterior(cls.all_trees, cls.X, cls.m, rng=rng, shape=shape).squeeze().T
+            return _sample_posterior(cls.all_trees, cls.X, rng=rng, shape=shape).squeeze().T
 
 
 bart = BARTRV()

pymc_bart/pgbart.py

@@ -491,8 +491,7 @@ def draw_leaf_value(
         if response == "constant":
             mu_mean = fast_mean(y_mu_pred) / m
         if response == "linear":
-            y_fit, linear_params = fast_linear_fit(x=x_mu, y=y_mu_pred)
-            mu_mean = y_fit / m
+            mu_mean, linear_params = fast_linear_fit(x=x_mu, y=y_mu_pred, m=m)
 
     draw = norm + mu_mean
     return draw, linear_params
@@ -518,9 +517,12 @@ def fast_mean(ari: npt.NDArray[np.float_]) -> Union[float, npt.NDArray[np.float_
 
 @njit
 def fast_linear_fit(
-    x: npt.NDArray[np.float_], y: npt.NDArray[np.float_]
+    x: npt.NDArray[np.float_],
+    y: npt.NDArray[np.float_],
+    m: int,
 ) -> Tuple[npt.NDArray[np.float_], List[npt.NDArray[np.float_]]]:
     n = len(x)
+    y = y / m
 
     xbar = np.sum(x) / n
     ybar = np.sum(y, axis=1) / n

pymc_bart/tree.py

@@ -221,7 +221,6 @@ def _predict(self) -> npt.NDArray[np.float_]:
     def predict(
         self,
         x: npt.NDArray[np.float_],
-        m: int,
         excluded: Optional[List[int]] = None,
         shape: int = 1,
     ) -> npt.NDArray[np.float_]:
@@ -232,8 +231,6 @@ 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
 
@@ -244,12 +241,11 @@ def predict(
         """
         if excluded is None:
             excluded = []
-        return self._traverse_tree(x=x, m=m, excluded=excluded, shape=shape)
+        return self._traverse_tree(x=x, excluded=excluded, shape=shape)
 
     def _traverse_tree(
         self,
         x: npt.NDArray[np.float_],
-        m: int,
         excluded: Optional[List[int]] = None,
         shape: int = 1,
     ) -> npt.NDArray[np.float_]:
@@ -260,8 +256,6 @@ def _traverse_tree(
         ----------
         x : npt.NDArray[np.float_]
             (Un)observed point
-        m : int
-            Number of trees
         node_index : int
             Index of the node to start the traversal from
         split_variable : int
@@ -274,7 +268,7 @@ def _traverse_tree(
         npt.NDArray[np.float_]
             Leaf node value or mean of leaf node values
         """
-        stack = [(0, 1.0)]  # (node_index, prop) initial state
+        stack = [(0, 1.0)]  # (node_index, weight) initial state
         p_d = np.zeros(shape)
         while stack:
             node_index, weight = stack.pop()
@@ -285,7 +279,7 @@ def _traverse_tree(
                     p_d += weight * node.value
                 else:
                     # this produce nonsensical results
-                    p_d += weight * ((params[0] + params[1] * x[node.idx_split_variable]) / m)
+                    p_d += weight * (params[0] + params[1] * x[node.idx_split_variable])
                     # this produce reasonable result
                     # p_d += weight * node.value.mean()
             else:

pymc_bart/utils.py

@@ -21,7 +21,6 @@
 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[List[int]] = None,
@@ -37,8 +36,6 @@ 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.
@@ -66,7 +63,7 @@ def _sample_posterior(
 
     for ind, p in enumerate(pred):
         for tree in stacked_trees[idx[ind]]:
-            p += np.vstack([tree.predict(x=x, m=m, excluded=excluded, shape=shape) for x in X])
+            p += np.vstack([tree.predict(x=x, excluded=excluded, shape=shape) for x in X])
     pred.reshape((*size_iter, shape, -1))
     return pred
 
@@ -239,7 +236,6 @@ def plot_ice(
     axes: matplotlib axes
     """
     all_trees = bartrv.owner.op.all_trees
-    m: int = bartrv.owner.op.m
     rng = np.random.default_rng(random_seed)
 
     if func is None:
@@ -271,7 +267,7 @@ def plot_ice(
             fake_X[:, indices_mi] = X[:, indices_mi][instance]
             y_pred.append(
                 np.mean(
-                    _sample_posterior(all_trees, X=fake_X, m=m, rng=rng, size=samples, shape=shape),
+                    _sample_posterior(all_trees, X=fake_X, rng=rng, size=samples, shape=shape),
                     0,
                 )
             )
@@ -386,7 +382,6 @@ def plot_pdp(
     axes: matplotlib axes
     """
     all_trees: list = bartrv.owner.op.all_trees
-    m: int = bartrv.owner.op.m
     rng = np.random.default_rng(random_seed)
 
     if func is None:
@@ -411,7 +406,7 @@ def plot_pdp(
         excluded.remove(var)
         fake_X, new_x = _create_pdp_data(X, xs_interval, var, xs_values, var_discrete)
         p_d = _sample_posterior(
-            all_trees, X=fake_X, m=m, rng=rng, size=samples, excluded=excluded, shape=shape
+            all_trees, X=fake_X, rng=rng, size=samples, excluded=excluded, shape=shape
         )
 
         for s_i in range(shape):
@@ -738,8 +733,6 @@ def plot_variable_importance(
     """
     _, axes = plt.subplots(2, 1, figsize=figsize)
 
-    m: int = bartrv.owner.op.m
-
     if bartrv.ndim == 1:  # type: ignore
         shape = 1
     else:
@@ -775,7 +768,7 @@ def plot_variable_importance(
     all_trees = bartrv.owner.op.all_trees
 
     predicted_all = _sample_posterior(
-        all_trees, X=X, m=m, rng=rng, size=samples, excluded=None, shape=shape
+        all_trees, X=X, rng=rng, size=samples, excluded=None, shape=shape
     )
 
     ev_mean = np.zeros(len(var_imp))
@@ -784,7 +777,6 @@ def plot_variable_importance(
         predicted_subset = _sample_posterior(
             all_trees=all_trees,
             X=X,
-            m=m,
             rng=rng,
             size=samples,
             excluded=subset,

tests/test_bart.py

@@ -139,13 +139,9 @@ class TestUtils:
     def test_sample_posterior(self):
         all_trees = self.mu.owner.op.all_trees
         rng = np.random.default_rng(3)
-        pred_all = pmb.utils._sample_posterior(
-            all_trees, X=self.X, m=self.mu.owner.op.m, rng=rng, size=2
-        )
+        pred_all = pmb.utils._sample_posterior(all_trees, X=self.X, rng=rng, size=2)
         rng = np.random.default_rng(3)
-        pred_first = pmb.utils._sample_posterior(
-            all_trees, X=self.X[:10], m=self.mu.owner.op.m, rng=rng
-        )
+        pred_first = pmb.utils._sample_posterior(all_trees, X=self.X[:10], rng=rng)
 
         assert_almost_equal(pred_first[0], pred_all[0, :10], decimal=4)
         assert pred_all.shape == (2, 50, 1)

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)
+    y_fit, linear_params = fast_linear_fit(x, y, m=1)
     assert linear_params[0] == a_expected
     assert linear_params[1] == b_expected
     np.testing.assert_almost_equal(