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.

Fixes #732
Fixes #733
Fixes #734
Fixes #735

[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).

[shawcharles]

I have opened a focused dependent PR for #733: gorgeousfish#1. It targets this branch and preserves the overall ATT and inference already computed by fit() rather than reconstructing them from marginal cohort effects.

It also adds a slow strict-xfail bootstrap concordance test for the separate covariance-aware inference follow-up, #735. The default not-yet-treated identification defect in #734 remains the merge-blocking item; I have not attempted to fold that larger (g, t) redesign into this bounded fix.

[gorgeousfish]

This round addresses the four issues @shawcharles filed against this branch. Rebased onto upstream/main at 686bed30; branch head is now 10c900ae.

Commits

Commit	Issue
e17cc53f refactor(lwdid): put LWDiDResults on the shared results contract	#733, #732
5d988b2c fix(lwdid): build staggered cells per (g,t) and aggregate them jointly	#734, #735
ac9667d5 test(lwdid): cover cohort-time cells, joint inference and the aggregate contract	tests for all four
10c900ae docs(lwdid): register the staggered module and the never-treated replication note	docs

#734 — staggered control eligibility (merge-blocking)

The old path applied control eligibility as a unit-level filter and then averaged whatever calendar window each cohort happened to retain, so the cohort means were taken over different calendar periods and a common time trend entered the contrast directly. On a panel whose only signal is that trend it returned ATT ≈ 1.0 with SE ≈ 5e-16.

Estimation now runs one cell per (g, t) in a new module diff_diff/lwdid_staggered.py, each on its own cross-section, with control eligibility evaluated at t rather than once per cohort: never-treated plus {G > max(g, t)}, which is the Appendix D.3 rule. The trend-only panel now returns ATT = 0, and adding a constant shift to every unit-period leaves the ATT unchanged. Cells that cannot be identified — no treated or no control units left after the eligibility filter, too few pre-periods for the transformation — are recorded with a reason rather than averaged over, and a cohort with no supported cell is dropped instead of contributing a NaN.

#735 — overall staggered standard error

Cohort effects that share control units are correlated, so the previous sqrt(sum w_g^2 se_g^2) was not valid. Every estimator now returns a per-observation influence function through _dispatch_estimator; cell influence functions are accumulated to units, combined with the treated-mass weights, and the overall SE is read off the joint function.

The basis is now reported explicitly in LWDiDResults.inference_basis and printed by summary() rather than left implicit:

• composite_regression — the gate control_group='never_treated' and estimator='ra' and vce='classical' and no covariates. This is where eq. 7.18/7.19 applies, and it is the path that reproduces the paper's castle numbers; it keeps _composite_regression_aggregation, which exists for exactly this case.
• joint_influence_function — every other combination, i.e. the general path.
• unavailable_matching — estimator='psm' has no influence-function representation. The overall SE is NaN and warns, pointing at ipwra.
• unavailable_degenerate_cells — some cell SE is degenerate or non-finite. The overall SE is NaN and warns, naming the cohorts.

The two NaN cases fail closed rather than falling back to an independence assumption.

#733 — aggregate() recomputing the ATT

LWDiDResults now inherits the shared AggregationMixin, so aggregate() follows the library contract. aggregate("simple") reports the estimand fit() already computed, exactly, including df_inference. The old aggregate(by="overall") rebuilt the ATT from the marginal cohort effects under cohort independence, which gave a narrower SE than fit() for the same estimand and dropped df_inference, so the aggregated object reported a different significance level than the fit it came from. "overall" is no longer accepted; the supported set is 'simple', 'event_study', 'group', and unknown types, a weights selector, balance_e off event-study, and common-timing fits all raise.

#732 — event-study output

LWDiDResults inherits BaseResults, and aggregate("event_study") returns EventStudyResults on EVENT_STUDY_SCHEMA instead of a private dict. The anchor period is carried as an explicit is_reference row at att = 0 rather than dropped, and the sup-t band and its critical value thread through. to_dict() now serialises the event study, the reference periods, the band metadata and inference_basis, so a round-tripped result no longer loses them.

Paper replication

All replication goldens still pass:

• Prop 99 Table 3 — average ATT and SE, per-period ATTs, exact-inference and randomization-inference p-values, both demean and detrend.
• Prop 99 Table 4 (southern donor pool) and Table A1 (midwest donor pool), both transformations.
• Castle tau_omega — demean point and usual OLS SE, detrend point, plus the composite-regression reference.
• Walmart event study — all six point-estimate columns exact.

The registry note in docs/methodology/REGISTRY.md records why the staggered replication goldens pass control_group='never_treated' explicitly: the implementation default is 'not_yet_treated', matching OVLS (eq. 4.10) as the text states, but the printed staggered results are computed against the never-treated pool only. A default-pool fit gives different and equally valid estimates because the (g, t) cells then draw on a strictly larger control sample.

Tests and lint

tests/test_lwdid*.py, tests/test_methodology_lwdid.py, tests/test_aggregate_contract.py, tests/test_results_serialization.py, tests/test_v4_matrix.py: 625 passed, 57 skipped, 5 xfailed, 1 xpassed.

The xfails are all six parametrizations of TestEventStudySpec::test_walmart_eventstudy_se_goldens, non-strict, which is why one of them currently xpasses. I want to be explicit that this changed in this round rather than being pre-existing: on the pre-change baseline, detrend-ra/a5_wholesale passed. Changing which controls enter each cell moves the bootstrap SEs slightly in both directions — that column goes 0.0570 → 0.0581 against a printed 0.057, while demean-ipwra/a4_retail moves the other way into agreement. At B = 999 the Monte Carlo error on a bootstrap SD is roughly 1/sqrt(2B) ≈ 2.2%, larger than the 1.9% gap, so which individual columns land inside three-decimal precision is not a stable property to encode; the marker now covers the whole parametrization instead of an enumerated subset. PRINTED_ATOL is unchanged at 1e-3 and the point-estimate goldens remain exact equality checks.

Lint Gate reproduced with the pinned versions from lint.yml (mypy 2.3.0, numpy 2.4.5, pandas 3.0.3, scipy 1.17.1, ruff 0.15.21, black 26.5.1): mypy diff_diff clean across all 98 source files, ruff check diff_diff tests and black --check diff_diff tests both clean.

I also ran the whole suite. The 24 failures I see are confined to test_rust_backend.py, test_estimators_vcov_type.py::TestDfConvention and one test_doc_snippets.py case, and they are artifacts of my local checkout rather than of this change: the first two come from a stale locally-built _rust_backend extension and disappear when it is removed, and the snippet failure is a NumPy/numba version clash from an unrelated package on my path. Applying this branch's diff to a clean worktree leaves all of them passing, so CI is the authority here.

Credit

test_simple_preserves_the_fitted_result and test_matches_unit_cluster_bootstrap in ac9667d5 are adapted from @shawcharles's patch in gorgeousfish#1, which independently identified the aggregate() SE discrepancy behind #733. He is credited as co-author on that commit. His patch targeted the _aggregate_legacy code path, which this rewrite removes, so the diff itself no longer applies — but both of his tests are preserved.

[gorgeousfish]

@igerber — first off, apologies for the wait; this round took longer than it
should have. Status update on the review checklist: the branch has been rebased
onto current upstream/main (post-#752), and every item from your 7/17 review is
now addressed on top of the v4 API. Heads-up: this was a force-push
(--force-with-lease), so the commit anchors in the older review comments no
longer resolve; the mirror PR #690 will need a re-sync to the new HEAD before
the AI review can re-run.

Review checklist reconciliation

Status tags: Done (this round), Done earlier (7/30) (re-verified green on
the rebased branch), Maintainer-side, Needs your decision.

§1 IPW influence-function centering — Done earlier (7/30). Re-verified on
the rebase; translation-invariance tests green.

§2 Staggered estimation — Done earlier (7/30). Cells are built per (g, t)
and aggregated with joint influence-function inference (a6ca7b9, e73b56e).
Still green after the rebase.

§3 Event study — Point-estimate goldens (LW 2025 Appendix F, Tables A4/A5)
are green. The event-study output is now aligned with the shared
EventStudyResults schema including the provenance fields (01217a6). What
remains are the 6 Walmart SE goldens, kept as non-strict xfail: the paper's
printed SEs come from a B=999 multiplier bootstrap, and a re-seeded run sits
near the printed-precision tolerance boundary (this run: 5 xfailed, 1 xpassed —
consistent with MC noise, not a systematic gap). Needs your decision:
accept the non-strict xfails as-is, or have us raise B to push the MC noise
below the tolerance.

§4 Randomization inference — The test validates against the paper value
(RI p = 0.020, Table 3 detrend) with atol = 0.015 covering permutation noise at
1,000 draws. Needs your decision: please confirm this matches the intended
convention of Stata lwdid, ri (p-value definition / tie handling), so we can
tighten or re-anchor the golden if needed.

§5 Silent-failure warnings — Done earlier (7/30). All seven sites warn
explicitly; re-verified green.

§5b Design-validation gaps — Done this round. A single vectorized
_check_treatment_design() now enforces all three contracts — absorbing
treatment, unique onset for the common-timing interface, and staggered
consistency D_it = 1[t >= g_i] (0666b34). Transformation diagnostics use
per-cohort pre-periods instead of a pooled window (0ea2c61), and the API docs
gained an explicit "Input Contract" section (b58787d).

§6 Exceptions — Done, with one caveat. All exception types are now plain
stdlib aliases (no custom hierarchy). The lwdid_exceptions.py file itself is
kept as a shim because lwdid_wild_bootstrap.py — which is in the §8 scope you
are taking over — still imports NumericalWarning from it. Suggestion: delete
the shim in the same change that moves WCB onto the house machinery.

§7 API surface — Done. Top-level exports are LWDiD, LWDiDResults, and
the LW alias only. test_parallel_trends has been made private as
_placebo_pre_trends (internal dependency of recommend_transformation;
also removes the pytest mis-collection hazard of a public test_* name), and a
guard test asserts the retired names do not reappear (11c1f73).

§8 Wild cluster bootstrap — Maintainer-side. Untouched per agreement;
waiting for your takeover.

§9 Diagnostics triage — Done. Sensitivity is scoped down to T0-robustness
plus no-anticipation (946 → 682 lines, 75bc2d2); the clustering module is
deleted with its advisory moved into the REGISTRY (e8e73f1); the standalone
trend pre-tests are retired in favor of the house placebo machinery
(run_placebo_test and friends, 11c1f73).

§10 Tutorial and lint — Done. Tutorial 27 is fully re-executed with
committed outputs (05a616a; 30 code cells, zero errors), now demonstrating the
placebo test and the scoped sensitivity analyses. mypy/ruff/black are clean on
our files; note local mypy chokes on the numpy stubs' 3.12-only syntax, so CI
is authoritative there.

Rebase adaptation (new work outside the checklist)

The v4 rebase required a few items your review predates: LWDiD now inherits
the BaseEstimator mixin and drops the hand-written get_params/set_params
(fc43596); the estimator surface is registered with the naming guard
(921d4fe); the tutorial card and homepage estimator row are registered for the
docs-IA guards (62c08a8); the choosing_estimator example was switched to the
staggered interface after the new design validation correctly rejected it
(ed948bc); and pre-period identification now works on the native time scale
(datetime64 etc.), matching the main estimator (b58787d).

Test evidence

The full lwdid suite plus the naming and docs-IA guards: 333 passed,
43 skipped, 0 failed (plus the 6 non-strict Walmart SE goldens: 5 xfailed,
1 xpassed this run). The whole repository suite is at 11,800+ passed with only
pre-existing items unrelated to this PR: the StackedDiD sklearn-clone failure
(reproduces on upstream/main) and one pyfixest/numba third-party environment
incompatibility. After rebuilding the Rust backend against the rebased tree,
all 151 backend tests pass.

Requests

1. Re-sync the mirror [Mirror] LWDiD estimator - CI-review mirror of PR #588 (do not merge) #690 to the new HEAD so the AI review can re-run against
   the actual diff.
2. Apply the ready-for-ci label so the real CI matrix runs on the rebased
   branch.
3. Decisions on the two open items above: §3 (keep the non-strict SE xfails vs.
   raise B) and §4 (RI p-value convention of Stata lwdid, ri).

[shawcharles]

I did a narrow pass over the new b58787d5 head. The main #732/#733/#734/#735 replacement work looks directionally aligned with the prior findings, but I found two bounded follow-ups worth fixing before CI/review hardens this branch.

1. Native time-scale staggered cohorts still fail in the estimator path

b58787d5 advertises native time-scale pre-period handling, but the staggered estimator still treats cohort values as numeric in several places:

• diff_diff/lwdid.py _check_treatment_design(): cohort_by_unit == 0, cohort_by_unit > 0, and cohort_by_unit <= max_time.
• diff_diff/lwdid_staggered.py: treated-cohort discovery and not-yet-treated eligibility also use value > 0 / cohort_by_unit > threshold.
• event-time construction uses relative_time = t - g and later casts labels with int(relative_time).

Minimal reproduction on the PR head:

import pandas as pd
from diff_diff.lwdid import LWDiD

periods = pd.date_range("2020-01-01", periods=6, freq="YS")
rows = []
for unit, cohort in [("t1", periods[2]), ("t2", periods[2]), ("c1", pd.NaT), ("c2", pd.NaT)]:
    for i, period in enumerate(periods):
        rows.append({
            "unit": unit,
            "time": period,
            "cohort": cohort,
            "treat": int(pd.notna(cohort) and period >= cohort),
            "y": float(i),
        })

df = pd.DataFrame(rows)
LWDiD(rolling="demean", estimator="ra", control_group="never_treated").fit(
    df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="cohort"
)

This raises TypeError: Invalid comparison between dtype=datetime64[ns] and int. The same pattern fails with Period cohorts.

Concrete fix: normalise the ordered time support once, map cohort/time labels to integer positions for design checks, event-time labels, reference-period exclusion, and not-yet-treated eligibility. Keep the original calendar labels in metadata if useful, but do not compare datetime/period cohorts to numeric 0. For non-numeric time scales, never-treated should be represented by missing cohort values; numeric 0 can remain a numeric-time convenience.

2. Public LWDiD examples are stale after the stricter first_treat contract

Two examples now fail if copied:

• diff_diff/lwdid.py class docstring generates staggered data and calls fit(..., treatment="treated") without first_treat, which now correctly raises heterogeneous-onset ValueError.
• docs/api/lwdid.rst does the same in the basic and robustness examples, and the staggered example passes cohort="first_treat" even though the public fit keyword is first_treat.

I reproduced the docstring/basic example failure on the PR head:

ValueError: Treated units have heterogeneous first-treatment periods [3, 5, 7] but no cohort column was given. Common-timing LWDiD requires a single treatment onset; pass first_treat= to use the staggered (cohort) interface.

and the RST staggered example raises:

TypeError: LWDiD.fit() got an unexpected keyword argument 'cohort'

Concrete fix: either use a true common-timing fixture in the common-timing examples, or pass first_treat="first_treat" consistently for generated staggered data. The latter is probably simpler and aligns with the new input-contract section.

Decision note on the Walmart SE goldens

I would avoid merging with non-strict xfails as the final representation of the SE check. If the printed B=999 bootstrap SEs are Monte-Carlo-limited, the cleaner contract is a passing test with an explicit MC-informed tolerance or a slow higher-B calibration test, not an xfail that can fail or xpass without changing merge status.

[shawcharles]

Full follow-up review of b58787d5

Conclusion: this is not ready for ready-for-ci or merge yet. The per-(g, t) rewrite, joint-inference work, and shared result-object integration materially improve the branch and address the original concerns in #732-#735. A full merge-readiness review nevertheless found five correctness blockers and several bounded contract/documentation defects.

[High] vce="classical" manufactures cross-cell covariance

_ols_treatment_influence() returns sigma * basis, rather than the residual-based contribution stated in its own docstring and in the Lee-Wooldridge Algorithm 1 representation. The vector is calibrated to reproduce each marginal classical OLS SE, but lwdid_staggered.py then treats it as a unit influence function when constructing covariance across overlapping cells.

A deterministic shared-control panel (seed=101, 60 never-treated units, 30 units in each of G={5,7,9}, t=1..12, y_i,t = alpha_i + 0.2t + u_i + eps_i,t + 1.5D_i,t) with control_group="not_yet_treated", RA and demeaning gives:

                    ATT       analytic SE   unit-bootstrap SE (B=400)
classical       1.500387       0.113235             0.057876
hc1             1.500387       0.062159             0.057876

The existing reconciliation test checks only sqrt(sum(psi**2)) == marginal_se; it cannot validate off-diagonal covariance. The existing bootstrap concordance test covers HC1 only.

Required: use the residual-based RA influence contribution for joint and event-study covariance. Keep any finite-sample classical marginal SE adjustment separate from that influence function. Add a cross-cell covariance or unit-bootstrap oracle for the classical path, including simultaneous event-study inference.

[High] randomisation inference rejects a null distribution made entirely of ties

_compute_pvalue() uses strict >. With a constant outcome, every randomised statistic equals the observed statistic, so the valid two-sided p-value is 1. Instead:

y = np.ones(20)
d = np.r_[np.ones(5), np.zeros(15)]
r = randomization_inference(y, d, n_reps=999, method="permutation", seed=42)
print(r.att_observed, np.unique(r.att_distribution), r.pvalue)
# 0.0 [0.] 0.001

The +1 correction does not make strict exclusion of all tied statistics valid.

Required: count statistics at least as extreme as observed, document whether the implementation is Monte Carlo or exact, and add all-tie and duplicated-statistic regression tests. The separate Prop 99 convention question can then be reconciled against that declared rule.

[High] all-eventually-treated support is claimed but the identified final-period cell is dropped

For not-yet-treated controls, the implementation requires G > max(g, t). In the Section 4.3 all-eventually-treated design, the last cohort is the reference cohort at t=T; its own ATT is not identified, but it remains the control for earlier cohorts at T.

On a balanced no-effect panel with cohorts 3 and 5 and T=5, the branch marks (3,5) as zero_treated_control, omits event time 2, and retains only cohort 3. The methodology registry still describes Section 4.3 support.

Required: either implement an explicit all-eventually-treated mode using the final cohort as the reference while excluding its own ATT, with complete cell-set and aggregation tests, or reject such panels until support exists. A silently truncated estimand is not acceptable.

[High] failed sensitivity fits are reported as highly_robust

_fit_single_spec() catches every exception and returns NaNs. _compute_sensitivity_ratio() maps one or fewer finite estimates to zero, which is classified as highly_robust; the insufficient-pre-period branch does the same explicitly.

Both robustness_pre_periods(..., outcome="missing_outcome") and sensitivity_no_anticipation(..., outcome="missing_outcome") return baseline_att=nan, sensitivity_ratio=0.0, robustness_level="highly_robust" rather than raising or reporting non-estimability.

Required: validate columns and arguments before fitting; do not swallow broad exceptions; require a finite baseline and a documented minimum number of successful alternatives. Failed or insufficient analyses must be inconclusive/not_estimable, never robust. Given the current PR size, removing these helpers from this contribution is the smallest safe resolution.

[High] staggered estimation can adjust for post-treatment covariates

The common-timing path takes the first covariate row per unit, while the staggered path takes cell[controls] at each calendar time. There is no within-unit constancy check. Replacing only treated post-period values of an otherwise unit-level covariate changed the same panel's ATT from 1.162507 to 1.919646 (+0.757139), without an error or warning.

This permits post-treatment adjustment and gives common and staggered paths different covariate semantics.

Required: enforce unit-constant covariates, or define and implement a pre-treatment baseline snapshot rule. Add a test that mutating post-treatment covariate values is rejected or cannot change the estimate.

[Medium] recommend_transformation() ignores cohort and overstates what pre-tests establish

The public function accepts cohort, but _safe_lwdid_fit() never passes first_treat; the placebo routine collapses all ever-treated units into one pseudo-treated group. Passing cohort="cohort" and cohort="not_a_column" produces the same recommendation and p-value. The joint statistic also sums squared marginal t-statistics under an independence assumption, despite overlapping estimates from the same units, and non-rejection is converted into a high-confidence recommendation.

Required: preferably remove this recommender and rely on the library's existing event-study/placebo machinery. If retained, validate and use cohort, estimate the joint covariance, and describe non-rejection as inconclusive rather than evidence that parallel trends holds.

[Medium] to_dict() is not JSON-serialisable as documented

to_dict() copies nested cohort dictionaries without normalising NumPy scalars. On an ordinary integer staggered panel, json.dumps(result.to_dict()) raises TypeError: Object of type int64 is not JSON serializable; cohort, time, and relative_time retain np.int64 values.

Required: recursively convert arrays and NumPy scalars to built-in JSON types, and test json.dumps(result.to_dict()) for both common and staggered/event-study results.

[Medium] the strict documentation build fails, and the earlier public examples/time-scale defects remain

python -m sphinx -W -b html docs <tmp-output> fails with:

diff_diff/lwdid.py:docstring of diff_diff.lwdid.LWDiD:91: ERROR: Malformed table.
docs/tutorials/27_lwdid.ipynb:1315: WARNING: Title underline too short.
build finished with problems, 3 warnings (with warnings treated as errors).

The table failure comes from columns that exceed the declared grid width. The datetime/Period staggered failures and stale public examples reported in the earlier review comment are also still present on this head.

Required: make the warning-as-error documentation build pass, execute the public examples as tests, and support native ordered time scales or reject them with a clear validation error.

Verification performed

• Focused non-slow/non-realdata LWDiD suite: 246 passed, 43 skipped, 31 deselected.
• Slow LWDiD tests: 2 passed.
• Real-data LWDiD tests: 23 passed, 5 xfailed, 1 xpassed; these include 2,284 propensity/separation warnings.
• Ruff 0.16.0, Black 26.5.1, mypy 2.3.0, and git diff --check: clean on the touched LWDiD source/test files.
• Independent numerical/API probes: failures reproduced as described above.

My recommended scope is: fix the first, second, third, and fifth findings in the estimator core; remove the sensitivity and transformation-recommendation helpers from this PR unless they receive their own identification and failure-contract work; then close the result serialisation and documentation defects. That keeps the contribution substantive without asking this already large PR to establish two additional diagnostic APIs at the same time.

[gorgeousfish]

@igerber @shawcharles — all items from the Aug 8 review rounds are now resolved on the branch, which has also been rebased onto latest main (post-#759, 54df914). Point-by-point below.

High

1. classical cross-cell covariance — _ols_treatment_influence() now returns a residual-direction influence vector rescaled so its norm still reproduces the textbook classical SE for the marginal cell. Joint covariance therefore keeps unit-level residual dependence instead of the spurious sigma*basis correlation. On a shared-control staggered panel the classical overall SE went from ~1.84x the unit bootstrap to within ~6.5% of it; single-cell classical SEs, Prop 99 goldens, and all existing test values are unchanged. New tests: joint influence basis, bootstrap agreement (B=400, slow-marked), sup-t band sanity.

2. randomization inference under ties — _compute_pvalue() now counts permutations at least as extreme (>=) under the existing Phipson-Smyth plus-one convention; an all-tie distribution returns exactly p = 1.0. Tests cover all-tie (end-to-end and unit), half-tie, and a small discrete design against the exact value. This fixes only the unambiguous degenerate case; the broader Stata lwdid, ri tie-handling question from your §4 remains yours to call.

3. all-eventually-treated designs — adopted the reject option: control_group='not_yet_treated' with no never-treated unit now raises ValueError instead of silently dropping final-period cells. Docstring and the Input Contract in docs/api/lwdid.rst state the requirement; the datetime path is covered by the same guard. Designs with never-treated units are unaffected (tested: full event-time coverage).

4. sensitivity failed fits — non-finite baselines or <2 finite alternative estimates now yield a NaN ratio classified as a new "not_estimable" level with a SensitivityWarning; missing columns raise ValueError before the try block, and the except scope is narrowed to the fit itself. We kept the helpers (rather than removing them) because the §9 triage retained them deliberately — happy to drop them instead if you prefer the smaller surface.

5. post-treatment covariates in staggered designs — covariates must now be unit-constant; time-varying columns raise ValueError (same contract as lwdid-py's validation, https://github.com/gorgeousfish/lwdid-py). One disclosure: the common-timing path takes each unit's first-row covariate value and does not validate constancy, whereas lwdid-py raises there too — flagging as a follow-up candidate rather than bundling it here.

Medium

6. recommend_transformation now forwards the cohort/gvar column through the placebo machinery (pseudo-cohort staggered fits) and raises ValueError on unknown columns; a spy-based test verifies first_treat actually reaches the internal fits. The joint statistic's independence assumption is untouched — your call whether to revisit it.

7. LWDiDResults.to_dict() is fully JSON-native: numpy scalars/arrays and dict keys are converted recursively, and datetime/Period labels serialize as ISO-8601 / period strings (NaT -> None), so json.dumps works for common-timing, staggered, event-study, and datetime panels. Round-trip tests added.

8. Sphinx / datetime / stale examples — the malformed docstring table is fixed and a full sphinx -W -b html build now succeeds. Staggered fits natively accept datetime64 and Period time scales via an ordered-support integer-position encoding with labels mapped back afterwards; numeric panels bypass the encoding entirely (pointwise-equality tested). The class docstring example and all four lwdid.rst staggered examples now use the first_treat= contract and were re-executed.

Additional fixes found while validating

• cluster equal to the unit column crashed the staggered path (KeyError after set_index); fixed, with a regression test asserting identical results to an explicit duplicate cluster column.
• Equivalence-suite aggregation basis: with lwdid-py 0.2.2 installed, the staggered SE comparisons were exercising lwdid-py's default aggregate='cohort', which assumes independence across cohort estimates and understates the overall SE when control units are shared. LW 2026 eq. 7.19 recommends the pooled cross-section regression precisely because "it automatically accounts for the correlations among the tau_g" — the Stata implementation and lwdid-py's aggregate='overall' both follow that basis, and diff-diff's joint influence-function SE matches it (0.02509 vs 0.02495; panel bootstrap ~0.0248). The suite now requests aggregate="overall" and compares SEs at rtol 0.01 (RA) / 0.05 (IPW/IPWRA, logit optimization path differences), with ATT assertions unchanged at strict tolerances. All 43 equivalence tests pass when run against lwdid-py 0.2.2.

Contract disclosures

The Input Contract in docs/api/lwdid.rst now lists five enforced requirements (adding the all-eventually-treated rejection and unit-constant covariates), plus notes on recommend_transformation validation and the JSON-native to_dict output. CHANGELOG has a consolidated entry for the round under Unreleased.

Validation

• Rebased onto main @ 54df914; branch is 35 commits ahead, review-round work is the top 15 commits; conflicts were limited to README.md and guides/llms.txt catalog lines.
• Full lwdid suite: 351 passed, 0 failed (including all 43 lwdid-py equivalence tests executed for real); docs IA tests 19 passed; sphinx -W clean.
• Lint at CI-pinned versions: ruff 0.16.0 and black 26.5.1 clean; mypy reports zero new errors relative to main.

Standing asks

Unchanged from before: your decision on the Walmart SE goldens (§3 — agreed that xfail should not be the final representation), the RI p-value convention against Stata (§4 — degenerate all-tie case fixed above, convention choice still open), §8 wild-bootstrap integration stays untouched on my side, and when you get a chance: the ready-for-ci label plus a re-sync of mirror #690.