Adding commits for ETWFE extension to staggered DiD - #1083
Conversation
Signed-off-by: Nathaniel <NathanielF@users.noreply.github.com>
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
# Conflicts: # causalpy/experiments/staggered_did.py # causalpy/maketables_adapters.py # causalpy/pymc_models.py # causalpy/tests/test_pymc_models.py # docs/source/notebooks/index.md # docs/source/references.bib
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1083 +/- ##
==========================================
+ Coverage 95.98% 96.17% +0.19%
==========================================
Files 104 106 +2
Lines 16269 18081 +1812
Branches 912 1074 +162
==========================================
+ Hits 15615 17389 +1774
- Misses 488 507 +19
- Partials 166 185 +19 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
The ETWFE estimator was written before causalpy/experiments/model_adapter.py existed, so it dispatched on isinstance(self.model, ...) and called model.fit/predict directly. ARCHITECTURE.md is explicit that backend dispatch is centralized in the adapter, that experiments call self._model_backend.fit(...) unconditionally, and that capability is discovered via ModelAdapter rather than by probing for attributes. - Widen ModelAdapter.fit with a **fit_kwargs passthrough so models whose fit needs more than X/y (ETWFERegression's panel index arrays) can still be reached through the adapter instead of past it. - Route _fit_model_etwfe through the backend for fit, predict and coefficients; drop the hand-built canonical prediction array in favour of the adapter's. - Replace all 12 isinstance(self.model, ...) checks in ETWFE methods with the adapter's is_bayesian / is_ols properties. The one remaining isinstance is a deliberate concrete-class check (warn when the OLS sandwich is applied to a non-plain-OLS fitter). - _etwfe_idata now delegates to require_idata() rather than getattr(self.model, 'idata', None). Suite 1369 passed, 6 skipped; Tier-1 exactness still recovers tau[g,k] and the ATT to 1e-14. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Per-commit hooks only lint changed files, so these sat unnoticed: - Notebook: C408 (dict() -> literal, x2) and B905 (zip without strict=). Main removed the ^docs/ exclusion from the ruff hook, so notebooks are now linted. Notebook re-executed; all reported numbers unchanged. - ETWFERegression class docstring: dropped the Parameters section (numpydoc validates class params against the signature, and the constructor is inherited from PyMCModel -> PR02), folding the prior-key documentation into Notes; renamed 'Example' to 'Examples' so numpydoc recognises the section and GL07 ordering holds. - ETWFERegression.fit: gave X, y and coords real descriptions (PR07). - ETWFERegression.predict: documented **kwargs (PR01). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
daimon-pymclabs
left a comment
daimon-pymclabs
left a comment
There was a problem hiding this comment.
Review
I reviewed this across staggered_did.py, pymc_models.py, reporting.py, the data generator, tests, and the adapters. No MAJOR correctness bugs. Summary below, with a couple of things worth addressing before merge (none blocking).
The ATT aggregation change is correct
The imputation ATT correction lives in reporting.py (_treated_weighted_mean, reporting.py:350-388). n_obs per event-time cell is int(e_mask.sum()), the count of treated observations at that event time. Because each per-event-time att is itself the mean of tau over the observations in that cell, an n_obs-weighted mean over event times collapses exactly to the simple mean of tau over all treated observations, i.e. the Callaway-Sant'Anna / Wooldridge "average over the treated" estimand. Pre- and post-treatment rows are split and weighted separately, so placebo rows never leak into the headline. The old unweighted mean was biased upward under staggered adoption exactly as the PR description states.
The three backends are internally consistent: the ETWFE in-model att = (tau * att_weights).sum() uses treated-cell shares N_gk / ΣN_gk, and the OLS scalar ATT uses the same weights.
What checks out
- Indexing: the non-negative-index-plus-
effect_indicatorscheme avoids-1aliasing onto a real lead column. Top-binning, lead-window/reference handling, and unbalanced/non-contiguous panel re-indexing are all sound. _ols_error_bar_seis a genuine fix: ETWFE'satt_stdis already√(w'Vw)and must not be divided by√n_obsthe way the imputation sample-SD is.- Backward-compat for the imputation path is explicitly guarded (
_validate_imputation_argumentshard-errors if any ETWFE-only arg is set). The only behavior change for existing users is the summary ATT switching to the weighted value, which is the intended fix. - The
g_tnon-identification under Mundlak is correctly isolated from the ATT (it entersmu_linadditively and never touchestauor the weights). generate_staggered_did_datais bit-identical when no new argument is passed (covariates are drawn after the main loop, so the default RNG stream is untouched).
Worth addressing before merge
1. No fit-time guard on att_weights (robustness, not a live bug). The in-model att Deterministic (pymc_models.py:2920) sums over the entire (cohorts, ev) surface, so its correctness rests entirely on att_weights being zero on every non-treated / unidentified cell. By construction they are (weights come from treated-cell counts only), so it's correct today, but given the in-model ATT is the headline feature, an assertion that att_weights are nonnegative, sum to 1, and are zero off the identified support would be cheap insurance against a future edit silently leaking prior-only tau into the ATT.
2. The balanced validation panels don't actually discriminate weighted vs unweighted. On the Tier-1/Tier-2 panels every treated (g,k) cell has equal size, so test_tier1_att_equals_treated_cell_mean_tau and the noisy-recovery tests would pass even with the wrong (unweighted) aggregation, despite docstrings implying they exercise the weighting. The real discrimination comes from test_tier1_exact_on_unbalanced_panel, test_tier1_att_weights_are_treated_cell_shares, and test_staggered_did.py (test_staggered_did_effect_summary_reports_treated_weighted_att, which asserts weighted ≈ true within 1e-8 while unweighted is off by >0.1). The claim holds across the suite, but making one balanced Tier-1 panel use unequal cohort sizes would make the headline test self-contained.
Minor
- Priors scale to
sd_yrather than rescaling inputs. Reasonable response to the too-narrow-intervals criticism, but covariates enterbetaunscaled and assume O(1) inputs; worth a note or a standardization step. betacarries a vestigial length-1treated_unitsdim; onlybeta[0]is ever used. Consider dropping the dim or documenting the length-1 assumption.- Tier-3/4 real-MCMC tests are
@pytest.mark.slow, so they'd drop out of PR-time coverage if CI ever adds-m "not slow". A comment pinning that expectation would help, since they're the only genuine-sampler coverage of the ATT deterministic.
Nice work overall, the backend-contract consistency across PyMC / OLS / event-time surfaces is the part that could most easily have gone wrong, and it's coherent.
Two items from review, plus three minor notes.
1. Fit-time guard on att_weights. The in-model 'att' deterministic sums
tau * att_weights over the entire (cohorts, ev) surface, so its
correctness rests on the weights being zero wherever no treated
observation sits. That holds by construction, but nothing in the model
would notice if a future change broke it, and it is the single assumption
the headline estimand depends on. _validate_att_weights now checks
non-negativity, sum-to-one, and zero mass off the occupied support, once
per fit.
2. Validation panels now use unequal cohort sizes. With {4: 8, 8: 8} every
treated (g, k) cell was the same size, the weight matrix was uniform, and
weighted and unweighted aggregation were numerically identical -- so the
Tier-1 exactness tests would have passed with the weighting dropped
entirely. With {4: 10, 8: 6} the weighted ATT is exact (0.0e+00) while an
unweighted one misses by 1.6e-2, six orders of magnitude outside the 1e-8
tolerance. test_tier1_att_weights_are_treated_cell_shares now asserts the
weights are non-uniform, so this discriminating power cannot silently
regress.
Minor: documented that priors scale to sd_y only and covariates are assumed
roughly unit-scale; documented the length-1 treated_units dim on beta; pinned
the expectation that the slow-marked Tier 3/4 tests are the only real-sampler
coverage of the ATT deterministic.
Suite 1372 passed, 6 skipped.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
drbenvincent
added
the
review:high
Automated triageRecommendation: Why:
Review focus:
Confidence: high |
juanitorduz
left a comment
juanitorduz
left a comment
There was a problem hiding this comment.
This looks very interesting @NathanielF ! I love the notebook! From my side, I approve (maybe @drbenvincent has still some small comments, but for me we can iterate over them and address them if they are easy :) )
|
Thanks @juanitorduz . FWIW @drbenvincent the fact that the PR is also a replication test of the existing Staggered DiD should give some confidence in the implementation. But i guess the bigger question might be around architecture and code bloat within the Staggered DiD class? There is definitely some work i could do to pull out the supporting functions into a separate module but let me know what you think. |
IMO, we can leave this for another iteration (please create an issue). At least I like merging MVP PRs, getting feedback, and then doing a more thorough design and refactor after (but this is just me XD) |
derwells
added
the
needs:maintainer-decision
5 participants
Relates to: #1056 (comment)
Add Wooldridge/Mundlak ETWFE estimator to StaggeredDifferenceInDifferences
Summary
Adds an estimator="etwfe" switch to the existing staggered DiD class, implementing Wooldridge's extended two-way fixed effects estimator via the Mundlak device.
The existing BJS imputation estimator is a fit–predict–subtract pipeline: the treatment effect is a post-hoc residual, and its uncertainty is assembled by differencing posterior-predictive draws outside the model. ETWFE is the structural phrasing of the same estimand — the cohort × event-time surface τ[g,k] is a model parameter, and the aggregated ATT ⟨W, T⟩ is a pm.Deterministic inside the model, so it arrives with its own posterior. It also replaces one free dummy per unit with partially pooled intercepts plus treatment means, which is what lets it scale.
Both estimators sit behind one entry point — users reach for "staggered DiD", not for a paper name — so switching costs one keyword.
API
New: ETWFERegression (PyMC), att_ (in-model ATT draws / float), att_se_, tau_surface_, att_weights_, event_time_grid_, etwfe_formula_, plot_tau_surface(). att_event_time_ / att_group_time_ / data_ / hdi_prob_ keep identical schemas, so effect_summary, PreTreatmentPlaceboCheck and the maketables adapters work unchanged.
generate_staggered_did_data gains a callable treatment_effects, cohort_effect_scale, and covariates — backwards compatible, verified bit-identical output when no new argument is passed.
_effect_summary_staggered_did previously reported an unweighted mean over event-time rows as the headline ATT. That is not the ATT. Under staggered adoption only the earliest cohorts reach the largest event times, so late rows rest on fewer treated observations; when effects grow with event time those thin rows are also the large ones, and the summary is biased upward.
On noise-free data where the truth is exactly recoverable:
estimate error
unweighted mean (old) 3.4200 0.2178
n_obs-weighted (new) 3.2022 0.0000
Now weighted by n_obs on both estimator paths. This changes a user-visible number, downward in the typical case. The per-event-time ATTs were never wrong — only the aggregation.
Validation
Correctness is proved algebraically rather than by MCMC, because mock_pymc_sample makes posterior recovery impossible on the default test path:
Tier 1 — on noise-free data the saturated OLS fit recovers the generator's τ[g,k] surface and the treated-cell ATT to 1e-14 or better, including with covariates, with leads, and on an unbalanced panel with a non-contiguous index. Any error in ev_idx, cohort_idx, effect_indicator, top-binning or W shows up here.
Tier 2 — noisy recovery, plus the assertion that motivates the estimator: naive single-δ TWFE misses by more than ETWFE does.
Tier 3/4 — real MCMC (not mocked): true ATT inside the 94% HDI, r̂ = 1.0023; Bayesian and OLS agree to 0.007 against a 0.25 threshold. Combined runtime 23s.
Notebook headline (σ=0.3, cohort-heterogeneous growing effects, true ATT 3.2022): naive TWFE 2.8219 (off by 4.1 SE), ETWFE 3.1876.
953 passed, 5 skipped · 37 doctests · ruff clean · interrogate 87.4% · notebook executes under the papermill runner.
Known limitations
Covariates enter additively. Wooldridge's centred-covariate × (g,k) interactions are not implemented; documented in the class Notes and notebook.
conditioning="mundlak" raises for sklearn models — D̄_unit is exactly collinear with free unit dummies, so pinv would silently drop it and return the dummy fit.
Under Mundlak, g_t is identified only by its prior (D̄_time lies in the span of the time effects) and must not be interpreted. The ATT is unaffected. Documented in the class Notes and flagged in the notebook.