From 3728b5a9bb4e8b46e755d320370ac6f07c4747d4 Mon Sep 17 00:00:00 2001 From: Michael Ellis Date: Wed, 25 Mar 2026 19:41:18 -0500 Subject: [PATCH 1/2] Fix DRTester to accept pandas DataFrames Signed-off-by: Michael Ellis --- econml/tests/test_drtester.py | 31 +++++++++++++++++++++++++++++++ econml/validate/drtester.py | 10 ++++++++++ 2 files changed, 41 insertions(+) diff --git a/econml/tests/test_drtester.py b/econml/tests/test_drtester.py index 3396f0fe2..32dd404c0 100644 --- a/econml/tests/test_drtester.py +++ b/econml/tests/test_drtester.py @@ -286,3 +286,34 @@ def test_exceptions(self): autoc_res = my_dr_tester.evaluate_uplift(Xval, Xtrain, metric='toc') self.assertLess(autoc_res.pvals[0], 0.05) + + def test_dataframe_input(self): + Xtrain, Dtrain, Ytrain, Xval, Dval, Yval = self._get_data(num_treatments=1) + + reg_t = RandomForestClassifier(random_state=0) + reg_y = GradientBoostingRegressor(random_state=0) + + cate = DML( + model_y=reg_y, + model_t=reg_t, + model_final=reg_y, + discrete_treatment=True + ).fit(Y=Ytrain, T=Dtrain, X=Xtrain) + + # Fit with numpy arrays as baseline + dr_numpy = DRTester( + model_regression=reg_y, + model_propensity=reg_t, + cate=cate + ).fit_nuisance(Xval, Dval, Yval, Xtrain, Dtrain, Ytrain) + + # Fit with DataFrames/Series + dr_pandas = DRTester( + model_regression=reg_y, + model_propensity=reg_t, + cate=cate + ).fit_nuisance( + pd.DataFrame(Xval), pd.Series(Dval), pd.Series(Yval), + pd.DataFrame(Xtrain), pd.Series(Dtrain), pd.Series(Ytrain) + ) + np.testing.assert_array_equal(dr_pandas.dr_val_, dr_numpy.dr_val_) diff --git a/econml/validate/drtester.py b/econml/validate/drtester.py index c79330218..2d2becac1 100644 --- a/econml/validate/drtester.py +++ b/econml/validate/drtester.py @@ -221,6 +221,16 @@ def fit_nuisance( If training data is provided, also adds attributes for the doubly robust outcomes for the training set (dr_train) and the training treatments (Dtrain) """ + Xval = np.asarray(Xval) + Dval = np.asarray(Dval) + yval = np.asarray(yval) + if Xtrain is not None: + Xtrain = np.asarray(Xtrain) + if Dtrain is not None: + Dtrain = np.asarray(Dtrain) + if ytrain is not None: + ytrain = np.asarray(ytrain) + self.Dval = Dval # Unique treatments (ordered, includes control) From fdcc2f568d589a6de818427c0ebc0bc328ea08b5 Mon Sep 17 00:00:00 2001 From: Michael Ellis Date: Wed, 25 Mar 2026 20:05:29 -0500 Subject: [PATCH 2/2] Use check_input_arrays for input conversion Signed-off-by: Michael Ellis --- econml/validate/drtester.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/econml/validate/drtester.py b/econml/validate/drtester.py index 2d2becac1..60cf3fb2a 100644 --- a/econml/validate/drtester.py +++ b/econml/validate/drtester.py @@ -8,7 +8,7 @@ from statsmodels.api import OLS from statsmodels.tools import add_constant -from econml.utilities import deprecated +from econml.utilities import check_input_arrays, deprecated from .results import CalibrationEvaluationResults, BLPEvaluationResults, UpliftEvaluationResults, EvaluationResults from .utils import calculate_dr_outcomes, calc_uplift @@ -221,15 +221,8 @@ def fit_nuisance( If training data is provided, also adds attributes for the doubly robust outcomes for the training set (dr_train) and the training treatments (Dtrain) """ - Xval = np.asarray(Xval) - Dval = np.asarray(Dval) - yval = np.asarray(yval) - if Xtrain is not None: - Xtrain = np.asarray(Xtrain) - if Dtrain is not None: - Dtrain = np.asarray(Dtrain) - if ytrain is not None: - ytrain = np.asarray(ytrain) + Xval, Dval, yval = check_input_arrays(Xval, Dval, yval) + Xtrain, Dtrain, ytrain = check_input_arrays(Xtrain, Dtrain, ytrain) self.Dval = Dval