updated missing changes in #776 per @paullo's feedback

uber · May 26, 2024 · ef3a41c · ef3a41c
1 parent 3bbce11
commit ef3a41c
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 4 deletions.
diff --git a/causalml/inference/tree/_tree/_classes.py b/causalml/inference/tree/_tree/_classes.py
@@ -89,7 +89,13 @@ def __init__(
 
     @abstractmethod
     def fit(
-        self, X, y, sample_weight=None, check_input=True, X_idx_sorted="deprecated"
+        self,
+        X,
+        treatment,
+        y,
+        sample_weight=None,
+        check_input=True,
+        X_idx_sorted="deprecated",
     ):
         pass
 

diff --git a/causalml/inference/tree/causal/_tree.py b/causalml/inference/tree/causal/_tree.py
@@ -44,8 +44,8 @@ def _support_missing_values(self, X) -> bool:
     def fit(
         self,
         X,
-        y,
         treatment,
+        y,
         sample_weight=None,
         check_input=True,
         X_idx_sorted="deprecated",

diff --git a/causalml/inference/tree/causal/causaltree.py b/causalml/inference/tree/causal/causaltree.py
@@ -164,7 +164,7 @@ def fit(
             treatment (np.ndarray): treatment vector
             y (np.ndarray): outcome vector
             sample_weight (np.ndarray): sample_weight
-            check_input (bool, optional), default=False
+            check_input (bool, optional): default=False
         Returns:
             self
         """

diff --git a/tests/test_causal_trees.py b/tests/test_causal_trees.py
@@ -115,7 +115,7 @@ def test_fit_predict(
     def test_predict(self, generate_regression_data):
         y, X, treatment, tau, b, e = generate_regression_data(mode=2)
         ctree = self.prepare_model()
-        ctree.fit(X=X, y=y, treatment=treatment)
+        ctree.fit(X=X, treatment=treatment, y=y)
         y_pred = ctree.predict(X[:1, :])
         y_pred_with_outcomes = ctree.predict(X[:1, :], with_outcomes=True)
         assert y_pred.shape == (1,)