Update test_dml_data.py

DoubleML · Jan 13, 2025 · 5bd1d3f · 5bd1d3f
1 parent d1a4915
commit 5bd1d3f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/doubleml/tests/test_dml_data.py b/doubleml/tests/test_dml_data.py
@@ -156,7 +156,8 @@ def test_add_vars_in_df():
     np.random.seed(3141)
     df = make_plr_CCDDHNR2018(n_obs=100, return_type="DataFrame")
     dml_data_full_df = DoubleMLData(df, "y", "d", ["X1", "X11", "X13"])
-    dml_data_subset = DoubleMLData(df[["X1", "X11", "X13"] + ["y", "d"]], "y", "d", ["X1", "X11", "X13"])
+    subset_columns = ["X1", "X11", "X13", "y", "d"]
+    dml_data_subset = DoubleMLData(df[subset_columns], "y", "d", ["X1", "X11", "X13"])
     dml_plr_full_df = DoubleMLPLR(dml_data_full_df, Lasso(), Lasso())
     dml_plr_subset = DoubleMLPLR(dml_data_subset, Lasso(), Lasso(), draw_sample_splitting=False)
     dml_plr_subset.set_sample_splitting(dml_plr_full_df.smpls)