feat: add LWDiD estimator (Lee & Wooldridge 2025, 2026)

[gorgeousfish]

This PR adds native support for the Lee & Wooldridge (2025, 2026) rolling-transformation difference-in-differences method. The approach works by applying unit-specific time-series transformations (demeaning or detrending) to panel outcomes before treatment, converting the panel DiD problem into a standard cross-sectional one. Once transformed, any treatment-effect estimator — regression adjustment, inverse probability weighting, doubly robust, or propensity-score matching — can be applied directly to the cross-section. The method handles both common-timing and staggered-adoption designs with flexible control-group selection.

The implementation lives entirely within diff_diff/lwdid*.py (9 modules) and introduces zero new runtime dependencies — it reuses the existing solve_ols, solve_logit, and safe_inference infrastructure. The IPW and IPWRA standard errors use the full semiparametric influence function with propensity-score and outcome-model correction terms, matching the variance formula in the authors' Stata package.

Numerical correctness has been validated against the reference lwdid Python package across all supported configurations. The RA path achieves machine-precision agreement (≤1e-10), and IPW/IPWRA paths agree to within 1%. The California Proposition 99 results from Table 3 of LW (2026) are reproduced exactly: demeaning ATT = −0.422, detrending ATT = −0.227.

Beyond core estimation, the PR includes wild cluster bootstrap inference (Rademacher/Mammen/Webb), Fisher randomization inference, parallel-trends pre-testing, sensitivity analysis, clustering-level diagnostics, and visualization methods. A tutorial notebook walks through the full workflow on the papers' empirical datasets (California smoking data, Castle Doctrine laws, Walmart county-level entry).

Methodology references (required if estimator / math changes)

• Method name(s): LWDiD (Lee & Wooldridge rolling-transformation DiD)
• Paper / source link(s):
  • Lee, S. J. & Wooldridge, J. M. (2025). https://ssrn.com/abstract=4516518
  • Lee, S. J. & Wooldridge, J. M. (2026). https://ssrn.com/abstract=5325686
• Any intentional deviations from the source (and why): None. Implementation follows Procedures 2.1, 3.1, and 4.1 exactly. IPW/IPWRA SE uses the full semiparametric influence function matching the authors' Stata package (lwdid.ado).

Validation

• Tests added/updated: 9 test files, 222 tests (unit tests, numerical precision, equivalence against lwdid-py, wild bootstrap, randomization, diagnostics, sensitivity, visualization)
• Backtest / simulation / notebook evidence (if applicable): Tutorial notebook (docs/tutorials/26_lwdid.ipynb) reproduces Tables 3–4 from LW (2026) on California Proposition 99 data. ATT values match published results to 0.04% precision.

Security / privacy

• Confirm no secrets/PII in this PR: Confirmed. No secrets, tokens, or PII. Datasets included (smoking.csv, walmart.csv, castle.csv) are publicly available research data from the referenced papers.

[igerber]

@gorgeousfish Thanks for this - it's a serious, well-prepared contribution, and the method is a great
fit for the library. I'd like to move it forward.

That said, before I do a full review, one thing to clear first: licensing. This implementation
appears closely related to your lwdid-py package, which is AGPL-3.0, while diff-diff
is MIT. Could you confirm:

1. You hold the copyright to all code in this PR (i.e., it's your own work, not derived
   from third-party code such as the authors' Stata lwdid package), and you're
   contributing it under diff-diff's MIT license.
2. The bundled datasets (smoking.csv, castle.csv, walmart.csv, .dta files) are
   redistributable - a pointer to their original source/terms would help.

[gorgeousfish]

Thanks for the quick response and the positive signal.

On licensing:

Yes, I hold the copyright to all code in this PR. It's my own independent implementation, not derived from Lee & Wooldridge's Stata lwdid package or any third-party source. I'm contributing it under diff-diff's MIT license.

On datasets:

smoking.csv: Abadie, Diamond & Hainmueller (2010), California tobacco control program. Publicly available, widely redistributed in academic packages.

castle.csv: Cheng & Hoekstra (2013) / Cunningham (2021), Castle Doctrine laws. Publicly available.

walmart.csv: county-level panel from Brown & Butts (2025, Journal of Econometrics), constructed from County Business Patterns (CBP) data. Publicly available government statistical data.

[igerber]

@gorgeousfish Thanks for the licensing confirmation - that closes the question. Since then we've completed a full evaluation on our side: fresh reviews of both papers against their current SSRN revisions, an independent replication of your headline results, and a code review against the papers and this library's conventions.

The short version: this is a strong contribution and we want to merge it. We reproduced your Prop 99 numbers ourselves (demeaning ATT -0.4222, SE 0.1208; detrending -0.2270, SE 0.0941, with the exact-inference p-value matching the paper), and the core is exactly right: the transformation reproduces our DifferenceInDifferences estimator to machine precision where the theory says it must, and the control-group logic matches the current paper revision precisely.

Getting it merged takes real work on both sides, so below is the full plan - split, sequenced, and complete (no surprise rounds later). None of it starts until you confirm you're on board. If any part looks wrong to you, push back and we'll discuss.

First, one decision from you: the event study

The current revision's Appendix D event study (placebo/dynamic WATT(r)) plus the Algorithm 1 influence-function multiplier bootstrap (simultaneous sup-t bands) is the one substantial piece of the papers not yet implemented - and it's also this estimator's proper pre-trends diagnostic. Your call:

• Include it in this PR -> LWDiD ships as a standard estimator.
• Defer it to a follow-up PR -> we merge this PR as an experimental preview (as BR/DR are today), promoted to standard when the event study lands.

Both options are genuinely fine with us.

The plan

Step 0 - you, now. Reply confirming (a) you're good with this plan and (b) your event-study choice above. Nothing below starts until then.

Step 1 - us. We open a maintainer PR to main with (a) our methodology notes for both papers, written against the current SSRN revisions (June 8 and February 3, 2026 - both papers were revised after you wrote your implementation notes, and the June revision reworked the event-study and inference appendices), and (b) diff_diff.datasets loaders for the smoking and Walmart data (the library's existing download+cache mechanism). Those notes become the canonical spec this estimator is validated against. We then push directly to your branch (maintainer-edit is enabled): the rebase onto current main (~85 PRs of drift - version strings, CHANGELOG, tutorial renumbering to 27), removal of the committed data files in favor of the loaders (castle.csv and both .dta files are referenced by nothing and just go - this also moots any dataset redistribution questions), and a suite of independent validation tests: paper-table goldens (Prop 99 Table 3, Tables 4/A1, the castle-laws staggered targets), from-scratch reference implementations of the core procedures, cross-estimator equivalence pins, and property tests. Tests the current code doesn't pass will be xfail-marked so CI stays green - they are your acceptance criteria. Alongside this we'll leave a detailed review on this PR with file:line comments for every item in step 2, so each ask is anchored in code, not prose.

Step 2 - you. Work the fix list below, flipping xfails as you go. One consequence of step 1 to fold into your work: our methodology notes supersede the two review docs currently in this PR - please drop those during the rebase and treat the maintainer versions as canonical (the file paths collide anyway).

Step 3 - us. Wild-bootstrap dedupe into the existing diff_diff.utils machinery (one public bootstrap, not two), the REGISTRY entry and remaining docs surfaces, final review, merge.

The step 2 fix list

1. Fix the IPW standard error (must-fix bug). In the IPW variance path (lwdid.py, roughly lines 2061-2088 on your current branch) the influence-function terms enter un-centered - raw weighted outcomes where the formula needs deviations from their weighted means. The observable symptom: adding a constant to all post-period outcomes leaves the ATT unchanged (as it must) but shifts the IPW SE. Your IPWRA path (roughly lines 2404-2466) centers correctly and is the internal template for the fix. Our xfail'd invariance test is both the repro and the acceptance criterion, and the review will mark the exact lines. (Worth checking whether lwdid-py inherits the same issue.)
2. Remove silent failures. Dropped units/periods, PSM-unmatched treated units, and discarded failed bootstrap replicates all need explicit warnings, and lwdid() must not silently ignore unrecognized kwargs. We found seven sites; the review will list each one.
3. Fold the custom exceptions into house conventions (ValueError/ImportError plus standard warning categories). As written, validation errors don't derive from ValueError, so downstream except ValueError handlers miss them.
4. Trim top-level exports to LWDiD, LWDiDResults, and the LW alias; everything else stays importable at module level but leaves diff_diff.__init__. test_parallel_trends needs a rename regardless - it's pytest-collectable in downstream codebases.
5. Adopt the house test conventions: ci_params for iteration scaling, assert_nan_inference, @pytest.mark.slow on the heavy bootstrap/RI tests. Keep the lwdid-py equivalence suite as an optional upstream-drift check - but drop the lwdid entry from the dev extra in pyproject.toml. Reference implementations aren't dependencies here (R isn't one either): the tests stay importorskip-gated and run only where lwdid is manually installed, with a note in the test-file docstring on how to enable them.
6. Source or fold the diagnostics modules. Keep the paper-sourced sensitivity checks (varying-T0 robustness per Sec 8.1; no-anticipation) with citations. lwdid_clustering.py becomes docstring/tutorial guidance - Sec 8.2 is advice, not an algorithm. lwdid_trend_diagnostics.py is superseded by the Appendix D event study, and its generic slope test duplicates the existing check_parallel_trends.
7. mypy to zero for the new modules (currently +43 errors against a zero baseline for comparable estimators).

Thanks again - the core here is exactly right, which is the hard part. The rest is the normal cost of moving a standalone package into a library with strong invariants, and we're glad to carry our share of it.

[gorgeousfish]

Thanks for the thorough plan - I'm on board.

On the event study: Option A - I'd like to include the Appendix D event study (WATT(r) + Algorithm 1 multiplier bootstrap) in this PR so LWDiD ships as a standard estimator.

Ready for Step 1 whenever you are. I'll start looking into the IPW centering issue in the meantime.

[igerber]

@gorgeousfish - as promised, here is the detailed file:line review. Three pieces of news first:

1. Your branch is rebased onto current main (single commit, your authorship preserved). The rebase already took care of the mechanical items so you don't need to: committed datasets removed (the tutorial now uses the checksummed load_prop99() / load_walmart() loaders that landed on main; your lw_walmart.dta was byte-different from the SSC file but content-identical - only column order differed, so all your tutorial numbers stand), the lwdid dev dependency removed from pyproject (reference implementations stay environmental, importorskip-gated - your equivalence tests still run when the package is present), tutorial renumbered 26 -> 27, and ruff/black conformance fixes to your test files under the pinned linters that now gate every PR (Lint Gate required check). Your estimator modules are untouched.
2. The maintainer validation suite is live on your branch: tests/test_methodology_lwdid.py now collects against your code - currently 30 pass / 18 xfail. The xfail markers are the acceptance criteria for the work below: each reason string names the item it gates, and strict=True markers will force their own removal in the commit that fixes them (an XPASS fails the run). Green-with-no-xfails = done.
3. A mirror PR ([Mirror] LWDiD estimator - CI-review mirror of PR #588 (do not merge) #690) now exists in the main repo carrying this branch, purely so our CI AI reviewer can run (it cannot fire on fork PRs for security reasons). Ignore it otherwise - this PR remains the merge vehicle and your authorship stands.

What replicates today, verified first-hand on your rebased head: all three Prop 99 donor-pool tables (Table 3: demean -0.4222/0.1208, detrend -0.2270/0.0941; Table 4 Southern; Table A1 Midwestern), per-period effects, and the exact-t p-value (0.0209 vs the paper's printed 0.021). The transformation core is genuinely solid. The items below are ordered by how much they matter.

---

1. IPW influence function is un-centered (confirmed bug)

diff_diff/lwdid.py:2058-2059:

psi_ht[treat_mask] = (y[treat_mask] - att) / p_bar
psi_ht[ctrl_mask] = -w_ctrl * y[ctrl_mask] / p_bar

The control term is linear in raw y, so the variance of the IF - and hence the SE - changes under a constant shift of the outcome (verified: adding +100 to all post-period outcomes moves the IPW SE from 0.063385 to 0.063882; ATT is invariant as expected). Your IPWRA path at lwdid.py:2401-2402 does this correctly - the control term there is centered (resid_ctrl - control_term). Centering the Hajek control term the same way fixes it.

Acceptance: TestTranslationInvariance::test_ipw_se_translation_invariant (strict xfail; RA and IPWRA already pass the same parametrized test).

2. Staggered path: per-(g, r) cells + the composite-outcome regression (7.18)/(7.19)

This is the deepest item, with two coupled halves.

(a) Cell construction. _fit_staggered (diff_diff/lwdid.py:689-899) builds one cross-section per cohort by averaging each unit's post-treatment transformed outcomes into _ydot_avg, then estimates a single per-cohort ATT. The papers' estimand is per-(g, r) cells with calendar-period-specific control eligibility: A_{g,t} = {G = g} union {G = 0} union {G > max(g, t)} - under not-yet-treated controls a later cohort is a valid control for period r only while cohort > r, which a per-cohort averaged window cannot express (your filter is by g, and later-treated controls contribute unequal calendar windows inside the average). Under never-treated controls the pools are period-invariant, which is why your point estimates still replicate the paper's tables - the construction and the paper coincide exactly there, and diverge for not-yet-treated.

(b) Aggregation and SE. _aggregate_cohort_effects (diff_diff/lwdid.py:995-1042) computes SE = sqrt(sum of w_g^2 se_g^2), assuming independence across cohorts - not valid when cohorts reuse control units, and in neither paper. It measurably diverges on the paper's own application: on the castle-doctrine data (LW 2026 Section 7.2), the paper's tau_omega-hat via the composite-outcome regression is 0.0917 with OLS SE 0.0571 (demeaning; detrending 0.0666) - a from-scratch implementation of (7.18)/(7.19) in our suite reproduces those numbers to printed precision - while your aggregation gives SE 0.0512. The paper's exact-t theory is stated for the composite regression, not for the independence formula; adopting (7.18)/(7.19) resolves (b) and forces the per-(g, r) structure of (a) at the same time.

Please also add a hand-built staggered test in your test files asserting period-specific control eligibility directly (e.g. a panel where one later cohort is a valid control at r = 3 but not at r = 5, with the cell samples asserted explicitly) - our golden-based suite catches the SE divergence but not the eligibility semantics on its own.

Two sub-points discovered while calibrating:

• The paper's (7.10) staggered target is never-treated-controlled; your control_group default is 'not_yet_treated' (lwdid.py:61,160), under which the castle point estimate is 0.074 rather than 0.092. Both control groups are legitimate options - but the docs/REGISTRY should be explicit that replication of the paper's numbers requires control_group='never_treated'.
• With any per-cohort SE non-finite, your aggregate SE silently becomes NaN alone (lwdid.py:1035-1040); with the composite regression this case disappears structurally.

Acceptance: TestCastleTauOmegaAdjudicator::test_demean_tau_omega_ols_se (strict xfail; the point-estimate tests already pass under never-treated controls).

3. Event study (your Option A): Appendix D + Algorithm 1

The API spec we agreed to is written as normative docstrings in TestEventStudySpec (tests/test_methodology_lwdid.py): invocation fit(..., aggregate="event_study"), results exposing event_study_effects: Dict[int, {effect, se, t_stat, p_value, conf_int, cband_conf_int}] plus result-level cband_method / cband_crit_value / cband_n_bootstrap; WATT(r) weights per E.1 (contributing treated units at event time r - reduces to cohort-size weights only in balanced panels); anchor-period exclusions (r = -1 for demeaning; r = -2, -1 for detrending); Algorithm 1 unit-level Rademacher multiplier bootstrap, B = 999, sup-t bands.

Full per-period goldens for both Walmart outcomes are committed at benchmarks/data/lwdid_walmart_eventstudy_golden.json (Tables A4/A5 of the current June 8, 2026 revision, r = 0..13, with provenance and PDF SHA-256; note these are the current-revision numbers - your tutorial's Walmart figures came from an older PDF revision and should be re-checked against these when you get there).

Related guard: N_infinity >= 2 for the never-treated-only control strategy (LW 2026 p. 26) does not exist yet - TestExactSmallSampleInference::test_never_treated_pool_of_one_is_rejected (strict xfail).

Acceptance: 4 TestEventStudySpec spec tests + 6 point-golden tests (strict), 6 SE-golden tests (non-strict - printed bootstrap draws; we will re-calibrate tolerance together when it lands).

4. Randomization inference convention (discussion item, not necessarily a bug)

diff_diff/lwdid_randomization.py:253 (randomization_inference) produces the seed-stable exact permutation atom ~2/39 ≈ 0.051 for Prop 99 (N1 = 1 among 39 states) - arguably the standard exact answer - while the paper prints 0.020, and its permutation scheme is under-documented. Please check what Stata lwdid, ri does on the same data; whichever convention wins should be documented in the docstring. Tracked as a non-strict xfail (TestProp99Table3Goldens::test_detrend_randomization_inference_p_value).

5. Silent-failure sites (library policy: never drop/alter user data without a warning)

• lwdid.py:537 (and the bootstrap copies at 2723, 2814; staggered analogue 2642): units whose transformation yields NaN are dropped by dropna; a warning fires only when ALL units drop. Partial drops need a warning with the count.
• lwdid.py:600-601: a user-supplied cluster= is silently ignored unless vce == "cluster". Either warn or (better) treat cluster= as implying cluster vce, matching the rest of the library.
• lwdid.py:2922-2923: staggered (g, t) cells with no treated or no control units are skipped via continue with no record; report skipped cells (cf. the loud handling you already do at 877).
• lwdid.py:3175 (lwdid() wrapper): **{k: v for k, v in kwargs.items() if k in LWDiD().get_params()} silently discards typo'd kwargs - raise on unknown keys instead.
• lwdid.py:2016 / 2204 / 2338: propensity scores silently clipped at trim_threshold; warn with the affected count (the CS estimator's trimming warning is the house pattern).
• lwdid.py:2239-2243: PSM raises only when ALL matches exceed the caliper; partially dropped treated units (NaN matches) go uncounted.
• Bootstrap replications that fail are excluded from the distribution without a reported count (both in lwdid.py bootstrap loops and lwdid_randomization.py degenerate draws) - report n_failed.

5b. Design validation gaps

_validate_inputs (diff_diff/lwdid.py:220-281 region) checks column presence, missingness, binary treatment, and duplicate unit-time rows, but not the design assumptions the estimator relies on: treatment absorbing within unit, a single first-treatment time when fitting common timing (cohort=None), cohort constant within unit and equal to first treated period, and consistency between treatment and cohort (in the staggered path the treatment column is unused after validation, so an inconsistent pair silently estimates a different design than the user described). Please add these checks - loud ValueErrors, matching how the rest of the library treats design violations.

Related: get_transformation_diagnostics(..., cohort=...) (lwdid.py:319-326) defines the pre-period globally as time < earliest_cohort, which understates available pre-periods for later cohorts; diagnostics should be cohort-specific (time < g).

6. Exceptions: fold into house style

diff_diff/lwdid_exceptions.py:12-131 defines 13 stateless exception/warning classes; no other estimator has a custom hierarchy, and validation errors that are not ValueError subclasses break user except ValueError handling that works everywhere else in the library. Validation -> ValueError, warnings -> UserWarning (module retired).

7. API surface: trim top-level exports

Keep LWDiD, LWDiDResults, LW top-level; everything else module-level. Two naming collisions to resolve regardless: test_parallel_trends (lwdid_trend_diagnostics.py:401) is pytest-collectable by name and shadows the concept behind the existing check_parallel_trends; wild_cluster_bootstrap / randomization_inference / sensitivity_analysis are generic names claiming library-wide meaning for estimator-specific machinery.

8. Wild cluster bootstrap: we will fold this into house machinery (our take-on)

FYI only: lwdid_wild_bootstrap.py:723-726 fits the restricted model as intercept-only even when controls are supplied (the CGM null-imposed model should regress y on intercept + controls). Since the house wild bootstrap already handles this, the dedupe lands on our side in step 3 - no action needed from you beyond not building further on that module.

9. Diagnostics triage (as agreed)

• lwdid_sensitivity.py: keep the paper-sourced parts (Section 8.1 T0-robustness, no-anticipation) with citations; drop the rest.
• lwdid_clustering.py: advisory content -> docs; not estimator API.
• lwdid_trend_diagnostics.py: superseded by the Appendix D placebo machinery once the event study lands (pre-period placebo WATT(r) IS the trend diagnostic).

10. Typing and tests

• mypy: 43 errors, all in lwdid.py + lwdid_results.py (the required Lint Gate enforces zero, so both this PR and the mirror will show that check red until this lands; ruff/black are already clean after the rebase).
• Test conventions: adopt ci_params for bootstrap/grid scaling, assert_nan_inference() for NaN-tuple checks, behavioral assertions over no-exception checks. Your equivalence suite vs lwdid-py stays, demoted to an importorskip'd upstream-drift watchdog.
• Tutorial: the notebook is committed without outputs and our Sphinx config never executes (nbsphinx_execute = "never") - before merge it needs one full execution committed, at which point please re-check the Walmart section against the current-revision goldens (see item 3).

---

Sequencing from here is as agreed in the plan comment: you work down this list on the rebased branch (small commits are easier for us to sync into the mirror for CI review rounds); the xfail suite tracks progress objectively; once it is green-with-no-xfails we do the final REGISTRY/experimental-surface pass and the merge mechanics on our side. The two of the seven step-2 items already taken off your plate: packaging (done in the rebase) and the wild-bootstrap dedupe (our step 3).