From 2003ff0dbac46f7aa04a89577c54a8f6f74592ed Mon Sep 17 00:00:00 2001 From: Asad Gudzevic Date: Mon, 2 Nov 2026 16:15:41 +0100 Subject: [PATCH 1/2] Keep never-treated units in make_did_CS2021 with float time and anticipation --- doubleml/did/datasets/dgp_did_CS2021.py | 7 +++---- doubleml/did/tests/test_datasets.py | 9 +++++++++ 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/doubleml/did/datasets/dgp_did_CS2021.py b/doubleml/did/datasets/dgp_did_CS2021.py index 1756d085..dae11483 100644 --- a/doubleml/did/datasets/dgp_did_CS2021.py +++ b/doubleml/did/datasets/dgp_did_CS2021.py @@ -313,16 +313,15 @@ def make_did_CS2021(n_obs=1000, dgp_type=1, include_never_treated=True, time_typ if anticipation_periods > 0: # filter time periods df = df[df["t"] >= time_periods[anticipation_periods]] - # filter treatment after anticipation periods - df = df[(df["d"] <= time_periods[-(anticipation_periods + 1)]) | pd.isna(df["d"])] + # filter treatment after anticipation periods (never treated: NaT for datetime, inf for float) + never_treated = pd.isna(df["d"]) if time_type == "datetime" else np.isinf(df["d"]) + df = df[(df["d"] <= time_periods[-(anticipation_periods + 1)]) | never_treated] # update time periods by subtracting time delta if time_type == "datetime": - df = df[(df["d"] <= time_periods[-(anticipation_periods + 1)]) | pd.isna(df["d"])] df["t"] = df["t"].apply(lambda x: x - pd.DateOffset(months=anticipation_periods)) else: assert time_type == "float" - df = df[(df["d"] <= time_periods[-(anticipation_periods + 1)]) | np.isinf(df["d"])] df["t"] = df["t"] - anticipation_periods return df diff --git a/doubleml/did/tests/test_datasets.py b/doubleml/did/tests/test_datasets.py index 8e079e9a..4e94c9fd 100644 --- a/doubleml/did/tests/test_datasets.py +++ b/doubleml/did/tests/test_datasets.py @@ -72,6 +72,15 @@ def test_make_did_CS2021_return_types(dgp_type, include_never_treated, time_type assert isinstance(df, pd.DataFrame) +@pytest.mark.ci +@pytest.mark.parametrize("anticipation_periods", [0, 1, 2]) +def test_make_did_CS2021_keeps_never_treated(time_type, anticipation_periods): + np.random.seed(3141) + df = make_did_CS2021(n_obs=500, time_type=time_type, anticipation_periods=anticipation_periods) + never_treated = pd.isna(df["d"]) if time_type == "datetime" else np.isinf(df["d"]) + assert never_treated.any() + + @pytest.mark.ci def test_make_did_CS2021_exceptions(): msg = r"time_type must be one of \('datetime', 'float'\). Got 2." From 07a91848c1b057ab7b9cbe39cf9c692ed1374b23 Mon Sep 17 00:00:00 2001 From: Sven Klaassen <47529404+SvenKlaassen@users.noreply.github.com> Date: Mon, 2 Nov 2026 19:14:55 +0100 Subject: [PATCH 2/2] update time type test --- doubleml/did/tests/test_datasets.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/doubleml/did/tests/test_datasets.py b/doubleml/did/tests/test_datasets.py index 4e94c9fd..3198a8ba 100644 --- a/doubleml/did/tests/test_datasets.py +++ b/doubleml/did/tests/test_datasets.py @@ -74,10 +74,16 @@ def test_make_did_CS2021_return_types(dgp_type, include_never_treated, time_type @pytest.mark.ci @pytest.mark.parametrize("anticipation_periods", [0, 1, 2]) -def test_make_did_CS2021_keeps_never_treated(time_type, anticipation_periods): +def test_make_did_CS2021_time_types_equivalent(anticipation_periods): + """Datetime and float time types simulate the same panel, including all never-treated units.""" np.random.seed(3141) - df = make_did_CS2021(n_obs=500, time_type=time_type, anticipation_periods=anticipation_periods) - never_treated = pd.isna(df["d"]) if time_type == "datetime" else np.isinf(df["d"]) + df_dt = make_did_CS2021(n_obs=200, time_type="datetime", anticipation_periods=anticipation_periods) + np.random.seed(3141) + df_fl = make_did_CS2021(n_obs=200, time_type="float", anticipation_periods=anticipation_periods) + + pd.testing.assert_frame_equal(df_dt.drop(columns=["d", "t"]), df_fl.drop(columns=["d", "t"])) + never_treated = np.isinf(df_fl["d"]) + np.testing.assert_array_equal(pd.isna(df_dt["d"]), never_treated) assert never_treated.any()