add universal option to gt_combinations

Author: SvenKlaassen

doubleml/did/tests/test_did_multi_exceptions.py

@@ -100,7 +100,7 @@ def test_exception_learners():
 
 @pytest.mark.ci
 def test_exception_gt_combinations():
-    msg = r"gt_combinations must be one of \['standard', 'all'\]. test was passed."
+    msg = r"gt_combinations must be one of \['standard', 'all', 'universal'\]. test was passed."
     with pytest.raises(ValueError, match=msg):
         invalid_arguments = {"gt_combinations": "test"}
         _ = dml.did.DoubleMLDIDMulti(**(valid_arguments | invalid_arguments))

doubleml/did/tests/test_did_multi_plot.py

@@ -23,8 +23,13 @@ def n_rep(request):
     return request.param
 
 
+@pytest.fixture(scope="module", params=["standard", "all", "universal"])
+def gt_comb(request):
+    return request.param
+
+
 @pytest.fixture(scope="module")
-def doubleml_did_fixture(did_score, panel, n_rep):
+def doubleml_did_fixture(did_score, panel, n_rep, gt_comb):
     n_obs = 1000
     dgp = 5  # has to be experimental (for experimental score to be valid)
     np.random.seed(42)
@@ -35,7 +40,7 @@ def doubleml_did_fixture(did_score, panel, n_rep):
         "obj_dml_data": dml_data,
         "ml_g": LinearRegression(),
         "ml_m": LogisticRegression(),
-        "gt_combinations": "all",
+        "gt_combinations": gt_comb,
         "score": did_score,
         "panel": panel,
         "n_rep": n_rep,

doubleml/did/utils/_did_utils.py

@@ -84,12 +84,6 @@ def _check_gt_combination(gt_combination, g_values, t_values, never_treated_valu
     if t_value_pre == t_value_eval:
         raise ValueError(f"The pre-treatment and evaluation period must be different. Got {t_value_pre} for both.")
 
-    if t_value_pre > t_value_eval:
-        raise ValueError(
-            "The pre-treatment period must be before the evaluation period. "
-            f"Got t_value_pre {t_value_pre} and t_value_eval {t_value_eval}."
-        )
-
     # get t_value equal to g_value and adjust for anticipation periods
     maximal_t_pre = t_values[max(np.where(t_values == g_value)[0] - anticipation_periods, 0)]
     if t_value_pre >= maximal_t_pre:
@@ -128,7 +122,7 @@ def _construct_gt_combinations(setting, g_values, t_values, never_treated_value,
     Returns:
         list: List of (g_val, t_pre, t_eval) tuples
     """
-    valid_settings = ["standard", "all"]
+    valid_settings = ["standard", "all", "universal"]
     if setting not in valid_settings:
         raise ValueError(f"gt_combinations must be one of {valid_settings}. {setting} was passed.")
 
@@ -163,6 +157,15 @@ def _construct_gt_combinations(setting, g_values, t_values, never_treated_value,
                     for t_pre in valid_t_pre_values:
                         gt_combinations.append((g_val, t_pre, t_eval))
 
+    if setting == "universal":
+        for g_val in treatment_groups:
+            t_values_before_g = t_values[t_values < g_val]
+            if len(t_values_before_g) > anticipation_periods:
+                base_period = g_val - anticipation_periods - 1
+                for t_eval in t_values:
+                    if t_eval != base_period:
+                        gt_combinations.append((g_val, base_period, t_eval))
+
     if len(gt_combinations) == 0:
         raise ValueError(
             "No valid group-time combinations found. "

doubleml/did/utils/tests/test_did_utils.py

@@ -99,11 +99,6 @@ def test_check_gt_combination():
             ValueError,
             "The pre-treatment and evaluation period must be different. Got 1 for both.",
         ),
-        (
-            {"gt_combination": (1, 1, 0)},
-            ValueError,
-            "The pre-treatment period must be before the evaluation period. Got t_value_pre 1 and t_value_eval 0.",
-        ),
         (
             {"gt_combination": (-1, 0, 1)},
             ValueError,
@@ -171,7 +166,7 @@ def test_input_check_gt_values():
 
 @pytest.mark.ci
 def test_construct_gt_combinations():
-    msg = r"gt_combinations must be one of \['standard', 'all'\]. test was passed."
+    msg = r"gt_combinations must be one of \['standard', 'all', 'universal'\]. test was passed."
     with pytest.raises(ValueError, match=msg):
         _construct_gt_combinations(
             setting="test",
@@ -256,6 +251,24 @@ def test_construct_gt_combinations():
     ]
     assert all_combinations == expected_all
 
+    # Test universal setting
+    universal_combinations = _construct_gt_combinations(
+        setting="universal",
+        g_values=np.array([2, 3]),
+        t_values=np.array([0, 1, 2, 3]),
+        never_treated_value=np.inf,
+        anticipation_periods=0,
+    )
+    expected_universal = [
+        (2, 1, 0),  # g=2, pre=1, eval=0
+        (2, 1, 2),  # g=2, pre=1, eval=2
+        (2, 1, 3),  # g=2, pre=1, eval=3
+        (3, 2, 0),  # g=3, pre=2, eval=0
+        (3, 2, 1),  # g=3, pre=2, eval=1
+        (3, 2, 3),  # g=3, pre=2, eval=3
+    ]
+    assert universal_combinations == expected_universal
+
     # Test standard setting with anticipation periods
     standard_combinations_anticipation = _construct_gt_combinations(
         setting="standard",
@@ -282,6 +295,21 @@ def test_construct_gt_combinations():
     ]
     assert all_combinations_anticipation == expected_all_anticipation
 
+    # Test universal setting with anticipation periods
+    universal_combinations_anticipation = _construct_gt_combinations(
+        setting="universal",
+        g_values=np.array([2, 3]),
+        t_values=np.array([0, 1, 2, 3]),
+        never_treated_value=np.inf,
+        anticipation_periods=2,
+    )
+    expected_universal_anticipation = [
+        (3, 0, 1),  # g=3, pre=0, eval=1 with anticipation 2
+        (3, 0, 2),
+        (3, 0, 3),
+    ]
+    assert universal_combinations_anticipation == expected_universal_anticipation
+
 
 @pytest.mark.ci
 def test_construct_gt_index():