diff --git a/CHANGELOG.md b/CHANGELOG.md index f1fa6e1a..2d584854 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **`SpilloverDiD(survey_design=SurveyDesign.subpopulation(...))` full-design retention via zero-pad scores (Wave E.3).** Closes the Wave E.1/E.2/follow-up documented limitation at `REGISTRY.md:3249`: `SurveyDesign.subpopulation()`-derived designs AND warn-and-drop fits now preserve the full-domain resolved survey design — `n_psu` / `n_strata` / `df_survey` / Binder TSL per-stratum centering reflect the FULL domain rather than the post-`finite_mask` fit sample. **Documented synthesis (library-convention adoption, NOT new methodology):** Wave E.3 adopts the canonical "zero-pad scores to full panel + retain full-design resolved survey" pattern from R `survey::svyrecvar(subset())` (Lumley 2010 §2.5) already established in `diff_diff/imputation.py:2175-2183` (PreTrendsImputation lead regression — Omega_0 scores zero-padded back to full panel length) and `diff_diff/prep.py:1401-1432` (DCDH cell variance — IF zero-padded outside the cell). Wave E.3 propagates the same convention to SpilloverDiD's Wave E.1 Binder TSL × Wave D Gardner GMM × Wave E.2/follow-up stratified-Conley + serial Bartlett meat. **Mechanical realization (one new `_compute_gmm_corrected_meat` kwarg):** the gamma_hat / Psi build stays on SURVEY-FINITE-MASK inputs (`X_1_sparse_fit`, `X_10_sparse_fit`, `eps_10_fit` built on `survey_finite_mask = finite_mask & survey_weights > 0`; `X_2_kept_gamma`, `eps_2_fit_gamma`, `survey_weights_fit_gamma` projected from the fit-sample frame down to survey_finite_mask) so the drop-first stage-1 FE column space is bit-identical to the pre-E.3 path. `_compute_gmm_corrected_meat` gains a new optional kwarg `score_pad_mask: Optional[np.ndarray] = None`: when supplied, the helper zero-pads the fit-sample `Psi` to full panel length AFTER construction but BEFORE kernel dispatch via `Psi_padded[score_pad_mask] = Psi`. Kernel-dispatch arrays (`cluster_ids`, `conley_coords`, `conley_time`, `conley_unit`, `resolved_survey`) are passed at FULL length so the meat helpers (Binder TSL / stratified-Conley / serial Bartlett) see the full-domain PSU / strata / centroid / time geometry. The `_validate_conley_kwargs` call inside the helper reads `n_for_conley = len(score_pad_mask)` when the kwarg is set so the Conley shape checks see the full-length geometry. **`gamma_hat` invariance:** the gamma_hat solve operates on fit-sample inputs throughout — bit-identical to the pre-E.3 path (critical for the case where `_build_butts_fe_design_csr`'s `pd.factorize` re-compaction would drop a different unit's column under a full-length FE build than under a fit-length one). **Bread invariance:** `A_22 = X_2_kept' W X_2_kept` at `spillover.py:3187-3214` still uses fit-length `X_2_kept` because `A_22_full = X_2_full' W_full X_2_full` equals `A_22_kept` when zero-weight rows contribute zero. **A2 invariant:** warn-and-drop and `SurveyDesign.subpopulation()` drops are treated identically — both apply the zero-pad mechanism. The "both mechanisms compose cleanly" case (subpop-excluded row that is ALSO warn-and-dropped) produces `Psi = 0` from either cause; the PSU still counts toward `n_psu_full`. Hand-computation methodology anchor at `_scratch/wave_e3_smoke.py` codifies the A2 invariant on 4 PSU × 4 period × 3 obs synthetic. **Subpopulation parity vs upstream-subset:** `df_survey` matches the full domain regardless of how many rows the subpopulation mask excludes (mirrors R `svyglm(design=subset(d, mask))` vs `svyglm(design=svydesign(data=data[mask], ...))`). SE may differ by design — subpopulation retains zero-padded PSU geometry; upstream-subset drops PSUs entirely. **Pre-E.3 baseline parity:** when `finite_mask.all() == True` AND all weights `> 0`, the Wave E.3 zero-pad is a no-op — ATT + SE + n_psu + df_survey match pre-E.3 baseline values via FIXED GOLDEN values at `test_c` (`rtol=1e-12, atol=1e-12`). **Cross-surface n_psu consistency:** top-level `res.n_psu` reads from `len(resolved_survey_fit.weights)` on the implicit-PSU branch (was `int(finite_mask.sum())` pre-codex-R1-P2-fix); this keeps `res.n_psu == res.survey_metadata.n_psu` on weights-only / strata-only survey designs under warn-and-drop. Regression at `test_c2`. **Restrictions inherited:** replicate-weight variance + subpopulation continues to raise `NotImplementedError` at the Wave E.1 gate. TwoStageDiD's analogous `finite_mask + design-subset` pattern at `two_stage.py:567-601` is NOT yet adopted to Wave E.3 — separate parity follow-up tracked in `TODO.md` (an expected-divergence test was attempted but TwoStageDiD's always-treated handling at `two_stage.py:294-336` differs from SpilloverDiD's per-unit Omega_0 check, so the divergence didn't materialize on the standard fixture; the parity follow-up should add its own targeted regression). **Implementation:** `spillover.py:2845-2896` design-subset block deleted; `survey_weights_fit = survey_weights[finite_mask]` retained for the stage-2 OLS solve which still operates on the fit sample; `cluster_ids_full[finite_mask]` subset dropped on the survey path. `_compute_gmm_corrected_meat` call at `spillover.py:3163` now receives FIT-LENGTH gamma_hat-construction inputs (unchanged) plus FULL-LENGTH kernel-dispatch arrays (`cluster_ids_for_meat`, `conley_*_for_meat`, `resolved_survey_fit`) plus the new `score_pad_mask=survey_finite_mask` kwarg; no-survey path passes `score_pad_mask=None` and uses fit-length variables throughout (bit-identical to pre-E.3). `_compute_gmm_corrected_meat` at `two_stage.py:62-80` adds one new optional kwarg `score_pad_mask: Optional[np.ndarray] = None` and one post-Psi-construction zero-pad block; the `_validate_conley_kwargs` call uses `n_for_conley = len(score_pad_mask)` when the kwarg is set. Within-unit-constancy validator at `spillover.py:2913` updated to operate on full-length unit array. Second `compute_survey_metadata` recompute at `spillover.py:2954-2959` uses full-length `raw_w`. No `_compute_stratified_meat_from_psu_scores` / `_compute_stratified_conley_meat` / `_compute_stratified_serial_bartlett_meat` signature changes. **Tests:** new `TestSpilloverDiDWaveE3SubpopulationFullDesign` and `TestSpilloverDiDWaveE3SubpopulationFullDesignEventStudy` classes in `tests/test_spillover.py` (19 tests: pre-E.3 baseline parity via pinned goldens, n_psu cross-surface consistency on implicit-PSU branch, A2 invariant (zero-pad mechanics via mock-spy), subpopulation × explicit-PSU parity, conley + lag>0 + subpopulation × explicit-PSU / cluster-injection / weights-only branches, cluster-as-PSU + subpopulation parity, unit with BOTH zero weight AND no Omega_0 support, gamma_hat-build sample excludes zero-weight rows, n_obs / n_treated / n_control / n_far_away_obs reflect count_mask, warn-drop SE drift golden, ATT bit-equality under PSU-last-sort exclusion, exact event-study n_obs propagation, event-study on both is_staggered branches with analytical + conley+lag variants). Pre-existing Wave E.1 `test_p2_finite_mask_forces_drop_under_survey` assertion flipped from `n_psu=8` (subset) to `n_psu=10` (full domain) to reflect the new contract. - **ChaisemartinDHaultfoeuille (DCDH) methodology-review-tracker promotion.** Tracker row flipped **In Progress** → **Complete** with full Verified Components / Test Coverage / Corrections Made / Deviations / Outstanding Concerns structure mirroring the HAD precedent (PR #473) and ContinuousDiD precedent (PR #476). REGISTRY `## ChaisemartinDHaultfoeuille` gains a formal `### Deviations from the paper / from R / library extensions` block consolidating 7 documented deviations into a single AI-review-recognized labeled surface (per CLAUDE.md "Documenting Deviations (AI Review Compatibility)"): (D1) equal-cell weighting (deviation from BOTH AER 2020 Equation 3 AND R `DIDmultiplegtDYN`); (D2) period-based vs cohort-based stable controls; (D3) balanced-baseline panel + interior-gap drops + terminal-missingness retention + cell-period-allocator targeted `ValueError`; (D4) SE normalization `N_l` vs R `G` (~4% smaller analytical SE); (D5) singleton-cohort degeneracy → NaN with `UserWarning`; (D6) `<50%` switcher warning at far horizons (library extension citing Favara-Imbs application, footnote 14 of NBER WP 29873); (D7) Phase 3 `DID^X` covariate first-stage equal-cell weights. R cross-language coverage holds at documented tolerance bands in `tests/test_chaisemartin_dhaultfoeuille_parity.py` (`POINT_RTOL = 1e-4` on pure-direction point estimates, `MIXED_POINT_RTOL = 0.025` on mixed-direction, `PURE_DIRECTION_SE_RTOL = 0.05` on pure-direction SE, `SE_RTOL = 0.10` on multi-horizon SE, `se_rtol=0.15` on the long-panel `L_max=5` joiners-only scenario where cell-count-weighting compounds). No source code changes, no new tests, no new docstrings — consolidation only against the existing 12 methodology tests (`tests/test_methodology_chaisemartin_dhaultfoeuille.py`), 26 R-parity tests (`tests/test_chaisemartin_dhaultfoeuille_parity.py`), 352 unit tests (`tests/test_chaisemartin_dhaultfoeuille.py`), survey suites (`tests/test_survey_dcdh.py`, `tests/test_survey_dcdh_replicate_psu.py`, three cell-period coverage suites), and two primary-source DCDH paper reviews on disk (2020 AER + 2022/2023 NBER WP 29873 via PR #478; the `dechaisemartin-2026-review.md` on disk is HAD's primary source, not DCDH's, and is referenced as adjacent context only). The REGISTRY Deviations block uses semantic section-name anchors (rather than fragile line numbers) for back-references to other parts of the DCDH section — an intentional divergence from the PR #476 ContinuousDiD precedent reflecting PR-A wording-drift CI feedback that flagged line-number cross-references as drift-prone in long sections. `METHODOLOGY_REVIEW.md` DCDH row promoted **In Progress** → **Complete**; L27 In Progress example paragraph re-pointed to WooldridgeDiD; L1289 priority-order queue item #6 (DCDH) removed and items #7-#11 renumbered to #6-#10. +### Changed +- **Internal refactor: dedup serial Bartlett kernel construction and PSD guard between Conley no-survey and TwoStage survey paths.** Extracts `_serial_bartlett_kernel_matrix(t_codes, L)` and `_validate_meat_psd(M, *, error_msg, warning_template, stacklevel=3)` to `diff_diff/conley.py`. Replaces three inline kernel constructions (`conley.py` panel-block branch, `two_stage.py` survey singleton-adjust branch, `two_stage.py` survey multi-PSU branch) and two inline finite-plus-eigvalsh guards (`conley.py::_compute_conley_meat`, `two_stage.py` survey panel-block orchestrator) with helper calls. No behavior change — methodology anchor at `tests/test_spillover.py::TestSpilloverDiDWaveE2FollowupConleySurveyLagCutoff` (21 tests including hand-computed serial Bartlett HAC at L=1) and existing PSD-warning monkey-patch tests at `tests/test_conley_vcov.py::TestConleyDirectHelper::{test_uniform_kernel_negative_eigenvalue_warns, test_indefinite_meat_warning_fires_for_bartlett}` still pass unchanged (substring `"bartlett"` / `"uniform"` / `"negative eigenvalue"` in warning messages preserved byte-for-byte). New `TestSerialBartlettKernelMatrix`-grouped tests in `TestConleyKernels` (5 tests: hand-computed L=2 / L=1 / L=0 degenerate / single-element / int-vs-float bit-equality contract) and new `TestValidateMeatPsd` class (4 tests: non-finite raises with caller's `error_msg`, negative-eigenvalue warns with `{eigval:.2e}` substituted, PSD matrix silent, threshold boundary at -5e-13 silent). Closes `TODO.md` Bartlett-dedup row. + ## [3.4.1] - 2026-05-21 ### Added diff --git a/TODO.md b/TODO.md index fdbc2c26..dce7b041 100644 --- a/TODO.md +++ b/TODO.md @@ -155,7 +155,6 @@ Deferred items from PR reviews that were not addressed before merge. | `SyntheticDiD(vcov_type="conley")` support. Currently raises `TypeError` at `__init__` because SyntheticDiD uses `variance_method ∈ {bootstrap, jackknife, placebo}` rather than the analytical sandwich that Conley plugs into. Wiring would require either reimplementing an analytical sandwich path for SyntheticDiD or designing a spatial-block bootstrap (new methodology, Politis-Romano 1994 territory). | `synthetic_did.py::SyntheticDiD` | follow-up (spillover-conley) | Low | | `SpilloverDiD(survey_design=...)` replicate-weight variance (BRR / Fay / JK1 / JKn / SDR). Wave E.1 ships Taylor-linearization only. Per Gerber (2026) Appendix A, the IF-reweighting shortcut does NOT apply to TwoStageDiD-class estimators because `gamma_hat` is weight-sensitive; correct support requires per-replicate full re-fit of stage 1 and stage 2 (200+ LoC of test surface beyond E.1). | `spillover.py::SpilloverDiD.fit`, `survey.py::compute_replicate_refit_variance` | follow-up | Low | | `compute_survey_metadata(resolved_survey, raw_w_for_meta)` helper extraction. Wave E.1/E.3 contain two near-duplicate `raw_w_for_meta` constructions (upstream + post-cluster-injection) that differ only in which point of the resolution pipeline they fire at. Factor out a shared helper that takes `(survey_design, data, [finite_mask])` and returns `(resolved_survey, raw_w_for_meta)` to reduce drift risk between the two paths. Cosmetic; behaviour unchanged. | `spillover.py::SpilloverDiD.fit` | follow-up | Low | -| Serial Bartlett kernel logic duplicated between `diff_diff/two_stage.py::_compute_stratified_serial_bartlett_meat` (survey path) and `diff_diff/conley.py::_compute_conley_meat` panel-block branch (no-survey path). Both compute `K[t,s] = (1 - |t-s|/(L+1)) * 1{|t-s| <= L, t != s}` over dense panel-period codes. Factor out a shared `_serial_bartlett_kernel_matrix(t_codes, L)` helper and a shared post-meat finite + PSD-warning guard so the survey and no-survey paths can't drift on diagnostics or kernel weights. Cosmetic; refactor doesn't change behavior. | `two_stage.py::_compute_stratified_serial_bartlett_meat`, `conley.py::_compute_conley_meat` | follow-up | Low | | `SpilloverDiD(vcov_type="conley", conley_lag_cutoff > 0, survey_design=...)` no-effective-PSU serial Bartlett HAC. Wave E.2 follow-up ships the panel-block composition when an effective PSU exists (explicit `survey_design.psu` OR injected via `cluster=` per `_inject_cluster_as_psu`). Weights-only / strata-only survey designs WITHOUT a cluster fallback raise `NotImplementedError` at `SpilloverDiD.fit` post-resolution because under the pseudo-PSU = obs-index fallback each pseudo-PSU appears in exactly one period — the per-PSU serial cross-period loop would silently contribute zero. Fix would either derive a unit-level serial fallback for no-PSU designs (mixes IF allocators with the pseudo-PSU spatial term — needs methodology work) or route the serial loop through `conley_unit` with explicit documentation of the IF-allocator asymmetry. Regression goldens vs the effective-PSU shipped path. | `spillover.py::SpilloverDiD.fit`, `two_stage.py::_compute_stratified_serial_bartlett_meat` | follow-up (Wave E.2 follow-up tail) | Low | | `SpilloverDiD(ring_method="count")` extension. Currently only the nearest-treated-ring specification is exposed. Count-of-treated-in-ring (paper Section 3.2 end) is methodologically supported by Butts but re-introduces functional-form dependence; expose with an explicit kwarg gate and documentation warning. | `spillover.py::SpilloverDiD.fit` | follow-up | Low | | `SpilloverDiD` data-driven `d_bar` selection (Butts 2021b / Butts 2023 JUE Insight cross-validation). | `spillover.py::SpilloverDiD` | follow-up | Low | diff --git a/diff_diff/conley.py b/diff_diff/conley.py index 1df54697..4deceff2 100644 --- a/diff_diff/conley.py +++ b/diff_diff/conley.py @@ -365,6 +365,61 @@ def _uniform_kernel(u: np.ndarray) -> np.ndarray: return (np.abs(u) <= 1.0).astype(np.float64) +def _serial_bartlett_kernel_matrix(t_codes: np.ndarray, L: int) -> np.ndarray: + """Within-unit Newey-West (1987) Bartlett HAC kernel matrix for serial + correlation in panel data, indexed by panel-wide dense time codes. + + Returns the K matrix with ``K[i, j] = 1 - |t_i - t_j| / (L + 1)`` for + ``0 < |t_i - t_j| <= L``, else 0. The lag-0 diagonal is excluded so + callers can add this to a spatial within-period meat without + double-counting the diagonal. + + Uses the 1-D radial pairwise form (matches conleyreg::time_dist), NOT + Conley 1999 Eq 3.14's 2-D separable product window — see the methodology + lock at :func:`_compute_conley_meat` for context. + """ + t = t_codes.astype(np.float64, copy=False) + lag_mat = np.abs(t[:, None] - t[None, :]) + return ((lag_mat <= L) & (lag_mat != 0)).astype(np.float64) * (1.0 - lag_mat / (L + 1.0)) + + +def _validate_meat_psd( + M: np.ndarray, + *, + error_msg: str, + warning_template: str, + stacklevel: int = 3, +) -> None: + """Finite + PSD guard for sandwich meat matrices. Raises ``ValueError`` + on non-finite entries; warns ``UserWarning`` when ``min(eigvalsh(M)) < + -1e-12``. + + Parameters + ---------- + error_msg + Message passed to ``ValueError`` on non-finite entries. + warning_template + Format string for the negative-eigenvalue warning. May contain an + ``{eigval}`` placeholder; the caller embeds ``{eigval:.2e}`` directly + in the template so the helper formats the minimum eigenvalue with + scientific notation. + stacklevel + Frame count from inside ``_validate_meat_psd``: ``stacklevel=N`` + attributes the warning to the Nth caller above the helper itself. + Default 3 covers a single intermediate frame (the helper's direct + caller's caller); pass an explicit value matching call-site depth. + """ + if not np.all(np.isfinite(M)): + raise ValueError(error_msg) + eigvals = np.linalg.eigvalsh(M) + if eigvals.size and eigvals.min() < -1e-12: + warnings.warn( + warning_template.format(eigval=eigvals.min()), + UserWarning, + stacklevel=stacklevel, + ) + + def _compute_spatial_bartlett_meat_sparse( S: np.ndarray, coords: np.ndarray, @@ -957,37 +1012,33 @@ def _spatial_meat_for_mask(mask: Optional[np.ndarray] = None) -> np.ndarray: mask_u = unit_arr == u_val scores_u = scores[mask_u] # Use dense panel-period codes (NOT raw labels) for lag math. - t_u = time_codes[mask_u].astype(np.float64) - lag_mat = np.abs(t_u[:, None] - t_u[None, :]) - K_u = ((lag_mat <= L) & (lag_mat != 0)).astype(np.float64) * ( - 1.0 - lag_mat / (L + 1.0) - ) + K_u = _serial_bartlett_kernel_matrix(time_codes[mask_u], L) meat += scores_u.T @ K_u @ scores_u - if not np.all(np.isfinite(meat)): - raise ValueError( - "Conley meat contains non-finite values; check residuals and " - "score matrix for NaN/Inf." - ) - # PSD guard. Neither the uniform kernel (Conley 1999 fn 11) nor the # radial 1-D Bartlett specialization is formally PSD-guaranteed — # Conley's explicit PSD Bartlett formula (Eq 3.14) is the 2-D separable # product window, not the 1-D radial pairwise form that R `conleyreg`, # Stata `acreg`, and this implementation use. Check both kernels. - eigvals = np.linalg.eigvalsh(meat) - if eigvals.size and eigvals.min() < -1e-12: - warnings.warn( + # ``{eigval:.2e}`` is a literal placeholder for ``_validate_meat_psd``; + # only ``{kernel!r}`` is interpolated by the f-string here. + _validate_meat_psd( + meat, + error_msg=( + "Conley meat contains non-finite values; check residuals and " + "score matrix for NaN/Inf." + ), + warning_template=( f"Conley meat with conley_kernel={kernel!r} has a materially " - f"negative eigenvalue ({eigvals.min():.2e}); the variance " + "negative eigenvalue ({eigval:.2e}); the variance " "estimator is not guaranteed PSD on this design. Both " "supported kernels (radial bartlett and uniform) are " "practitioner specializations of Conley 1999 and are not " "formally PSD-guaranteed; consider varying conley_cutoff_km " "or reviewing the design for collinearity / degenerate " - "residual structure.", - UserWarning, - stacklevel=3, - ) + "residual structure." + ), + stacklevel=4, + ) return meat diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index ceadb3ff..c5ee0170 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -32,7 +32,9 @@ from diff_diff.conley import ( ConleyMetric, _compute_conley_meat, + _serial_bartlett_kernel_matrix, _validate_conley_kwargs, + _validate_meat_psd, ) from diff_diff.linalg import solve_ols from diff_diff.two_stage_bootstrap import TwoStageDiDBootstrapMixin @@ -904,34 +906,35 @@ def _compute_stratified_conley_meat( return np.full((p_2, p_2), np.nan) # Finite + PSD guards on the COMBINED survey meat (spatial + serial). - # Mirrors :func:`diff_diff.conley._compute_conley_meat` L966-990 so the - # survey panel-block path has the same diagnostic surface as the + # Shares ``_validate_meat_psd`` with :func:`diff_diff.conley._compute_conley_meat` + # so the survey panel-block path has the same diagnostic surface as the # no-survey path. The radial 1-D Bartlett spatial kernel and the # Newey-West Bartlett serial kernel are both practitioner # specializations that are NOT formally PSD-guaranteed; adding two # non-PSD-guaranteed terms can produce a more indefinite combined # meat, so the check matters most on the panel-block path. CI codex # R1 P2 fix. - if not np.all(np.isfinite(meat)): - raise ValueError( + # ``{eigval:.2e}`` is a literal placeholder for ``_validate_meat_psd``; + # only ``{conley_kernel!r}`` is interpolated by the f-string here. + _validate_meat_psd( + meat, + error_msg=( "SpilloverDiD Wave E.2 stratified-Conley meat contains non-finite " "values; check Psi for NaN/Inf upstream of the sandwich." - ) - eigvals = np.linalg.eigvalsh(meat) - if eigvals.size and eigvals.min() < -1e-12: - warnings.warn( + ), + warning_template=( f"SpilloverDiD Wave E.2 stratified-Conley meat with conley_kernel=" f"{conley_kernel!r} has a materially negative eigenvalue " - f"({eigvals.min():.2e}); the variance estimator is not guaranteed " + "({eigval:.2e}); the variance estimator is not guaranteed " "PSD on this design. Both supported kernels (radial bartlett and " "uniform spatial) plus the hardcoded serial Bartlett term are " "practitioner specializations of Conley 1999 / Newey-West 1987 " "and are not formally PSD-guaranteed; consider varying " "conley_cutoff_km / conley_lag_cutoff, or reviewing the design " - "for collinearity / degenerate residual structure.", - UserWarning, - stacklevel=2, - ) + "for collinearity / degenerate residual structure." + ), + stacklevel=3, + ) return meat @@ -1164,10 +1167,7 @@ def _compute_stratified_serial_bartlett_meat( if int(present_g.sum()) < 2: continue t_g = t_codes_full[present_g] - lag_mat = np.abs(t_g[:, None] - t_g[None, :]) - K_g = ((lag_mat <= L) & (lag_mat != 0)).astype(np.float64) * ( - 1.0 - lag_mat / (L + 1.0) - ) + K_g = _serial_bartlett_kernel_matrix(t_g, L) S_g_centered = S_psu_panel[g, present_g] - _global_psu_mean with np.errstate(invalid="ignore", over="ignore"): meat += S_g_centered.T @ K_g @ S_g_centered @@ -1194,9 +1194,8 @@ def _compute_stratified_serial_bartlett_meat( t_g = t_codes_full[present_g] # PANEL-WIDE dense time codes for the serial kernel (NOT per-PSU # positional encoding). See test (g) in TestSpilloverDiDWaveE2Followup - # for the methodology lock; matches conley.py:940 R-deviation. - lag_mat = np.abs(t_g[:, None] - t_g[None, :]) - K_g = ((lag_mat <= L) & (lag_mat != 0)).astype(np.float64) * (1.0 - lag_mat / (L + 1.0)) + # for the methodology lock; matches conley.py R-deviation. + K_g = _serial_bartlett_kernel_matrix(t_g, L) S_g_centered = S_centered[g, present_g] with np.errstate(invalid="ignore", over="ignore"): M_h_serial += S_g_centered.T @ K_g @ S_g_centered diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index f007c88b..6c7a7e27 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -3426,6 +3426,8 @@ Degrees of freedom for the t-distribution lookup use `ResolvedSurveyDesign.df_su **Implementation:** new sibling helper `_compute_stratified_serial_bartlett_meat` in `diff_diff/two_stage.py` (parallel to the Wave E.2 spatial orchestrator; ~200 LoC; three-mode singleton-stratum branching with FPC scaling inside the multi-PSU branch to avoid divide-by-zero; panel-wide dense time codes for the lag math). Orchestrator `_compute_stratified_conley_meat` signature extended with `conley_lag_cutoff: Optional[int] = None`; spatial loop unchanged; serial helper called after spatial loop when `conley_lag_cutoff > 0`; saturation NaN-fail accounting merges both terms' `(variance_computed, legitimate_zero)` flags. Dispatch in `_compute_gmm_corrected_meat` conley branch threads `conley_lag_cutoff` through to the orchestrator. Spillover-side gate at `spillover.py:2210` deleted (Wave E.2 era `NotImplementedError` for lag>0 + survey). Stage-1 weighted FE solver, `finite_mask` survey-array subsetting, `df_survey` threading, bread weighting, and `SpilloverDiDResults` survey metadata are all inherited UNCHANGED — Psi construction is bit-identical regardless of vcov_type or lag. Tests: `TestSpilloverDiDWaveE2FollowupConleySurveyLagCutoff` + `TestSpilloverDiDWaveE2FollowupConleySurveyLagCutoffEventStudy` at `tests/test_spillover.py`. +**Note (shared helpers):** the within-PSU serial Bartlett kernel matrix construction and the post-meat finite/PSD guard are shared with the no-survey reference path via `_serial_bartlett_kernel_matrix` and `_validate_meat_psd` in `conley.py`. The survey orchestrator's per-PSU serial loop (`_compute_stratified_serial_bartlett_meat`) and the orchestrator's combined-meat PSD guard call the same helpers as `_compute_conley_meat`, so survey and no-survey paths cannot drift on kernel weights or PSD threshold (`< -1e-12`). + ### Variance (Wave E.3 — `SurveyDesign.subpopulation()` / warn-and-drop full-design retention via zero-pad scores) `SurveyDesign.subpopulation()`-derived designs AND warn-and-drop fits now preserve the full-domain resolved survey design for inference bookkeeping. `n_psu`, `n_strata`, `df_survey`, and the Binder TSL per-stratum centering reflect the FULL domain rather than the post-`finite_mask` fit sample. SHIPPED in Wave E.3. @@ -3584,6 +3586,8 @@ robust default; pass `time` as a dense integer index for bit-exact R parity. **Note (deviation / source specialization):** Conley 1999's explicitly PSD-guaranteed Bartlett formula (Eq 3.14, page 12) is the **2-D separable product window** `K(j, k) = (1 - |j|/L_M)(1 - |k|/L_N)` indexed on a lattice. The 1-D radial form on pairwise distance that diff-diff implements (matching R `conleyreg`) is a practitioner specialization that is not explicitly written in the paper and is therefore **not formally PSD-guaranteed**. We apply the same indefiniteness check to both kernels: a `UserWarning` is emitted if any meat eigenvalue is materially negative (< `-1e-12`). +**Note (shared helpers):** the within-unit serial Bartlett kernel matrix construction and the finite/PSD guard are shared with the SpilloverDiD Wave E.2 follow-up survey path via `_serial_bartlett_kernel_matrix` and `_validate_meat_psd` in `conley.py`. The no-survey reference path (`_compute_conley_meat`) and the survey panel-block path (`_compute_stratified_serial_bartlett_meat`'s caller in `two_stage.py`) exercise the same kernel weights and PSD threshold (`< -1e-12`); see the parallel note in the SpilloverDiD Wave E.2 follow-up section below. + **Distance metrics:** - `conley_metric="haversine"` (default): great-circle in km using Earth's mean radius (6371.01 km, matching R `conleyreg`). Validates `lat ∈ [-90, 90]`, `lon ∈ [-180, 180]`. - `conley_metric="euclidean"`: Euclidean from projected coords. Skips lat/lon range checks (user owns the projection's units). diff --git a/tests/test_conley_vcov.py b/tests/test_conley_vcov.py index cdd5d71f..2843c064 100644 --- a/tests/test_conley_vcov.py +++ b/tests/test_conley_vcov.py @@ -19,8 +19,10 @@ _compute_spatial_bartlett_meat_sparse, _haversine_km, _pairwise_distance_matrix, + _serial_bartlett_kernel_matrix, _uniform_kernel, _validate_conley_kwargs, + _validate_meat_psd, ) from diff_diff.linalg import ( LinearRegression, @@ -107,6 +109,59 @@ def test_bartlett_kernel_finite_and_in_unit_interval(self): assert np.all(K <= 1.0) np.testing.assert_allclose(K, K.T, atol=1e-15) + def test_serial_bartlett_kernel_matrix_basic(self): + """Hand-computed Bartlett HAC kernel for t=[0,1,2,3], L=2. + + K[i,j] = (1 - |i-j|/(L+1)) for 0 < |i-j| <= L, else 0. With L=2 + and (L+1)=3: lag 1 -> 2/3, lag 2 -> 1/3, lag 3 -> 0 (out of band). + """ + K = _serial_bartlett_kernel_matrix(np.array([0, 1, 2, 3]), L=2) + expected = np.array( + [ + [0.0, 2.0 / 3.0, 1.0 / 3.0, 0.0], + [2.0 / 3.0, 0.0, 2.0 / 3.0, 1.0 / 3.0], + [1.0 / 3.0, 2.0 / 3.0, 0.0, 2.0 / 3.0], + [0.0, 1.0 / 3.0, 2.0 / 3.0, 0.0], + ] + ) + np.testing.assert_allclose(K, expected, atol=1e-15) + + def test_serial_bartlett_kernel_matrix_l_one(self): + """L=1: only adjacent lags survive, with weight (1 - 1/2) = 0.5.""" + K = _serial_bartlett_kernel_matrix(np.array([0, 1, 2]), L=1) + expected = np.array( + [ + [0.0, 0.5, 0.0], + [0.5, 0.0, 0.5], + [0.0, 0.5, 0.0], + ] + ) + np.testing.assert_allclose(K, expected, atol=1e-15) + + def test_serial_bartlett_kernel_matrix_l_zero_returns_zero(self): + """L=0 is degenerate-but-callable: every off-diagonal lag fails + ``lag <= 0`` (since `lag != 0`), so K is the zero matrix. Callers + guard externally (``conley.py`` skips the loop when L == 0).""" + K = _serial_bartlett_kernel_matrix(np.array([0, 1, 2]), L=0) + np.testing.assert_array_equal(K, np.zeros((3, 3))) + + def test_serial_bartlett_kernel_matrix_single_element(self): + """Single-element input yields a 1x1 zero matrix (no off-diagonal + lags exist).""" + K = _serial_bartlett_kernel_matrix(np.array([7]), L=2) + np.testing.assert_array_equal(K, np.zeros((1, 1))) + + def test_serial_bartlett_kernel_matrix_int_input_bit_equal_to_float(self): + """Contract test: int64 and float64 inputs must yield bit-equal + matrices. The helper does ``astype(np.float64, copy=False)`` and one + of the three current call sites (``conley.py`` panel-block branch) + passes the result of array slicing on int time codes.""" + K_int = _serial_bartlett_kernel_matrix(np.array([0, 1, 2, 3], dtype=np.int64), L=2) + K_float = _serial_bartlett_kernel_matrix( + np.array([0.0, 1.0, 2.0, 3.0], dtype=np.float64), L=2 + ) + np.testing.assert_array_equal(K_int, K_float) + # --------------------------------------------------------------------------- # TestConleyDistanceMetrics — haversine, euclidean, callable @@ -604,6 +659,309 @@ def test_panel_unit_pd_na_object_raises(self): ) +# --------------------------------------------------------------------------- +# TestValidateMeatPsd — _validate_meat_psd guard helper +# --------------------------------------------------------------------------- + + +class TestValidateMeatPsd: + def test_nonfinite_raises(self): + """Non-finite meat must raise ValueError with the caller's exact + ``error_msg`` so site-specific guidance reaches the user.""" + M = np.array([[1.0, np.nan], [np.nan, 1.0]]) + with pytest.raises(ValueError, match="custom guidance for caller XYZ"): + _validate_meat_psd( + M, + error_msg="custom guidance for caller XYZ", + warning_template="unused-here ({eigval:.2e})", + ) + + def test_negative_eigenvalue_warns_with_template_substitution(self): + """An indefinite meat triggers UserWarning with ``{eigval}`` + substituted in scientific notation.""" + # Symmetric matrix with eigenvalues {2, -1}: aggressively indefinite. + M = np.array([[0.5, 1.5], [1.5, 0.5]]) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _validate_meat_psd( + M, + error_msg="not used on this path", + warning_template="SITEX meat: min eigenvalue = {eigval:.2e}", + ) + psd = [ + msg + for msg in w + if issubclass(msg.category, UserWarning) and "SITEX" in str(msg.message) + ] + assert len(psd) == 1, f"Expected one PSD warning; got {[str(m.message) for m in w]}" + # Min eigenvalue is -1.0; verify scientific-notation substitution. + assert "-1.00e+00" in str(psd[0].message) + + def test_psd_matrix_silent(self): + """A PSD meat (identity) must not emit any warning.""" + M = np.eye(3) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _validate_meat_psd( + M, + error_msg="not used", + warning_template="not used ({eigval:.2e})", + ) + psd = [msg for msg in w if issubclass(msg.category, UserWarning)] + assert psd == [], f"Expected no warnings; got {[str(m.message) for m in psd]}" + + def test_threshold_boundary_above_threshold_silent(self): + """An eigenvalue just above the -1e-12 threshold (at -5e-13) is + absorbed as numerical noise and must NOT warn.""" + # diag(-5e-13, 1.0): symmetric, eigenvalues exactly {-5e-13, 1.0}. + M = np.diag([-5e-13, 1.0]) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _validate_meat_psd( + M, + error_msg="not used", + warning_template="not used ({eigval:.2e})", + ) + psd = [msg for msg in w if issubclass(msg.category, UserWarning)] + assert psd == [], ( + f"Eigenvalue -5e-13 is above -1e-12 threshold; expected no warning. " + f"Got: {[str(m.message) for m in psd]}" + ) + + def test_stacklevel_attributes_to_caller_frame(self): + """Stacklevel parameter must attribute the warning to the caller's + frame (or higher), NOT to the helper's own frame. + + Locks the contract that the helper itself is invisible in warning + attribution. The helper extraction added one frame to the stack, so + call sites had to bump their stacklevel by +1 (conley.py 3→4, + two_stage.py 2→3). Defaulting to ``stacklevel=3`` from inside the + helper attributes the warning to the helper's direct caller's caller + (one intermediate frame between caller and outer-caller).""" + M = np.array([[0.5, 1.5], [1.5, 0.5]]) # eigenvalues {2, -1} + + def inner_caller(): + _validate_meat_psd( + M, + error_msg="x", + warning_template="attr-check ({eigval:.2e})", + stacklevel=3, + ) + + def outer_caller(): + inner_caller() + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + outer_caller() + + psd = [msg for msg in w if "attr-check" in str(msg.message)] + assert len(psd) == 1, f"Expected one PSD warning; got {[str(m.message) for m in w]}" + # With stacklevel=3 from inside _validate_meat_psd, the warning is + # attributed to the caller of inner_caller, which is outer_caller. + # Both inner_caller and outer_caller are defined in this test file, + # so the attribution must land here (NOT in conley.py inside the + # helper body). + assert psd[0].filename.endswith("test_conley_vcov.py"), ( + f"Expected attribution to test_conley_vcov.py (outer_caller " + f"frame); got {psd[0].filename!r}:{psd[0].lineno}. If " + f"attribution landed in conley.py the helper's stacklevel " + f"contract has regressed." + ) + + def test_no_survey_path_attributes_warning_to_user_code(self): + """End-to-end warning-capture test on ``_compute_conley_vcov``: when + the indefinite-meat path triggers the PSD warning via the shared + helper, attribution must bubble all the way through the three + internal frames (``_validate_meat_psd`` → ``_compute_conley_meat`` → + ``_compute_conley_vcov``) to land at user code (this test's frame). + + Locks the stacklevel=4 contract at the no-survey call site that + compensates for the +1 frame the helper extraction added. Pre- + extraction the inline warn used ``stacklevel=3`` which already + attributed to user code; preserving that behavior is the whole + point of the +1 bump.""" + from diff_diff import conley as conley_mod + + rng = np.random.default_rng(seed=11) + n = 6 + coords = rng.uniform(0, 1, size=(n, 2)) + X = np.column_stack([np.ones(n), rng.standard_normal(n)]) + eps = np.ones(n) + bread = X.T @ X + + # Monkey-patch the bartlett kernel to force an indefinite meat + # (mirrors the existing test_indefinite_meat_warning_fires_for_bartlett + # fixture pattern). + original = conley_mod._bartlett_kernel + + def _indefinite(u: np.ndarray) -> np.ndarray: + base = np.eye(u.shape[0]) + for i in range(u.shape[0]): + for j in range(u.shape[0]): + if i != j: + base[i, j] = -10.0 + return base + + try: + conley_mod._bartlett_kernel = _indefinite + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + conley_mod._compute_conley_vcov( + X, + eps, + coords, + cutoff=10.0, + metric="euclidean", + kernel="bartlett", + bread_matrix=bread, + ) + psd = [ + msg + for msg in w + if issubclass(msg.category, UserWarning) + and "negative eigenvalue" in str(msg.message) + ] + finally: + conley_mod._bartlett_kernel = original + + assert len(psd) >= 1, "Expected a PSD UserWarning from the indefinite meat path" + msg = psd[0] + # Attribution must be in this test file (user code), NOT any of the + # three internal production frames. If a future refactor regresses + # the stacklevel to 3, attribution would stick at _compute_conley_vcov + # in conley.py; to 2, at _compute_conley_meat; to 1, inside the helper. + assert msg.filename.endswith("test_conley_vcov.py"), ( + f"Expected attribution to user code (test_conley_vcov.py); got " + f"{msg.filename!r}:{msg.lineno}. The stacklevel=4 contract at " + f"the conley.py call site has regressed." + ) + + def test_survey_call_site_passes_stacklevel_3(self): + """Static source check: the survey orchestrator + ``_compute_stratified_conley_meat`` in two_stage.py must pass + ``stacklevel=3`` to ``_validate_meat_psd``. The pre-extraction + inline warn used ``stacklevel=2``; after the helper extraction the + +1 frame shift means the call site must pass ``stacklevel=3`` to + attribute the warning to the same outer caller. Pairs with the + runtime test ``test_survey_path_attributes_warning_to_user_code`` + which exercises the actual frame walk; this static check pins the + literal kwarg to surface bare-text regressions even if the runtime + test's fixture changes.""" + import inspect + + from diff_diff.two_stage import _compute_stratified_conley_meat + + src = inspect.getsource(_compute_stratified_conley_meat) + # Find the _validate_meat_psd call and verify the kwarg block + # includes stacklevel=3 (not 2, not 4, not missing). + assert "_validate_meat_psd(" in src, ( + "_compute_stratified_conley_meat no longer calls " + "_validate_meat_psd; survey-side PSD guard regressed." + ) + assert "stacklevel=3" in src, ( + "_compute_stratified_conley_meat does not pass stacklevel=3 " + "to _validate_meat_psd. The pre-extraction inline warn used " + "stacklevel=2; the +1 frame shift from extracting the helper " + "requires stacklevel=3 to preserve attribution." + ) + + def test_survey_path_attributes_warning_to_user_code(self): + """Runtime warning-capture test on the SURVEY orchestrator + ``_compute_stratified_conley_meat``: when the panel-block path + produces an indefinite combined meat the PSD warning must bubble + through the orchestrator frame to land at user code (this test). + Locks the stacklevel=3 contract end-to-end (not just by source + substring), addressing CI codex R1 P3. + + Mirrors the no-survey test's monkey-patch pattern: replaces the + serial-Bartlett kernel helper bound inside ``two_stage.py`` with + an aggressively-negative-off-diagonal stub so the serial meat is + indefinite and the combined meat's min eigenvalue drops below + the -1e-12 PSD threshold. Uses the minimal 4-PSU x 2-period x + 3-obs survey fixture from + ``tests/test_spillover.py::TestSpilloverDiDWaveE2Followup``.""" + from diff_diff import two_stage as two_stage_mod + from diff_diff.survey import ResolvedSurveyDesign + from diff_diff.two_stage import _compute_stratified_conley_meat + + rng = np.random.default_rng(seed=29) + n_obs, T, G, p_2 = 24, 2, 4, 3 + obs_per_psu_period = 3 + psu_id = np.repeat(np.arange(G), obs_per_psu_period * T) + time_arr = np.tile(np.repeat(np.arange(T), obs_per_psu_period), G) + Psi = rng.standard_normal((n_obs, p_2)) + psu_centroids = np.array([[40.0, -120.0], [40.1, -120.0], [40.2, -120.0], [40.3, -120.0]]) + coords = psu_centroids[psu_id] + psu_strata = np.array([0, 0, 1, 1]) + resolved = ResolvedSurveyDesign( + weights=np.ones(n_obs), + weight_type="pweight", + strata=np.repeat(psu_strata, obs_per_psu_period * T), + psu=psu_id, + fpc=np.full(n_obs, 20.0), + n_strata=2, + n_psu=4, + lonely_psu="remove", + ) + + # Monkey-patch the serial Bartlett kernel helper as bound inside + # diff_diff.two_stage (the `from diff_diff.conley import ...` + # rebind at module load time) so the serial meat is indefinite. + # The combined meat = spatial + indefinite_serial then drops + # below the -1e-12 PSD threshold. + original = two_stage_mod._serial_bartlett_kernel_matrix + + def _indefinite(t_codes: np.ndarray, L: int) -> np.ndarray: + n = t_codes.shape[0] + K = np.eye(n) + for i in range(n): + for j in range(n): + if i != j: + K[i, j] = -10.0 + return K + + try: + two_stage_mod._serial_bartlett_kernel_matrix = _indefinite + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _compute_stratified_conley_meat( + Psi, + conley_coords=coords, + conley_cutoff_km=0.30, + conley_metric="euclidean", + conley_kernel="bartlett", + resolved_survey=resolved, + conley_time=time_arr, + conley_lag_cutoff=1, + ) + finally: + two_stage_mod._serial_bartlett_kernel_matrix = original + + psd = [ + msg + for msg in w + if issubclass(msg.category, UserWarning) and "negative eigenvalue" in str(msg.message) + ] + assert len(psd) >= 1, ( + f"Expected a PSD UserWarning from the indefinite combined " + f"survey meat. Got: {[str(m.message) for m in w]}" + ) + msg = psd[0] + # Attribution must be in this test file (user code), proving the + # warning bubbled through both the helper (_validate_meat_psd in + # conley.py) and the survey orchestrator + # (_compute_stratified_conley_meat in two_stage.py). A regression + # of the call-site stacklevel from 3 to 2 would stick the + # attribution inside two_stage.py; to 1 inside the helper itself. + assert msg.filename.endswith("test_conley_vcov.py"), ( + f"Expected attribution to user code (test_conley_vcov.py); got " + f"{msg.filename!r}:{msg.lineno}. The stacklevel=3 contract at " + f"two_stage.py's _compute_stratified_conley_meat call site has " + f"regressed." + ) + + # --------------------------------------------------------------------------- # TestConleyDirectHelper — _compute_conley_vcov correctness # ---------------------------------------------------------------------------