add did binary models to set nuisance params test

Author: SvenKlaassen

doubleml/tests/test_set_ml_nuisance_params_models.py

@@ -3,6 +3,9 @@
 from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
 
 from doubleml import DoubleMLAPO, DoubleMLCVAR, DoubleMLIIVM, DoubleMLIRM, DoubleMLLPQ, DoubleMLPLIV, DoubleMLPLR, DoubleMLPQ
+from doubleml.data import DoubleMLPanelData
+from doubleml.did import DoubleMLDIDBinary, DoubleMLDIDCSBinary
+from doubleml.did.datasets import make_did_CS2021
 from doubleml.irm.datasets import make_iivm_data, make_irm_data
 from doubleml.plm.datasets import make_pliv_CHS2015, make_plr_CCDDHNR2018
 
@@ -18,6 +21,10 @@
 dml_data_irm = make_irm_data(n_obs=1000)
 dml_data_iivm = make_iivm_data(n_obs=2000)
 
+# Create DID data
+df_did = make_did_CS2021(n_obs=500, dgp_type=1, n_pre_treat_periods=2, n_periods=4, time_type="float")
+dml_data_did = DoubleMLPanelData(df_did, y_col="y", d_cols="d", id_col="id", t_col="t", x_cols=["Z1", "Z2", "Z3", "Z4"])
+
 reg_learner = RandomForestRegressor(max_depth=2)
 class_learner = RandomForestClassifier(max_depth=2)
 
@@ -43,6 +50,35 @@
 dml_cvar.fit(store_models=True)
 dml_apo.fit(store_models=True)
 
+# DID models
+dml_did_binary = DoubleMLDIDBinary(
+    obj_dml_data=dml_data_did,
+    ml_g=reg_learner,
+    ml_m=class_learner,
+    g_value=2,
+    t_value_pre=0,
+    t_value_eval=1,
+    score="observational",
+    n_folds=n_folds,
+)
+
+dml_did_cs_binary = DoubleMLDIDCSBinary(
+    obj_dml_data=dml_data_did,
+    ml_g=reg_learner,
+    ml_m=class_learner,
+    g_value=2,
+    t_value_pre=0,
+    t_value_eval=1,
+    score="observational",
+    n_folds=n_folds,
+)
+
+dml_did_binary.set_ml_nuisance_params("ml_g0", "d", {"n_estimators": n_est_test})
+dml_did_cs_binary.set_ml_nuisance_params("ml_g_d0_t0", "d", {"n_estimators": n_est_test})
+
+dml_did_binary.fit(store_models=True)
+dml_did_cs_binary.fit(store_models=True)
+
 # nonlinear models
 dml_pq = DoubleMLPQ(dml_data_irm, ml_g=class_learner, ml_m=class_learner, n_folds=n_folds)
 dml_lpq = DoubleMLLPQ(dml_data_iivm, ml_g=class_learner, ml_m=class_learner, n_folds=n_folds)
@@ -108,3 +144,13 @@ def test_lpq_params():
     assert all(param == n_est_default for param in param_list_2)
     param_list_2 = [dml_lpq.models["ml_g_du_z1"]["d"][0][fold].n_estimators for fold in range(n_folds)]
     assert all(param == n_est_default for param in param_list_2)
+
+
+@pytest.mark.ci
+def test_did_binary_params():
+    _assert_nuisance_params(dml_did_binary, "ml_g0", "ml_g1")
+
+
+@pytest.mark.ci
+def test_did_cs_binary_params():
+    _assert_nuisance_params(dml_did_cs_binary, "ml_g_d0_t0", "ml_g_d1_t0")