feat(continuous-did): covariate support (reg + dr) under conditional parallel trends (#616)

Author: igerber

CHANGELOG.md

@@ -8,6 +8,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- **`ContinuousDiD` covariate support** (`covariates=`, `estimation_method ∈ {"reg", "dr"}`) for
+  dose-response estimation under **conditional** parallel trends
+  (`E[ΔY(0) | D=d, X] = E[ΔY(0) | D=0, X]`). `reg` uses an outcome-regression control counterfactual;
+  `dr` (default) is doubly-robust (DRDID `drdid_panel`). The scalar `overall_att` + standard error
+  match `DRDID::reg_did_panel` / `drdid_panel` to ~1e-8; analytical, multiplier-bootstrap, and
+  event-study inference all compose with covariates. `reg` and `dr` share the dose-response *shape*
+  and `ACRT(d)`, differing only in the `overall_att` / ATT(d) level and the doubly-robust SE.
+  `estimation_method="ipw"` with covariates raises `NotImplementedError` (pure IPW's covariate
+  adjustment is a scalar level shift and cannot adjust the curve shape); `covariates=` +
+  `survey_design=` is deferred (`NotImplementedError`).
 - **`HeterogeneousAdoptionDiD` event-study `cluster=` support** (both designs). On
   `aggregate="event_study"`, `cluster=` now produces cluster-robust per-horizon pointwise
   confidence intervals AND a cluster-robust simultaneous sup-t confidence band (`cband=True`),

TODO.md

@@ -29,7 +29,7 @@ The `Origin` column (Actionable tables) and the `PR` column (Deferred tables) bo
 | Issue | Location | Origin | Effort | Priority |
 |-------|----------|--------|--------|----------|
 | `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 |
-| `ContinuousDiD` CGBS-2024 extensions (matches R `contdid` v0.1.0 deferral set): (a) `covariates=` kwarg; (b) discrete-treatment saturated regression (integer dose currently warned, not routed to per-level coefficients); (c) lowest-dose-as-control per Remark 3.1 when `P(D=0)=0`. REGISTRY `## ContinuousDiD` → Implementation Checklist marks these `[ ]`. | `continuous_did.py` | CGBS-2024 | Heavy | Low |
+| `ContinuousDiD` CGBS-2024 extensions. (a) `covariates=` kwarg — **DONE (reg/dr)**; remaining: (b) discrete-treatment saturated regression (integer dose currently warned, not routed to per-level coefficients); (c) lowest-dose-as-control per Remark 3.1 when `P(D=0)=0`. Also deferred from the covariate work: `estimation_method="ipw"` on the dose curve (scalar-adjustment / degenerate — documented `NotImplementedError`), and `covariates=` × `survey_design=` (weighted OR + weighted nuisance IF). | `continuous_did.py` | CGBS-2024 | Heavy | Low |
 | `ImputationDiD` LOO conservative-variance refinement (BJS 2024 Supp. Appendix A.9) — a finite-sample improvement to the auxiliary-model residuals reducing overfit of `tau_tilde_g` to `epsilon`. Asymptotic Theorem-3 variance is implemented and matches R `didimputation` (which also omits LOO by default). | `imputation.py` | imputation-validation | Mid | Low |
 | `TwoWayFixedEffects(vcov_type in {hc2, hc2_bm})` with replicate-weight designs raises `NotImplementedError` (`twfe.py:~233`). The replicate path re-demeans per replicate, which doesn't compose with the full-dummy HC2/HC2-BM build — a correct impl needs per-replicate full-dummy refit. Workaround: `hc1` for replicate-weight CR1. | `twfe.py::fit` | follow-up | Heavy | Low |
 | TWFE's HC2/HC2-BM inline full-dummy build (`twfe.py:280-315`) duplicates the dummy-construction logic in `DifferenceInDifferences(fixed_effects=...)` (`estimators.py:478-486`). Extract a shared helper, or delegate TWFE's HC2/HC2-BM path to DiD's `fixed_effects=` branch (with TWFE-specific cluster-default threading), to reduce drift risk on FE naming / survey behavior / result-surface conventions. Substantive refactor — touches both estimators. | `twfe.py::fit`, `estimators.py::DifferenceInDifferences.fit` | follow-up | Heavy | Low |

diff_diff/continuous_did.py

@@ -139,6 +139,14 @@ class ContinuousDiDResults:
     bootstrap_weights: str = "rademacher"
     seed: Optional[int] = None
     rank_deficient_action: str = "warn"
+    # Covariate adjustment (conditional parallel trends). ``covariates`` is None
+    # for the unconditional path; ``estimation_method`` is only meaningful when
+    # covariates are used (``"reg"`` or ``"dr"``).
+    covariates: Optional[List[str]] = field(default=None)
+    estimation_method: str = "dr"
+    pscore_trim: float = 0.01
+    epv_threshold: float = 10.0
+    pscore_fallback: str = "error"
     event_study_effects: Optional[Dict[int, Dict[str, Any]]] = field(default=None)
     # Survey design metadata (SurveyMetadata instance from diff_diff.survey)
     survey_metadata: Optional[Any] = field(default=None)
@@ -224,8 +232,11 @@ def summary(self, alpha: Optional[float] = None) -> str:
             f"{'Interior knots:':<30} {self.num_knots:>10}",
             f"{'Base period:':<30} {self.base_period:>10}",
             f"{'Anticipation:':<30} {self.anticipation:>10}",
-            "",
         ]
+        if self.covariates:
+            lines.append(f"{'Covariates:':<30} {', '.join(self.covariates):>10}")
+            lines.append(f"{'Estimation method:':<30} {self.estimation_method:>10}")
+        lines.append("")
 
         # Add survey design info
         if self.survey_metadata is not None:

diff_diff/continuous_did_results.py

@@ -696,9 +696,21 @@ ContinuousDiD(
     bootstrap_weights: str = "rademacher",
     seed: int | None = None,
     rank_deficient_action: str = "warn",
+    covariates: list[str] | None = None,      # Conditional parallel trends (X). None = unconditional
+    estimation_method: str = "dr",            # "reg" or "dr" (used only with covariates); "ipw" not
+                                              #   supported on the dose curve (raises NotImplementedError)
+    pscore_trim: float = 0.01,                # dr propensity trimming
+    epv_threshold: float = 10.0,              # dr propensity events-per-variable
+    pscore_fallback: str = "error",           # dr propensity-failure action: "error" (fail-closed
+                                              #   default) or "unconditional" (opt-in reg-like fallback)
 )
 ```
 
+Covariates give conditional parallel trends: each (g,t) cell's control counterfactual becomes a
+covariate-adjusted prediction. `reg`/`dr` share the ATT(d) *shape* and ACRT(d); `dr` (doubly-robust,
+default) differs only in the overall_att/ATT(d) level and SE. `overall_att`+SE match DRDID
+reg_did_panel/drdid_panel. `covariates=` + `survey_design=` is not yet supported.
+
 **Alias:** `CDiD`
 
 **fit() parameters:**

diff_diff/guides/llms-full.txt

@@ -889,6 +889,12 @@ No selection into dose groups on the basis of treatment effects.
 Implies `ATT(d|d) = ATT(d)` for all d.
 Additionally identifies: `ATT(d)`, `ACRT(d)`, `ACRT^{glob}`, and cross-dose comparisons.
 
+**Conditional Parallel Trends (with covariates).** When `covariates=` is passed the PT/SPT
+assumptions are conditional on covariates `X`:
+`E[Y_t(0) - Y_{t-1}(0) | D = d, X] = E[Y_t(0) - Y_{t-1}(0) | D = 0, X]`. The per-`(g,t)` cell's
+control counterfactual becomes a covariate-adjusted prediction instead of the unconditional control
+mean (see Key Equations and the covariate Note below).
+
 See `docs/methodology/continuous-did.md` Section 4 for full details.
 
 ### Key Equations
@@ -907,6 +913,17 @@ See `docs/methodology/continuous-did.md` Section 4 for full details.
 3. OLS: `beta = (Psi'Psi)^{-1} Psi' Delta_tilde_Y`
 4. `ATT(d) = Psi(d)' beta`, `ACRT(d) = dPsi(d)/dd' beta`
 
+**Covariate-adjusted `Delta_tilde_Y` (conditional PT).** With `covariates=`, step 1's scalar control
+mean is replaced by a per-treated-unit covariate-adjusted counterfactual (`X_i` includes an intercept):
+- `reg` (outcome regression): fit `gamma_hat = (X_C'X_C)^{-1} X_C' Delta_Y_C` on controls;
+  `Delta_tilde_Y_i = Delta_Y_i - X_i' gamma_hat`.
+- `dr` (doubly-robust, DRDID `drdid_panel`): same OLS `gamma_hat`, plus a propensity model and a scalar
+  augmentation `eta_cont = odds_weighted_mean_C(Delta_Y - X' gamma_hat)`;
+  `Delta_tilde_Y_i = Delta_Y_i - X_i' gamma_hat - eta_cont`.
+Steps 2-4 are unchanged. Because the augmentation is a constant, `reg` and `dr` share the same
+`ACRT(d)` (a constant only shifts the B-spline intercept, which `dPsi` annihilates); they differ only
+in the `ATT(d)` / `ATT^{glob}` level (by `-eta_cont`) and in the doubly-robust SE.
+
 ### Edge Cases
 
 - **No untreated group**: Remark 3.1 (lowest-dose-as-control) not implemented; requires P(D=0) > 0.
@@ -939,6 +956,7 @@ labels.*
 2. **Note:** `bspline_derivative_design_matrix` derivative-failure `UserWarning` — Phase 2 axis-C #12 silent-failures audit fix. No R correspondence; `contdid` v0.1.0 does not implement an equivalent warning. Cross-references the § Edge Cases `**Note:**` bullet above (`bspline_derivative_design_matrix` entry) and `METHODOLOGY_REVIEW.md` § ContinuousDiD Deviations #2. Locked in `tests/test_continuous_did.py::TestBSplineDerivativeDegenerateBasis` (3 tests); source-level aggregate-warning block at `diff_diff/continuous_did_bspline.py:150-187`.
 3. **Note:** `+inf` → `0` never-treated recoding emits `UserWarning` reporting the affected row count; negative `first_treat` (including `-inf`) raises `ValueError`. Axis-E silent-coercion fix per Phase 2 audit. No R correspondence; `contdid` v0.1.0 silently absorbs `+inf` without a signal. Cross-references the § Implementation Checklist `**Note:**` below and `METHODOLOGY_REVIEW.md` § ContinuousDiD Deviations #3.
 4. **Note:** Zero-`first_treat` rows with nonzero `dose` are force-zeroed with `UserWarning` reporting the affected row count (axis-E silent-coercion). No R correspondence; `contdid` v0.1.0 has the same `first_treat = 0` → `D = 0` invariant but silently coerces without a warning. Cross-references the § Implementation Checklist `**Note:**` below and `METHODOLOGY_REVIEW.md` § ContinuousDiD Deviations #4.
+5. **Note (covariate support — library extension beyond `contdid` v0.1.0):** `covariates=` with `estimation_method ∈ {reg, dr}` adds conditional-parallel-trends adjustment. This is a **library extension**: `contdid` v0.1.0 hard-stops on any covariate (`stop("covariates not currently supported…")`), so there is **no external R anchor for the covariate-adjusted dose *curve***. Validation instead: (a) the **scalar `overall_att` + SE** map *exactly* onto `DRDID::reg_did_panel` (reg) / `DRDID::drdid_panel` (dr) — a tight (~1e-8) component anchor, skip-guarded since DRDID is not in CI (`tests/test_methodology_continuous_did.py::TestCovariateReg`); (b) an **R-free NumPy reconstruction** of the reg/dr `att`+SE runs *in CI* at p≥2 (`test_dr_reg_numpy_crosscheck_p2`) — the guard the p=1 reduction cannot provide (at p=1 the intercept-only propensity is constant, so `eta_cont ≡ 0` and dr collapses to reg); (c) DGP recovery + MC coverage (reg 96%, dr 95%). **`ipw` restricted:** `estimation_method="ipw"` with covariates raises `NotImplementedError` — pure IPW's covariate adjustment is a single scalar (a propensity-reweighted control mean) that shifts only the ATT(d) level and leaves `ACRT(d)` identical to the unconditional fit, so it cannot adjust the dose-response *shape*. **Deviations from DRDID:** unit weights are 1 (unweighted; `covariates=` + `survey_design=` raises `NotImplementedError`, deferred); propensity trimming uses clip semantics (`pscore_trim`) rather than DRDID's drop-trimming — identical on moderate-overlap data (the anchor regime), diverging only at extreme propensities. **Fail-closed policies (no-silent-failures):** (i) missing/non-finite covariate values raise `ValueError` up front — a per-cell fallback to unconditional estimation would silently mix conditional-PT and unconditional-PT cells in the aggregate; (ii) `dr` propensity-estimation failure raises by default (`pscore_fallback="error"`) so a `dr` fit never silently degrades to a non-DR estimate — `pscore_fallback="unconditional"` opts into the graceful (warned, reg-like) fallback. Cross-references `docs/methodology/continuous-did.md` § Covariates.
 
 ### Implementation Checklist
 
@@ -948,7 +966,7 @@ labels.*
 - [x] Multiplier bootstrap for inference
 - [x] Analytical SEs via influence functions
 - [x] Equation verification tests (linear, quadratic, multi-period)
-- [ ] Covariate support (deferred, matching R v0.1.0)
+- [x] Covariate support (reg / dr) — conditional parallel trends. **ipw restricted** (see Note below); survey × covariate deferred; discrete-saturated & Remark 3.1 still deferred.
 - [ ] Discrete treatment saturated regression
 - [ ] Lowest-dose-as-control (Remark 3.1)
 - [x] Survey design support (Phase 3): weighted B-spline OLS, TSL on influence functions; bootstrap+survey supported (Phase 6)

docs/methodology/REGISTRY.md

@@ -434,7 +434,14 @@ can be meaningfully aggregated.
 ### Phase 3 (Advanced)
 8. CCK nonparametric estimation
 9. Uniform confidence bands
-10. Covariates support (DR/IPW/OR)
+10. Covariates support — **implemented** for outcome-regression (`reg`) and doubly-robust (`dr`)
+   under conditional parallel trends (`covariates=`, `estimation_method=`). Each `(g,t)` cell replaces
+   the unconditional control mean with a covariate-adjusted counterfactual (`reg`:
+   `ΔY_i − X_i'γ̂`; `dr`: additionally minus the DRDID augmentation `η̄_cont`); the B-spline dose
+   layer is unchanged. Scalar `overall_att` + SE match `DRDID::reg_did_panel` / `drdid_panel`. `ipw`
+   is not offered on the dose curve (its covariate adjustment is a scalar level shift → `ACRT(d)`
+   unchanged), and `covariates=` + `survey_design=` is deferred. See REGISTRY § ContinuousDiD
+   Note #5.
 
 ### Defer
 - Discrete treatment (saturated regression — simpler, add later)