Fix sensitivity_benchmark() of DID models

doubleml/did/did_binary.py

@@ -186,7 +186,6 @@ def __init__(
         self._n_treated_subset = self._data_subset["G_indicator"].sum()
 
         # Save x and y for later ML estimation
-        self._x_data_subset = self._data_subset.loc[:, self._dml_data.x_cols].values
         self._y_data_subset = self._data_subset.loc[:, "y_diff"].values
         self._g_data_subset = self._data_subset.loc[:, "G_indicator"].values
 
@@ -308,6 +307,11 @@ def anticipation_periods(self):
         """
         return self._anticipation_periods
 
+    @property
+    def _x_data_subset(self):
+        # Read with the current covariates, so that sensitivity_benchmark() can drop some of them.
+        return self._data_subset.loc[:, self._dml_data.x_cols].values
+
     @property
     def data_subset(self):
         """

doubleml/did/did_cs_binary.py

@@ -174,7 +174,6 @@ def __init__(
         self._n_obs_subset = self.data_subset.shape[0]  # Effective sample size used for resampling
 
         # Save x and y for later ML estimation
-        self._x_data_subset = self.data_subset.loc[:, self._dml_data.x_cols].values
         self._y_data_subset = self.data_subset.loc[:, self._dml_data.y_col].values
         self._g_data_subset = self.data_subset.loc[:, "G_indicator"].values
         self._t_data_subset = self.data_subset.loc[:, "t_indicator"].values
@@ -299,6 +298,11 @@ def anticipation_periods(self):
         """
         return self._anticipation_periods
 
+    @property
+    def _x_data_subset(self):
+        # Read with the current covariates, so that sensitivity_benchmark() can drop some of them.
+        return self._data_subset.loc[:, self._dml_data.x_cols].values
+
     @property
     def data_subset(self):
         """

doubleml/did/did_multi.py

@@ -670,6 +670,9 @@ def fit(self, n_jobs_models=None, n_jobs_cv=None, store_predictions=True, store_
         framework_list = [None] * self.n_gt_atts
 
         for i_gt in range(self.n_gt_atts):
+            # parallel workers return copies with their own data object; share the data again as in serial fitting,
+            # so that changes of the covariates (e.g. in sensitivity_benchmark()) reach all group-time models
+            fitted_models[i_gt]._dml_data = self._dml_data
             self._modellist[i_gt] = fitted_models[i_gt]
             framework_list[i_gt] = self._modellist[i_gt].framework
 

doubleml/did/tests/test_did_sensitivity_benchmark.py

@@ -0,0 +1,50 @@
+import numpy as np
+import pytest
+from sklearn.linear_model import LinearRegression, LogisticRegression
+
+import doubleml as dml
+from doubleml.did.datasets import make_did_CS2021, make_did_cs_CS2021
+
+
+def _panel_data(panel, x_cols):
+    np.random.seed(3141)
+    if panel:
+        df = make_did_CS2021(n_obs=500, dgp_type=1, time_type="float")
+    else:
+        df = make_did_cs_CS2021(n_obs=1000, dgp_type=1, time_type="float", lambda_t=0.5)
+    return dml.data.DoubleMLPanelData(df, y_col="y", d_cols="d", id_col="id", t_col="t", x_cols=x_cols)
+
+
+def _fit(model_class, dml_data, fit_args=None, **kwargs):
+    np.random.seed(42)
+    model = model_class(dml_data, ml_g=LinearRegression(), ml_m=LogisticRegression(), **kwargs)
+    return model.fit(**(fit_args or {}))
+
+
+@pytest.mark.ci
+@pytest.mark.parametrize("panel", [True, False])
+@pytest.mark.parametrize("n_jobs_models", [None, 2])
+def test_did_multi_sensitivity_benchmark_drops_covariates(panel, n_jobs_models):
+    x_long = ["Z1", "Z2", "Z3", "Z4"]
+    x_short = ["Z2", "Z3", "Z4"]
+    kwargs = {"gt_combinations": "standard", "panel": panel}
+    fit_args = {"n_jobs_models": n_jobs_models}
+    dml_long = _fit(dml.did.DoubleMLDIDMulti, _panel_data(panel, x_long), fit_args=fit_args, **kwargs)
+    dml_short = _fit(dml.did.DoubleMLDIDMulti, _panel_data(panel, x_short), **kwargs)
+
+    benchmark = dml_long.sensitivity_benchmark(benchmarking_set=["Z1"])
+    np.testing.assert_allclose(benchmark["delta_theta"].values, dml_short.coef - dml_long.coef, rtol=1e-9, atol=1e-12)
+    assert np.all(benchmark["delta_theta"].values != 0)
+
+
+@pytest.mark.ci
+@pytest.mark.parametrize("panel", [True, False])
+def test_did_binary_sensitivity_benchmark_drops_covariates(panel):
+    model_class = dml.did.DoubleMLDIDBinary if panel else dml.did.DoubleMLDIDCSBinary
+    kwargs = {"g_value": 2.0, "t_value_pre": 1.0, "t_value_eval": 2.0}
+    dml_long = _fit(model_class, _panel_data(panel, ["Z1", "Z2", "Z3", "Z4"]), **kwargs)
+    dml_short = _fit(model_class, _panel_data(panel, ["Z2", "Z3", "Z4"]), **kwargs)
+
+    benchmark = dml_long.sensitivity_benchmark(benchmarking_set=["Z1"])
+    np.testing.assert_allclose(benchmark["delta_theta"].values, dml_short.coef - dml_long.coef, rtol=1e-9, atol=1e-12)
+    assert np.all(benchmark["delta_theta"].values != 0)