fix incorrect effect_learner_objective in XGBRRegressor (#504)

uber · Apr 30, 2022 · 1cd2906 · 1cd2906
1 parent 3d59609
commit 1cd2906
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 10 deletions.
diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py
@@ -516,7 +516,7 @@ def __init__(
         early_stopping=True,
         test_size=0.3,
         early_stopping_rounds=30,
-        effect_learner_objective="rank:pairwise",
+        effect_learner_objective="reg:squarederror",
         effect_learner_n_estimators=500,
         random_state=42,
         *args,
@@ -531,7 +531,7 @@ def __init__(
             early_stopping_rounds (int, optional): validation metric needs to improve at least once in every
                                                    early_stopping_rounds round(s) to continue training
             effect_learner_objective (str, optional): the learning objective for the effect learner
-                                                      (default = 'rank:pairwise')
+                                                      (default = 'reg:squarederror')
             effect_learner_n_estimators (int, optional): number of trees to fit for the effect learner (default = 500)
         """
 

diff --git a/tests/test_meta_learners.py b/tests/test_meta_learners.py
@@ -840,18 +840,24 @@ def test_BaseRClassifier_with_sample_weights(generate_classification_data):
     # higher than it would be under random targeting
     assert cumgain["tau_pred"].sum() > cumgain["Random"].sum()
 
+
+def test_XGBRegressor_with_sample_weights(generate_regression_data):
+    y, X, treatment, tau, b, e = generate_regression_data()
+
+    weights = np.random.rand(y.shape[0])
+
     # Check if XGBRRegressor successfully produces treatment effect estimation
     # when sample_weight is passed
     uplift_model = XGBRRegressor()
     uplift_model.fit(
-        X=df_train[x_names].values,
-        p=df_train["propensity_score"].values,
-        treatment=df_train["treatment_group_key"].values,
-        y=df_train[CONVERSION].values,
-        sample_weight=df_train["sample_weights"],
-    )
-    tau_pred = uplift_model.predict(X=df_test[x_names].values)
-    assert len(tau_pred) == len(df_test["sample_weights"].values)
+        X=X,
+        p=e,
+        treatment=treatment,
+        y=y,
+        sample_weight=weights,
+    )
+    tau_pred = uplift_model.predict(X=X)
+    assert len(tau_pred) == len(weights)
 
 
 def test_pandas_input(generate_regression_data):