fix(staggered): uniform no-IF NaN cells + cover covariate paths (review #582)

igerber · claude · igerber · commit fa149d55eabc · 2026-06-29T09:12:35.000-04:00
- The covariate-regression non-finite cell now materializes via _nan_gt_entry
  with NO influence-function entry, matching the other paths and the documented
  REGISTRY/helper contract (previously it wrote a zero-IF entry and ran batch
  inference). Aggregates and SEs are unchanged -- the cell was finite-masked /
  IF-membership-filtered out either way; now the "NaN cells carry no IF entry"
  invariant holds uniformly across all paths.
- Extend the materialization test to cover covariate IPW/DR (panel + RCS) paths.
- Remove the now-implemented CallawaySantAnna NaN-cell row from TODO.md.

Co-Authored-By: Claude Opus 4.8 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/TODO.md b/TODO.md
@@ -29,7 +29,6 @@ The `Origin` column (Actionable tables) and the `PR` column (Deferred tables) bo
 | Issue | Location | Origin | Effort | Priority |
 |-------|----------|--------|--------|----------|
 | `SyntheticControl` cv: thread an `"infeasible"` reason-code from `_outer_solve_V_cv()` / `_placebo_fit_unit()` so `in_space_placebo()` / `leave_one_out()` distinguish a structural cv-refit exclusion (donor-indistinguishable re-aggregated window) from a genuine inner-solver non-convergence — mirror the split `in_time_placebo()` already emits. Warnings already distinguish the two causes; only the machine-readable status/count is missing. | `synthetic_control.py`, `synthetic_control_results.py` | follow-up | Mid | Low |
-| `CallawaySantAnna`: materialize NaN entries for non-estimable `(g,t)` cells in `group_time_effects` (currently omitted with a consolidated warning); requires updating downstream consumers (event study, `balance_e`, aggregation). | `staggered.py` | #256 | Mid | Low |
 | Survey-design resolution / collapse patterns are inconsistent across panel estimators — `ContinuousDiD` rebuilds unit-level design in SE code, `EfficientDiD` builds once in `fit()`, `StackedDiD` re-resolves on stacked data. Extract shared helpers for panel-to-unit collapse, post-filter re-resolution, and metadata recomputation. | `continuous_did.py`, `efficient_did.py`, `stacked_did.py` | #226 | Mid | Low |
 | `SyntheticControl` remaining ADH-2015 §4 items: the regression-weight `W^reg = X_0'(X_0 X_0')^{-1} X_1` extrapolation diagnostic (flag implied OLS weights outside `[0,1]`) and sparse-SC subset search (`l < J`, holding `V` fixed). LOO, in-time placebo, CV `V`-selection, and inverse-variance `V` have landed; these two are the deferred tail. | `synthetic_control.py`, `synthetic_control_results.py` | ADH-2015 | Mid | Low |
 | `SyntheticControl` conformal (CWZ 2021) extensions: (a) one-sided / signed-`t` variants (§7); (b) covariates in the conformal proxy (`X_jt`, eqs 4/6 — current proxy is outcomes-only); (c) AR / innovation-permutation path (Lemmas 5-7) for time-series proxies. The joint test, pointwise CIs, and average-effect CI have landed. | `conformal.py`, `synthetic_control_results.py` | CWZ-2021 | Heavy | Low |
diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py
@@ -1522,16 +1522,26 @@ def _compute_all_att_gt_covariate_reg(
 
                     if nan_cell:
                         att = np.nan
-                        se = np.nan
-                        inf_treated = np.zeros(n_t)
-                        inf_control = np.zeros(n_c)
                     else:
                         var_t = float(np.var(treated_residuals, ddof=1)) if n_t > 1 else 0.0
                         var_c = float(np.var(residuals, ddof=1)) if pair_n_c > 1 else 0.0
                         se = float(np.sqrt(var_t / n_t + var_c / pair_n_c))
                         inf_treated = (treated_residuals - np.mean(treated_residuals)) / n_t
                         inf_control = -residuals / pair_n_c
 
+                # Non-estimable (non-finite regression) cell: materialize NaN with
+                # NO influence-function entry — uniform with the missing-period /
+                # zero-control / zero-weight paths, so every NaN cell is excluded
+                # from aggregation and the bootstrap (the IF-membership filter and
+                # the finite-mask both drop it). Skip batch inference (already NaN).
+                if not np.isfinite(att):
+                    group_time_effects[(g, t)] = _nan_gt_entry(
+                        n_treated=n_t,
+                        n_control=n_c,
+                        skip_reason="non_finite_regression",
+                    )
+                    continue
+
                 group_time_effects[(g, t)] = {
                     "effect": att,
                     "se": se,
@@ -1540,7 +1550,7 @@ def _compute_all_att_gt_covariate_reg(
                     "conf_int": (np.nan, np.nan),
                     "n_treated": n_t,
                     "n_control": n_c,
-                    "skip_reason": None if np.isfinite(att) else "non_finite_regression",
+                    "skip_reason": None,
                 }
 
                 all_units = precomputed["all_units"]
diff --git a/tests/test_staggered.py b/tests/test_staggered.py
@@ -1298,6 +1298,7 @@ def _cs_nonestimable_data(panel: bool, n: int = 25, seed: int = 0) -> pd.DataFra
     uid = 0
     for g in (2, 3, 4):
         for _ in range(n):
+            x = rng.normal(0, 1)  # unit-level covariate (for IPW/DR/reg covariate paths)
             if panel:
                 fe = rng.normal(0, 1)
                 for t in range(1, 5):
@@ -1306,20 +1307,23 @@ def _cs_nonestimable_data(panel: bool, n: int = 25, seed: int = 0) -> pd.DataFra
                         {
                             "unit": uid,
                             "time": t,
-                            "outcome": fe + 0.3 * t + 1.5 * post + rng.normal(0, 0.5),
+                            "outcome": fe + 0.3 * t + 1.5 * post + 0.5 * x + rng.normal(0, 0.5),
                             "first_treat": g,
+                            "x": x,
                         }
                     )
                 uid += 1
             else:
                 for t in range(1, 5):
                     post = 1.0 if t >= g else 0.0
+                    x_t = rng.normal(0, 1)
                     rows.append(
                         {
                             "unit": uid,
                             "time": t,
-                            "outcome": rng.normal(0, 1) + 0.3 * t + 1.5 * post,
+                            "outcome": rng.normal(0, 1) + 0.3 * t + 1.5 * post + 0.5 * x_t,
                             "first_treat": g,
+                            "x": x_t,
                         }
                     )
                     uid += 1
@@ -1341,10 +1345,19 @@ class TestCallawaySantAnnaNonEstimableMaterialization:
     }
 
     @pytest.mark.parametrize(
-        "method,panel",
-        [("reg", True), ("ipw", True), ("dr", True), ("reg", False)],
+        "method,panel,covariates",
+        [
+            ("reg", True, None),  # no-covariate vectorized path
+            ("ipw", True, None),  # general path, no covariates
+            ("dr", True, None),  # general path, no covariates
+            ("reg", False, None),  # repeated cross-section path
+            ("reg", True, ["x"]),  # covariate-regression vectorized path
+            ("ipw", True, ["x"]),  # general path, covariate IPW
+            ("dr", True, ["x"]),  # general path, covariate DR
+            ("dr", False, ["x"]),  # repeated cross-section, covariate DR
+        ],
     )
-    def test_materializes_nan_cell_with_skip_reason(self, method, panel):
+    def test_materializes_nan_cell_with_skip_reason(self, method, panel, covariates):
         """Each previously-omitting path now stores the non-estimable cell as NaN."""
         data = _cs_nonestimable_data(panel=panel, seed=0)
         cs = CallawaySantAnna(
@@ -1356,7 +1369,12 @@ def test_materializes_nan_cell_with_skip_reason(self, method, panel):
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
             results = cs.fit(
-                data, outcome="outcome", unit="unit", time="time", first_treat="first_treat"
+                data,
+                outcome="outcome",
+                unit="unit",
+                time="time",
+                first_treat="first_treat",
+                covariates=covariates,
             )
 
         # The (g=4, t=4) cell has no not-yet-treated controls -> materialized, not omitted.
