Add pre-period event study coefficients to ImputationDiD and TwoStageDiD

[igerber]

Summary

• Add pretrends: bool = False parameter to both ImputationDiD and TwoStageDiD
• When enabled, event study includes pre-treatment horizons (should be ~0 under parallel trends) for visual pre-trends assessment
• ImputationDiD: computes tau_hat for pre-treatment observations of eventually-treated units, aggregates by relative time, uses extended Theorem 3 variance (pre-period targets contribute both directly and via FE correction)
• TwoStageDiD: extends Stage 2 design matrix with pre-period relative-time dummies; GMM sandwich variance handles naturally
• Bootstrap inference extended for both estimators
• Addresses AI code review: fixed P0 bootstrap balance_e mismatch and P1 survey-weighted variance regression

Methodology references (required if estimator / math changes)

• Method name(s): ImputationDiD (Borusyak-Jaravel-Spiess 2024), TwoStageDiD (Gardner 2022)
• Paper / source link(s): BJS Section 4.2 (pre-period event study coefficients), Gardner Theorem 1 (GMM sandwich validity)
• Any intentional deviations from the source (and why): None — pre-period variance uses Theorem 3 extension where v_it = w_it + FE_correction for observations in both omega_0 and the aggregation target

Validation

• Tests added/updated: tests/test_pretrends_event_study.py (22 tests covering both estimators, cross-estimator consistency, edge cases, and regression tests for survey/covariates/bootstrap paths)
• Backtest / simulation / notebook evidence (if applicable): N/A

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

Overall Assessment

⛔ Blocker

The conceptual pretrends=True extensions for ImputationDiD and TwoStageDiD are documented in the registry, and I did not find another unmitigated paper/registry mismatch. The blocker is an implementation bug in the new ImputationDiD pre-period inference path: it uses DataFrame index labels as positional indices when mapping a horizon back into Omega_0, which can silently produce wrong pre-period SEs/CI/p-values or raise on non-default indices.

Executive Summary

• ImputationDiD(pretrends=True) has a P0 in the new pre-period SE/bootstrap path: horizon targets are mapped back into Omega_0 with label-based indexing, so non-default DataFrame indices can corrupt inference or crash.
• The conceptual methodology changes themselves are documented in docs/methodology/REGISTRY.md:L839-L850 and docs/methodology/REGISTRY.md:L941-L946; I did not find an additional undocumented deviation there.
• The bug affects pre-period inference in ImputationDiD, not the pre-period point estimates.
• TwoStageDiD’s analytical and bootstrap horizon construction looks internally consistent on static inspection.
• The new tests miss non-default-index regression cases, which is exactly how the blocker above slips through.

Methodology

• Severity: P0. Impact: The new pre-period ImputationDiD inference code builds a positional boolean array and then marks it with df_pre.index[h_mask_pre]. Because fit() preserves the caller’s original index instead of resetting it, gapped indices can raise IndexError, and permuted integer indices can silently mark the wrong untreated rows. That corrupts preperiod_weights_h, so pre-period SEs, p-values, CIs, and bootstrap inference can be wrong with no warning, even when the point estimates are correct. Concrete fix: derive the pre-period target mask strictly from positional masks (omega_pre_mask / omega_0_mask) or reset the working index once at fit start and keep all later mappings positional; apply the same fix in both the analytical and bootstrap paths. See diff_diff/imputation.py:L239-L240, diff_diff/imputation.py:L1838-L1840, and diff_diff/imputation_bootstrap.py:L231-L233.

No other unmitigated methodology deviations found. The intended BJS/Gardner extensions are documented in docs/methodology/REGISTRY.md:L839-L850 and docs/methodology/REGISTRY.md:L941-L946.

Code Quality

No additional findings beyond the blocker above.

Performance

No new performance findings in the changed code.

Maintainability

• Severity: P3. Impact: The analytical and bootstrap ImputationDiD pre-period paths duplicate the same horizon-to-Omega_0 weight-construction logic, which is why the same indexing defect exists in both places. Concrete fix: extract a shared helper for building pre-period target masks/weights and reuse it from both callers. See diff_diff/imputation.py:L1831-L1867 and diff_diff/imputation_bootstrap.py:L229-L268.

Tech Debt

No deferrable finding mitigates the blocker. This issue is not tracked in TODO.md, and per the review policy it would not be eligible for downgrade anyway because it affects statistical output.

Security

No findings.

Documentation/Tests

• Severity: P3. Impact: The new test file exercises TwoStageDiD pretrends only analytically, and ImputationDiD bootstrap only on default indices, so the non-default-index failure mode and some new bootstrap/survey parameter interactions are unguarded. Concrete fix: add regression tests for permuted and gapped indices under ImputationDiD(pretrends=True) with and without bootstrap, plus at least one TwoStageDiD(pretrends=True, n_bootstrap>0) case and one survey-weighted event-study case. See tests/test_pretrends_event_study.py:L297-L425 and tests/test_pretrends_event_study.py:L503-L612.

Execution note: I could not run the test file in this environment because numpy and pytest are not installed.

Path to Approval

1. Fix the ImputationDiD pre-period target-to-Omega_0 mapping so it is purely positional in both diff_diff/imputation.py:L1838-L1840 and diff_diff/imputation_bootstrap.py:L231-L233.
2. Add regression tests that rerun ImputationDiD(pretrends=True, aggregate="event_study") on the same data with RangeIndex, permuted integer index, and nonconsecutive filtered index, both with n_bootstrap=0 and n_bootstrap>0, and assert identical pre-period inference after reset_index(drop=True).

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 64dc042953500fa43145de082b4378a1a17c5236

---

Overall Assessment

⚠️ Needs changes

The prior ImputationDiD non-default-index blocker appears fixed, but this re-review still has one unmitigated P1: ImputationDiD(pretrends=True) does not implement the BJS pre-trend coefficients it cites. The new path computes placebo-style averages of pre-period residuals and Theorem 3-based SEs in diff_diff/imputation.py:L458, diff_diff/imputation.py:L1694, and diff_diff/imputation.py:L1088, while the existing paper-consistent implementation in diff_diff/imputation.py:L2138 estimates untreated-only lead coefficients from Test 1 / Equation 9. BJS describes event-study plots as combining ATT-by-horizon estimates with pre-trend coefficients from Test 1 and treats placebo estimates as a separate approach with drawbacks; Gardner’s two-stage event-study extension does support lead coefficients, so I did not find the same source-material problem in TwoStageDiD.

Executive Summary

• The previous P0 on ImputationDiD pre-period index mapping looks addressed: both analytical and bootstrap paths now use positional masks, and regression coverage was added in tests/test_pretrends_event_study.py:L620.
• Severity: P1. ImputationDiD(pretrends=True) is an undocumented methodology deviation from BJS. It outputs placebo-style pre-period residual averages instead of the Test 1 pre-trend coefficients BJS says should appear in event-study plots.
• The new REGISTRY entry in docs/methodology/REGISTRY.md:L839 documents that placebo construction as if it were BJS-conformant, but it is not labeled as a deviation, so it does not mitigate the issue.
• TwoStageDiD’s pre-period lead handling in diff_diff/two_stage.py:L994 is directionally consistent with Gardner’s event-study extension.
• The new tests reinforce the wrong ImputationDiD contract by asserting cross-estimator equality of pre-period coefficients in tests/test_pretrends_event_study.py:L446 instead of checking ImputationDiD against its existing pretrend_test() semantics.
• I could not execute the test suite here because pytest and numpy are not installed in this environment.

Methodology

• Severity: P1. Impact: ImputationDiD(pretrends=True) is not following the cited BJS source. The new code imputes/stores tau_hat on pre-treatment observations and averages those residuals by negative horizon in diff_diff/imputation.py:L458 and diff_diff/imputation.py:L1796, with SEs from the new Theorem-3-style path in diff_diff/imputation.py:L1088. But BJS says event-study plots should use pre-trend coefficients from Test 1, while “placebo estimates” are a different approach with stated drawbacks; this library already has that Test 1 implementation in diff_diff/imputation.py:L2138. Because the new REGISTRY entry in docs/methodology/REGISTRY.md:L839 presents the placebo path as BJS Section 4.2 with no **Note:**/deviation label, this is an unmitigated undocumented methodology deviation. Concrete fix: either reimplement ImputationDiD negative horizons by reusing the existing Test 1 lead-regression / cluster-robust inference machinery, or explicitly mark the placebo construction as a deviation in the registry and stop presenting it as BJS-conformant.

Code Quality

• No additional findings beyond the methodology issue above.

Performance

• No new performance findings in the changed code.

Maintainability

• No separate maintainability finding worth blocking or downgrading beyond the methodology issue.

Tech Debt

• The P1 above is not tracked under “Tech Debt from Code Reviews” in TODO.md:L47, so it is not mitigated.

Security

• No findings.

Documentation/Tests

• Severity: P2. Impact: the new test file locks in the same wrong ImputationDiD contract it should be checking. tests/test_pretrends_event_study.py:L446 asserts that ImputationDiD and TwoStageDiD pre-period point estimates should match, which is only true under the PR’s placebo-style Imputation implementation. A source-consistent ImputationDiD fix would instead be expected to match pretrend_test() / Test 1, not TwoStage’s lead regression. Concrete fix: replace the cross-estimator pre-period parity test with an ImputationDiD contract test against results.pretrend_test().lead_coefficients, and keep TwoStage pre-period tests separate.
• I could not run tests/test_pretrends_event_study.py locally because this environment lacks pytest and numpy.

Path to Approval

1. Change ImputationDiD negative-horizon event-study outputs to use the existing untreated-only lead regression from diff_diff/imputation.py:L2138, so pre-period coefficients/SEs reflect BJS Test 1 rather than placebo residual averages.
2. If placebo-style pre-period coefficients are intentional, add an explicit **Note:** deviation in docs/methodology/REGISTRY.md:L839 and update the docstrings so the feature no longer claims BJS Section 4.2 equivalence.
3. Update tests/test_pretrends_event_study.py:L446 to validate the chosen contract instead of asserting ImputationDiD/TwoStageDiD pre-period equality.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 2cdc050f58af92b779d8d332df6adc913f8de2aa

---

Overall Assessment

⛔ Blocker

The previous source-material concern appears resolved: ImputationDiD(pretrends=True) now pulls pre-period coefficients from the Test 1 / Equation 9 lead regression rather than the earlier placebo-style path, and I did not find a new paper/registry mismatch in the TwoStageDiD changes. But the PR still introduces one silent correctness bug and one event-study contract bug in ImputationDiD, so it is not ready for approval.

Executive Summary

• The prior BJS methodology issue looks fixed: the new Imputation negative-horizon path now reuses the same lead-regression machinery as pretrend_test(), which is consistent with the new registry text in docs/methodology/REGISTRY.md:L839-L850 and the shared helper in diff_diff/imputation.py:L1784-L1983.
• Severity P0: ImputationDiD(pretrends=True, survey_design=...) silently emits non-survey-aware pre-period coefficients and SEs. The estimator’s main fit path is survey-weighted, but the new pre-period branch calls an unweighted helper that pretrend_test() itself explicitly says is not survey-aware. See diff_diff/imputation.py:L397-L400, diff_diff/imputation.py:L1526-L1544, diff_diff/imputation.py:L1824-L1867, and diff_diff/imputation.py:L1892-L1897.
• Severity P1: ImputationDiD(pretrends=True, balance_e=...) does not apply the balanced-cohort restriction to the new pre-period coefficients, so negative horizons can come from a different cohort set than the post-treatment event study. When no cohort qualifies, Imputation can still return negative-horizon estimates instead of only the reference period. See diff_diff/imputation.py:L210-L218, diff_diff/imputation.py:L1492-L1544, diff_diff/imputation.py:L1654-L1671, and docs/methodology/REGISTRY.md:L862-L868.
• The added tests do not cover the two failing Imputation interactions above, and they also do not exercise the changed TwoStageDiD bootstrap path. The file adds horizon_max/anticipation checks and an unweighted Imputation bootstrap check, but no balance_e cases and no pretrends + survey_design case. See tests/test_pretrends_event_study.py:L220-L252, tests/test_pretrends_event_study.py:L520-L622, tests/test_pretrends_event_study.py:L297-L438, and diff_diff/two_stage_bootstrap.py:L321-L420.
• I could not run the test file in this environment because pytest and numpy are unavailable here.

Methodology

• Severity: P0. Impact: ImputationDiD(pretrends=True, survey_design=...) produces silently wrong pre-period statistical output. The estimator otherwise fits the untreated model with survey weights (diff_diff/imputation.py:L397-L400, diff_diff/imputation.py:L960-L999), but the new pre-period event-study branch always routes through _compute_lead_coefficients() (diff_diff/imputation.py:L1526-L1544), and that helper uses unweighted demeaning and unweighted solve_ols() (diff_diff/imputation.py:L1824-L1851). The same machinery is explicitly rejected for survey use in pretrend_test() because it “does not account for survey weights” (diff_diff/imputation.py:L1892-L1897). Concrete fix: either implement a survey-aware lead regression for pre-period coefficients (weighted demeaning, weighted OLS, and survey df/inference consistent with the estimator’s survey path) or raise NotImplementedError whenever pretrends=True is combined with survey_design.
• Severity: P1. Impact: balance_e is not propagated to Imputation’s new pre-period coefficients. The public contract says event study should restrict to cohorts observed at all relative times in [-balance_e, max_h] (diff_diff/imputation.py:L210-L212, docs/methodology/REGISTRY.md:L862-L868), and post-treatment horizons do use the balanced cohort mask (diff_diff/imputation.py:L1492-L1502). But the new pre-period branch computes coefficients on all untreated observations before any balanced-cohort restriction (diff_diff/imputation.py:L1526-L1544). That means negative horizons can be estimated on a different population than nonnegative horizons, and when no balanced cohort exists the empty-result warning logic can be bypassed because those unfiltered pre-period entries keep real_effects nonempty (diff_diff/imputation.py:L1654-L1671). Concrete fix: compute the balanced cohort set first and restrict the pre-period lead regression to those cohorts as well; if no cohort qualifies, return only the reference period with the same warning behavior as the post-treatment path.

Code Quality

• No additional findings beyond the methodological propagation bugs above.

Performance

• No findings in the changed code.

Maintainability

• No separate maintainability finding. The main risk is semantic inconsistency in the new Imputation event-study branches, not code organization.

Tech Debt

• No mitigating entry exists for these issues in TODO.md:L47-L71, so they are not downgraded by the project’s deferred-work policy.

Security

• No findings.

Documentation/Tests

• Severity: P1. Impact: the new parameter interaction coverage is incomplete in the exact places the changed code is fragile. The new tests cover horizon_max, anticipation, and an unweighted Imputation bootstrap path (tests/test_pretrends_event_study.py:L220-L252, tests/test_pretrends_event_study.py:L579-L735), and they include a survey event-study test only for pretrends=False (tests/test_pretrends_event_study.py:L520-L548). There is no balance_e case in the new file, no ImputationDiD(pretrends=True, survey_design=...) case, and no TwoStageDiD(pretrends=True, n_bootstrap>0) case even though diff_diff/two_stage_bootstrap.py:L321-L420 was modified. Concrete fix: add targeted regression tests for ImputationDiD(pretrends=True, survey_design=...), ImputationDiD(pretrends=True, balance_e=...) including the no-qualifying-cohort case, and TwoStageDiD(pretrends=True, n_bootstrap>0) negative-horizon bootstrap inference.
• I could not execute tests/test_pretrends_event_study.py here because this environment does not have pytest or numpy installed.

Path to Approval

1. Make ImputationDiD(pretrends=True) survey-safe: either implement a weighted/survey-aware lead regression for the pre-period coefficients or explicitly reject survey_design for that combination.
2. Thread balance_e through the Imputation pre-period path so negative and nonnegative event-study horizons use the same balanced cohort set; when no cohort qualifies, return only the reference period with a warning.
3. Add regression tests for pretrends + survey_design, pretrends + balance_e (including zero qualifying cohorts), and TwoStageDiD pretrends + bootstrap.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 38b707e026057f6e3accb47f503d91eb7c5e50b2

---

Overall Assessment

⚠️ Needs changes

Executive Summary

• The previous silent pretrends + survey_design correctness issue in ImputationDiD is now fail-fast instead of silent, which is the right direction.
• Severity P1: ImputationDiD(pretrends=True, balance_e=...) still applies balance_e incorrectly in the new pre-period path by shrinking the entire untreated lead-regression sample to balanced cohorts, which drops never-treated controls and changes the comparison set.
• Severity P1: the new survey rejection for ImputationDiD(pretrends=True) is not documented in the Methodology Registry, and the new error message points users to pretrend_test() even though that path is also survey-unsupported.
• TwoStageDiD looks methodologically consistent in this diff; I did not find a new source-material mismatch there.
• Test coverage improved, but the new Imputation balance_e test only checks that the call succeeds and would not catch the remaining control-sample bug.

Methodology

At a high level, the estimator choices are aligned with the cited sources: BJS Test 1 estimates Equation 9 by OLS on untreated observations only and explicitly says those coefficients can be used for event-study plots, while Gardner’s event-study extension amends the second stage to regress residualized outcomes on lead and duration indicators and keeps GMM-based asymptotics for valid SEs. I did not find a new paper-level mismatch in the TwoStageDiD changes. (academic.oup.com)

• Severity: P1. Impact: ImputationDiD(pretrends=True, balance_e=...) enforces balance by subsetting the entire untreated regression sample to balanced_cohorts in diff_diff/imputation.py:L1535-L1550. That drops never-treated controls and untreated observations from unbalanced cohorts before estimating Equation 9, even though BJS Test 1 is an OLS on untreated observations only, and the registry says balance_e analytical aggregation should use the shared cohort-mask contract in docs/methodology/REGISTRY.md:L862-L868. This means the previous balance_e finding is only partially fixed: negative-horizon coefficients are now restricted to balanced cohorts, but with the wrong untreated comparison set. Concrete fix: compute balanced_cohorts once, keep the full Omega_0 sample, and only activate/report lead indicators for balanced cohorts inside _compute_lead_coefficients() instead of filtering df_0. (academic.oup.com)
• Severity: P1. Impact: the PR now intentionally rejects ImputationDiD(pretrends=True, survey_design=...) in diff_diff/imputation.py:L239-L245, and the test suite encodes that behavior in tests/test_pretrends_event_study.py:L737-L757, but the new registry text still describes Imputation pre-period coefficients and survey support as if this combination is available in docs/methodology/REGISTRY.md:L839-L850 and docs/methodology/REGISTRY.md:L868-L868. Under the review rubric, that is an undocumented methodology/capability mismatch. The new error text is also misleading because pretrend_test() is itself survey-unsupported. Concrete fix: add a labeled **Note:** in the Imputation registry section documenting that pretrends=True with survey_design raises NotImplementedError, and change the error text to remove the pretrend_test() suggestion.

Code Quality

• No separate code-quality findings beyond the methodological issues above.

Performance

• No findings in the changed code.

Maintainability

• No separate maintainability findings.

Tech Debt

• No mitigating entry in TODO.md:L47-L71 covers the P1 issues above, so they remain unmitigated.

Security

• No findings.

Documentation/Tests

• Severity: P2. Impact: the new Imputation balance_e regression test in tests/test_pretrends_event_study.py:L759-L787 only asserts that the call succeeds. It would not fail if never-treated controls were incorrectly removed from the Equation 9 lead regression, which is exactly the remaining bug. Concrete fix: add a regression test that compares pretrends=True, balance_e=... against a manual Equation 9 fit that keeps never-treated controls in Omega_0, plus a no-qualifying-cohort case that expects only the reference period and a warning.
• I could not execute tests/test_pretrends_event_study.py here because this environment does not have pytest or numpy installed.

Path to Approval

1. In diff_diff/imputation.py, change the balance_e handling for pre-period coefficients so the Equation 9 regression keeps the full untreated sample and restricts only the lead targets to balanced cohorts.
2. In docs/methodology/REGISTRY.md, add a labeled **Note:** that ImputationDiD(pretrends=True, survey_design=...) is currently unsupported, and update the NotImplementedError text in diff_diff/imputation.py:L239-L245 so it does not recommend pretrend_test().
3. In tests/test_pretrends_event_study.py, add a regression that fails if never-treated controls are dropped from the balance_e pre-period lead regression, and add the zero-qualifying-cohort warning case.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 0684ad4c39adf966fdb29e12283c1132cfcbff80

---

Overall Assessment

⚠️ Needs changes

Executive Summary

• The previous ImputationDiD comparison-set problem is fixed: the new lead-regression path keeps the full untreated Omega_0 sample and only restricts the lead indicators themselves, which matches BJS Test 1’s untreated-only OLS setup. (academic.oup.com)
• The previous survey-support/documentation issue is also fixed: pretrends=True now fails fast with survey_design, and the registry documents that limitation in a labeled note at diff_diff/imputation.py:L239-L245 and docs/methodology/REGISTRY.md:L839-L853.
• Severity P1: ImputationDiD(pretrends=True, balance_e=...) still mishandles the zero-qualifying-cohort branch. When no cohort satisfies balance_e, the code resets balanced_cohorts to None, so pre-period coefficients are still estimated from all untreated eventual-treated observations instead of returning only the reference period.
• TwoStageDiD looks methodologically consistent in this diff: the changed Stage 2 path now includes pre-period relative-time dummies when requested, which fits Gardner’s untreated-first-stage / event-study extension. (arxiv.org)
• Test coverage is substantially better, but the new Imputation balance_e regression at tests/test_pretrends_event_study.py:L761-L789 only checks that the call succeeds and would not catch the remaining empty-sample bug.

Methodology

BJS Test 1 estimates lead coefficients by OLS on untreated observations only and explicitly notes that those coefficients can be used for event-study plots; Gardner’s two-stage approach estimates fixed effects on untreated observations first and extends to event-study analyses of pre-trends and duration-specific effects. Against that benchmark, I do not see a new paper-level mismatch in TwoStageDiD, and the earlier ImputationDiD control-sample issue is resolved. (academic.oup.com)

• Severity: P1. Impact: in diff_diff/imputation.py:L1541-L1565, balanced_cohorts is reset to None when no cohort qualifies for balance_e, so rel_time_for_leads falls back to all untreated eventual-treated observations. Because the empty-result guard at diff_diff/imputation.py:L1689-L1706 keys off n_obs, the estimator can return pre-period coefficients from an empty balanced sample instead of the documented reference-period-only result. This violates the registry’s stated balance_e contract at docs/methodology/REGISTRY.md:L851-L864. Concrete fix: keep the empty balanced-cohort set as “no qualifying cohorts” and skip pre-period coefficient generation entirely in that branch so the existing empty-result warning fires; while there, compute pre-period n_obs from the same balanced-cohort restriction used to build the lead indicators instead of the unrestricted count at diff_diff/imputation.py:L1927-L1938.

Code Quality

• No findings.

Performance

• No findings.

Maintainability

• No findings.

Tech Debt

• No separate findings. The relevant issue is not tracked in TODO.md:L47-L88, so the P1 above is unmitigated.

Security

• No findings.

Documentation/Tests

• Severity: P2. Impact: tests/test_pretrends_event_study.py:L761-L789 only asserts that ImputationDiD(pretrends=True, balance_e=1) runs. It would not fail if the zero-qualifying-cohort branch incorrectly returned pre-period coefficients, and it does not verify that pre-period n_obs reflects the actually restricted sample. Concrete fix: add a regression with balance_e larger than feasible support that expects the warning and a reference-period-only event-study result, plus an assertion on pre-period n_obs in a mixed balanced/unbalanced panel.

I could not execute tests/test_pretrends_event_study.py here because this environment does not have pytest or numpy installed.

Path to Approval

1. In diff_diff/imputation.py:L1541-L1578, treat an empty balanced_cohorts set as “no qualifying cohorts” rather than falling back to None, so pretrends=True, balance_e=... cannot emit lead coefficients from the unrestricted sample.
2. In diff_diff/imputation.py:L1927-L1938, compute pre-period n_obs using the same balanced-cohort mask used for the lead dummies.
3. In tests/test_pretrends_event_study.py:L761-L789, add a no-qualifying-cohort regression that asserts the warning and reference-period-only output for ImputationDiD(pretrends=True).

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 0b65daa6038df51a9217c86ed78263757c7d9d95

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings in the changed estimator, variance, or bootstrap paths.

Executive Summary

• Previous P1 resolved: ImputationDiD no longer falls back to unrestricted pre-period leads when balance_e leaves no qualifying cohorts. The new guard skips pre-period generation at diff_diff/imputation.py:L1543, the existing empty-result warning path at diff_diff/imputation.py:L1705 now handles that branch, and pre-period n_obs is taken from the restricted lead indicator at diff_diff/imputation.py:L1945.
• Previous survey-support concern is now explicit and mitigated: pretrends=True with survey_design fails fast at diff_diff/imputation.py:L239, and the limitation is documented in docs/methodology/REGISTRY.md:L853.
• Methodologically, the ImputationDiD pre-period path at diff_diff/imputation.py:L1534 and diff_diff/imputation.py:L1834 matches BJS Test 1’s untreated-only lead regression, and the TwoStageDiD event-study path at diff_diff/two_stage.py:L1025 is consistent with Gardner’s event-study formulation and asymptotics. citeturn2view0turn2view1turn2view2turn1view1turn4view3
• I did not find any unmitigated P0/P1 issues in the changed control-group logic, weighting, variance, NaN handling, or parameter propagation.
• Remaining non-blocking gap: tests/test_pretrends_event_study.py:L761 still does not assert the zero-qualifying-cohort or restricted-n_obs cases from the prior review.

Methodology

• Severity: P3 informational. Impact: the only in-scope methodology limitation is the new ImputationDiD(pretrends=True, survey_design=...) rejection. It is explicitly enforced at diff_diff/imputation.py:L239 and documented with a labeled note at docs/methodology/REGISTRY.md:L853, so this is mitigated rather than a defect. Concrete fix: none required for approval.

Code Quality

• No findings.

Performance

• No findings.

Maintainability

• No findings.

Tech Debt

• No findings.

Security

• No findings.

Documentation/Tests

• Severity: P2. Impact: tests/test_pretrends_event_study.py:L761 still only verifies that ImputationDiD(pretrends=True, balance_e=1) runs. It does not exercise the zero-qualifying-cohort branch fixed at diff_diff/imputation.py:L1543, and it does not assert that pre-period n_obs follows the restricted lead indicator at diff_diff/imputation.py:L1945. That leaves the specific regression from the last review vulnerable to reintroduction without a failing test. Concrete fix: add one case with balance_e larger than feasible support that expects the warning and reference-period-only output, and one mixed-support case that asserts pre-period n_obs falls when an unbalanced cohort is excluded.
• I could not execute tests/test_pretrends_event_study.py here because this environment does not have pytest, numpy, pandas, or scipy installed.