Warn when staggered DiD ATTs lack untreated-period support (#949)

* Warn when staggered DiD ATTs lack untreated-period support

Detect calendar periods with no untreated observations, emit a clear
UserWarning, and mark non-identified ATT(g, t) and ATT(e) cells with
identified=False and NaN estimates instead of silently returning biased
values when never-treated units are absent.

Closes #938

Co-authored-by: Cursor <cursoragent@cursor.com>

* Mask non-identified cells in get_plot_data_bayesian recompute path

When hdi_prob differs from the stored aggregation value, pass the
recomputed event-time DataFrame through _mark_non_identified_att_rows()
so non-identified ATTs stay NaN. Add regression test for the non-default
HDI path on a no-never-treated panel.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Add tests for non-identified ATT coverage gaps

Cover effect_summary filtering and _mark_non_identified_att_rows edge
cases so Codecov patch coverage passes on the identification changes.

Related: #950 (local patch coverage in prek)
Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: drbenvincent

causalpy/experiments/staggered_did.py

@@ -19,6 +19,7 @@
 different units receive treatment at different times.
 """
 
+import warnings
 from typing import Any, Literal
 
 import numpy as np
@@ -58,6 +59,13 @@ class StaggeredDifferenceInDifferences(BaseExperiment):
        units would have followed parallel outcome trajectories.
     3. **No anticipation**: Units do not change their behavior in anticipation
        of future treatment.
+    4. **Untreated support at each calendar period**: The time fixed effect
+       :math:`\\gamma_t` for calendar period :math:`t` is identified only if at
+       least one unit is untreated in that period. Without never-treated units,
+       post-treatment effects for the last-treated cohort (and any calendar
+       periods where every unit is already treated) are not identified. CausalPy
+       warns when this condition fails and marks the affected ``ATT(g, t)`` and
+       ``ATT(e)`` cells as non-identified in the output tables.
 
     Parameters
     ----------
@@ -94,8 +102,14 @@ class StaggeredDifferenceInDifferences(BaseExperiment):
         and tau_hat (treatment effect) columns.
     att_group_time_ : pd.DataFrame
         Group-time ATT estimates: ATT(g, t) for each cohort g and calendar time t.
+        Includes an ``identified`` column; non-identified cells have ``NaN`` estimates.
     att_event_time_ : pd.DataFrame
         Event-time ATT estimates: ATT(e) for each event-time e = t - G.
+        Includes an ``identified`` column; non-identified cells have ``NaN`` estimates.
+    non_identified_periods_ : set
+        Calendar periods with no untreated observations.
+    non_identified_cohorts_ : set
+        Treatment cohorts with at least one non-identified post-treatment ATT(g, t).
 
     Notes
     -----
@@ -185,6 +199,9 @@ def __init__(
         # Step 3: Identify untreated observations (training set)
         self._identify_untreated_observations()
 
+        # Step 3b: Check calendar-period identification support
+        self._check_att_identification()
+
         # Step 4: Build design matrices
         self._build_design_matrices()
 
@@ -319,6 +336,92 @@ def _identify_untreated_observations(self) -> None:
                 "Ensure there are never-treated units or pre-treatment periods."
             )
 
+    def _get_periods_without_untreated_support(self) -> set[Any]:
+        """Return calendar periods with zero untreated observations."""
+        untreated_periods = set(
+            self.data.loc[self.data["_is_untreated"], self.time_variable_name].unique()
+        )
+        all_periods = set(self.data[self.time_variable_name].unique())
+        return all_periods - untreated_periods
+
+    def _get_non_identified_cohorts(self, periods: set[Any]) -> set[Any]:
+        """Return cohorts with post-treatment cells in non-identified periods."""
+        non_identified_cohorts: set[Any] = set()
+        for cohort in self.cohorts:
+            for period in periods:
+                if period >= cohort:
+                    non_identified_cohorts.add(cohort)
+                    break
+        return non_identified_cohorts
+
+    def _check_att_identification(self) -> None:
+        """Detect non-identified ATT cells and warn when untreated support is missing."""
+        self.non_identified_periods_ = self._get_periods_without_untreated_support()
+        self.non_identified_cohorts_ = self._get_non_identified_cohorts(
+            self.non_identified_periods_
+        )
+
+        if not self.non_identified_periods_:
+            return
+
+        periods_str = ", ".join(str(p) for p in sorted(self.non_identified_periods_))
+        cohorts_str = ", ".join(str(c) for c in sorted(self.non_identified_cohorts_))
+        warnings.warn(
+            "No untreated observations in calendar period(s) "
+            f"{{{periods_str}}}; treatment effects for cohort(s) "
+            f"{{{cohorts_str}}} are not identified at the affected post-treatment "
+            "cells. Provide never-treated units or restrict the event window. "
+            "Non-identified ATT(g, t) and ATT(e) cells are marked in the output "
+            "tables (identified=False) with NaN estimates.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    def _is_calendar_period_identified(self, period: Any) -> bool:
+        """Return whether calendar period ``period`` has untreated support."""
+        return period not in self.non_identified_periods_
+
+    def _is_event_time_att_identified(self, event_time: int) -> bool:
+        """Return whether aggregated ATT(e) is identified."""
+        for cohort in self.cohorts:
+            period = cohort + event_time
+            has_contributing_obs = (
+                (self.data["G"] == cohort)
+                & (self.data[self.time_variable_name] == period)
+                & (self.data["event_time"] == event_time)
+            ).any()
+            if has_contributing_obs and not self._is_calendar_period_identified(period):
+                return False
+        return True
+
+    def _mark_non_identified_att_rows(self, att_df: pd.DataFrame) -> pd.DataFrame:
+        """Add ``identified`` column and mask non-identified point estimates."""
+        if len(att_df) == 0:
+            att_df = att_df.copy()
+            att_df["identified"] = pd.Series(dtype=bool)
+            return att_df
+
+        att_df = att_df.copy()
+        if "cohort" in att_df.columns and "time" in att_df.columns:
+            att_df["identified"] = att_df["time"].map(
+                self._is_calendar_period_identified
+            )
+        elif "event_time" in att_df.columns:
+            att_df["identified"] = att_df["event_time"].apply(
+                lambda e: self._is_event_time_att_identified(int(e))
+            )
+        else:
+            att_df["identified"] = True
+
+        value_columns = [
+            col
+            for col in ("att", "att_lower", "att_upper", "att_std")
+            if col in att_df.columns
+        ]
+        for col in value_columns:
+            att_df.loc[~att_df["identified"], col] = np.nan
+        return att_df
+
     def _build_design_matrices(self) -> None:
         """Build design matrices using patsy."""
         # Build design matrix for the full data
@@ -499,7 +602,9 @@ def _aggregate_effects_bayesian(
                     "att_upper": float(np.percentile(tau_gt, upper_pct)),
                 }
             )
-        self.att_group_time_ = pd.DataFrame(att_gt_rows)
+        self.att_group_time_ = self._mark_non_identified_att_rows(
+            pd.DataFrame(att_gt_rows)
+        )
 
         # --- Event-time ATTs (including pre-treatment placebo) ---
         att_et_rows: list[dict] = []
@@ -563,7 +668,9 @@ def _aggregate_effects_bayesian(
                 }
             )
 
-        self.att_event_time_ = pd.DataFrame(att_et_rows)
+        self.att_event_time_ = self._mark_non_identified_att_rows(
+            pd.DataFrame(att_et_rows)
+        )
 
     def _aggregate_effects_ols(
         self, treated_data: pd.DataFrame, pretreatment_data: pd.DataFrame
@@ -585,7 +692,7 @@ def _aggregate_effects_ols(
             .reset_index()
         )
         att_gt.columns = ["cohort", "time", "att", "att_std", "n_obs"]
-        self.att_group_time_ = att_gt
+        self.att_group_time_ = self._mark_non_identified_att_rows(att_gt)
 
         # --- Event-time ATTs (including pre-treatment placebo) ---
         # Compute tau_hat for pre-treatment observations (residuals)
@@ -613,7 +720,7 @@ def _aggregate_effects_ols(
         )
         att_et.columns = ["event_time", "att", "att_std", "n_obs"]
         att_et["event_time"] = att_et["event_time"].astype(int)
-        self.att_event_time_ = att_et
+        self.att_event_time_ = self._mark_non_identified_att_rows(att_et)
 
     def summary(
         self, round_to: int | None = 2, include_group_time: bool = False
@@ -1506,7 +1613,7 @@ def get_plot_data_bayesian(self, hdi_prob: float = HDI_PROB) -> pd.DataFrame:
                 }
             )
 
-        return pd.DataFrame(att_et_rows)
+        return self._mark_non_identified_att_rows(pd.DataFrame(att_et_rows))
 
     def get_plot_data_ols(self) -> pd.DataFrame:
         """Get plotting data for OLS model.

causalpy/reporting.py

@@ -382,6 +382,10 @@ def _effect_summary_staggered_did(
     # Separate pre-treatment (placebo) and post-treatment effects
     pre_treatment = att_et[att_et["event_time"] < 0]
     post_treatment = att_et[att_et["event_time"] >= 0]
+    if "identified" in post_treatment.columns:
+        post_treatment = post_treatment[post_treatment["identified"]]
+    if "identified" in pre_treatment.columns:
+        pre_treatment = pre_treatment[pre_treatment["identified"]]
 
     # Build summary table with all event-time effects
     table = att_et.copy()

causalpy/tests/test_staggered_did.py

@@ -15,6 +15,8 @@
 Tests for StaggeredDifferenceInDifferences experiment class.
 """
 
+import warnings
+
 import numpy as np
 import pandas as pd
 import pytest
@@ -703,6 +705,183 @@ def test_staggered_did_group_time_att_structure():
     assert len(result.att_group_time_["cohort"].unique()) >= 2
 
 
+def _no_never_treated_staggered_did_df() -> pd.DataFrame:
+    """Return a staggered panel where every unit is eventually treated."""
+    return generate_staggered_did_data(
+        n_units=20,
+        n_time_periods=10,
+        treatment_cohorts={3: 10, 7: 10},
+        seed=42,
+    )
+
+
+@pytest.mark.parametrize(
+    "model_factory",
+    [
+        pytest.param(lambda: LinearRegression(), id="ols"),
+        pytest.param(
+            lambda: cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+            id="pymc",
+        ),
+    ],
+)
+def test_staggered_did_warns_non_identified_without_never_treated(
+    model_factory, mock_pymc_sample
+):
+    """No never-treated units should warn about non-identified post-treatment ATTs."""
+    df = _no_never_treated_staggered_did_df()
+
+    with pytest.warns(
+        UserWarning, match="No untreated observations in calendar period"
+    ):
+        result = cp.StaggeredDifferenceInDifferences(
+            df,
+            formula="y ~ 1 + C(unit) + C(time)",
+            unit_variable_name="unit",
+            time_variable_name="time",
+            treated_variable_name="treated",
+            treatment_time_variable_name="treatment_time",
+            model=model_factory(),
+        )
+
+    assert result.non_identified_periods_
+    assert 7 in result.non_identified_periods_
+    assert 7 in result.non_identified_cohorts_
+
+    assert "identified" in result.att_group_time_.columns
+    non_identified_gt = result.att_group_time_[~result.att_group_time_["identified"]]
+    assert len(non_identified_gt) > 0
+    assert non_identified_gt["att"].isna().all()
+
+    assert "identified" in result.att_event_time_.columns
+    non_identified_et = result.att_event_time_[
+        (result.att_event_time_["event_time"] >= 0)
+        & ~result.att_event_time_["identified"]
+    ]
+    assert len(non_identified_et) > 0
+    assert non_identified_et["att"].isna().all()
+
+
+def test_staggered_did_get_plot_data_bayesian_masks_non_identified_on_recompute(
+    mock_pymc_sample,
+):
+    """Non-default hdi_prob recompute path must still mask non-identified cells."""
+    df = _no_never_treated_staggered_did_df()
+
+    with pytest.warns(
+        UserWarning, match="No untreated observations in calendar period"
+    ):
+        result = cp.StaggeredDifferenceInDifferences(
+            df,
+            formula="y ~ 1 + C(unit) + C(time)",
+            unit_variable_name="unit",
+            time_variable_name="time",
+            treated_variable_name="treated",
+            treatment_time_variable_name="treatment_time",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        )
+
+    plot_data = result.get_plot_data_bayesian(hdi_prob=0.80)
+
+    assert "identified" in plot_data.columns
+    non_identified_post = plot_data[
+        (plot_data["event_time"] >= 0) & ~plot_data["identified"]
+    ]
+    assert len(non_identified_post) > 0
+    assert non_identified_post["att"].isna().all()
+    assert non_identified_post["att_lower"].isna().all()
+    assert non_identified_post["att_upper"].isna().all()
+
+
+def test_staggered_did_effect_summary_excludes_non_identified_cells():
+    """effect_summary prose should average only identified post-treatment ATTs."""
+    df = _no_never_treated_staggered_did_df()
+
+    with pytest.warns(
+        UserWarning, match="No untreated observations in calendar period"
+    ):
+        result = cp.StaggeredDifferenceInDifferences(
+            df,
+            formula="y ~ 1 + C(unit) + C(time)",
+            unit_variable_name="unit",
+            time_variable_name="time",
+            treated_variable_name="treated",
+            treatment_time_variable_name="treatment_time",
+            model=LinearRegression(),
+        )
+
+    post_treatment = result.att_event_time_[result.att_event_time_["event_time"] >= 0]
+    assert (~post_treatment["identified"]).any()
+
+    summary = result.effect_summary()
+    assert "Staggered DiD" in summary.text
+    assert isinstance(summary.table, pd.DataFrame)
+
+
+def test_staggered_did_mark_non_identified_att_rows_edge_cases():
+    """Cover empty and unknown-column paths in _mark_non_identified_att_rows."""
+    df = _no_never_treated_staggered_did_df()
+
+    with pytest.warns(
+        UserWarning, match="No untreated observations in calendar period"
+    ):
+        result = cp.StaggeredDifferenceInDifferences(
+            df,
+            formula="y ~ 1 + C(unit) + C(time)",
+            unit_variable_name="unit",
+            time_variable_name="time",
+            treated_variable_name="treated",
+            treatment_time_variable_name="treatment_time",
+            model=LinearRegression(),
+        )
+
+    empty = result._mark_non_identified_att_rows(pd.DataFrame())
+    assert list(empty.columns) == ["identified"]
+    assert len(empty) == 0
+
+    unknown_cols = result._mark_non_identified_att_rows(pd.DataFrame({"x": [1]}))
+    assert unknown_cols["identified"].all()
+
+
+@pytest.mark.parametrize(
+    "model_factory",
+    [
+        pytest.param(lambda: LinearRegression(), id="ols"),
+        pytest.param(
+            lambda: cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+            id="pymc",
+        ),
+    ],
+)
+def test_staggered_did_all_identified_with_never_treated(
+    model_factory, mock_pymc_sample
+):
+    """Never-treated units should identify all post-treatment ATT cells."""
+    df = generate_staggered_did_data(
+        n_units=30,
+        n_time_periods=10,
+        treatment_cohorts={3: 10, 7: 10},
+        seed=42,
+    )
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", UserWarning)
+        result = cp.StaggeredDifferenceInDifferences(
+            df,
+            formula="y ~ 1 + C(unit) + C(time)",
+            unit_variable_name="unit",
+            time_variable_name="time",
+            treated_variable_name="treated",
+            treatment_time_variable_name="treatment_time",
+            model=model_factory(),
+        )
+
+    assert result.non_identified_periods_ == set()
+    assert result.att_group_time_["identified"].all()
+    post_treatment = result.att_event_time_[result.att_event_time_["event_time"] >= 0]
+    assert post_treatment["identified"].all()
+
+
 def test_staggered_did_no_untreated_observations():
     """Test that having no untreated observations raises DataException."""
     # All units treated from time 0