diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py index 2fea90eb..e2a64209 100644 --- 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 index f9f7122a..c9c6e34f 100644 --- 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):