diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c662621..920bc0c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - **HAD pretest workflow: stratified survey-design support (Phase 4.5 C continuation).** Lifts the `NotImplementedError` gate on `SurveyDesign(strata=...)` in `stute_test` (`had_pretests.py:1927-1940` pre-PR) and `stute_joint_pretest` (`:3259-3271` pre-PR), and by inheritance in `joint_pretrends_test`, `joint_homogeneity_test`, and `did_had_pretest_workflow` (the wrappers delegate to the joint Stute helper). Implements a documented synthesis of clustered-wild-bootstrap ingredients (Cameron-Gelbach-Miller 2008 cluster-level multipliers; Davidson-Flachaire 2008 wild-bootstrap centering; Djogbenou-MacKinnon-Nielsen 2019 cluster-wild consistency for nonlinear functionals; Kreiss-Lahiri 2012 within-block centering analogy; Wu 1986 / Liu 1988 Bessel small-sample correction) — no single paper covers the exact composition for the stratified Stute CvM functional. The recipe: within-stratum demean + `sqrt(n_h/(n_h-1))` Bessel rescale applied to the PSU multipliers `psu_mults` BEFORE the per-obs broadcast `eta_obs = psu_mults[b, psu_col_idx]` in the wild-residual loop. Bootstrap CvM variance matches the analytical Binder-TSL stratified target `V_S = sum_h (1 - f_h) (n_h / (n_h - 1)) sum_j (psi_hj - psi_h_bar)²` exactly (the `(1 - f_h)` FPC factor was already baked in by `generate_survey_multiplier_weights_batch`; this PR bakes the remaining `(n_h / (n_h - 1))` factor and enforces within-stratum-mean-zero centering). New shared helper `bootstrap_utils.apply_stratum_centering(psu_mults, resolved_survey, psu_ids, psu_axis=...)` is called from both the new Stute path (psu_axis=1 on the multiplier matrix) AND the existing HAD sup-t event-study cband bootstrap (psu_axis=0 on the PSU-aggregated influence tensor; refactored bit-exactly from the inline block previously at `had.py:2172-2204`). Locks the algebraic identity architecturally instead of leaving parallel code blocks to drift. MC oracle consistency validated under a 4-stratum × 6-PSU/stratum stratified null DGP with weights+strata+PSU (200 seeded draws, empirical Type I at α=0.05 in `[0, 0.10]` — 3σ band; the FPC bake-in is covered separately by the helper-unit test `test_fpc_baked_in_helper_is_fpc_agnostic`); MC power validated under a known-alternative stratified DGP (rejection > 0.50). HAD sup-t event-study cband bit-parity preserved (`atol=1e-14, rtol=1e-14` on the refactored helper output + 29 existing cband tests passing post-refactor; that helper-level bit-parity test locks the axis-0 algebra). A separate wired-in regression at `tests/test_had_pretests.py::TestStuteStratifiedSurveyBootstrap::test_stute_call_sites_invoke_apply_stratum_centering` monkey-patches the helper and asserts both Stute call sites (`stute_test` at `had_pretests.py:1985` and `stute_joint_pretest` at `:3312`) invoke it with `psu_axis=1` — that test fails if either call site is disconnected (the axis-0 helper-parity test alone does not catch that case). See `docs/methodology/REGISTRY.md` § HeterogeneousAdoptionDiD — "Note (Stute stratified survey-bootstrap calibration)" for the full derivation. Remaining deferrals: `lonely_psu='adjust'` + singleton-strata (same pseudo-stratum centering gap as the HAD sup-t deviation at REGISTRY:2382) and replicate-weight designs (BRR/Fay/JK1/JKn/SDR — separate Rao-Wu / JKn bootstrap composition). Unblocks the realistic survey-weighted HAD workflow on BRFSS/CPS/NHANES/ACS-shaped designs. -- **Conley (1999) spatial-HAC standard errors via `vcov_type="conley"`** on cross-sectional `LinearRegression` / `compute_robust_vcov` plus panel `MultiPeriodDiD` / `TwoWayFixedEffects` (Phases 1 and 2 of the spillover-conley initiative). **Cross-sectional contract:** `conley_coords` (n × 2 array of lat/lon or projected coords), `conley_cutoff_km=` (positive finite bandwidth in km for haversine, or coord units for euclidean — REQUIRED, no default per the no-silent-failures contract), `conley_metric="haversine"|"euclidean"|callable` (default `"haversine"`; great-circle uses Earth's mean radius 6371.01 km matching R `conleyreg`), `conley_kernel="bartlett"|"uniform"` (default `"bartlett"`; both kernels emit a `UserWarning` if the resulting meat has a materially negative eigenvalue — neither the radial 1-D Bartlett nor the uniform kernel is formally PSD-guaranteed; Conley 1999's explicit PSD formula is the 2-D separable lattice product window at Eq 3.14). Cross-sectional variance estimator `Var̂(β) = (X'X)^{-1} · ( Σ_{i,j} K(d_ij/h) · X_i ε_i ε_j X_j' ) · (X'X)^{-1}` (Conley 1999 Eq 4.2). **Panel contract (Phase 2, new):** Three new co-required kwargs `conley_time` (n-length array), `conley_unit` (n-length array), and `conley_lag_cutoff=` (non-negative; 0 means within-period spatial only, no serial component) switch into the **block-decomposed panel sandwich** that matches R `conleyreg` with `lag_cutoff > 0`: `XeeX_total = Σ_t (within-period spatial sandwich) + Σ_u (within-unit Bartlett temporal sandwich, lag ∈ {1..L}, same-time excluded)`. This is NOT a multiplicative product kernel — verified empirically against `conleyreg::time_dist` and `XeeXhC` at ~1e-14 on the panel parity fixtures. The temporal kernel is hardcoded Bartlett `(1 - |lag|/(L+1))` regardless of `conley_kernel`, mirroring `conleyreg::time_dist.cpp`; documented as a `Note (deviation from R-symmetric API)` in REGISTRY. **Panel estimator wire-up (Phase 2):** `MultiPeriodDiD(vcov_type="conley", conley_lag_cutoff=...).fit(..., unit=...)` and `TwoWayFixedEffects(vcov_type="conley", conley_lag_cutoff=...).fit(..., unit=...)` lift the Phase 1 fit-time rejection; the `conley_time` and `conley_unit` arrays are auto-derived from the existing `time` and `unit` column-name arguments at fit-time. **`DifferenceInDifferences(vcov_type="conley")` continues to raise `NotImplementedError`** because `DiD.fit()` has no `unit` column declaration; the error message redirects users to `MultiPeriodDiD` / `TwoWayFixedEffects`. **Other constraints (Phase 1, unchanged):** `SyntheticDiD(vcov_type="conley")` raises `TypeError` (uses bootstrap variance, not analytical sandwich); `set_params` mirrors the constructor rejection. `vcov_type="conley"` + `cluster_ids=` / `weights=` / `survey_design=` raises `NotImplementedError` (combined product kernel + Bertanha-Imbens 2014 weighted-Conley deferred to follow-up PRs). TWFE's default auto-cluster on the Conley path is silently dropped (the explicit `cluster=` raises with a clear redirect). `inference="wild_bootstrap"` + Conley raises (incompatible inference modes). `n > 20_000` emits a `UserWarning` about the dense O(n²) distance-matrix memory; sparse k-d-tree fast path is a follow-up. **Implementation:** Helpers live in `diff_diff/conley.py` (`_haversine_km`, `_pairwise_distance_matrix`, `_bartlett_kernel`, `_uniform_kernel`, `_validate_conley_kwargs`, `_compute_conley_vcov` — the validator and sandwich helper now accept keyword-only `time` / `unit` / `lag_cutoff` for the panel path); `compute_robust_vcov` in `diff_diff/linalg.py` threads the new kwargs through. **R `conleyreg` parity (Düsterhöft 2021, CRAN v0.1.9)** on **six** benchmark fixtures (`benchmarks/data/r_conleyreg_conley_golden.json`, regenerable via `benchmarks/R/generate_conley_golden.R`): 3 cross-sectional (Phase 1) + 3 new panel fixtures (`panel_haversine_lag1`, `panel_haversine_lag2`, `panel_lat_lon_realistic_lag1`; n_units × T = 60×3, 80×5, 100×4 at lag={1,2,1}); observed max abs diff ~5.7e-16. Earth radius 6371.01 km matches `conleyreg::haversine_dist`. Test file `tests/test_conley_vcov.py` skips parity cleanly when the JSON is absent. New REGISTRY section `## ConleySpatialHAC`. Subsequent phases of the spillover-conley initiative (sparse k-d-tree fast path for `n > 20_000`; Conley + `cluster_ids` combined kernel; ring-indicator spillover-aware DiD per Butts 2021; survey-design / replicate-weight support; `SyntheticDiD` Conley path) are tracked in `TODO.md` under "Tech Debt from Code Reviews" → spillover-conley rows. +- **Conley (1999) Wave A mechanical extensions** on top of the Phase 1+2 sandwich (`diff_diff/conley.py`, `diff_diff/linalg.py`, `diff_diff/estimators.py`, `diff_diff/twfe.py`). **(1) DiD support (#118):** `DifferenceInDifferences(vcov_type="conley").fit(..., unit="")` is now supported. `unit` is a fit-time kwarg (NOT on `__init__`; unused unless Conley is set; not part of `get_params()` / `set_params()`) mirroring `MultiPeriodDiD.fit(unit=...)` / `TwoWayFixedEffects.fit(unit=...)`. DiD inherits the same panel block-decomposed sandwich as MPD/TWFE; on a 2-period panel it matches `MultiPeriodDiD(...).fit(..., post_periods=[1], reference_period=0)` bit-exactly. Missing `unit=`/`conley_lag_cutoff`/`conley_coords`/`conley_cutoff_km` raise `ValueError`; `survey_design=` + Conley raises `NotImplementedError` (Bertanha-Imbens 2014 follow-up); `inference="wild_bootstrap"` + Conley raises `NotImplementedError`. **(2) Combined spatial + cluster product kernel (#119):** `compute_robust_vcov(vcov_type="conley", cluster_ids=...)` / `LinearRegression(vcov_type="conley", cluster_ids=...)` / `TwoWayFixedEffects(vcov_type="conley", cluster="")` / `MultiPeriodDiD(vcov_type="conley", cluster="")` / `DifferenceInDifferences(vcov_type="conley", cluster="")` apply `K_total[i, j] = K_space(d_ij/h) · 1{c_i = c_j}`. On the panel block-decomposed path the cluster indicator multiplies BOTH the spatial sandwich AND the serial sandwich; the validator enforces that `cluster_ids` is constant within each unit across periods (the within-unit serial mask is then trivially all-ones; cross-sectional path has no such constraint). TWFE's default auto-cluster on the Conley path remains silently dropped (combining with unit-level clusters would zero out all between-unit pairs and defeat the spatial pooling); users must pass an explicit above-unit cluster (e.g. region) to opt in. DiD has no auto-cluster — the choice is fully explicit. Two limit fixtures anchor correctness (no R parity — R `conleyreg` does not support combined kernels): all-unique-clusters reduces to HC0; huge-cutoff reduces to pure within-cluster CR1. The huge-cutoff reduction is EXACT only for `conley_kernel="uniform"` (`K(u) = 1` for `|u| ≤ 1`); for `conley_kernel="bartlett"` the identity is asymptotic since `K_bartlett(u) = 1 - |u| < 1` for `u > 0`. The fixture anchor uses uniform for an exact identity check. Per-slice mask construction (NOT full n×n) preserves memory on panel paths. **(3) Sparse k-d-tree fast path (#120):** auto-activates for the spatial Bartlett meat when `n > 5_000` AND metric is `"haversine"` or `"euclidean"` AND kernel is `"bartlett"`. Builds a CSR sparse kernel matrix via `scipy.spatial.cKDTree.query_ball_tree` instead of materializing the full n×n distance matrix; haversine projects to a 3-D unit-sphere chord representation with the exact great-circle recomputed for in-range neighbors only. Bit-identity parity vs the dense path at `atol=1e-10`; R parity at `atol=1e-6` is preserved on the existing 3 panel R fixtures with the sparse path force-enabled. The bartlett-only gate is for boundary correctness — bartlett at `u=1` is exactly 0, so the sparse path safely drops at-cutoff pairs; uniform at `u=1` is 1 and would require a closed-interval query semantic that haversine chord projection cannot reliably preserve. Constants: `_CONLEY_SPARSE_N_THRESHOLD = 5_000` (auto-toggle); `_CONLEY_DENSE_WARN_N` renamed `_CONLEY_DENSE_OOM_WARN_N = 20_000` (memory exhaustion threshold for the dense fallback — independent of the sparse threshold). Private `_conley_sparse: Optional[bool]` kwarg on `_compute_conley_vcov` controls the toggle (`None` = auto, `True` = force, `False` = force dense; `True` with an unsupported kernel/metric raises). The serial component (within-unit Bartlett over time) remains dense regardless — per-unit slices are small. **(4) Callable `conley_metric` validation (#123):** result must satisfy shape `(n, n)`, finite, non-negative, symmetric to `atol=1e-10`, AND zero on the diagonal (`|d(i, i)| ≤ 1e-10`); each failure raises a targeted `ValueError` naming the violated invariant. The zero-diagonal contract is load-bearing for the Conley sandwich: the `i = j` term must reduce to the HC0 diagonal `X_i ε_i² X_i'` via `K(0) = 1`; positive self-distance would silently attenuate the HC0 contribution by `K(d_ii / h) < 1`. Built-in metrics (`"haversine"`, `"euclidean"`) satisfy this by construction. Previously, malformed callables produced opaque BLAS errors deep in the pipeline. **Tests:** `tests/test_conley_vcov.py::TestConleySparse` (12), `::TestConleySparseRParityForced` (3), `::TestConleyCluster` (10), `::TestConleyDistanceMetrics` extended (7 new); existing rejection tests flipped to behavioral; `test_did_conley_matches_mpd_post_periods_1` locks the DiD-vs-MPD bit-exact agreement. **Docs:** REGISTRY `## ConleySpatialHAC` updates: new "Combined spatial + cluster product kernel" + "Performance / scale" subsections, DiD-vs-TWFE cluster asymmetry paragraph, updated panel-API restrictions table. TODO rows 118 / 119 / 120 / 123 removed; rows 121 (Conley + survey_design / weights, Bertanha-Imbens 2014) and 122 (`SyntheticDiD(vcov_type="conley")`, spatial-block bootstrap per Politis-Romano 1994) retained for future waves. +- **Conley (1999) spatial-HAC standard errors via `vcov_type="conley"`** on cross-sectional `LinearRegression` / `compute_robust_vcov` plus panel `MultiPeriodDiD` / `TwoWayFixedEffects` (Phases 1 and 2 of the spillover-conley initiative). **Cross-sectional contract:** `conley_coords` (n × 2 array of lat/lon or projected coords), `conley_cutoff_km=` (positive finite bandwidth in km for haversine, or coord units for euclidean — REQUIRED, no default per the no-silent-failures contract), `conley_metric="haversine"|"euclidean"|callable` (default `"haversine"`; great-circle uses Earth's mean radius 6371.01 km matching R `conleyreg`), `conley_kernel="bartlett"|"uniform"` (default `"bartlett"`; both kernels emit a `UserWarning` if the resulting meat has a materially negative eigenvalue — neither the radial 1-D Bartlett nor the uniform kernel is formally PSD-guaranteed; Conley 1999's explicit PSD formula is the 2-D separable lattice product window at Eq 3.14). Cross-sectional variance estimator `Var̂(β) = (X'X)^{-1} · ( Σ_{i,j} K(d_ij/h) · X_i ε_i ε_j X_j' ) · (X'X)^{-1}` (Conley 1999 Eq 4.2). **Panel contract (Phase 2, new):** Three new co-required kwargs `conley_time` (n-length array), `conley_unit` (n-length array), and `conley_lag_cutoff=` (non-negative; 0 means within-period spatial only, no serial component) switch into the **block-decomposed panel sandwich** that matches R `conleyreg` with `lag_cutoff > 0`: `XeeX_total = Σ_t (within-period spatial sandwich) + Σ_u (within-unit Bartlett temporal sandwich, lag ∈ {1..L}, same-time excluded)`. This is NOT a multiplicative product kernel — verified empirically against `conleyreg::time_dist` and `XeeXhC` at ~1e-14 on the panel parity fixtures. The temporal kernel is hardcoded Bartlett `(1 - |lag|/(L+1))` regardless of `conley_kernel`, mirroring `conleyreg::time_dist.cpp`; documented as a `Note (deviation from R-symmetric API)` in REGISTRY. **Panel estimator wire-up (Phase 2):** `MultiPeriodDiD(vcov_type="conley", conley_lag_cutoff=...).fit(..., unit=...)` and `TwoWayFixedEffects(vcov_type="conley", conley_lag_cutoff=...).fit(..., unit=...)` lift the Phase 1 fit-time rejection; the `conley_time` and `conley_unit` arrays are auto-derived from the existing `time` and `unit` column-name arguments at fit-time. `DifferenceInDifferences(vcov_type="conley")` is also supported (Wave A #118 in this release; see the Wave A entry above) — pass `unit=` as a fit-time kwarg to `DiD.fit(...)`. **Other constraints (Phase 1, unchanged):** `SyntheticDiD(vcov_type="conley")` raises `TypeError` (uses bootstrap variance, not analytical sandwich); `set_params` mirrors the constructor rejection. `vcov_type="conley"` + `weights=` / `survey_design=` raises `NotImplementedError` (Bertanha-Imbens 2014 weighted-Conley deferred to a follow-up PR). `vcov_type="conley"` + explicit `cluster_ids=` is supported via the combined spatial + cluster product kernel (Wave A #119; see the Wave A entry above). TWFE's default auto-cluster on the Conley path is silently dropped (combining with unit-level clusters would defeat the spatial pooling); users opt into the combined kernel by passing an explicit above-unit cluster. `inference="wild_bootstrap"` + Conley raises (incompatible inference modes). A sparse k-d-tree fast path auto-activates for the spatial Bartlett meat when `n > 5_000` with bartlett kernel and haversine/euclidean metric (Wave A #120); the dense fallback still emits an OOM `UserWarning` at `n > 20_000`. **Implementation:** Helpers live in `diff_diff/conley.py` (`_haversine_km`, `_pairwise_distance_matrix`, `_bartlett_kernel`, `_uniform_kernel`, `_validate_conley_kwargs`, `_compute_conley_vcov` — the validator and sandwich helper now accept keyword-only `time` / `unit` / `lag_cutoff` for the panel path); `compute_robust_vcov` in `diff_diff/linalg.py` threads the new kwargs through. **R `conleyreg` parity (Düsterhöft 2021, CRAN v0.1.9)** on **six** benchmark fixtures (`benchmarks/data/r_conleyreg_conley_golden.json`, regenerable via `benchmarks/R/generate_conley_golden.R`): 3 cross-sectional (Phase 1) + 3 new panel fixtures (`panel_haversine_lag1`, `panel_haversine_lag2`, `panel_lat_lon_realistic_lag1`; n_units × T = 60×3, 80×5, 100×4 at lag={1,2,1}); observed max abs diff ~5.7e-16. Earth radius 6371.01 km matches `conleyreg::haversine_dist`. Test file `tests/test_conley_vcov.py` skips parity cleanly when the JSON is absent. New REGISTRY section `## ConleySpatialHAC`. Subsequent phases of the spillover-conley initiative (ring-indicator spillover-aware DiD per Butts 2021; survey-design / replicate-weight support; `SyntheticDiD` Conley path) are tracked in `TODO.md` under "Tech Debt from Code Reviews" → spillover-conley rows. - **Tutorial 21: HAD Pre-test Workflow** (`docs/tutorials/21_had_pretest_workflow.ipynb`) — composite pre-test walkthrough for `HeterogeneousAdoptionDiD` building on Tutorial 20's brand-campaign framing. Uses a 60-DMA × 8-week panel close in shape to T20's but with the dose distribution drawn from `Uniform[$0.01K, $50K]` (vs T20's `[$5K, $50K]`); the true support is strictly positive but very near zero, chosen so the QUG step in `did_had_pretest_workflow` fails-to-reject `H0: d_lower = 0` in this finite sample and the verdict text fires the load-bearing "Assumption 7 deferred" pivot for the upgrade-arc narrative. (HAD's `design="auto"` selector — a separate min/median heuristic at `had.py::_detect_design`, NOT the QUG p-value — independently lands on the `continuous_at_zero` identification path with target `WAS` on this panel because `d.min() < 0.01 * median(|d|)`. The QUG test and the design selector are independent rules that point to the same identification path here.) Walks through three surfaces: (a) `did_had_pretest_workflow(aggregate="overall")` on a two-period collapse, where the verdict explicitly flags Step 2 (Assumption 7 pre-trends) as not run because a single pre-period structurally cannot support a pre-trends test, and the structural fields `pretrends_joint` / `homogeneity_joint` are both `None`; (b) `did_had_pretest_workflow(aggregate="event_study")` on the full multi-period panel, where the verdict reads "TWFE admissible under Section 4 assumptions" because all three testable diagnostics (QUG + joint pre-trends Stute over 3 horizons + joint homogeneity Stute over 4 horizons) fail-to-reject — non-rejection evidence under finite-sample power and test specification, not proof that the identifying assumptions hold; and (c) a side panel exercising both `yatchew_hr_test` null modes — `null="linearity"` (default, paper Theorem 7) vs `null="mean_independence"` (Phase 4 R-parity with R `YatchewTest::yatchew_test(order=0)`) — on the within-pre-period first-difference paired with post-period dose, illustrating the stricter null's larger residual variance (`sigma2_lin` 7.01 vs 6.53) and smaller p-value (0.29 vs 0.49). Companion drift-test file `tests/test_t21_had_pretest_workflow_drift.py` (16 tests pinning panel composition, both verdict pivots, structural anchors on both paths, deterministic QUG / Yatchew statistics, bootstrap p-value tolerance bands per `feedback_bootstrap_drift_tests_need_backend_tolerance`, and `HAD(design="auto")` resolution to `continuous_at_zero` on this panel). T20's "Composite pretest workflow" Extensions bullet updated with a forward-pointer to T21. T22 weighted/survey HAD tutorial remains queued as a separate notebook PR. - **`ChaisemartinDHaultfoeuille.by_path` and `paths_of_interest` now compose with `survey_design`** for analytical Binder TSL SE and replicate-weight bootstrap variance. The `NotImplementedError` gate at `chaisemartin_dhaultfoeuille.py:1233-1239` is replaced by a per-path multiplier-bootstrap-only gate (`survey_design + n_bootstrap > 0` under by_path / paths_of_interest still raises, since the survey-aware perturbation pivot for path-restricted IFs is methodologically underived). Per-path SE routes through the existing `_survey_se_from_group_if` cell-period allocator: the per-period IF (`U_pp_l_path`) is built with non-path switcher-side contributions skipped (control contributions are unchanged, matching the joiners/leavers IF convention; preserves the row-sum identity `U_pp.sum(axis=1) == U`), cohort-recentered via `_cohort_recenter_per_period`, then expanded to observations as `psi_i = U_pp[g_i, t_i] · (w_i / W_{g_i, t_i})`. Replicate-weight designs unconditionally use the cell allocator (Class A contract from PR #323). New `_refresh_path_inference` helper post-call refreshes `safe_inference` on every populated entry across `multi_horizon_inference`, `placebo_horizon_inference`, `path_effects`, and `path_placebos` so all four surfaces use the same final `df_survey` after per-path replicate fits append `n_valid` to the shared accumulator. Path-enumeration ranking under `survey_design` remains unweighted (group-cardinality, not population-weight mass). Lonely-PSU policy stays sample-wide, not per-path. Telescope invariant: on a single-path panel, per-path SE matches the global non-by_path survey SE bit-exactly. **No R parity** — R `did_multiplegt_dyn` does not support survey weighting; this is a Python-only methodology extension. The global non-by_path TSL multiplier-bootstrap path is unaffected (anti-regression test `tests/test_chaisemartin_dhaultfoeuille.py::TestByPathSurveyDesignAnalytical::test_global_survey_plus_n_bootstrap_still_works` locks the per-path-only scope of the new gate). Cross-surface invariants regression-tested at `TestByPathSurveyDesignAnalytical` (~17 tests across gate / dispatch / analytical SE / replicate-weight SE / per-path placebos / `trends_linear` composition / unobserved-path warnings / final-df refresh regressions) and `TestByPathSurveyDesignTelescope`. See `docs/methodology/REGISTRY.md` §`ChaisemartinDHaultfoeuille` `Note (Phase 3 by_path ...)` → "Per-path survey-design SE" for the full contract. - **Inference-field aliases on staggered result classes** for adapter / external-consumer compatibility. Read-only `@property` aliases expose the flat `att` / `se` / `conf_int` / `p_value` / `t_stat` names (matching `DiDResults` / `TROPResults` / `SyntheticDiDResults` / `TripleDifferenceResults` / `HeterogeneousAdoptionDiDResults`) on every result class that previously only carried prefixed canonical fields: `CallawaySantAnnaResults`, `StackedDiDResults`, `EfficientDiDResults`, `ChaisemartinDHaultfoeuilleResults`, `StaggeredTripleDiffResults`, `WooldridgeDiDResults`, `SunAbrahamResults`, `ImputationDiDResults`, `TwoStageDiDResults` (mapping to `overall_*`); `ContinuousDiDResults` (mapping to `overall_att_*`, ATT-side as the headline, ACRT-side accessible unchanged via `overall_acrt_*`); `MultiPeriodDiDResults` (mapping to `avg_*`). `ContinuousDiDResults` additionally exposes `overall_se` / `overall_conf_int` / `overall_p_value` / `overall_t_stat` aliases for naming consistency with the rest of the staggered family. Aliases are pure read-throughs over the canonical fields — no recomputation, no behavior change — so the `safe_inference()` joint-NaN contract (per CLAUDE.md "Inference computation") is inherited automatically (NaN canonical → NaN alias, locked at `tests/test_result_aliases.py::test_pattern_b_aliases_propagate_nan`). The native `overall_*` / `overall_att_*` / `avg_*` fields remain canonical for documentation and computation. Motivated by the `balance.interop.diff_diff.as_balance_diagnostic()` adapter (`facebookresearch/balance` PR #465) which calls `getattr(res, "se", None)` / `getattr(res, "conf_int", None)` without a fallback chain — pre-alias, every staggered result class returned `None` on those keys, silently dropping `se` and `conf_int` from the adapter's diagnostic dict. 23 alias-mechanic + balance-adapter regression tests at `tests/test_result_aliases.py`. Patch-level (additive on stable surfaces). diff --git a/README.md b/README.md index b174711b..b54d8aa0 100644 --- a/README.md +++ b/README.md @@ -124,7 +124,7 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`. - [Honest DiD](https://diff-diff.readthedocs.io/en/stable/api/honest_did.html) - Rambachan & Roth (2023) sensitivity analysis: robust CI under PT violations, breakdown values - [Pre-Trends Power Analysis](https://diff-diff.readthedocs.io/en/stable/api/pretrends.html) - Roth (2022) minimum detectable violation and power curves - [Power Analysis](https://diff-diff.readthedocs.io/en/stable/api/power.html) - analytical and simulation-based MDE, sample size, power curves for study design -- Conley spatial HAC SE (`vcov_type="conley"`) on cross-sectional `LinearRegression` / `compute_robust_vcov` plus panel `MultiPeriodDiD` / `TwoWayFixedEffects` (with `conley_lag_cutoff` for within-unit Bartlett temporal HAC) - Conley (1999) spatial-correlation-aware SEs with parity vs R `conleyreg` on cross-sectional + panel fixtures +- Conley spatial HAC SE (`vcov_type="conley"`) on cross-sectional `LinearRegression` / `compute_robust_vcov` plus panel `DifferenceInDifferences` / `MultiPeriodDiD` / `TwoWayFixedEffects` (with `conley_lag_cutoff` for within-unit Bartlett temporal HAC) - Conley (1999) spatial-correlation-aware SEs with parity vs R `conleyreg` on cross-sectional + panel fixtures, optional combined spatial + cluster product kernel via explicit `cluster=`, auto-activating sparse k-d-tree fast path for `n > 5_000` ## Survey Support diff --git a/TODO.md b/TODO.md index 5aaf2001..aaae5dfb 100644 --- a/TODO.md +++ b/TODO.md @@ -116,12 +116,8 @@ Deferred items from PR reviews that were not addressed before merge. | `HeterogeneousAdoptionDiD` time-varying dose on event study: Phase 2b REJECTS panels where `D_{g,t}` varies within a unit for `t >= F` (the aggregation uses `D_{g, F}` as the single regressor for all horizons, paper Appendix B.2 constant-dose convention). A follow-up PR could add a time-varying-dose estimator for these panels; current behavior is front-door rejection with a redirect to `ChaisemartinDHaultfoeuille`. | `diff_diff/had.py::_validate_had_panel_event_study` | Phase 2b | Low | | `HeterogeneousAdoptionDiD` repeated-cross-section support: paper Section 2 defines HAD on panel OR repeated cross-section, but Phase 2a is panel-only. RCS inputs (disjoint unit IDs between periods) are rejected by the balanced-panel validator with the generic "unit(s) do not appear in both periods" error. A follow-up PR will add an RCS identification path based on pre/post cell means (rather than unit-level first differences), with its own validator and a distinct `data_mode` / API surface. | `diff_diff/had.py::_validate_had_panel`, `diff_diff/had.py::_aggregate_first_difference` | Phase 2a | Medium | | SyntheticDiD: bootstrap cross-language parity anchor against R's default `synthdid::vcov(method="bootstrap")` (refit; rebinds `opts` per draw) or Julia `Synthdid.jl::src/vcov.jl::bootstrap_se` (refit by construction). Same-library validation (placebo-SE tracking, AER §6.3 MC truth) is in place; a cross-language anchor is desirable to bolster the methodology contract. Julia is the cleanest target — minimal wrapping work and refit-native vcov. Tolerance target: 1e-6 on Monte Carlo samples (different BLAS + RNG paths preclude 1e-10). The R-parity fixture from the previous release was deleted because it pinned the now-removed fixed-weight path. | `benchmarks/R/`, `benchmarks/julia/`, `tests/` | follow-up | Low | -| `DifferenceInDifferences(vcov_type="conley")` support: Phase 2 lifted the panel rejection on `MultiPeriodDiD` / `TwoWayFixedEffects` but kept DiD rejected because `DiD.fit()` has no `unit` column declaration. A follow-up could either add a `unit` kwarg to `DiD.__init__` / `.fit()` and wire the block-decomposed sandwich, or document the redirect to `MultiPeriodDiD` permanently. | `diff_diff/estimators.py::DifferenceInDifferences.fit` | follow-up (spillover-conley) | Low | -| Conley + cluster_ids combined product kernel `K_space(d_ij/h) · 1{cluster_i = cluster_j}`. Deferred to a follow-up PR per the Phase 2 scope decision (methodology PR only); tracked here for the next Conley wave. Currently raises `NotImplementedError` at the linalg validator (cross-sectional Conley + cluster) and at `TwoWayFixedEffects.fit` when the user sets `cluster=` explicitly. | `linalg.py::_validate_vcov_args`, `twfe.py::TwoWayFixedEffects.fit` | follow-up (spillover-conley) | Medium | -| Conley sparse k-d-tree fast path via `scipy.spatial.cKDTree.query_ball_tree` for `n > 20_000`; lifts the dense O(n²) `UserWarning` that fires at the validator. Kernel must be compact-support (bartlett or uniform); callable metric not supported in the fast path. Performance-only; semantics unchanged. | `diff_diff/conley.py::_pairwise_distance_matrix`, `_compute_conley_vcov` | follow-up (spillover-conley) | Low | | Conley + survey weights / `survey_design`. Score-reweighted meat `s_i = w_i · X_i · ε_i` is mechanical, but PSU clustering interaction with the spatial kernel and replicate-weights variance under spatial correlation are non-trivial (Bertanha-Imbens 2014 covers cluster-sample but not the explicit Conley case). Phase 5 of the spillover-conley initiative; paper review prerequisite. Currently raises `NotImplementedError` at the linalg validator. | `linalg.py::_validate_vcov_args` | Phase 5 (spillover-conley) | Medium | | `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 | -| Validate user-supplied callable `conley_metric` for shape `(n, n)`, finiteness, non-negativity, and symmetry. Currently `np.asarray(metric(coords, coords))` is accepted unchecked; a malformed callable produces opaque matmul errors and a non-symmetric distance matrix produces a non-symmetric vcov. CI reviewer flagged as P2 M3 in PR #(spillover-conley). | `diff_diff/conley.py::_pairwise_distance_matrix`, `_compute_conley_vcov` | follow-up (spillover-conley) | Low | #### Performance diff --git a/diff_diff/conley.py b/diff_diff/conley.py index 0fac96e2..812a8f16 100644 --- a/diff_diff/conley.py +++ b/diff_diff/conley.py @@ -28,14 +28,26 @@ ``conleyreg::haversine_dist`` (Düsterhöft 2021, CRAN v0.1.9). See ``benchmarks/R/README.md`` for the cross-language parity convention. -A sparse k-d-tree fast path for ``n > 20_000`` is a follow-up; the dense -O(n²) distance matrix is emitted with a ``UserWarning`` for now. +A sparse k-d-tree fast path auto-activates for the spatial component when +``n > _CONLEY_SPARSE_N_THRESHOLD`` (5_000), ``conley_metric`` is one of +``{"haversine", "euclidean"}``, and ``conley_kernel`` is ``"bartlett"``. +The bartlett-only gate is for boundary correctness — bartlett at +``u = 1.0`` returns exactly ``0.0``, so pairs at exactly the cutoff +contribute zero and can be safely dropped by the sparse path. The +uniform kernel returns ``1.0`` at the cutoff and is methodologically +incompatible with the chord-projection roundoff in the haversine query; +it falls back to the dense path. The serial component (within-unit +Bartlett HAC over time) is always dense regardless of n. The +``_CONLEY_DENSE_OOM_WARN_N`` constant (20_000) is a separate +memory-exhaustion warning, retained even when the sparse path activates +because dense fallback still applies for non-bartlett kernels and for +callable metrics. """ from __future__ import annotations import warnings -from typing import Callable, Literal, Optional, Union +from typing import Any, Callable, Literal, Optional, Union, cast import numpy as np @@ -58,8 +70,34 @@ # at many more digits) but matters for the 1e-6 cross-language parity bound. _CONLEY_EARTH_RADIUS_KM = 6371.01 -# Empirical threshold for warning about dense O(n²) distance matrix memory -_CONLEY_DENSE_WARN_N = 20_000 +# Empirical threshold above which the dense O(n²) distance matrix is expected +# to risk OOM on typical commodity workstations (64 GB float64 at n=89_443). +# Renamed from the original ``_CONLEY_DENSE_WARN_N`` to disambiguate it from +# ``_CONLEY_SPARSE_N_THRESHOLD``: this constant is about memory exhaustion; +# the sparse-path threshold is about compute. The two are independent because +# the sparse fast path requires ``conley_kernel='bartlett'`` and a non-callable +# metric — callable metrics and uniform kernels still take the dense path +# at any n. +_CONLEY_DENSE_OOM_WARN_N = 20_000 + +# Total-n threshold above which the sparse k-d-tree fast path auto-activates +# for the spatial Bartlett meat. Subject to two additional gates: the metric +# must be ``"haversine"`` or ``"euclidean"`` (not callable), and the kernel +# must be ``"bartlett"`` (not ``"uniform"``). Crossed in either direction, +# the dense path is used regardless of n. +_CONLEY_SPARSE_N_THRESHOLD = 5_000 + +# Density gate: fraction of n*n pairs within cutoff above which the sparse +# CSR matrix's storage overhead (~12 bytes/nnz: data + indices + indptr) +# loses its memory advantage over a dense float64 (8 bytes/cell). The +# break-even is at ~67% density; we gate at 30% (well below break-even) +# to give a comfortable safety margin: at 30% nnz, CSR uses ~45% of +# dense memory. Above 30%, fall back to dense + emit a UserWarning so +# users with large cutoffs aren't surprised by the "sparse" path +# materializing a near-dense matrix. Computed exactly via +# ``cKDTree.count_neighbors`` (O(n log n + nnz), shares tree traversal +# with the subsequent ``query_ball_tree``). +_CONLEY_SPARSE_DENSITY_THRESHOLD = 0.3 def _haversine_km( @@ -88,11 +126,200 @@ def _haversine_km( return _CONLEY_EARTH_RADIUS_KM * 2.0 * np.arcsin(np.sqrt(a)) +def _validate_callable_metric_result(result: object, n: int) -> np.ndarray: + """Validate the output of a user-supplied callable ``conley_metric``. + + A user-supplied distance callable must return an ``(n, n)`` matrix of + finite, non-negative, symmetric values with **zero diagonal**; + otherwise downstream code produces opaque BLAS errors or silently- + wrong vcov estimates. This helper performs all six checks at the + boundary and produces a targeted :class:`ValueError` for each + failure. + + The zero-diagonal invariant is load-bearing for the Conley sandwich: + the ``i = j`` term contributes ``K(d_ii / h) · X_i ε_i² X_i'``, which + must reduce to the HC0 diagonal ``X_i ε_i² X_i'`` (i.e., ``K(0) = 1``). + A callable with positive self-distance would attenuate the HC0 + contribution by ``K(d_ii / h) < 1`` and silently misstate Conley SEs. + Built-in metrics (``"haversine"``, ``"euclidean"``) satisfy this by + construction. + + Returns + ------- + arr : ndarray of shape (n, n), float64 + The validated distance matrix, ready for kernel evaluation. + + Raises + ------ + ValueError + Result cannot cast to ``float64``; shape is not ``(n, n)``; + contains NaN/inf; contains negative entries; is not symmetric + within ``atol=1e-10``; or has any nonzero diagonal entry + ``|d(i, i)| > 1e-10``. + """ + try: + arr = np.asarray(result, dtype=np.float64) + except (TypeError, ValueError) as exc: + raise ValueError( + "conley_metric callable returned a value that cannot be cast to " + f"a float64 array: {exc}." + ) from exc + if arr.shape != (n, n): + raise ValueError( + "conley_metric callable must return a (n, n) distance matrix; " + f"got shape {arr.shape}, expected ({n}, {n})." + ) + if not np.isfinite(arr).all(): + raise ValueError( + "conley_metric callable returned non-finite entries (NaN or inf); " + "all pairwise distances must be finite." + ) + if (arr < 0.0).any(): + raise ValueError( + "conley_metric callable returned negative entries; all pairwise " + "distances must be non-negative." + ) + asymmetry = float(np.max(np.abs(arr - arr.T))) if arr.size else 0.0 + if asymmetry > 1e-10: + raise ValueError( + "conley_metric callable returned an asymmetric matrix; the " + "distance matrix must satisfy d(i, j) = d(j, i). Max |D - D.T| = " + f"{asymmetry:.2e}, tolerance 1e-10." + ) + diag_max = float(np.max(np.abs(np.diag(arr)))) if arr.size else 0.0 + if diag_max > 1e-10: + raise ValueError( + "conley_metric callable returned a matrix with nonzero self-" + "distance on the diagonal; the Conley sandwich requires " + "d(i, i) = 0 so the kernel reduces to K(0) = 1 on the HC0 " + f"diagonal contribution. Max |diag(D)| = {diag_max:.2e}, " + "tolerance 1e-10." + ) + return arr + + +def _validate_conley_estimator_inputs( + *, + estimator_name: str, + data: "Any", + unit: Optional[str], + conley_coords: object, + conley_cutoff_km: Optional[float], + conley_lag_cutoff: Optional[int], + survey_design: object, + inference: str, + cluster: Optional[str] = None, +) -> None: + """Shared front-door validation for ``vcov_type='conley'`` on the + estimator entry points (``DifferenceInDifferences``, ``MultiPeriodDiD``, + ``TwoWayFixedEffects``). + + Each estimator's ``fit()`` calls this BEFORE building Conley arrays or + threading them into the variance computation. The eight checks below + are the union of what each estimator needs; estimator-specific bits + (e.g. building the array from a column name) remain inline at the + caller. Centralizing this in one place prevents the validation-class + drift that surfaced repeatedly across Wave A CI rounds (front-door + guards for `unit`, `conley_coords`, and `cluster` were each missing + on at least one estimator surface and required separate fixes). + + Parameters + ---------- + estimator_name : str + Class name surfaced in error messages (e.g. + ``"DifferenceInDifferences"``). + data : pandas.DataFrame + The dataset passed to ``fit()``. Used to check column existence + for ``unit``, ``conley_coords``, and ``cluster``. Typed as + ``Any`` here to avoid importing pandas at module load time. + unit : str or None + Name of the unit identifier column (required for Conley). + conley_coords : tuple/list/None + Pair of column names ``(lat_col, lon_col)``. + conley_cutoff_km : float or None + Positive finite bandwidth. + conley_lag_cutoff : int or None + Non-negative integer lag for the within-unit Bartlett serial sum. + survey_design : object + ``SurveyDesign`` instance or ``None``. Survey + Conley is deferred. + inference : str + Estimator inference mode. ``"wild_bootstrap"`` + Conley is rejected. + cluster : str or None, default None + Name of the cluster column if the user opted into the combined + spatial + cluster product kernel (Wave A #119). When non-None, + the column must exist in ``data``; otherwise the downstream + ``data[cluster]`` access would raise an opaque pandas + ``KeyError``. Pass ``None`` (the default) to skip this check. + + Raises + ------ + ValueError + Missing/malformed conley_coords, missing conley_cutoff_km, + missing/unknown unit, missing conley_lag_cutoff, coord column + not in ``data``, or ``cluster`` set to a name not in ``data``. + NotImplementedError + ``survey_design`` is non-None, or ``inference == "wild_bootstrap"``. + """ + if conley_coords is None or conley_cutoff_km is None: + raise ValueError( + f"{estimator_name}(vcov_type='conley') requires " + "conley_coords=(lat_col, lon_col) and conley_cutoff_km " + "on the constructor." + ) + if ( + not isinstance(conley_coords, (tuple, list)) + or len(conley_coords) != 2 + or not all(isinstance(c, str) for c in conley_coords) + ): + raise ValueError( + "conley_coords must be a 2-element tuple/list of column " + f"names (lat_col, lon_col); got {conley_coords!r}." + ) + for _coord_col in conley_coords: + if _coord_col not in data.columns: + raise ValueError(f"conley_coords column '{_coord_col}' not found in data.") + if unit is None: + raise ValueError( + f"{estimator_name}(vcov_type='conley') requires `unit=` " + "— the panel block-decomposed Conley sandwich needs the unit " + "identifier to compute the per-unit serial sum." + ) + if unit not in data.columns: + raise ValueError(f"Unit column '{unit}' not found in data") + if conley_lag_cutoff is None: + raise ValueError( + f"{estimator_name}(vcov_type='conley') requires conley_lag_cutoff " + "(non-negative int; 0 means spatial-within-period only, no serial " + "component). See R conleyreg's `lag_cutoff` argument for the convention." + ) + if cluster is not None and cluster not in data.columns: + raise ValueError(f"Cluster column '{cluster}' not found in data") + if survey_design is not None: + raise NotImplementedError( + f"{estimator_name}(vcov_type='conley') + survey_design is a " + "follow-up (Bertanha-Imbens 2014 weighted-Conley). Drop " + "survey_design for cross-sectional Conley, or use vcov_type='hc1' " + "for survey-aware cluster-robust without spatial HAC." + ) + if inference == "wild_bootstrap": + raise NotImplementedError( + f"{estimator_name}(vcov_type='conley', inference='wild_bootstrap') " + "is not supported: the wild bootstrap is a separate inference path " + "that does not consume the analytical Conley sandwich. Use " + "inference='analytical' for Conley SEs." + ) + + def _pairwise_distance_matrix(coords: np.ndarray, metric: ConleyMetric) -> np.ndarray: """Build the dense n×n pairwise distance matrix. ``metric`` is one of ``"haversine"`` (lat/lon in degrees, distance in km), ``"euclidean"`` (any units), or a callable ``f(coords1, coords2) -> n×n``. + + For the callable branch, the result is validated via + :func:`_validate_callable_metric_result`: shape ``(n, n)``, finite, + non-negative, symmetric within ``atol=1e-10``. Each failure raises + :class:`ValueError` naming the violated invariant. """ if metric == "haversine": lats = coords[:, 0] @@ -104,7 +331,7 @@ def _pairwise_distance_matrix(coords: np.ndarray, metric: ConleyMetric) -> np.nd diff = coords[:, None, :] - coords[None, :, :] return np.sqrt(np.sum(diff * diff, axis=-1)) if callable(metric): - return np.asarray(metric(coords, coords), dtype=np.float64) + return _validate_callable_metric_result(metric(coords, coords), coords.shape[0]) raise ValueError( f"conley_metric must be 'haversine', 'euclidean', or callable; got {metric!r}." ) @@ -138,6 +365,200 @@ def _uniform_kernel(u: np.ndarray) -> np.ndarray: return (np.abs(u) <= 1.0).astype(np.float64) +def _compute_spatial_bartlett_meat_sparse( + S: np.ndarray, + coords: np.ndarray, + cutoff: float, + metric: str, + *, + cluster_codes: Optional[np.ndarray] = None, + density_threshold: float = _CONLEY_SPARSE_DENSITY_THRESHOLD, +) -> Optional[np.ndarray]: + """Sparse k-d-tree-based spatial Bartlett meat: ``S.T @ K_bartlett @ S``. + + Used by :func:`_compute_conley_vcov` when ``_conley_sparse`` is True or + when the auto-toggle fires (n above the threshold, bartlett kernel, + non-callable metric). Avoids materializing the full n×n distance matrix + by using ``scipy.spatial.cKDTree.query_ball_tree`` to find in-range + neighbor pairs and assembling a CSR sparse kernel matrix on top of those. + + Parameters + ---------- + S : (n, k) array + Score matrix ``X * residuals[:, None]``. + coords : (n, 2) array of float64 + ``[lat, lon]`` (haversine) or arbitrary 2-D coordinates (euclidean). + cutoff : float + Conley bandwidth (km for haversine; same units as ``coords`` for + euclidean). + metric : {"haversine", "euclidean"} + Distance metric. Callable metrics fall back to the dense path and + are not supported here — the caller is expected to enforce that. + cluster_codes : (n,) int array, optional, keyword-only + Integer cluster codes (one per row of ``S``). When supplied, the + combined kernel ``K_total[i, j] = K_space(d_ij/h) · 1{cluster_i = + cluster_j}`` is applied: neighbors are filtered to same-cluster + pairs before the kernel evaluation. + + Returns + ------- + meat : (k, k) array + + Notes + ----- + For haversine, lat/lon is projected to a 3-D unit-sphere Cartesian + point cloud (``x = cos(lat) cos(lon), y = cos(lat) sin(lon), z = sin(lat)``) + so that the kd-tree's euclidean radius query corresponds to a chord + distance on the unit sphere. The chord radius is + ``2 sin(cutoff_km / (2 R_earth))`` (the chord on the unit sphere + matching the great-circle arc of length ``cutoff_km``). A small + relative epsilon (``1 + 1e-12``) is added to the query radius to + absorb chord-projection roundoff; after the query, the exact + great-circle distance is recomputed via :func:`_haversine_km` and + the Bartlett kernel is applied. Pairs at exactly the cutoff distance + contribute zero to the kernel (bartlett at ``u = 1.0`` is exactly + ``0.0``) and are filtered out, which is why the sparse path is + bartlett-only — uniform at ``u = 1.0`` returns ``1.0`` and would + require querying with a strict superset and applying the closed- + interval kernel by hand. + + Numerical tolerance vs the dense path on the same inputs is + typically ~1e-12 in absolute terms; the haversine→chord projection + introduces sub-eps error that propagates through the matmul. + """ + from scipy.sparse import csr_matrix + from scipy.spatial import cKDTree + + n = coords.shape[0] + k = S.shape[1] + + if metric == "haversine": + lat_r = np.radians(coords[:, 0]) + lon_r = np.radians(coords[:, 1]) + xyz = np.column_stack( + [ + np.cos(lat_r) * np.cos(lon_r), + np.cos(lat_r) * np.sin(lon_r), + np.sin(lat_r), + ] + ) + # Chord on the unit sphere matching arc-length cutoff_km. Clamp the + # arc to π radians (half-Earth circumference) before computing the + # chord: the function `arc -> 2 sin(arc/2)` is monotonically + # increasing only on [0, π] and reaches its global max of 2 (the + # unit-sphere diameter) at arc = π. Without the clamp, cutoffs + # above π·R (~20,015 km) would compute a SMALLER chord radius than + # at cutoff = π·R and silently drop antipodal pairs that still have + # positive Bartlett weight. The dense path saturates at π·R via + # `_haversine_km`'s `np.clip(a, 0, 1)` (arcsin caps at π/2), so the + # clamp mirrors that saturation. The 1+1e-12 epsilon then absorbs + # chord-projection float roundoff at sub-cutoff distances. + arc_radians = min(cutoff / _CONLEY_EARTH_RADIUS_KM, np.pi) + query_r = 2.0 * np.sin(arc_radians / 2.0) + query_r *= 1.0 + 1e-12 + tree = cKDTree(xyz) + tree_coords = xyz # used for per-cluster sub-trees in the density gate + elif metric == "euclidean": + # Small relative epsilon for symmetry with the haversine branch. + # Bartlett's <=-vs-< boundary is moot since kernel is exactly 0 at u=1. + query_r = cutoff * (1.0 + 1e-12) + tree = cKDTree(coords) + tree_coords = coords + else: + raise ValueError( + "sparse Conley path requires metric in {'haversine', 'euclidean'}; " + f"got {metric!r}. (Callable metrics fall back to the dense path.)" + ) + + # Density gate: count in-range pairs cheaply via the same tree (shares + # traversal cost with query_ball_tree, no extra allocation). If density + # exceeds the threshold, return None so the caller falls back to dense + # (avoids materializing a near-dense CSR matrix that would use MORE + # memory than dense float64). Codex CI R6 P2. + n_pairs_in_range = int(tree.count_neighbors(tree, r=query_r, p=2.0)) + density = n_pairs_in_range / float(n * n) + if density > density_threshold and cluster_codes is not None: + # Cluster-aware refinement: the actual CSR nnz reflects within- + # cluster pairs only (the inner loop drops cross-cluster neighbors + # before adding to the sparse matrix). On large clustered panels, + # spatial density can be ~100% while within-cluster density is + # tiny — without refinement we'd spuriously fall back to dense + # exactly when sparse helps most. Refine by counting per-cluster. + # Codex CI R8 P1: the unrefined density check was cluster-blind. + refined_pairs = 0 + for c in np.unique(cluster_codes): + in_c = cluster_codes == c + n_c = int(in_c.sum()) + if n_c <= 1: + # Singleton cluster contributes only the diagonal self-pair. + refined_pairs += n_c + continue + sub_tree = cKDTree(tree_coords[in_c]) + refined_pairs += int(sub_tree.count_neighbors(sub_tree, r=query_r, p=2.0)) + density = refined_pairs / float(n * n) + if density > density_threshold: + warnings.warn( + f"Conley sparse path: neighbor density {density:.1%} exceeds " + f"threshold {density_threshold:.1%}; falling back to dense " + "(CSR storage would use more memory than the dense float64 " + "matrix at this density). Consider a smaller conley_cutoff_km " + "for genuine memory savings.", + UserWarning, + stacklevel=3, + ) + return None + + neighbors = tree.query_ball_tree(tree, r=query_r, p=2.0) + + rows_list: list[np.ndarray] = [] + cols_list: list[np.ndarray] = [] + data_list: list[np.ndarray] = [] + for i, neigh in enumerate(neighbors): + if not neigh: + # Should not happen — i is always in its own ball — but guard. + continue + neigh_arr = np.asarray(neigh, dtype=np.intp) + # Combined spatial + cluster product kernel: drop neighbors that + # don't share i's cluster. This is the sparse-path analog of the + # dense path's Hadamard-with-cluster-mask step. + if cluster_codes is not None: + same_cluster = cluster_codes[neigh_arr] == cluster_codes[i] + neigh_arr = neigh_arr[same_cluster] + if neigh_arr.size == 0: + continue + # Recompute exact distance for the in-range neighbors (NOT chord + # for haversine — we apply the kernel to actual great-circle km). + if metric == "haversine": + dist = _haversine_km( + np.asarray(coords[i, 0]), + np.asarray(coords[i, 1]), + coords[neigh_arr, 0], + coords[neigh_arr, 1], + ) + else: + diff = coords[neigh_arr] - coords[i] + dist = np.sqrt(np.sum(diff * diff, axis=-1)) + kernel_vals = _bartlett_kernel(dist / cutoff) + nonzero_mask = kernel_vals > 0.0 + if not nonzero_mask.any(): + continue + nonzero_neighbors = neigh_arr[nonzero_mask] + nonzero_kernels = kernel_vals[nonzero_mask] + rows_list.append(np.full(nonzero_neighbors.shape, i, dtype=np.intp)) + cols_list.append(nonzero_neighbors) + data_list.append(nonzero_kernels) + + if not data_list: + return np.zeros((k, k)) + + K = csr_matrix( + (np.concatenate(data_list), (np.concatenate(rows_list), np.concatenate(cols_list))), + shape=(n, n), + ) + # scipy sparse @ dense returns dense; do K @ S first then S.T @ that. + return S.T @ (K @ S) + + def _validate_conley_kwargs( coords: Optional[np.ndarray], cutoff: Optional[float], @@ -148,19 +569,27 @@ def _validate_conley_kwargs( time: Optional[np.ndarray] = None, unit: Optional[np.ndarray] = None, lag_cutoff: Optional[int] = None, + cluster_ids: Optional[np.ndarray] = None, ) -> None: """Validate the Conley kwargs against the design's row count. The first five positional args define the cross-sectional contract (Phase 1). The three keyword-only ``time`` / ``unit`` / ``lag_cutoff`` args define the panel block-decomposed contract (Phase 2); they are three-way co-required. + The ``cluster_ids`` keyword-only arg enables the combined spatial + + cluster product kernel (Wave A item #119); on the panel block-decomposed + path it must be **constant within each unit across periods** (otherwise + the within-unit serial mask would zero out adjacent-time pairs and + produce a methodologically-muddled meat). Raises ------ ValueError Missing/malformed coords or cutoff; lat/lon out of range under haversine; unknown kernel/metric; non-finite or non-positive cutoff; - partial panel arg set (must pass all three of time/unit/lag_cutoff or none). + partial panel arg set (must pass all three of time/unit/lag_cutoff or none); + ``cluster_ids`` with wrong shape or NaN; ``cluster_ids`` not constant + within each unit on the panel path. Warnings -------- @@ -254,15 +683,86 @@ def _validate_conley_kwargs( f"conley_lag_cutoff must be a non-negative integer; got {lag_cutoff!r}." ) - if n > _CONLEY_DENSE_WARN_N: + # Combined spatial + cluster product kernel (Wave A #119). The validator + # checks shape, missing values, and — on the panel block-decomposed path + # only — that cluster membership is constant within each unit across + # periods. Cross-sectional path has no time dimension, so no + # invariance constraint applies. + if cluster_ids is not None: + cluster_arr = np.asarray(cluster_ids) + if cluster_arr.ndim != 1 or cluster_arr.shape[0] != n: + raise ValueError( + f"cluster_ids must be a 1-D array of length {n} when combined " + f"with vcov_type='conley'; got shape {cluster_arr.shape}." + ) + import pandas as _pd + + if _pd.isna(cluster_arr).any(): + raise ValueError( + "cluster_ids contains NaN or missing values when combined " + "with vcov_type='conley'." + ) + if n_panel_set == 3: + # Panel path: enforce time-invariance within unit. If a unit's + # cluster changes across periods, the within-unit serial mask + # would zero out adjacent-time pairs that should contribute. + cluster_codes_v, _ = _pd.factorize(cluster_arr) + unit_arr_local = np.asarray(unit) + unit_codes_v, _ = _pd.factorize(unit_arr_local) + # Count distinct cluster codes per unit. >1 means the cluster + # assignment varies across time within at least one unit. + df_check = _pd.DataFrame({"u": unit_codes_v, "c": cluster_codes_v}) + unit_n_clusters = df_check.groupby("u")["c"].nunique() + violator_mask = unit_n_clusters.to_numpy() > 1 + if bool(violator_mask.any()): + # Report up to 5 violating unit labels to help debugging. + violator_codes = np.asarray(unit_n_clusters.index)[violator_mask] + uniques_unit = _pd.unique(unit_arr_local) + shown = [str(uniques_unit[int(c)]) for c in violator_codes[:5]] + total = int(violator_mask.sum()) + more = "" if total <= 5 else f" (+{total - 5} more)" + raise ValueError( + "cluster_ids combined with the panel block-decomposed " + "Conley path (conley_lag_cutoff set) must be constant " + "within each unit across periods; the within-unit serial " + "Bartlett HAC's mask is trivially all-ones only under " + "this contract. Violating units: " + f"{', '.join(shown)}{more}. Either pass a " + "time-invariant cluster (e.g. unit-level region) or " + "drop cluster_ids." + ) + + if n > _CONLEY_DENSE_OOM_WARN_N: memory_gb = (n * n * 8) / 1e9 - warnings.warn( - f"vcov_type='conley' builds a dense {n}x{n} distance matrix " - f"(~{memory_gb:.1f} GB float64). The sparse k-d-tree fast path is " - "deferred to a follow-up PR.", - UserWarning, - stacklevel=3, - ) + # The sparse k-d-tree fast path will auto-activate for the spatial + # meat when n > _CONLEY_SPARSE_N_THRESHOLD AND metric is haversine + # or euclidean AND kernel is bartlett. Above the OOM warning + # threshold the dense fallback (callable metric, uniform kernel) + # is at material memory risk. + if ( + isinstance(metric, str) + and metric in ("haversine", "euclidean") + and kernel == "bartlett" + ): + warnings.warn( + f"vcov_type='conley': the sparse k-d-tree fast path will be " + f"used for the spatial Bartlett meat (n={n}, kernel='bartlett', " + f"metric={metric!r}); the per-unit serial Bartlett HAC remains " + f"dense. The dense {n}x{n} fallback would be ~{memory_gb:.1f} GB " + "float64.", + UserWarning, + stacklevel=3, + ) + else: + warnings.warn( + f"vcov_type='conley' builds a dense {n}x{n} distance matrix " + f"(~{memory_gb:.1f} GB float64). The sparse k-d-tree fast path " + "requires conley_kernel='bartlett' and conley_metric in " + "{'haversine', 'euclidean'}; " + f"current kernel={kernel!r}, metric={metric!r}.", + UserWarning, + stacklevel=3, + ) def _compute_conley_vcov( @@ -277,6 +777,8 @@ def _compute_conley_vcov( time: Optional[np.ndarray] = None, unit: Optional[np.ndarray] = None, lag_cutoff: Optional[int] = None, + cluster_ids: Optional[np.ndarray] = None, + _conley_sparse: Optional[bool] = None, ) -> np.ndarray: """Conley (1999) spatial HAC sandwich variance. @@ -300,6 +802,16 @@ def _compute_conley_vcov( exclusion avoids double-counting the diagonal already covered by the spatial component. + **Combined spatial + cluster product kernel** (``cluster_ids`` supplied): + + K_total[i, j] = K_space(d_ij/h) · 1{cluster_i = cluster_j} + + The cluster indicator multiplies the spatial kernel on the within-period + (or full cross-sectional) sandwich. On the panel path, the validator + enforces that ``cluster_ids`` is constant within each unit across + periods, which guarantees the within-unit serial mask is trivially + all-ones — no explicit per-unit-time mask needed in the serial loop. + Inputs are assumed already validated by :func:`_validate_conley_kwargs`; the helper only does the math. Caller is responsible for the validator. @@ -319,6 +831,15 @@ def _compute_conley_vcov( """ coords_arr = np.asarray(coords, dtype=np.float64) S = X * residuals[:, np.newaxis] + n = X.shape[0] + + # Factorize cluster_ids once if supplied, so per-slice mask construction + # below can use integer comparisons rather than re-factorizing per call. + cluster_codes: Optional[np.ndarray] = None + if cluster_ids is not None: + import pandas as _pd + + cluster_codes = np.asarray(_pd.factorize(np.asarray(cluster_ids))[0], dtype=np.int64) def _kernel_fn(u: np.ndarray) -> np.ndarray: if kernel == "bartlett": @@ -327,6 +848,72 @@ def _kernel_fn(u: np.ndarray) -> np.ndarray: return _uniform_kernel(u) raise ValueError(f"conley_kernel must be 'bartlett' or 'uniform'; got {kernel!r}.") + # Decide whether to use the sparse k-d-tree path for the spatial component. + # Three gates: a supported metric (haversine/euclidean, NOT callable), a + # supported kernel (bartlett — see module docstring for the boundary- + # semantics rationale), and either the auto-toggle threshold or an + # explicit override via the private _conley_sparse kwarg. When explicit + # forcing requests sparse on an unsupported metric/kernel, we raise + # rather than silently fall back so the caller sees the mismatch. + # Use direct equality (NOT isinstance(metric, str)) so pyright can keep + # the Literal narrowing on the dense fallback path below — checking + # `isinstance(metric, str)` widens metric from ConleyMetric to + # str | Callable and breaks downstream typing on `_pairwise_distance_matrix`. + metric_supports_sparse = metric == "haversine" or metric == "euclidean" + kernel_supports_sparse = kernel == "bartlett" + if _conley_sparse is True: + if not (metric_supports_sparse and kernel_supports_sparse): + raise ValueError( + "_conley_sparse=True requires conley_metric in " + "{'haversine', 'euclidean'} and conley_kernel='bartlett'; " + f"got metric={metric!r}, kernel={kernel!r}." + ) + use_sparse = True + elif _conley_sparse is False: + use_sparse = False + else: + use_sparse = ( + n > _CONLEY_SPARSE_N_THRESHOLD and metric_supports_sparse and kernel_supports_sparse + ) + + def _spatial_meat_for_mask(mask: Optional[np.ndarray] = None) -> np.ndarray: + """Compute the spatial meat ``S' K S`` for the given subset of rows. + + ``mask`` may be ``None`` (use all rows) or a boolean array of length n. + Dispatches to the sparse helper when ``use_sparse`` is True, otherwise + builds the dense n×n distance matrix. When ``cluster_codes`` is set, + the cluster indicator multiplies the spatial kernel on this slice + (per-slice mask, NOT full n×n — saves memory for panel paths). + """ + if mask is None: + S_sub = S + coords_sub = coords_arr + cluster_sub = cluster_codes + else: + S_sub = S[mask] + coords_sub = coords_arr[mask] + cluster_sub = cluster_codes[mask] if cluster_codes is not None else None + if use_sparse: + # The auto-toggle and explicit-True gate above both guarantee + # metric is a str (haversine/euclidean), so this cast is safe; + # using cast() avoids leaking the narrowing into the dense + # fallback below. The sparse helper returns None when the + # neighbor density exceeds the threshold (sparse CSR storage + # would use more memory than dense); fall through to dense in + # that case (the warning is already emitted by the helper). + sparse_meat = _compute_spatial_bartlett_meat_sparse( + S_sub, coords_sub, cutoff, cast(str, metric), cluster_codes=cluster_sub + ) + if sparse_meat is not None: + return sparse_meat + # Density-gated fallback: continue to the dense path below. + D = _pairwise_distance_matrix(coords_sub, metric) + K = _kernel_fn(D / cutoff) + if cluster_sub is not None: + cluster_mask = cluster_sub[:, None] == cluster_sub[None, :] + K = K * cluster_mask + return S_sub.T @ K @ S_sub + # Suppress spurious BLAS-level "divide by zero / overflow" warnings on # macOS Accelerate when K is sparse-ish (most off-diagonals are exactly # 0 outside the cutoff). The matmul result is mathematically correct; @@ -334,10 +921,8 @@ def _kernel_fn(u: np.ndarray) -> np.ndarray: # We verify finiteness immediately after. if time is None: # Phase 1 cross-sectional path: full n×n spatial sandwich. - D = _pairwise_distance_matrix(coords_arr, metric) - K = _kernel_fn(D / cutoff) with np.errstate(divide="ignore", over="ignore", invalid="ignore"): - meat = S.T @ K @ S + meat = _spatial_meat_for_mask(None) else: # Phase 2 panel block-decomposed path (matches R conleyreg). time_arr = np.asarray(time) @@ -358,13 +943,11 @@ def _kernel_fn(u: np.ndarray) -> np.ndarray: k = X.shape[1] meat = np.zeros((k, k)) # Spatial component: within-period sandwich, summed across periods. + # _spatial_meat_for_mask dispatches to sparse or dense per the toggle. with np.errstate(divide="ignore", over="ignore", invalid="ignore"): for t_code in np.unique(time_codes): mask_t = time_codes == t_code - S_t = S[mask_t] - D_t = _pairwise_distance_matrix(coords_arr[mask_t], metric) - K_t = _kernel_fn(D_t / cutoff) - meat += S_t.T @ K_t @ S_t + meat += _spatial_meat_for_mask(mask_t) # Serial component: within-unit Bartlett HAC for lag in {1..L}, # excluding lag=0 to avoid double-counting the spatial diagonal. # Bartlett form hardcoded (matches conleyreg::time_dist). diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index a9f4d8d0..01e8a377 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -69,15 +69,18 @@ class DifferenceInDifferences: with ``cluster=``, Pustejovsky-Tipton (2018) CR2 cluster-robust. (Note: ``MultiPeriodDiD`` does NOT yet support ``cluster=`` with ``"hc2_bm"`` — see ``MultiPeriodDiD`` docstring and REGISTRY.md.) - - ``"conley"``: Conley 1999 spatial-HAC sandwich. **Accepted by the - constructor for sklearn-style API symmetry but rejected at - fit-time on ``DifferenceInDifferences``** because DiD is - intrinsically a two-period panel design, and Phase 1's cross- - sectional Conley does not handle the time dimension. The - supported Phase 1 path for Conley is direct - ``compute_robust_vcov`` / ``LinearRegression`` on a single-period - regression. Phase 2 will add the space-time product kernel - (Driscoll-Kraay) and lift the rejection. + - ``"conley"``: Conley 1999 spatial-HAC sandwich. Pass + ``conley_coords=(lat_col, lon_col)``, ``conley_cutoff_km=``, + and ``conley_lag_cutoff=`` on the constructor; pass + ``unit=`` as a fit-time kwarg to :meth:`fit` (NOT on + ``__init__``; unused unless Conley is set; not part of + ``get_params()`` / ``set_params()``). The block-decomposed panel + sandwich (matches R ``conleyreg`` with ``lag_cutoff > 0``) sums + within-period spatial pairs plus within-unit Bartlett serial + pairs (lag=0 excluded). Explicit ``cluster=`` enables the + combined spatial + cluster product kernel; ``survey_design=`` + and ``inference='wild_bootstrap'`` both raise + ``NotImplementedError``. alpha : float, default=0.05 Significance level for confidence intervals. inference : str, default="analytical" @@ -97,13 +100,23 @@ class DifferenceInDifferences: - "warn": Issue warning and drop linearly dependent columns (default) - "error": Raise ValueError - "silent": Drop columns silently without warning - conley_coords, conley_cutoff_km, conley_metric, conley_kernel - Accepted by the constructor for sklearn-style API symmetry, but - ``vcov_type="conley"`` is rejected at fit-time on - ``DifferenceInDifferences`` (see ``vcov_type`` above). Use direct - ``compute_robust_vcov`` / ``LinearRegression`` on a single-period - regression for cross-sectional Conley in Phase 1; Phase 2 will lift - the panel rejection. + conley_coords, conley_cutoff_km, conley_metric, conley_kernel, conley_lag_cutoff + Conley (1999) spatial-HAC variance configuration. Pass + ``conley_coords=(lat_col, lon_col)``, ``conley_cutoff_km=``, + and ``conley_lag_cutoff=`` on the constructor; the ``unit`` + identifier is passed as a fit-time arg to ``fit(...)`` (NOT on + ``__init__``) — it is unused unless ``vcov_type="conley"`` and is + therefore not part of ``get_params()`` / ``set_params()`` (which + return constructor-arg dicts). The block-decomposed panel sandwich + (matching R ``conleyreg`` with ``lag_cutoff > 0``) sums within-period + spatial pairs plus within-unit Bartlett serial pairs (lag=0 excluded + to avoid double-counting). Explicit ``cluster=`` + Conley + enables the combined spatial + cluster product kernel; the cluster + must be constant within each unit across periods (validator-enforced). + DiD has no auto-cluster, so cluster is fully opt-in on the Conley + path — absent ``cluster=``, pure Conley spatial HAC applies. + ``survey_design=`` + Conley and ``inference='wild_bootstrap'`` + + Conley both raise ``NotImplementedError``. Attributes ---------- @@ -218,6 +231,7 @@ def fit( fixed_effects: Optional[List[str]] = None, absorb: Optional[List[str]] = None, survey_design=None, + unit: Optional[str] = None, ) -> DiDResults: """ Fit the Difference-in-Differences model. @@ -249,6 +263,16 @@ def fit( Survey design specification for design-based inference. When provided, uses Taylor Series Linearization for variance estimation and applies sampling weights to the regression. + unit : str, optional + Name of the unit identifier column. Required ONLY when + ``vcov_type="conley"`` — the panel block-decomposed Conley + sandwich (matching R ``conleyreg`` with ``lag_cutoff > 0``) + needs the unit identifier to compute the per-unit serial sum. + Mirrors :meth:`MultiPeriodDiD.fit(unit=...)` and + :meth:`TwoWayFixedEffects.fit(unit=...)`. Fit-time only — NOT + a constructor kwarg, so it is not part of ``get_params()`` / + ``set_params()`` (which return constructor-arg dicts). + Ignored when ``vcov_type`` is not ``"conley"``. Returns ------- @@ -358,22 +382,32 @@ def fit( "HC2/CR2-BM are computed on the full projection." ) - # Reject vcov_type='conley' on DifferenceInDifferences. The Phase 2 - # block-decomposed panel Conley sandwich (matching R conleyreg) - # requires the unit identifier to compute the per-unit serial sum, - # but DifferenceInDifferences.fit() has no `unit` column argument. - # Adding a `unit` arg is API churn for a path with a clean redirect: - # MultiPeriodDiD and TwoWayFixedEffects both take `unit` natively - # and support vcov_type="conley" via the same parity-tested helper. + # Validate vcov_type="conley" wire-up. DiD.fit() accepts `unit` + # as a fit-time arg (NOT on __init__) because cluster/unit + # semantics on DiD are opt-in rather than auto-derived (unlike + # MultiPeriodDiD / TwoWayFixedEffects which have a unit declaration + # at fit-time anyway). The panel block-decomposed Conley sandwich + # (matching R conleyreg with lag_cutoff > 0) needs unit/time/coords + # to assemble the within-period spatial and within-unit serial + # sums; we mirror MultiPeriodDiD's reject pattern for missing args + # and the survey/wild-bootstrap incompatibilities. if self.vcov_type == "conley": - raise NotImplementedError( - "DifferenceInDifferences(vcov_type='conley') is not supported " - "because DiD.fit() has no unit-column declaration; the panel " - "block-decomposed Conley sandwich requires per-unit serial " - "sums. Use MultiPeriodDiD(vcov_type='conley', " - "conley_lag_cutoff=...) or TwoWayFixedEffects(vcov_type='conley'," - " conley_lag_cutoff=...) — both take a `unit` argument in " - ".fit() and implement the same R conleyreg-parity sandwich." + # Shared front-door validation across DiD / MPD / TWFE entry + # points (Wave A holistic fix: replaces the inline drift that + # accumulated across CI R1/R2/R6 — same-class validation gaps + # mirrored across estimator surfaces). + from diff_diff.conley import _validate_conley_estimator_inputs + + _validate_conley_estimator_inputs( + estimator_name="DifferenceInDifferences", + data=data, + unit=unit, + conley_coords=self.conley_coords, + conley_cutoff_km=self.conley_cutoff_km, + conley_lag_cutoff=self.conley_lag_cutoff, + survey_design=survey_design, + inference=self.inference, + cluster=self.cluster, ) if absorb: @@ -466,6 +500,36 @@ def fit( # Remap implicit "classical" + cluster to CR1 for legacy-alias # backward compatibility (see `_resolve_effective_vcov_type`). _fit_vcov_type = self._resolve_effective_vcov_type(effective_cluster_ids) + + # Build Conley coord/time/unit arrays when applicable. CRITICAL: + # read from the ORIGINAL `data` frame, NOT `working_data` — `absorb` + # demeans `time` (and any column listed in `absorb`) in working_data, + # so reading `working_data[time]` would silently partition the + # within-period spatial sandwich on residualized floats instead of + # the true pre/post periods (Codex Wave A R1 P0). Coords are likewise + # read from raw `data` for symmetry with TwoWayFixedEffects + # (`twfe.py::TwoWayFixedEffects.fit`) which has the same FWL- + # composability contract: the meat is computed on demeaned scores + # but the kernel grid uses the original space (coords) and time/unit + # indexing. `_compute_conley_vcov` normalizes time labels to dense + # codes 0..T-1 internally, so non-numeric `time` labels (datetime64, + # pd.Period, strings) still work on the MultiPeriodDiD path; DiD's + # binary `time` column is integer 0/1 by convention and is unaffected + # by the normalization. + if _fit_vcov_type == "conley": + _conley_coords_arr: Optional[np.ndarray] = np.column_stack( + [ + data[self.conley_coords[0]].values.astype(np.float64), + data[self.conley_coords[1]].values.astype(np.float64), + ] + ) + _conley_time_arr: Optional[np.ndarray] = np.asarray(data[time].values) + _conley_unit_arr: Optional[np.ndarray] = data[unit].values + else: + _conley_coords_arr = None + _conley_time_arr = None + _conley_unit_arr = None + # Don't forward `robust=self.robust` when the vcov_type has been # remapped; `robust=False + vcov_type="hc1"` would otherwise trip # the conflict check inside `LinearRegression.__init__`. The @@ -479,6 +543,13 @@ def fit( weight_type=survey_weight_type, survey_design=_lr_survey, vcov_type=_fit_vcov_type, + conley_coords=_conley_coords_arr, + conley_cutoff_km=self.conley_cutoff_km, + conley_metric=self.conley_metric, + conley_kernel=self.conley_kernel, + conley_time=_conley_time_arr, + conley_unit=_conley_unit_arr, + conley_lag_cutoff=self.conley_lag_cutoff, ).fit(X, y, df_adjustment=n_absorbed_effects) coefficients = reg.coefficients_ @@ -1045,9 +1116,12 @@ class MultiPeriodDiD(DifferenceInDifferences): auto-derived from the ``time`` column at fit-time and normalized to dense panel-period codes ``0..T-1`` so ``conley_lag_cutoff`` always counts panel periods (works for int / datetime64 / - ``pd.Period`` / string encodings). Restrictions: ``cluster=``, - ``survey_design=``, and ``inference="wild_bootstrap"`` raise on - this path (Phase 5 / follow-up). + ``pd.Period`` / string encodings). Explicit ``cluster=`` + enables the combined spatial + cluster product kernel + (Wave A #119; cluster must be constant within each unit across + periods). Restrictions: ``survey_design=`` and + ``inference="wild_bootstrap"`` raise on this path + (Phase 5 / follow-up). alpha : float, default=0.05 Significance level for confidence intervals. conley_coords, conley_cutoff_km, conley_metric, conley_kernel, conley_lag_cutoff @@ -1175,9 +1249,14 @@ def fit( # type: ignore[override] If required parameters are missing or data validation fails. """ # Fall back to analytical inference if wild bootstrap requested - # (must happen before _resolve_survey_for_fit which rejects bootstrap+survey) + # (must happen before _resolve_survey_for_fit which rejects bootstrap+survey). + # SKIP the warning on the Conley path — the Conley validator below + # raises NotImplementedError for wild_bootstrap + Conley, so emitting + # the analytical-fallback warning first would produce contradictory + # guidance on the same call (warn "falling back" + raise "not + # supported"). The Conley raise takes precedence. Codex CI R11 P3. effective_inference = self.inference - if self.inference == "wild_bootstrap": + if self.inference == "wild_bootstrap" and self.vcov_type != "conley": warnings.warn( "Wild bootstrap inference is not yet supported for MultiPeriodDiD. " "Using analytical inference instead.", @@ -1384,47 +1463,24 @@ def fit( # type: ignore[override] ) # MultiPeriodDiD is intrinsically a multi-period panel estimator; - # Phase 2 panel block-decomposed Conley (matches R conleyreg): require - # `unit` and `conley_lag_cutoff` when vcov_type="conley". The actual - # array extraction (and conley_coords resolution from column names) - # happens just below at the solve_ols call. + # Phase 2 panel block-decomposed Conley (matches R conleyreg) needs + # `unit`, `conley_lag_cutoff`, and `conley_coords` at fit-time. The + # validation is shared with DiD / TWFE to avoid the validation-class + # drift that surfaced across Wave A CI R1/R2/R6. if self.vcov_type == "conley": - if self.conley_coords is None or self.conley_cutoff_km is None: - raise ValueError( - "MultiPeriodDiD(vcov_type='conley') requires " - "conley_coords=(lat_col, lon_col) and conley_cutoff_km " - "on the constructor." - ) - if unit is None: - raise ValueError( - "MultiPeriodDiD(vcov_type='conley') requires unit= at " - "fit-time (the panel block-decomposed sandwich computes " - "a per-unit serial sum, matching R conleyreg)." - ) - if self.conley_lag_cutoff is None: - raise ValueError( - "MultiPeriodDiD(vcov_type='conley') requires " - "conley_lag_cutoff (non-negative int; 0 means spatial-" - "within-period only, no serial component). See R " - "conleyreg's `lag_cutoff` argument for the convention." - ) - if survey_design is not None: - raise NotImplementedError( - "MultiPeriodDiD(vcov_type='conley', survey_design=...) " - "is not supported: Conley + survey weights / replicate " - "vcov is deferred to a follow-up PR (Bertanha-Imbens 2014 " - "territory). Use vcov_type='hc1' for survey-aware " - "cluster-robust without spatial HAC, or drop survey_design= " - "for panel Conley." - ) - if self.inference == "wild_bootstrap": - raise NotImplementedError( - "MultiPeriodDiD(vcov_type='conley', " - "inference='wild_bootstrap') is not supported: wild " - "bootstrap is a separate inference path that does not " - "consume the analytical Conley sandwich. Use " - "inference='analytical' for Conley SEs." - ) + from diff_diff.conley import _validate_conley_estimator_inputs + + _validate_conley_estimator_inputs( + estimator_name="MultiPeriodDiD", + data=data, + unit=unit, + conley_coords=self.conley_coords, + conley_cutoff_km=self.conley_cutoff_km, + conley_lag_cutoff=self.conley_lag_cutoff, + survey_design=survey_design, + inference=self.inference, + cluster=self.cluster, + ) # Pre-compute non_ref_periods (needed for absorb demeaning) non_ref_periods = [p for p in all_periods if p != reference_period] @@ -1558,21 +1614,30 @@ def fit( # type: ignore[override] _fit_vcov_type = self._resolve_effective_vcov_type(effective_cluster_ids) # Resolve Conley arrays from column names (init-time) plus the - # estimator's `time` / `unit` columns. When vcov_type != "conley", - # these are silently ignored downstream (Phase 1 / 2 convention). + # estimator's `time` / `unit` columns. CRITICAL: read from the + # ORIGINAL `data` frame, NOT `working_data` — if absorb is used + # with overlapping covariates (e.g. lat/lon or time listed in + # both `absorb` and `conley_coords`/`time`), `working_data` has + # those columns demeaned and the Conley helper would silently + # partition the spatial sandwich on residualized inputs. + # Mirrors the DiD/TWFE contract at `estimators.py::DifferenceInDifferences.fit` + # and `twfe.py::TwoWayFixedEffects.fit` (FWL composability: the meat + # is computed on demeaned scores but the kernel grid uses the raw + # coords + time/unit). When vcov_type != "conley", these are silently + # ignored downstream (Phase 1 / 2 convention). if _fit_vcov_type == "conley": _conley_coords_arr: Optional[np.ndarray] = np.column_stack( [ - working_data[self.conley_coords[0]].values.astype(np.float64), - working_data[self.conley_coords[1]].values.astype(np.float64), + data[self.conley_coords[0]].values.astype(np.float64), + data[self.conley_coords[1]].values.astype(np.float64), ] ) # Preserve the original time-label dtype (int, datetime64, pd.Period, # string). `_compute_conley_vcov` normalizes to dense 0..T-1 codes # internally; float coercion here would break datetime64 / Period / # string encodings before the normalizer runs. - _conley_time_arr: Optional[np.ndarray] = np.asarray(working_data[time].values) - _conley_unit_arr: Optional[np.ndarray] = working_data[unit].values + _conley_time_arr: Optional[np.ndarray] = np.asarray(data[time].values) + _conley_unit_arr: Optional[np.ndarray] = data[unit].values else: _conley_coords_arr = None _conley_time_arr = None diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index 6c4adbc7..27661bd1 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -1898,20 +1898,21 @@ modes: - **Cross-sectional:** Direct `compute_robust_vcov` / `LinearRegression` on a single-period design. - **Panel block-decomposed** (matches R `conleyreg` with `lag_cutoff > 0`): - `MultiPeriodDiD` / `TwoWayFixedEffects` with `vcov_type="conley"` and - `conley_lag_cutoff=`. The sandwich sums within-period spatial pairs - plus within-unit Bartlett serial pairs (excluding lag=0 to avoid - double-counting); NOT a multiplicative product kernel. + `DifferenceInDifferences` / `MultiPeriodDiD` / `TwoWayFixedEffects` with + `vcov_type="conley"` and `conley_lag_cutoff=`. The sandwich sums + within-period spatial pairs plus within-unit Bartlett serial pairs + (excluding lag=0 to avoid double-counting); NOT a multiplicative + product kernel. -**`DifferenceInDifferences(vcov_type="conley")` raises `NotImplementedError`** -with a redirect to `MultiPeriodDiD` / `TwoWayFixedEffects`: DiD.fit() has -no `unit` column declaration, but the block-decomposed sandwich requires -unit membership for the per-unit serial sum. +`DifferenceInDifferences(vcov_type="conley").fit(..., unit="")` is +supported (Wave A #118). `unit` is a fit-time kwarg (NOT on `__init__`; +unused unless Conley is set; not part of `get_params()` / `set_params()`) +mirroring `MultiPeriodDiD.fit(unit=...)` / `TwoWayFixedEffects.fit(unit=...)`. ```python import numpy as np from diff_diff.linalg import LinearRegression -from diff_diff import MultiPeriodDiD, TwoWayFixedEffects +from diff_diff import DifferenceInDifferences, MultiPeriodDiD, TwoWayFixedEffects # Cross-sectional design: 1 row per unit, n × 2 lat/lon coords. reg = LinearRegression( @@ -1949,6 +1950,30 @@ mp_res = MultiPeriodDiD( conley_lag_cutoff=2, # within-unit Bartlett up to 2 panel periods ).fit(data, outcome="y", treatment="treated", time="period", post_periods=[2, 3], unit="unit_id") + +# 2-period DiD on a panel: DiD.fit(unit="") opts into the Conley +# panel block-decomposed sandwich; on a 2-period design the ATT/SE match +# MultiPeriodDiD(...).fit(..., post_periods=[1], reference_period=0) bit-exactly. +did_res = DifferenceInDifferences( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=200.0, + conley_lag_cutoff=1, +).fit(data, outcome="y", treatment="treated", time="post", unit="unit_id") + +# Combined spatial + cluster product kernel: pass cluster= alongside +# Conley to apply K_total[i,j] = K_space(d_ij/h) · 1{c_i = c_j}. On the +# panel path the cluster must be constant within each unit across periods +# (e.g. an above-unit grouping like region); time-varying cluster raises +# ValueError. TWFE's default auto-cluster on the Conley path is silently +# dropped — users opt into the combined kernel explicitly. +combined_res = TwoWayFixedEffects( + vcov_type="conley", + cluster="region", # above-unit grouping; time-invariant within unit + conley_coords=("lat", "lon"), + conley_cutoff_km=500.0, + conley_lag_cutoff=1, +).fit(data, outcome="y", treatment="treated", time="post", unit="unit_id") ``` **Note on `conley_lag_cutoff` semantics:** the lag is counted in **panel @@ -1995,14 +2020,15 @@ recommends a sensitivity grid (e.g., 50, 100, 200, 500 km) and reporting the SE range. **Restrictions in this release:** -- `DifferenceInDifferences(vcov_type="conley")` → `NotImplementedError` (no `unit` declaration on `DiD.fit()`; use `MultiPeriodDiD` / `TwoWayFixedEffects`). -- `MultiPeriodDiD` / `TwoWayFixedEffects` `+ vcov_type="conley"` without `conley_lag_cutoff` → `ValueError` (no defensible default; explicit user choice required). -- `MultiPeriodDiD` `+ vcov_type="conley"` without `unit=` at fit-time → `ValueError`. -- `TwoWayFixedEffects(vcov_type="conley", cluster=...)` → `NotImplementedError` (combined spatial + cluster kernel deferred). Auto-cluster on the Conley path is silently dropped. -- `TwoWayFixedEffects(vcov_type="conley", inference="wild_bootstrap")` → `NotImplementedError`. +- `DifferenceInDifferences` / `MultiPeriodDiD` / `TwoWayFixedEffects` `+ vcov_type="conley"` without `conley_lag_cutoff` → `ValueError` (no defensible default; explicit user choice required). +- `DifferenceInDifferences` / `MultiPeriodDiD` `+ vcov_type="conley"` without `unit=` at fit-time → `ValueError`. +- Combining `vcov_type="conley"` with explicit `cluster=` applies the combined spatial + cluster product kernel (Wave A #119). On the panel path the cluster must be constant within each unit across periods (validator raises `ValueError` otherwise). TWFE's default auto-cluster on the Conley path is silently dropped; users opt into the combined kernel explicitly. +- `DifferenceInDifferences` / `MultiPeriodDiD` / `TwoWayFixedEffects` `(vcov_type="conley", inference="wild_bootstrap")` → `NotImplementedError` (wild bootstrap does not consume the analytical sandwich). +- `DifferenceInDifferences` / `MultiPeriodDiD` / `TwoWayFixedEffects` `(vcov_type="conley")` + `survey_design=` → `NotImplementedError` (Bertanha-Imbens 2014 follow-up). - `SyntheticDiD(vcov_type="conley")` → `TypeError` (uses bootstrap, not analytical sandwich). -- Any-mode `vcov_type="conley"` `+ cluster_ids=` / `weights=` / `survey_design=` on `LinearRegression` / `compute_robust_vcov` → `NotImplementedError` (combined kernel + Bertanha-Imbens 2014 weighted-Conley deferred to follow-up PRs). -- `n > 20_000` emits a `UserWarning` about O(n²) distance-matrix memory; sparse k-d-tree fast path is a follow-up. +- Any-mode `vcov_type="conley"` `+ weights=` / `survey_design=` on `LinearRegression` / `compute_robust_vcov` → `NotImplementedError` (Bertanha-Imbens 2014 weighted-Conley deferred to follow-up PR). +- A sparse k-d-tree fast path auto-activates for `n > 5_000` with `conley_kernel="bartlett"` AND `conley_metric` in `{"haversine", "euclidean"}` (Wave A #120); callable metrics and uniform kernel fall back to the dense path. `n > 20_000` with the dense fallback still emits a memory-OOM `UserWarning`. +- Callable `conley_metric` is validated at the boundary (shape `(n, n)`, finite, non-negative, symmetric to `atol=1e-10`, AND zero on the diagonal `|d(i, i)| ≤ 1e-10` so the kernel reduces to `K(0) = 1` on the HC0 contribution); each failure raises a targeted `ValueError` (Wave A #123). **Parity:** matches R `conleyreg` (Düsterhöft 2021, CRAN v0.1.9) to ≤ 1e-6 on six benchmark fixtures in diff --git a/diff_diff/guides/llms.txt b/diff_diff/guides/llms.txt index 40a7bc51..83e791d0 100644 --- a/diff_diff/guides/llms.txt +++ b/diff_diff/guides/llms.txt @@ -76,7 +76,7 @@ Full practitioner guide: call `diff_diff.get_llm_guide("practitioner")` - [Honest DiD](https://diff-diff.readthedocs.io/en/stable/api/honest_did.html): Rambachan & Roth (2023) sensitivity analysis — robust CI under parallel trends violations, breakdown values - [Pre-Trends Power Analysis](https://diff-diff.readthedocs.io/en/stable/api/pretrends.html): Roth (2022) minimum detectable violation and pre-trends test power curves - [Power Analysis](https://diff-diff.readthedocs.io/en/stable/api/power.html): Analytical and simulation-based power analysis — MDE, sample size, power curves for study design -- Conley spatial HAC SE (`vcov_type="conley"`) on cross-sectional `LinearRegression` / `compute_robust_vcov` PLUS panel `MultiPeriodDiD` / `TwoWayFixedEffects` (with `conley_lag_cutoff=` for within-unit Bartlett temporal HAC) — Conley (1999) spatial-correlation-aware SEs with haversine/euclidean/callable distance metric and Bartlett/uniform spatial kernel; panel path uses the R `conleyreg`-form block-decomposed sandwich (within-period spatial + within-unit Bartlett serial, same-time excluded); parity vs R `conleyreg` (Düsterhöft 2021) on cross-sectional AND panel `lag_cutoff > 0` fixtures. `DifferenceInDifferences(vcov_type="conley")` raises with a redirect to `MultiPeriodDiD` / `TwoWayFixedEffects` because DiD.fit() has no `unit` column declaration +- Conley spatial HAC SE (`vcov_type="conley"`) on cross-sectional `LinearRegression` / `compute_robust_vcov` PLUS panel `DifferenceInDifferences` / `MultiPeriodDiD` / `TwoWayFixedEffects` (with `conley_lag_cutoff=` for within-unit Bartlett temporal HAC) — Conley (1999) spatial-correlation-aware SEs with haversine/euclidean/callable distance metric and Bartlett/uniform spatial kernel; panel path uses the R `conleyreg`-form block-decomposed sandwich (within-period spatial + within-unit Bartlett serial, same-time excluded); parity vs R `conleyreg` (Düsterhöft 2021) on cross-sectional AND panel `lag_cutoff > 0` fixtures. Combining with explicit `cluster=` applies the combined spatial + cluster product kernel `K_total[i,j] = K_space · 1{c_i = c_j}` (cluster must be constant within each unit across periods on the panel path; validator-enforced). DiD takes `unit=` as a fit-time kwarg when `vcov_type="conley"` (not on `__init__`). Sparse k-d-tree fast path auto-activates for `n > 5_000` with bartlett kernel + haversine/euclidean metric ## Tutorials diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index 8bca8a83..793e1987 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -549,9 +549,11 @@ def solve_ols( design or pooled cross-section) and panel block-decomposed (matches R ``conleyreg`` with ``lag_cutoff > 0``); switch by passing the three co-required kwargs ``conley_time`` / ``conley_unit`` / - ``conley_lag_cutoff``. Combining with ``cluster_ids`` or ``weights`` - raises ``NotImplementedError`` (combined spatial-cluster kernel + - Bertanha-Imbens 2014 weighted-Conley deferred to follow-up PRs). + ``conley_lag_cutoff``. Combining with ``cluster_ids`` applies the + combined spatial + cluster product kernel (a diff-diff convention; + see :func:`compute_robust_vcov` for details). Combining with + ``weights`` raises ``NotImplementedError`` (Bertanha-Imbens 2014 + weighted-Conley deferred to a follow-up PR). conley_coords : ndarray of shape (n, 2), optional Required when ``vcov_type="conley"``. Two-column array of ``[lat, lon]`` (degrees, for ``conley_metric="haversine"``) or @@ -1170,17 +1172,12 @@ def _validate_vcov_args( "Bell-McCaffrey. Tracked in TODO.md." ) if vcov_type == "conley": - # Conley + cluster_ids (combined spatial-cluster kernel) and Conley + - # weights (Bertanha-Imbens 2014) are deferred to follow-up PRs. - # TwoWayFixedEffects has its own earlier raise at twfe.py with a - # TWFE-specific message; this is the fallback path for - # MultiPeriodDiD callers and direct compute_robust_vcov use. - if cluster_ids is not None: - raise NotImplementedError( - "vcov_type='conley' with cluster_ids is a follow-up " - "(combined spatial-cluster kernel). Use vcov_type='hc1' for " - "cluster-robust without spatial HAC, or drop cluster= for Conley." - ) + # Conley + cluster_ids is now supported (combined spatial + cluster + # product kernel; see ``docs/methodology/REGISTRY.md`` § ConleySpatialHAC + # → "Combined spatial + cluster product kernel"). Conley + weights + # (Bertanha-Imbens 2014) is still deferred to a follow-up PR — the + # design-based weighted-Conley methodology requires its own + # treatment. if weights is not None: raise NotImplementedError( "vcov_type='conley' with weights is a Phase 2+ follow-up " @@ -1293,9 +1290,13 @@ def compute_robust_vcov( bandwidth). Two operating modes: cross-sectional (default) and panel block-decomposed (pass the three co-required kwargs ``conley_time`` / ``conley_unit`` / ``conley_lag_cutoff``, matching R ``conleyreg`` with - ``lag_cutoff > 0``). Combining with ``cluster_ids`` or ``weights`` - raises ``NotImplementedError`` (combined spatial-cluster kernel and - Bertanha-Imbens weighted-Conley deferred to follow-up PRs). + ``lag_cutoff > 0``). Combining with ``cluster_ids`` applies the + combined spatial + cluster product kernel + ``K_total[i, j] = K_space(d_ij/h) · 1{c_i = c_j}`` (Wave A #119; on + the panel path the cluster must be constant within each unit across + periods). Combining with ``weights`` still raises + ``NotImplementedError`` (Bertanha-Imbens 2014 weighted-Conley + deferred to a follow-up PR). Parameters ---------- @@ -1763,7 +1764,9 @@ def _compute_robust_vcov_numpy( # is supplied; panel block-decomposed form (matches R conleyreg with # lag_cutoff > 0) when all three of conley_time / conley_unit / # conley_lag_cutoff are supplied. The validator above already raised - # on conley + cluster_ids and conley + weights. + # on conley + weights. ``cluster_ids`` is threaded through when set, + # combining the spatial kernel with the cluster indicator (combined + # spatial + cluster product kernel; see REGISTRY.md § ConleySpatialHAC). # ------------------------------------------------------------------ if vcov_type == "conley": _validate_conley_kwargs( @@ -1775,6 +1778,7 @@ def _compute_robust_vcov_numpy( time=conley_time, unit=conley_unit, lag_cutoff=conley_lag_cutoff, + cluster_ids=cluster_ids, ) vcov = _compute_conley_vcov( X, @@ -1787,6 +1791,7 @@ def _compute_robust_vcov_numpy( time=conley_time, unit=conley_unit, lag_cutoff=conley_lag_cutoff, + cluster_ids=cluster_ids, ) if return_dof: return vcov, np.full(k, n_eff - k, dtype=np.float64) @@ -2484,13 +2489,15 @@ class LinearRegression: pass the three co-required kwargs ``conley_time`` / ``conley_unit`` / ``conley_lag_cutoff``). Requires ``conley_coords`` (n × 2 array) and a positive ``conley_cutoff_km``. Combining ``vcov_type="conley"`` - with ``cluster_ids``, ``weights``, or ``survey_design`` raises - ``NotImplementedError`` (combined spatial-cluster kernel + - Bertanha-Imbens weighted-Conley deferred to follow-up PRs). The - panel DiD estimator rejects at fit-time (DiD.fit() has no unit - column declaration; use MultiPeriodDiD or TwoWayFixedEffects, which - thread their `unit`/`time` columns into the block-decomposed - sandwich via ``conley_lag_cutoff`` on the constructor). + with ``cluster_ids`` applies the combined spatial + cluster product + kernel (Wave A #119; cluster must be constant within each unit on + the panel path). Combining with ``weights`` / ``survey_design`` + still raises ``NotImplementedError`` (Bertanha-Imbens 2014 + weighted-Conley deferred to a follow-up PR). The DiD / MPD / TWFE + estimators all support panel Conley by passing ``unit`` at fit-time + (DiD as a fit-time kwarg; MPD/TWFE via the existing ``unit`` + argument), threading ``conley_time`` / ``conley_unit`` into the + block-decomposed sandwich. conley_coords : ndarray of shape (n, 2), optional Required when ``vcov_type="conley"``. Two-column array of ``[lat, lon]`` (degrees, for ``conley_metric="haversine"``) or diff --git a/diff_diff/results.py b/diff_diff/results.py index 8e54c666..e2270184 100644 --- a/diff_diff/results.py +++ b/diff_diff/results.py @@ -77,10 +77,15 @@ def _format_vcov_label( if vcov_type == "conley": # Cross-sectional Conley on direct LinearRegression / compute_robust_vcov, # or panel block-decomposed Conley (within-period spatial + within-unit - # Bartlett serial) on MultiPeriodDiD / TwoWayFixedEffects. - if conley_lag_cutoff is not None: - return f"Conley spatial HAC (1999), lag_cutoff={conley_lag_cutoff}" - return "Conley spatial HAC (1999)" + # Bartlett serial) on DifferenceInDifferences / MultiPeriodDiD / + # TwoWayFixedEffects. With an explicit cluster_name, the combined + # spatial + cluster product kernel applies (Wave A #119). + lag_suffix = f", lag_cutoff={conley_lag_cutoff}" if conley_lag_cutoff is not None else "" + if cluster_name: + return ( + f"Conley spatial HAC (1999) + cluster product kernel at {cluster_name}{lag_suffix}" + ) + return f"Conley spatial HAC (1999){lag_suffix}" return None @@ -135,12 +140,11 @@ class DiDResults: survey_metadata: Optional[Any] = field(default=None) # Variance-covariance family: "classical" | "hc1" | "hc2" | "hc2_bm" | # "conley". Plus cluster_name when cluster-robust. Used by summary() to - # label the SE family in the output. For vcov_type='conley' on - # MultiPeriodDiD / TwoWayFixedEffects, `conley_lag_cutoff` carries the - # within-unit Bartlett max lag (matches the constructor arg). On - # DifferenceInDifferences, vcov_type='conley' is rejected at fit-time - # (DiD.fit() has no unit column declaration); see MultiPeriodDiD / - # TwoWayFixedEffects for the panel block-decomposed path. + # label the SE family in the output. For vcov_type='conley' on the panel + # estimators (DifferenceInDifferences / MultiPeriodDiD / TwoWayFixedEffects), + # `conley_lag_cutoff` carries the within-unit Bartlett max lag (matches + # the constructor arg). DiD requires `unit=` as a fit-time kwarg + # when vcov_type='conley' (Wave A #118). vcov_type: Optional[str] = field(default=None) cluster_name: Optional[str] = field(default=None) conley_lag_cutoff: Optional[int] = field(default=None) diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index 121133d7..fefe8306 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -79,12 +79,16 @@ class TwoWayFixedEffects(DifferenceInDifferences): within-transformed scores ``S = X_demeaned * residuals_demeaned`` form the same meat as the full-dummy-expansion design. The temporal kernel is hardcoded Bartlett regardless of ``conley_kernel`` (matches - ``conleyreg::time_dist``). Restrictions: + ``conleyreg::time_dist``). Explicit ``cluster=`` + Conley + enables the combined spatial + cluster product kernel + ``K_total[i, j] = K_space(d_ij/h) · 1{cluster_i = cluster_j}``; + cluster membership must be constant within each unit across periods + (validator-enforced on the panel block-decomposed path). When + ``cluster=`` is unset, TWFE's default auto-cluster at the unit level + is silently dropped on the Conley path — Conley spatial HAC alone is + applied, not the combined kernel. Restrictions: ``inference="wild_bootstrap"`` + Conley raises (incompatible inference - modes); explicit ``cluster=...`` + Conley raises (combined spatial- - cluster kernel deferred to a follow-up PR; auto-cluster on the Conley - path is silently dropped); ``survey_design=`` + Conley raises (Phase 5 - follow-up). + modes); ``survey_design=`` + Conley raises (Phase 5 follow-up). Warning: TWFE can be biased with staggered treatment timing and heterogeneous treatment effects. Consider using @@ -169,36 +173,19 @@ def fit( # type: ignore[override] # with the original (un-demeaned) time / unit vectors and coords # yields the correct block-decomposed sandwich. if self.vcov_type == "conley": - if self.conley_lag_cutoff is None: - raise ValueError( - "TwoWayFixedEffects(vcov_type='conley') requires " - "conley_lag_cutoff (non-negative int; 0 means spatial-" - "within-period only, no serial component). See R " - "conleyreg's `lag_cutoff` argument for the convention." - ) - if self.conley_coords is None or self.conley_cutoff_km is None: - raise ValueError( - "TwoWayFixedEffects(vcov_type='conley') requires " - "conley_coords=(lat_col, lon_col) and " - "conley_cutoff_km on the constructor." - ) - if self.cluster is not None: - raise NotImplementedError( - "TwoWayFixedEffects(vcov_type='conley', cluster=...) " - "is not supported: the combined spatial-kernel + cluster " - "indicator product kernel is deferred to a follow-up PR " - "(see TODO.md). Use vcov_type='hc1' with cluster= for " - "cluster-robust without spatial HAC, or drop cluster= " - "for panel-Conley alone." - ) - if self.inference == "wild_bootstrap": - raise NotImplementedError( - "TwoWayFixedEffects(vcov_type='conley', " - "inference='wild_bootstrap') is not supported: the " - "wild bootstrap is a separate inference path that does " - "not consume the analytical Conley sandwich. Use " - "inference='analytical' for Conley SEs." - ) + from diff_diff.conley import _validate_conley_estimator_inputs + + _validate_conley_estimator_inputs( + estimator_name="TwoWayFixedEffects", + data=data, + unit=unit, + conley_coords=self.conley_coords, + conley_cutoff_km=self.conley_cutoff_km, + conley_lag_cutoff=self.conley_lag_cutoff, + survey_design=survey_design, + inference=self.inference, + cluster=self.cluster, + ) # Check for staggered treatment timing and warn if detected self._check_staggered_treatment(data, treatment, time, unit) @@ -374,11 +361,16 @@ def fit( # type: ignore[override] # string encodings before the normalizer runs. _conley_time_arr: Optional[np.ndarray] = np.asarray(data[time].values) _conley_unit_arr: Optional[np.ndarray] = data[unit].values - # vcov_type="conley" + cluster_ids raises at the linalg validator - # (combined kernel deferred). TWFE's auto-cluster would force that - # path; drop the cluster on the Conley path. The user's explicit - # cluster also conflicts and is rejected upstream. - _conley_cluster_override = None + # Combined spatial + cluster product kernel: thread the user's + # EXPLICIT cluster= through when set. TWFE's default auto- + # cluster at the unit level is silently dropped on the Conley + # path — combining Conley with the unit-cluster mask zeros out + # all between-unit pairs (defeating Conley's spatial pooling). + # Users who want the combined kernel must pass an above-unit + # cluster (e.g. region) explicitly. + _conley_cluster_override = ( + data[self.cluster].values if self.cluster is not None else None + ) else: _conley_coords_arr = None _conley_time_arr = None @@ -561,14 +553,16 @@ def _refit_twfe(w_r): # `self.cluster is None` AND the vcov family is cluster-compatible. # One-way families (`classical`, `hc2`) disable the auto-cluster (see # the `cluster_var` block above); report None so summary() labels the - # one-way family instead of "CR1 cluster-robust at unit". The Conley - # path drops the auto-cluster too (`_conley_cluster_override = None` - # above); mirror that here so the variance-provenance metadata - # (`res.cluster_name`, serialized `to_dict()["cluster_name"]`) - # accurately reflects the cluster IDs actually passed to - # `LinearRegression` — i.e. None on the Conley path. + # one-way family instead of "CR1 cluster-robust at unit". On the + # Conley path, the auto-cluster is silently dropped (combining with + # unit-level clusters would zero out all between-unit pairs and + # defeat the spatial pooling); when the user passes an explicit + # `cluster=` ALONGSIDE Conley, the combined spatial + cluster + # product kernel applies and the label tracks the user's column. + # Either way, the cluster_name reflects the cluster IDs actually + # passed to LinearRegression. if _fit_vcov_type == "conley": - _twfe_cluster_label: Optional[str] = None + _twfe_cluster_label: Optional[str] = self.cluster if self.cluster is not None else None elif self.cluster is not None: _twfe_cluster_label = self.cluster elif cluster_var is None: diff --git a/docs/api/_autosummary/diff_diff.DiDResults.rst b/docs/api/_autosummary/diff_diff.DiDResults.rst index 5178e915..0d95b394 100644 --- a/docs/api/_autosummary/diff_diff.DiDResults.rst +++ b/docs/api/_autosummary/diff_diff.DiDResults.rst @@ -29,6 +29,7 @@ ~DiDResults.cluster_name ~DiDResults.coef_var ~DiDResults.coefficients + ~DiDResults.conley_lag_cutoff ~DiDResults.fitted_values ~DiDResults.inference_method ~DiDResults.is_significant diff --git a/docs/api/_autosummary/diff_diff.MultiPeriodDiDResults.rst b/docs/api/_autosummary/diff_diff.MultiPeriodDiDResults.rst index b364264f..7f7dd215 100644 --- a/docs/api/_autosummary/diff_diff.MultiPeriodDiDResults.rst +++ b/docs/api/_autosummary/diff_diff.MultiPeriodDiDResults.rst @@ -31,6 +31,7 @@ ~MultiPeriodDiDResults.coef_var ~MultiPeriodDiDResults.coefficients ~MultiPeriodDiDResults.conf_int + ~MultiPeriodDiDResults.conley_lag_cutoff ~MultiPeriodDiDResults.fitted_values ~MultiPeriodDiDResults.inference_method ~MultiPeriodDiDResults.interaction_indices diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index af10e544..1df004b1 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -2962,19 +2962,33 @@ Newey-West-style Bartlett temporal HAC on panel data. estimator's `time` / `unit` column-name arguments; only `conley_lag_cutoff` is set on the constructor. -**Panel API restrictions (Phase 2):** -- `DifferenceInDifferences(vcov_type="conley")` continues to raise - `NotImplementedError`. DiD.fit() has no `unit` column declaration, and the - block-decomposed sandwich requires unit membership for the per-unit serial - sum. Use `MultiPeriodDiD` or `TwoWayFixedEffects` instead. +**Panel API restrictions (Phase 2 + Wave A):** +- `DifferenceInDifferences(vcov_type="conley")` is supported (Wave A #118): + pass `unit=` as a fit-time argument to `fit(...)` (NOT on `__init__`; + unused unless Conley is set). DiD inherits the same panel block-decomposed + sandwich as MPD/TWFE on the two-period design. - `SyntheticDiD(vcov_type="conley")` raises `TypeError`. SyntheticDiD uses bootstrap/jackknife/placebo variance, not the analytical sandwich. - TWFE's default auto-cluster on the Conley path is silently dropped (no - combined kernel). Explicit `cluster=...` with Conley raises (deferred - follow-up PR). + combined kernel from auto-cluster). Explicit `cluster=` + Conley + enables the **combined spatial + cluster product kernel** (Wave A #119; + see "Combined spatial + cluster product kernel" subsection below). On + the panel path the validator enforces that cluster membership is + constant within each unit across periods. - `inference="wild_bootstrap"` + Conley raises (wild bootstrap is a separate inference path that does not consume the analytical sandwich). +**Note (DiD vs TWFE cluster asymmetry on the Conley path):** TWFE auto-clusters +at the unit level by default, so combining with Conley silently drops the +auto-cluster (otherwise every between-unit pair would be zeroed out, defeating +the spatial pooling). To opt into the combined kernel, the user must pass an +explicit `cluster=` that is constant within each unit (typically an +above-unit grouping like region). DiD has no auto-cluster — combining with +Conley is fully opt-in: absent `cluster=`, pure Conley spatial HAC applies; +with `cluster=`, the combined kernel applies. This asymmetry preserves the +existing TWFE auto-cluster contract while making the cluster intent explicit +on the Conley path. + **Variance estimator — cross-sectional (Phase 1, Conley 1999 Eq 4.2 in pairwise-distance form, OLS specialization):** Var̂(β) = (X'X)^{-1} · ( Σ_{i,j} K(d_ij / h) · X_i ε_i ε_j X_j' ) · (X'X)^{-1} @@ -3054,19 +3068,133 @@ cross-sectional (Phase 1) plus three panel fixtures with `lag_cutoff > 0` `conleyreg::haversine_dist`. Regeneration: `cd benchmarks/R && Rscript generate_conley_golden.R`. +### Combined spatial + cluster product kernel (Wave A #119) + +When `cluster_ids` is supplied alongside `vcov_type="conley"`, the meat +applies the combined product kernel: + + K_total[i, j] = K_space(d_ij/h) · 1{cluster_i = cluster_j} + +On the panel block-decomposed path the cluster indicator multiplies BOTH +the within-period spatial sandwich AND the within-unit serial sandwich: + + XeeX_spatial = Σ_t Σ_{i,j∈units} K_space(d_ij/h) · 1{c_{i,t}=c_{j,t}} · X_{i,t} ε_{i,t} ε_{j,t} X_{j,t}' + XeeX_serial = Σ_u Σ_{|t-s|≤L, t≠s} (1 - |t-s|/(L+1)) · 1{c_{u,t}=c_{u,s}} · X_{u,t} ε_{u,t} ε_{u,s} X_{u,s}' + +**Cluster-time-invariance contract:** on the panel block-decomposed path +the validator REQUIRES that `cluster_ids` be constant within each unit +across periods. The within-unit serial sandwich's cluster mask is then +trivially all-ones, and the math simplifies to the bare serial Bartlett +HAC weighted by the spatial mask only. If a unit's cluster changes +across periods (e.g. a unit migrating between regions), the within-unit +mask would zero out adjacent-time pairs that should contribute, +producing a methodologically-muddled meat — the validator raises +`ValueError` naming the violating unit(s). The cross-sectional path +has no time dimension, so no invariance constraint applies. + +**Note:** R `conleyreg` does not support a combined spatial + cluster +product kernel; this is a diff-diff convention validated by two limit +fixtures rather than R parity: +1. **All-unique-clusters reduction:** when every observation is in its + own cluster, the cluster mask is the identity, and the meat reduces + to the diagonal HC0 contribution `Σ_i X_i ε_i² X_i'`. +2. **Huge-cutoff reduction:** when `conley_cutoff_km` is large enough + that `K_space = 1` on every pair, the meat reduces to the pure + within-cluster sum `Σ_g X_g' ε_g ε_g' X_g` (CR1 without the + Liang-Zeger small-sample correction). This exact reduction holds + only for `conley_kernel="uniform"` (`K_uniform(u) = 1` for `|u| ≤ 1`). + The Bartlett kernel gives `K_bartlett(u) = 1 - |u|`, which is + strictly less than 1 for `0 < |u| ≤ 1`, so the huge-cutoff limit + under Bartlett is asymptotic (`K → 1` as `cutoff → ∞` only at finite + off-diagonal distances), not exact at any finite cutoff. The fixture + anchor uses uniform for an exact identity check. + +The combined-kernel meat is well-defined either way; the two fixture +limits anchor the math, and the panel time-invariance contract +guarantees the serial component is unaffected by the cluster choice. + +### Performance / scale (Wave A #120) + +A sparse k-d-tree fast path auto-activates for the spatial Bartlett meat +when `n > _CONLEY_SPARSE_N_THRESHOLD` (default 5,000) AND `conley_metric` +is `"haversine"` or `"euclidean"` (NOT a callable) AND `conley_kernel` +is `"bartlett"`. The dense O(n²) distance matrix is replaced with a CSR +sparse kernel matrix built from `scipy.spatial.cKDTree.query_ball_tree` +neighbor queries. + +**Why bartlett-only:** Bartlett at `u = 1.0` returns exactly `0.0`, so +pairs at exactly the cutoff distance contribute zero to the meat — the +sparse path can safely drop them. Uniform at `u = 1.0` returns `1.0`, +which would require a closed-interval query semantic that the haversine +chord-projection roundoff cannot reliably preserve. The auto-toggle +falls back to the dense path for uniform regardless of n. Callable +metrics also fall back (kd-tree needs a vectorizable Minkowski distance). + +For haversine, the kd-tree operates on a 3-D unit-sphere projection +(`x = cos(lat)cos(lon), y = cos(lat)sin(lon), z = sin(lat)`) with the +chord radius matching the arc-length cutoff; the exact great-circle +distance is recomputed only for in-range neighbors before the kernel +evaluation. The numerical tolerance vs the dense path is typically +~1e-12 in absolute terms; R parity at `atol=1e-6` is preserved on the +existing fixtures under the auto-toggle. + +The `_CONLEY_DENSE_OOM_WARN_N = 20_000` constant remains as a separate +warning threshold for the dense fallback (callable metrics, uniform +kernel) where O(n²) memory is at material risk. The two thresholds are +independent — sparse auto-toggle at 5,000 is a compute optimization; +dense OOM warning at 20,000 is a memory caution. + +**Density gate (`_CONLEY_SPARSE_DENSITY_THRESHOLD = 0.3`):** the sparse +path's CSR storage carries ~12 bytes per non-zero (data + indices + +indptr) vs 8 bytes per cell for dense float64. The memory crossover +is at ~67% density, but at high density the CSR overhead loses its +advantage well before that. The sparse helper measures actual neighbor +density via `cKDTree.count_neighbors` (shares tree traversal with +`query_ball_tree`, no extra allocation) and falls back to the dense +path with a `UserWarning` when neighbor density exceeds 30%. This +prevents the "sparse" path from silently using MORE memory than dense +when cutoffs are large relative to the data span (e.g. cutoffs above +half-Earth circumference on a global panel, or unit-scale cutoffs on +a clustered dataset). Users see one line explaining the fallback so +they can either reduce `conley_cutoff_km` or accept the dense path. + +### Callable conley_metric validation (Wave A #123) + +When `conley_metric` is a user-supplied callable, the result is +validated at the boundary via `_validate_callable_metric_result`: + +1. Result casts to a float64 array (raises `ValueError` if not). +2. Shape is exactly `(n, n)` (raises if mismatched). +3. All entries are finite (NaN/inf raises). +4. All entries are non-negative (negative distances raise). +5. Symmetric to within `atol=1e-10` (asymmetric matrix raises). +6. Zero diagonal: `|d(i, i)| ≤ 1e-10` for all `i` (nonzero diagonal raises). + +Each failure produces a `ValueError` naming the violated invariant. +Sub-tolerance asymmetry (eps-level roundoff) is accepted. The zero- +diagonal invariant is load-bearing for the Conley sandwich: the +`i = j` term contributes `K(d_ii / h) · X_i ε_i² X_i'`, which must +reduce to the HC0 diagonal `X_i ε_i² X_i'` (i.e., `K(0) = 1`). A +callable with positive self-distance would attenuate the HC0 term +by `K(d_ii / h) < 1` and silently misstate Conley SEs. Built-in +metrics (`"haversine"`, `"euclidean"`) satisfy this by construction. + **Edge cases / restrictions:** -- `DifferenceInDifferences(vcov_type="conley")` → `NotImplementedError` at fit-time. DiD.fit() has no `unit` column declaration; redirect users to `MultiPeriodDiD` / `TwoWayFixedEffects` which take `unit` natively. -- `MultiPeriodDiD` / `TwoWayFixedEffects` `+ vcov_type="conley"` without `conley_lag_cutoff` → `ValueError` (no defensible default; explicit user choice required per Conley 1999 Section 5 sensitivity-grid recommendation). -- `MultiPeriodDiD(vcov_type="conley")` without `unit=` at fit-time → `ValueError`. -- `TwoWayFixedEffects(vcov_type="conley", cluster=...)` → `NotImplementedError` (combined spatial + cluster product kernel deferred to follow-up PR). TWFE's auto-cluster on the Conley path is silently dropped. -- `TwoWayFixedEffects(vcov_type="conley", inference="wild_bootstrap")` → `NotImplementedError` (wild bootstrap does not consume the analytical sandwich). -- `SyntheticDiD(vcov_type="conley")` → `TypeError` (SyntheticDiD uses bootstrap/jackknife/placebo variance, not the analytical sandwich; tracked in TODO.md) -- Cross-sectional `LinearRegression` / `compute_robust_vcov` `+ vcov_type="conley"` `+ cluster_ids=` → `NotImplementedError` (combined product kernel deferred to follow-up PR) -- Any-mode `vcov_type="conley"` `+ weights=` / `survey_design=` → `NotImplementedError` (Bertanha-Imbens 2014 territory; Phase 5 follow-up) -- Panel path: partial `conley_time` / `conley_unit` / `conley_lag_cutoff` (not all three set) → `ValueError` at validator -- `n > 20_000`: emits `UserWarning` about O(n²) distance-matrix memory (sparse k-d-tree fast path queued as follow-up) -- `conley_cutoff_km ≤ 0`, `nan`, or `inf`: rejected with `ValueError`. The HC0 reduction at h→0 is documented but not the sanctioned path; users should pass `vcov_type="hc1"` -- Identical coordinates (`d_ij = 0` for `i ≠ j`): `K(0) = 1`, contributing the full HC0 weight per Conley 1999 page 19. Documented behavior; no warning +- `DifferenceInDifferences(vcov_type="conley")` is supported (Wave A #118): pass `unit=` to `fit(...)` (NOT on `__init__`; unused unless Conley is set; not part of `get_params()` / `set_params()`). +- `MultiPeriodDiD` / `TwoWayFixedEffects` / `DifferenceInDifferences` `+ vcov_type="conley"` without `conley_lag_cutoff` → `ValueError`. +- `MultiPeriodDiD` / `DifferenceInDifferences` `(vcov_type="conley")` without `unit=` at fit-time → `ValueError`. +- `TwoWayFixedEffects(vcov_type="conley", cluster=)` is supported (Wave A #119): combined spatial + cluster product kernel applies. The cluster must be time-invariant within each unit on the panel path (validator-enforced). TWFE's default auto-cluster is silently dropped on the Conley path; explicit cluster is required to opt in. +- `DifferenceInDifferences(vcov_type="conley", cluster=)`: combined kernel applies; same time-invariance contract on the panel path. DiD has no auto-cluster, so the cluster choice is fully explicit. +- `DifferenceInDifferences` / `MultiPeriodDiD` / `TwoWayFixedEffects` `(vcov_type="conley", inference="wild_bootstrap")` → `NotImplementedError`. (MPD's pre-Conley analytical-fallback `UserWarning` is suppressed when `vcov_type="conley"` so the user gets one consistent error message.) +- `DifferenceInDifferences` / `MultiPeriodDiD` / `TwoWayFixedEffects` `(vcov_type="conley")` + `survey_design=` → `NotImplementedError` at the estimator level (deferred to Bertanha-Imbens 2014 weighted-Conley follow-up). +- `SyntheticDiD(vcov_type="conley")` → `TypeError` (SyntheticDiD uses bootstrap/jackknife/placebo variance, not the analytical sandwich; tracked in TODO.md). +- Any-mode `vcov_type="conley"` `+ weights=` / `survey_design=` → `NotImplementedError` (Bertanha-Imbens 2014 territory; Phase 5 follow-up). +- Panel path: partial `conley_time` / `conley_unit` / `conley_lag_cutoff` (not all three set) → `ValueError` at validator. +- Panel path with `cluster_ids` that vary across periods within a unit → `ValueError` (time-invariance contract). +- `n > 20_000` with `conley_kernel='uniform'` or callable metric: emits `UserWarning` about O(n²) distance-matrix memory (sparse fast path doesn't apply; consider switching to bartlett or projecting to euclidean for performance). +- `conley_cutoff_km ≤ 0`, `nan`, or `inf`: rejected with `ValueError`. The HC0 reduction at h→0 is documented but not the sanctioned path; users should pass `vcov_type="hc1"`. +- Identical coordinates (`d_ij = 0` for `i ≠ j`): `K(0) = 1`, contributing the full HC0 weight per Conley 1999 page 19. Documented behavior; no warning. +- Callable `conley_metric` returning a non-(n,n)/NaN/inf/negative/asymmetric/non-zero-diagonal matrix raises `ValueError` naming the violated invariant. The zero-diagonal contract (`|d(i, i)| ≤ 1e-10`) is load-bearing for the Conley sandwich's HC0 reduction `K(0) = 1`; see "Callable conley_metric validation" subsection above. **Reference implementations:** - R: `conleyreg::conleyreg(...)` (Düsterhöft 2021, CRAN v0.1.9) — **parity benchmark for diff-diff** diff --git a/tests/test_conley_vcov.py b/tests/test_conley_vcov.py index ea783829..cdd5d71f 100644 --- a/tests/test_conley_vcov.py +++ b/tests/test_conley_vcov.py @@ -13,8 +13,10 @@ from diff_diff.conley import ( _CONLEY_EARTH_RADIUS_KM, + _CONLEY_SPARSE_N_THRESHOLD, _bartlett_kernel, _compute_conley_vcov, + _compute_spatial_bartlett_meat_sparse, _haversine_km, _pairwise_distance_matrix, _uniform_kernel, @@ -162,22 +164,153 @@ def test_pairwise_distance_euclidean_matches_pdist(self): np.testing.assert_allclose(D, D_scipy, atol=1e-12) def test_pairwise_distance_callable(self): - """A user-supplied callable is dispatched and its output preserved.""" + """A user-supplied callable is dispatched and its output preserved. + Output must satisfy the validator's invariants (zero diagonal, finite, + non-negative, symmetric).""" coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) - def constant_metric(c1, c2): + def constant_offdiag_metric(c1, c2): n1 = len(c1) n2 = len(c2) - return np.full((n1, n2), 5.0) + out = np.full((n1, n2), 5.0) + np.fill_diagonal(out, 0.0) + return out - D = _pairwise_distance_matrix(coords, constant_metric) - np.testing.assert_allclose(D, np.full((3, 3), 5.0)) + D = _pairwise_distance_matrix(coords, constant_offdiag_metric) + expected = np.full((3, 3), 5.0) + np.fill_diagonal(expected, 0.0) + np.testing.assert_allclose(D, expected) def test_pairwise_distance_unknown_metric_raises(self): """Unknown metric strings raise ValueError from the dispatcher.""" with pytest.raises(ValueError, match="conley_metric"): _pairwise_distance_matrix(np.zeros((3, 2)), "manhattan") + def test_callable_metric_wrong_shape_raises(self): + """Callable returning a non-(n, n) matrix raises a targeted ValueError.""" + coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) + + def wrong_shape_metric(c1, c2): + return np.zeros((2, 5)) + + with pytest.raises(ValueError, match=r"\(n, n\) distance matrix"): + _pairwise_distance_matrix(coords, wrong_shape_metric) + + def test_callable_metric_returns_nan_raises(self): + """Callable returning a matrix with NaN raises a targeted ValueError.""" + coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) + + def nan_metric(c1, c2): + out = np.zeros((3, 3)) + out[0, 1] = np.nan + out[1, 0] = np.nan + return out + + with pytest.raises(ValueError, match="non-finite"): + _pairwise_distance_matrix(coords, nan_metric) + + def test_callable_metric_returns_inf_raises(self): + """Callable returning a matrix with inf raises (same branch as NaN).""" + coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) + + def inf_metric(c1, c2): + out = np.zeros((3, 3)) + out[0, 1] = np.inf + out[1, 0] = np.inf + return out + + with pytest.raises(ValueError, match="non-finite"): + _pairwise_distance_matrix(coords, inf_metric) + + def test_callable_metric_negative_entries_raise(self): + """Callable returning a negative distance raises (distances must be + non-negative).""" + coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) + + def negative_metric(c1, c2): + out = np.full((3, 3), 1.0) + out[0, 1] = -0.5 + out[1, 0] = -0.5 + np.fill_diagonal(out, 0.0) + return out + + with pytest.raises(ValueError, match="negative entries"): + _pairwise_distance_matrix(coords, negative_metric) + + def test_callable_metric_asymmetric_raises(self): + """Callable returning a non-symmetric matrix raises.""" + coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) + + def asymmetric_metric(c1, c2): + out = np.zeros((3, 3)) + out[0, 1] = 1.0 + out[1, 0] = 2.0 + return out + + with pytest.raises(ValueError, match="asymmetric matrix"): + _pairwise_distance_matrix(coords, asymmetric_metric) + + def test_callable_metric_nonzero_diagonal_raises(self): + """Callable returning a symmetric/finite/non-negative matrix with a + positive self-distance still raises, because the Conley sandwich + requires d(i, i) = 0 (K(0) = 1 reduces the i=j term to the HC0 + diagonal X_i ε_i² X_i'). A nonzero self-distance silently attenuates + the HC0 contribution by K(d_ii / h) < 1 and misstates Conley SEs. + Codex CI R4 P1.""" + coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) + + def positive_diagonal_metric(c1, c2): + # Symmetric, finite, non-negative — but d(i, i) = 0.5 > 0. + n1 = len(c1) + n2 = len(c2) + out = np.full((n1, n2), 5.0) + np.fill_diagonal(out, 0.5) + return out + + with pytest.raises(ValueError, match=r"nonzero self-distance"): + _pairwise_distance_matrix(coords, positive_diagonal_metric) + + def test_callable_metric_near_zero_diagonal_accepted(self): + """Sub-tolerance diagonal (roundoff scale) is accepted, mirroring + the symmetry-tolerance contract.""" + coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) + + def near_zero_diagonal_metric(c1, c2): + n1 = len(c1) + n2 = len(c2) + out = np.full((n1, n2), 5.0) + np.fill_diagonal(out, 0.0) + # Diagonal noise well below the 1e-10 tolerance + out[0, 0] = 1e-13 + return out + + D = _pairwise_distance_matrix(coords, near_zero_diagonal_metric) + assert D.shape == (3, 3) + + def test_callable_metric_non_array_result_raises(self): + """Callable returning a non-castable result raises a targeted ValueError.""" + coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) + + def non_array_metric(c1, c2): + return "not an array" + + with pytest.raises(ValueError, match="cannot be cast"): + _pairwise_distance_matrix(coords, non_array_metric) + + def test_callable_metric_near_symmetric_accepted(self): + """Sub-tolerance asymmetry (eps-level roundoff) is accepted.""" + coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) + + def near_symmetric_metric(c1, c2): + out = np.full((3, 3), 5.0) + np.fill_diagonal(out, 0.0) + # Asymmetry below the 1e-10 tolerance — round-off only + out[0, 1] += 1e-13 + return out + + D = _pairwise_distance_matrix(coords, near_symmetric_metric) + assert D.shape == (3, 3) + # --------------------------------------------------------------------------- # TestConleyValidatorHelpers — direct calls to _validate_conley_kwargs @@ -779,17 +912,31 @@ def test_conley_in_valid_set(self): assert "conley" in _VALID_VCOV_TYPES - def test_conley_with_cluster_raises(self, fit_inputs): + def test_conley_with_cluster_combined_kernel(self, fit_inputs): + """Conley + cluster_ids applies the combined spatial + cluster + product kernel; no longer raises. The shipped SE differs from + bare Conley because cross-cluster off-diagonals are zeroed out.""" X, residuals, coords = fit_inputs - with pytest.raises(NotImplementedError, match="conley.*cluster_ids"): - compute_robust_vcov( - X, - residuals, - cluster_ids=np.arange(len(X)) // 3, - vcov_type="conley", - conley_coords=coords, - conley_cutoff_km=100.0, - ) + cluster_ids = np.arange(len(X)) // 3 + V_combined = compute_robust_vcov( + X, + residuals, + cluster_ids=cluster_ids, + vcov_type="conley", + conley_coords=coords, + conley_cutoff_km=100.0, + ) + V_bare = compute_robust_vcov( + X, + residuals, + vcov_type="conley", + conley_coords=coords, + conley_cutoff_km=100.0, + ) + assert V_combined.shape == V_bare.shape + # Combined kernel zeros out off-cluster off-diagonals → the meat + # (and hence vcov) must differ from bare Conley on the same data. + assert not np.allclose(V_combined, V_bare, atol=1e-8) def test_conley_with_weights_raises(self, fit_inputs): X, residuals, coords = fit_inputs @@ -945,177 +1092,186 @@ def two_period_panel(self): return pd.DataFrame(rows) - def test_did_with_conley_raises(self, two_period_panel): - """DifferenceInDifferences + vcov_type='conley' is rejected - unconditionally. DiD is intrinsically a two-period panel; cross- - sectional Conley over (unit, t=0) ∪ (unit, t=1) rows would treat - same-unit cross-time pairs as d_ij=0 -> K=1, mishandling the space- - time HAC. Phase 2 will add the space-time product kernel; Phase 1's - supported Conley path is direct compute_robust_vcov on a single- - period design. Closes CI reviewer P1 #1. - """ + def test_did_with_conley_panel_finite_se(self, two_period_panel): + """DifferenceInDifferences + vcov_type='conley' + unit + lag_cutoff + produces a finite SE on a two-period panel (Wave A #118).""" from diff_diff import DifferenceInDifferences df = two_period_panel.copy() - with pytest.raises(NotImplementedError, match="DifferenceInDifferences.*conley"): + res = DifferenceInDifferences( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(df, outcome="y", treatment="treated", time="time", unit="unit") + assert np.isfinite(res.att) + assert np.isfinite(res.se) and res.se > 0 + assert res.vcov_type == "conley" + assert res.conley_lag_cutoff == 1 + + def test_did_conley_missing_unit_raises(self, two_period_panel): + """vcov_type='conley' without unit= at fit-time raises ValueError.""" + from diff_diff import DifferenceInDifferences + + with pytest.raises(ValueError, match=r"`unit=`"): DifferenceInDifferences( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, - ).fit(df, outcome="y", treatment="treated", time="time") + conley_lag_cutoff=1, + ).fit(two_period_panel, outcome="y", treatment="treated", time="time") - def test_did_with_conley_repeated_coords_raises(self, two_period_panel): - """Per CI reviewer P1 #1 recommendation: regression test where - coordinates repeat across multiple periods. The fit must reject - rather than silently produce wrong SE.""" + def test_did_conley_unknown_unit_column_raises(self, two_period_panel): + """vcov_type='conley' with `unit=` referring to an absent column + raises a clear estimator-level ValueError, NOT a raw pandas KeyError. + Front-door check mirrors MultiPeriodDiD / TwoWayFixedEffects. + Codex CI R1 P1 #1.""" from diff_diff import DifferenceInDifferences - # Confirm the fixture has time-invariant coords per unit. - coord_var = two_period_panel.groupby("unit")[["lat", "lon"]].nunique() - assert (coord_var.values == 1).all(), "Fixture coords must be time-invariant" - - with pytest.raises(NotImplementedError, match="conley"): + with pytest.raises(ValueError, match="Unit column 'missing_unit' not found"): DifferenceInDifferences( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, - ).fit(two_period_panel, outcome="y", treatment="treated", time="time") + conley_lag_cutoff=1, + ).fit( + two_period_panel, + outcome="y", + treatment="treated", + time="time", + unit="missing_unit", + ) - def test_multi_period_did_with_conley_panel(self): - """Phase 2 MultiPeriodDiD + vcov_type='conley' uses the block-decomposed - sandwich (matches R conleyreg). Verifies that finite SEs are produced - when conley_lag_cutoff and unit are both supplied.""" - from diff_diff import MultiPeriodDiD + def test_did_conley_unknown_coord_column_raises(self, two_period_panel): + """vcov_type='conley' with `conley_coords=(, )` raises + a clear estimator-level ValueError before downstream column access. + Codex CI R2 P1.""" + from diff_diff import DifferenceInDifferences - rng = np.random.default_rng(seed=13) - n_units = 30 - rows = [] - for u in range(n_units): - lat = rng.uniform(-30, 30) - lon = rng.uniform(-100, 100) - treated = u < 15 - for t in range(4): - y = 0.2 * t + (1.0 if (treated and t >= 2) else 0.0) + rng.normal(0, 0.5) - rows.append( - {"unit": u, "time": t, "y": y, "treated": int(treated), "lat": lat, "lon": lon} - ) - import pandas as pd + with pytest.raises(ValueError, match="conley_coords column 'missing_lat' not found"): + DifferenceInDifferences( + vcov_type="conley", + conley_coords=("missing_lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit( + two_period_panel, + outcome="y", + treatment="treated", + time="time", + unit="unit", + ) - df_mp = pd.DataFrame(rows) - res = MultiPeriodDiD( - vcov_type="conley", - conley_coords=("lat", "lon"), - conley_cutoff_km=2000.0, - conley_lag_cutoff=1, - ).fit( - df_mp, - outcome="y", - treatment="treated", - time="time", - post_periods=[2, 3], - unit="unit", - reference_period=1, - ) - assert np.all(np.isfinite(res.vcov)), "MPD+Conley vcov must be finite" + def test_did_conley_unknown_cluster_column_raises(self, two_period_panel): + """DiD + Conley + cluster= raises a clear estimator- + level ValueError before `data[self.cluster]` access (combined-kernel + path; codex CI R7 P1).""" + from diff_diff import DifferenceInDifferences + + with pytest.raises(ValueError, match="Cluster column 'missing_region' not found"): + DifferenceInDifferences( + vcov_type="conley", + cluster="missing_region", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit( + two_period_panel, + outcome="y", + treatment="treated", + time="time", + unit="unit", + ) + + def test_mpd_conley_unknown_cluster_column_raises(self): + """MPD + Conley + cluster= raises a clear estimator- + level ValueError before `data[self.cluster]` access (combined-kernel + path; codex CI R7 P1).""" + import pandas as _pd - def test_multi_period_did_conley_missing_unit_raises(self): - """MPD + vcov_type='conley' without unit= at fit-time raises ValueError.""" from diff_diff import MultiPeriodDiD - rng = np.random.default_rng(seed=13) - n_units = 20 + rng = np.random.default_rng(seed=83) rows = [] - for u in range(n_units): + for u in range(8): lat = rng.uniform(-30, 30) lon = rng.uniform(-100, 100) - treated = u < 10 for t in range(3): rows.append( { "unit": u, "time": t, - "y": rng.normal(), - "treated": int(treated), + "y": rng.standard_normal(), + "treated": int(u >= 4), "lat": lat, "lon": lon, } ) - import pandas as pd - - df_mp = pd.DataFrame(rows) - with pytest.raises(ValueError, match="unit="): + df = _pd.DataFrame(rows) + with pytest.raises(ValueError, match="Cluster column 'missing_region' not found"): MultiPeriodDiD( vcov_type="conley", + cluster="missing_region", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, conley_lag_cutoff=1, ).fit( - df_mp, + df, outcome="y", treatment="treated", time="time", - post_periods=[2], - reference_period=1, + unit="unit", + post_periods=[1, 2], + reference_period=0, ) - def test_multi_period_did_conley_missing_lag_cutoff_raises(self): - """MPD + vcov_type='conley' without conley_lag_cutoff raises ValueError - (no defensible default per Conley 1999 Section 5).""" - from diff_diff import MultiPeriodDiD + def test_twfe_conley_unknown_cluster_column_raises(self): + """TWFE + Conley + cluster= raises a clear estimator- + level ValueError before `data[self.cluster]` access (combined-kernel + path; codex CI R7 P1).""" + import pandas as _pd - rng = np.random.default_rng(seed=13) - n_units = 20 + from diff_diff import TwoWayFixedEffects + + rng = np.random.default_rng(seed=89) rows = [] - for u in range(n_units): - lat = rng.uniform(-30, 30) - lon = rng.uniform(-100, 100) - treated = u < 10 - for t in range(3): + for u in range(8): + lat = rng.uniform(-5, 5) + lon = rng.uniform(-5, 5) + for t in range(2): rows.append( { "unit": u, "time": t, - "y": rng.normal(), - "treated": int(treated), + "y": rng.standard_normal(), + "treated": int(u >= 4), "lat": lat, "lon": lon, } ) - import pandas as pd - - df_mp = pd.DataFrame(rows) - with pytest.raises(ValueError, match="conley_lag_cutoff"): - MultiPeriodDiD( + df = _pd.DataFrame(rows) + with pytest.raises(ValueError, match="Cluster column 'missing_region' not found"): + TwoWayFixedEffects( vcov_type="conley", + cluster="missing_region", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, - ).fit( - df_mp, - outcome="y", - treatment="treated", - time="time", - post_periods=[2], - unit="unit", - reference_period=1, - ) - - def test_multi_period_did_conley_with_survey_design_raises(self): - """MPD + vcov_type='conley' + survey_design raises NotImplementedError. + conley_lag_cutoff=1, + ).fit(df, outcome="y", treatment="treated", time="time", unit="unit") - Closes Codex P0: previously, MPD passed return_vcov=False to solve_ols - when _use_survey_vcov=True, bypassing the conley + weights guard, and - then overwrote vcov with compute_survey_vcov — silently returning - survey SEs under a Conley request. + def test_mpd_conley_wild_bootstrap_raises_without_warning(self): + """MPD + Conley + inference='wild_bootstrap' raises NotImplementedError + cleanly. The pre-Conley analytical-fallback UserWarning is suppressed + on this combination so the user gets one consistent error message + instead of "warn then raise". Codex CI R11 P3 #1. """ - import pandas as pd + import pandas as _pd from diff_diff import MultiPeriodDiD - from diff_diff.survey import SurveyDesign - rng = np.random.default_rng(seed=29) - n_units = 24 + rng = np.random.default_rng(seed=211) rows = [] - for u in range(n_units): + for u in range(8): lat = rng.uniform(-30, 30) lon = rng.uniform(-100, 100) for t in range(3): @@ -1123,406 +1279,837 @@ def test_multi_period_did_conley_with_survey_design_raises(self): { "unit": u, "time": t, - "y": rng.normal(), - "treated": int(u < 12), + "y": rng.standard_normal(), + "treated": int(u >= 4), "lat": lat, "lon": lon, - "weight": 1.0 + 0.1 * rng.random(), - "stratum": u % 4, - "psu": u // 6, } ) - df_mp = pd.DataFrame(rows) - # Pure pweight (no PSU / strata) — would route through analytical conley - # path; the guard must fire before solve_ols. - sd_tsl = SurveyDesign(weights="weight", weight_type="pweight") - with pytest.raises(NotImplementedError, match="survey_design"): - MultiPeriodDiD( + df = _pd.DataFrame(rows) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + with pytest.raises(NotImplementedError, match="wild_bootstrap"): + MultiPeriodDiD( + vcov_type="conley", + inference="wild_bootstrap", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit( + df, + outcome="y", + treatment="treated", + time="time", + unit="unit", + post_periods=[1, 2], + reference_period=0, + ) + fallback_warnings = [ + msg + for msg in w + if "falling back to analytical" in str(msg.message) + or "Wild bootstrap inference is not yet supported" in str(msg.message) + ] + assert len(fallback_warnings) == 0, ( + "Got the analytical-fallback warning on a Conley fit that will " + "raise NotImplementedError — contradictory guidance." + ) + + def test_did_conley_malformed_coord_tuple_raises(self, two_period_panel): + """vcov_type='conley' with a malformed conley_coords (wrong arity or + non-string elements) raises ValueError before downstream access. + Codex CI R2 P1.""" + from diff_diff import DifferenceInDifferences + + # Wrong arity (1-element tuple) + with pytest.raises(ValueError, match="2-element tuple/list of column"): + DifferenceInDifferences( vcov_type="conley", - conley_coords=("lat", "lon"), + conley_coords=("lat",), # type: ignore[arg-type] conley_cutoff_km=2000.0, conley_lag_cutoff=1, ).fit( - df_mp, + two_period_panel, outcome="y", treatment="treated", time="time", - post_periods=[2], unit="unit", - reference_period=1, - survey_design=sd_tsl, ) - # Stratified PSU survey design — would route through Taylor TSL path - # and was the canonical bypass case the codex reviewer flagged. - sd_psu = SurveyDesign( - weights="weight", - strata="stratum", - psu="psu", - weight_type="pweight", - nest=True, - ) - with pytest.raises(NotImplementedError, match="survey_design"): - MultiPeriodDiD( + # Non-string element + with pytest.raises(ValueError, match="2-element tuple/list of column"): + DifferenceInDifferences( vcov_type="conley", - conley_coords=("lat", "lon"), + conley_coords=("lat", 0), # type: ignore[arg-type] conley_cutoff_km=2000.0, conley_lag_cutoff=1, ).fit( - df_mp, + two_period_panel, outcome="y", treatment="treated", time="time", - post_periods=[2], unit="unit", - reference_period=1, - survey_design=sd_psu, ) - def test_multi_period_did_conley_with_datetime64_time(self): - """End-to-end MPD + vcov_type='conley' with datetime64 time labels. - Closes Codex re-review P1: the wrapper must NOT coerce time to float64 - before passing to _compute_conley_vcov; the helper normalizes to - dense codes internally. Verifies the SEs match an equivalent - dense-integer-coded fit. - """ - import pandas as pd - - from diff_diff import MultiPeriodDiD + def test_did_conley_missing_lag_cutoff_raises(self, two_period_panel): + """vcov_type='conley' without conley_lag_cutoff raises ValueError.""" + from diff_diff import DifferenceInDifferences - rng = np.random.default_rng(seed=37) - n_units = 12 - date_labels = pd.to_datetime(["2024-01-01", "2024-04-01", "2024-08-01"]) - rows = [] - for u in range(n_units): - lat = rng.uniform(-30, 30) - lon = rng.uniform(-100, 100) - for t_idx, dt in enumerate(date_labels): - treated = u < 6 - y = 0.2 * t_idx + (1.0 if (treated and t_idx >= 1) else 0.0) + rng.normal(0, 0.4) - rows.append( - { - "unit": u, - "time_dt": dt, - "time_int": t_idx, - "y": y, - "treated": int(treated), - "lat": lat, - "lon": lon, - } - ) - df_mp = pd.DataFrame(rows) + with pytest.raises(ValueError, match="conley_lag_cutoff"): + DifferenceInDifferences( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + ).fit( + two_period_panel, + outcome="y", + treatment="treated", + time="time", + unit="unit", + ) + + def test_did_conley_matches_mpd_post_periods_1(self, two_period_panel): + """DiD + Conley on a 2-period panel matches MultiPeriodDiD with + post_periods=[1], reference_period=0 on the same data (locks the + DiD wire-up correctness against the already-shipped MPD path).""" + from diff_diff import DifferenceInDifferences, MultiPeriodDiD + + df = two_period_panel.copy() kwargs = dict( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, conley_lag_cutoff=1, ) - res_int = MultiPeriodDiD(**kwargs).fit( - df_mp, + res_did = DifferenceInDifferences(**kwargs).fit( + df, outcome="y", treatment="treated", time="time", unit="unit" + ) + res_mpd = MultiPeriodDiD(**kwargs).fit( + df, outcome="y", treatment="treated", - time="time_int", - post_periods=[1, 2], + time="time", unit="unit", + post_periods=[1], reference_period=0, ) - res_dt = MultiPeriodDiD(**kwargs).fit( - df_mp, + # MPD reports ATT for the single post period (1) + np.testing.assert_allclose(res_did.att, res_mpd.att, atol=1e-10) + np.testing.assert_allclose(res_did.se, res_mpd.se, atol=1e-10) + + def test_did_conley_with_absorb_uses_raw_time_labels(self, two_period_panel, monkeypatch): + """DiD + Conley + absorb=[] must feed the Conley helper the + ORIGINAL time/unit/coord columns from `data`, not the absorb-demeaned + `working_data` (in which time has been residualized to floats). + Otherwise the within-period spatial sandwich silently partitions on + per-unit demeaned floats instead of the true pre/post periods. + Codex Wave A R1 P0 #2. + """ + import diff_diff.linalg as linalg_module + from diff_diff import DifferenceInDifferences + + df = two_period_panel.copy() + captured: dict = {"time_arg": None, "unit_arg": None} + orig = linalg_module._compute_conley_vcov + + def _spy(*args, **kwargs): + captured["time_arg"] = kwargs.get("time") + captured["unit_arg"] = kwargs.get("unit") + return orig(*args, **kwargs) + + monkeypatch.setattr(linalg_module, "_compute_conley_vcov", _spy) + DifferenceInDifferences( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit( + df, outcome="y", treatment="treated", - time="time_dt", - post_periods=[date_labels[1], date_labels[2]], + time="time", unit="unit", - reference_period=date_labels[0], + absorb=["unit"], ) - # Per-coefficient SE should match across the two encodings (dense - # codes normalize identically). MPD orders coefficients by the - # reference-vs-non-reference period split; with reference_period=0 - # and post_periods=[1,2] the coefficient ordering is bit-identical. - se_int = np.sqrt(np.diag(res_int.vcov)) - se_dt = np.sqrt(np.diag(res_dt.vcov)) - np.testing.assert_allclose(se_dt, se_int, atol=1e-10) - - def test_multi_period_did_conley_to_dict_carries_lag_cutoff(self): - """Closes Codex re-review round 4 P1 (Maintainability) on MPD: - serialized `to_dict()` must include `vcov_type` and - `conley_lag_cutoff` so downstream programmatic consumers can tell - which Conley variant produced the SEs.""" - import pandas as pd + assert captured["time_arg"] is not None + # Raw labels are integer 0/1 (the binary post-treatment indicator); + # demeaned values would be floats from absorb's within-unit + # demeaning. np.unique on raw labels yields exactly 2 distinct + # values; on demeaned floats it would yield ~n_units distinct. + time_arg = np.asarray(captured["time_arg"]) + uniques = np.unique(time_arg) + assert len(uniques) == 2, ( + f"Expected 2 unique time labels (raw 0/1), got {len(uniques)}: " + f"{uniques[:5]} — absorb is leaking demeaned time into the " + "Conley helper." + ) + assert set(uniques.tolist()) == {0, 1}, f"Expected raw integer labels 0/1, got {uniques}" + + def test_mpd_conley_with_absorb_uses_raw_coords_and_time(self, monkeypatch): + """MultiPeriodDiD + Conley + absorb=[] must feed the Conley + helper the ORIGINAL coords/time/unit columns from `data`, not the + absorb-demeaned `working_data`. If a user lists the `time` column + (or coord columns) in `absorb`, working_data has those demeaned and + the Conley helper would partition the within-period spatial sandwich + on residualized floats. Mirrors the DiD raw-time contract. + Codex CI R5 P0. + """ + import pandas as _pd + import diff_diff.linalg as linalg_module from diff_diff import MultiPeriodDiD - rng = np.random.default_rng(seed=41) + rng = np.random.default_rng(seed=71) n_units = 10 + T = 4 rows = [] for u in range(n_units): + treated = u >= n_units // 2 lat = rng.uniform(-30, 30) lon = rng.uniform(-100, 100) - for t in range(3): + for t in range(T): + effect = 1.0 if (treated and t >= 2) else 0.0 + yv = 0.2 * t + effect + rng.normal(0, 0.3) rows.append( { "unit": u, "time": t, - "y": rng.normal(), - "treated": int(u < 5), + "y": yv, + "treated": int(treated), "lat": lat, "lon": lon, } ) - df_mp = pd.DataFrame(rows) - res = MultiPeriodDiD( + df = _pd.DataFrame(rows) + + captured: dict = {"time_arg": None, "unit_arg": None, "coords_arg": None} + orig = linalg_module._compute_conley_vcov + + def _spy(*args, **kwargs): + captured["time_arg"] = kwargs.get("time") + captured["unit_arg"] = kwargs.get("unit") + # coords is the 3rd positional arg to _compute_conley_vcov + if len(args) >= 3: + captured["coords_arg"] = args[2] + return orig(*args, **kwargs) + + monkeypatch.setattr(linalg_module, "_compute_conley_vcov", _spy) + MultiPeriodDiD( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, - conley_lag_cutoff=2, + conley_lag_cutoff=1, ).fit( - df_mp, + df, outcome="y", treatment="treated", time="time", - post_periods=[1, 2], unit="unit", + post_periods=[2, 3], reference_period=0, + absorb=["unit"], + ) + # Raw time labels span exactly 0..T-1 = 4 distinct values; demeaned + # absorb would collapse to per-unit means → ~n_units distinct values. + time_arg = np.asarray(captured["time_arg"]) + uniques = np.unique(time_arg) + assert len(uniques) == T, ( + f"Expected {T} unique time labels (raw 0..{T - 1}), got {len(uniques)}: " + f"{uniques[:5]} — absorb is leaking demeaned time into the Conley helper." + ) + assert set(uniques.tolist()) == set(range(T)) + # Raw coords are time-invariant within unit and span n_units distinct + # (lat, lon) pairs. If absorb=["unit"] leaked demeaned coords, all + # within-unit coord values would collapse to 0 (per-unit mean), giving + # only 1 distinct row across all observations. + coords_arg = np.asarray(captured["coords_arg"]) + # Expect n_units distinct (lat, lon) pairs since each unit has its own + unique_coords = np.unique(coords_arr_view := coords_arg, axis=0) + del coords_arr_view + assert len(unique_coords) == n_units, ( + f"Expected {n_units} unique coord pairs, got {len(unique_coords)} — " + "absorb=['unit'] is leaking demeaned coords into the Conley helper." ) - d = res.to_dict() - assert d["vcov_type"] == "conley" - assert d["conley_lag_cutoff"] == 2 - def test_multi_period_did_conley_missing_coords_raises(self): - """MPD + vcov_type='conley' without conley_coords raises a clean - ValueError instead of a raw TypeError on `self.conley_coords[0]`. - Closes Codex P2 #1. - """ + def test_multi_period_did_with_conley_panel(self): + """Phase 2 MultiPeriodDiD + vcov_type='conley' uses the block-decomposed + sandwich (matches R conleyreg). Verifies that finite SEs are produced + when conley_lag_cutoff and unit are both supplied.""" + from diff_diff import MultiPeriodDiD + + rng = np.random.default_rng(seed=13) + n_units = 30 + rows = [] + for u in range(n_units): + lat = rng.uniform(-30, 30) + lon = rng.uniform(-100, 100) + treated = u < 15 + for t in range(4): + y = 0.2 * t + (1.0 if (treated and t >= 2) else 0.0) + rng.normal(0, 0.5) + rows.append( + {"unit": u, "time": t, "y": y, "treated": int(treated), "lat": lat, "lon": lon} + ) import pandas as pd + df_mp = pd.DataFrame(rows) + res = MultiPeriodDiD( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit( + df_mp, + outcome="y", + treatment="treated", + time="time", + post_periods=[2, 3], + unit="unit", + reference_period=1, + ) + assert np.all(np.isfinite(res.vcov)), "MPD+Conley vcov must be finite" + + def test_multi_period_did_conley_missing_unit_raises(self): + """MPD + vcov_type='conley' without unit= at fit-time raises ValueError.""" from diff_diff import MultiPeriodDiD - rng = np.random.default_rng(seed=31) - n_units = 10 + rng = np.random.default_rng(seed=13) + n_units = 20 rows = [] for u in range(n_units): - for t in range(2): + lat = rng.uniform(-30, 30) + lon = rng.uniform(-100, 100) + treated = u < 10 + for t in range(3): rows.append( { "unit": u, "time": t, "y": rng.normal(), - "treated": int(u < 5), + "treated": int(treated), + "lat": lat, + "lon": lon, } ) + import pandas as pd + df_mp = pd.DataFrame(rows) - with pytest.raises(ValueError, match="conley_coords.*conley_cutoff_km"): + with pytest.raises(ValueError, match="unit="): MultiPeriodDiD( vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, conley_lag_cutoff=1, ).fit( df_mp, outcome="y", treatment="treated", time="time", - post_periods=[1], - unit="unit", - reference_period=0, + post_periods=[2], + reference_period=1, ) + def test_multi_period_did_conley_missing_lag_cutoff_raises(self): + """MPD + vcov_type='conley' without conley_lag_cutoff raises ValueError + (no defensible default per Conley 1999 Section 5).""" + from diff_diff import MultiPeriodDiD -class TestConleyTWFE: - """TwoWayFixedEffects + vcov_type='conley' uses the Phase 2 block-decomposed - panel HAC (matches R conleyreg). The within-transformed scores feed the same - block-decomposed helper that LinearRegression uses; FWL composability - ensures the FE-residualized meat matches the full-dummy-expansion meat. - """ - - @pytest.fixture - def panel(self): - """Build a 2-period panel with geocoords for TWFE tests.""" - rng = np.random.default_rng(seed=17) - n_units = 30 + rng = np.random.default_rng(seed=13) + n_units = 20 rows = [] for u in range(n_units): lat = rng.uniform(-30, 30) lon = rng.uniform(-100, 100) - treated = u < 15 - unit_fe = rng.normal(0, 0.3) - for t in range(2): - time_fe = 0.5 if t == 1 else 0.0 - effect = 1.0 if (treated and t == 1) else 0.0 - y = unit_fe + time_fe + effect + rng.normal(0, 0.4) + treated = u < 10 + for t in range(3): rows.append( - {"unit": u, "time": t, "y": y, "treated": int(treated), "lat": lat, "lon": lon} + { + "unit": u, + "time": t, + "y": rng.normal(), + "treated": int(treated), + "lat": lat, + "lon": lon, + } ) import pandas as pd - return pd.DataFrame(rows) - - def test_twfe_conley_panel_finite_se(self, panel): - """TWFE + vcov_type='conley' on a balanced panel produces a finite SE.""" - from diff_diff import TwoWayFixedEffects - - res = TwoWayFixedEffects( - vcov_type="conley", - conley_coords=("lat", "lon"), - conley_cutoff_km=2000.0, - conley_lag_cutoff=1, - ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") - assert np.isfinite(res.att), "ATT must be finite" - assert np.isfinite(res.se) and res.se > 0, "SE must be positive and finite" - - def test_twfe_conley_with_explicit_cluster_raises(self, panel): - """TWFE + vcov_type='conley' + explicit cluster=... raises: the combined - spatial-kernel + cluster-indicator product kernel is deferred to a - follow-up PR. Auto-cluster on the Conley path is silently dropped.""" - from diff_diff import TwoWayFixedEffects - - with pytest.raises(NotImplementedError, match="cluster"): - TwoWayFixedEffects( + df_mp = pd.DataFrame(rows) + with pytest.raises(ValueError, match="conley_lag_cutoff"): + MultiPeriodDiD( vcov_type="conley", - cluster="unit", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, - conley_lag_cutoff=1, - ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + ).fit( + df_mp, + outcome="y", + treatment="treated", + time="time", + post_periods=[2], + unit="unit", + reference_period=1, + ) - def test_twfe_conley_with_wild_bootstrap_raises(self, panel): - """vcov_type='conley' + inference='wild_bootstrap' raises: wild bootstrap - does not consume the analytical Conley sandwich.""" - from diff_diff import TwoWayFixedEffects + def test_multi_period_did_conley_with_survey_design_raises(self): + """MPD + vcov_type='conley' + survey_design raises NotImplementedError. - with pytest.raises(NotImplementedError, match="wild_bootstrap"): - TwoWayFixedEffects( - vcov_type="conley", - inference="wild_bootstrap", - conley_coords=("lat", "lon"), - conley_cutoff_km=2000.0, - conley_lag_cutoff=1, - ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + Closes Codex P0: previously, MPD passed return_vcov=False to solve_ols + when _use_survey_vcov=True, bypassing the conley + weights guard, and + then overwrote vcov with compute_survey_vcov — silently returning + survey SEs under a Conley request. + """ + import pandas as pd - def test_twfe_conley_repeated_coords_panel_finite_se(self, panel): - """Phase 2 regression for the Phase-1 silent-bug case: each unit's - coords are time-invariant. The block-decomposed sandwich correctly - sums within-period (period 0 and period 1 separately) plus - within-unit serial (lag=1) so the same-unit cross-time pairs at - d_ij=0 do NOT inflate the meat.""" - from diff_diff import TwoWayFixedEffects + from diff_diff import MultiPeriodDiD + from diff_diff.survey import SurveyDesign - coord_var = panel.groupby("unit")[["lat", "lon"]].nunique() - assert (coord_var.values == 1).all(), "Fixture coords must be time-invariant" - res = TwoWayFixedEffects( - vcov_type="conley", - conley_coords=("lat", "lon"), - conley_cutoff_km=2000.0, - conley_lag_cutoff=1, - ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") - assert np.isfinite(res.se) and res.se > 0 - - def test_twfe_conley_missing_lag_cutoff_raises(self, panel): - """conley_lag_cutoff is required; no defensible default per Conley §5.""" - from diff_diff import TwoWayFixedEffects - - with pytest.raises(ValueError, match="conley_lag_cutoff"): - TwoWayFixedEffects( + rng = np.random.default_rng(seed=29) + n_units = 24 + rows = [] + for u in range(n_units): + lat = rng.uniform(-30, 30) + lon = rng.uniform(-100, 100) + for t in range(3): + rows.append( + { + "unit": u, + "time": t, + "y": rng.normal(), + "treated": int(u < 12), + "lat": lat, + "lon": lon, + "weight": 1.0 + 0.1 * rng.random(), + "stratum": u % 4, + "psu": u // 6, + } + ) + df_mp = pd.DataFrame(rows) + # Pure pweight (no PSU / strata) — would route through analytical conley + # path; the guard must fire before solve_ols. + sd_tsl = SurveyDesign(weights="weight", weight_type="pweight") + with pytest.raises(NotImplementedError, match="survey_design"): + MultiPeriodDiD( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, - ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + conley_lag_cutoff=1, + ).fit( + df_mp, + outcome="y", + treatment="treated", + time="time", + post_periods=[2], + unit="unit", + reference_period=1, + survey_design=sd_tsl, + ) + # Stratified PSU survey design — would route through Taylor TSL path + # and was the canonical bypass case the codex reviewer flagged. + sd_psu = SurveyDesign( + weights="weight", + strata="stratum", + psu="psu", + weight_type="pweight", + nest=True, + ) + with pytest.raises(NotImplementedError, match="survey_design"): + MultiPeriodDiD( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit( + df_mp, + outcome="y", + treatment="treated", + time="time", + post_periods=[2], + unit="unit", + reference_period=1, + survey_design=sd_psu, + ) - def test_twfe_conley_binary_post_label_normalization(self, panel): - """TWFE with binary `post` (values {0,1}) + `conley_lag_cutoff=1` - produces the same finite vcov as the equivalent dense-period-index - fit. Closes the Codex P1 example — the time-label normalization - means lag is counted in panel periods regardless of how `time` is - encoded (binary post indicator vs. dense period index). + def test_multi_period_did_conley_with_datetime64_time(self): + """End-to-end MPD + vcov_type='conley' with datetime64 time labels. + Closes Codex re-review P1: the wrapper must NOT coerce time to float64 + before passing to _compute_conley_vcov; the helper normalizes to + dense codes internally. Verifies the SEs match an equivalent + dense-integer-coded fit. """ - from diff_diff import TwoWayFixedEffects - - # `panel` fixture uses `time` in {0, 1}, identical to a binary post. - # Rename to `post` to make the test scenario explicit. - df_post = panel.rename(columns={"time": "post"}) - res = TwoWayFixedEffects( - vcov_type="conley", - conley_coords=("lat", "lon"), - conley_cutoff_km=2000.0, - conley_lag_cutoff=1, - ).fit(df_post, outcome="y", treatment="treated", time="post", unit="unit") - assert np.isfinite(res.se) and res.se > 0 + import pandas as pd - def test_twfe_conley_summary_emits_conley_label(self, panel): - """Panel result summary must label the variance family as Conley - spatial HAC and surface `lag_cutoff` so downstream consumers can tell - which Conley variant produced the SEs. Closes Codex P3 and the - re-review P1 (Maintainability). - """ - from diff_diff import TwoWayFixedEffects + from diff_diff import MultiPeriodDiD - res = TwoWayFixedEffects( + rng = np.random.default_rng(seed=37) + n_units = 12 + date_labels = pd.to_datetime(["2024-01-01", "2024-04-01", "2024-08-01"]) + rows = [] + for u in range(n_units): + lat = rng.uniform(-30, 30) + lon = rng.uniform(-100, 100) + for t_idx, dt in enumerate(date_labels): + treated = u < 6 + y = 0.2 * t_idx + (1.0 if (treated and t_idx >= 1) else 0.0) + rng.normal(0, 0.4) + rows.append( + { + "unit": u, + "time_dt": dt, + "time_int": t_idx, + "y": y, + "treated": int(treated), + "lat": lat, + "lon": lon, + } + ) + df_mp = pd.DataFrame(rows) + kwargs = dict( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, conley_lag_cutoff=1, - ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") - summary = res.summary() - assert "Conley spatial HAC" in summary - assert "lag_cutoff=1" in summary - # The result dataclass also carries the lag for programmatic access. - assert res.conley_lag_cutoff == 1 + ) + res_int = MultiPeriodDiD(**kwargs).fit( + df_mp, + outcome="y", + treatment="treated", + time="time_int", + post_periods=[1, 2], + unit="unit", + reference_period=0, + ) + res_dt = MultiPeriodDiD(**kwargs).fit( + df_mp, + outcome="y", + treatment="treated", + time="time_dt", + post_periods=[date_labels[1], date_labels[2]], + unit="unit", + reference_period=date_labels[0], + ) + # Per-coefficient SE should match across the two encodings (dense + # codes normalize identically). MPD orders coefficients by the + # reference-vs-non-reference period split; with reference_period=0 + # and post_periods=[1,2] the coefficient ordering is bit-identical. + se_int = np.sqrt(np.diag(res_int.vcov)) + se_dt = np.sqrt(np.diag(res_dt.vcov)) + np.testing.assert_allclose(se_dt, se_int, atol=1e-10) - def test_twfe_conley_to_dict_carries_lag_cutoff(self, panel): - """Closes Codex re-review round 4 P1 (Maintainability): the + def test_multi_period_did_conley_to_dict_carries_lag_cutoff(self): + """Closes Codex re-review round 4 P1 (Maintainability) on MPD: serialized `to_dict()` must include `vcov_type` and - `conley_lag_cutoff` so downstream programmatic consumers (notebooks, - adapters, pipelines) can tell which Conley variant produced the SEs - without re-deriving from the summary string.""" - from diff_diff import TwoWayFixedEffects + `conley_lag_cutoff` so downstream programmatic consumers can tell + which Conley variant produced the SEs.""" + import pandas as pd - res = TwoWayFixedEffects( + from diff_diff import MultiPeriodDiD + + rng = np.random.default_rng(seed=41) + n_units = 10 + rows = [] + for u in range(n_units): + lat = rng.uniform(-30, 30) + lon = rng.uniform(-100, 100) + for t in range(3): + rows.append( + { + "unit": u, + "time": t, + "y": rng.normal(), + "treated": int(u < 5), + "lat": lat, + "lon": lon, + } + ) + df_mp = pd.DataFrame(rows) + res = MultiPeriodDiD( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, - conley_lag_cutoff=1, - ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + conley_lag_cutoff=2, + ).fit( + df_mp, + outcome="y", + treatment="treated", + time="time", + post_periods=[1, 2], + unit="unit", + reference_period=0, + ) d = res.to_dict() assert d["vcov_type"] == "conley" - assert d["conley_lag_cutoff"] == 1 + assert d["conley_lag_cutoff"] == 2 - def test_twfe_conley_cluster_name_is_none(self, panel): - """Closes Codex re-review round 5 P1 (Maintainability): TWFE drops - its auto-unit-cluster on the Conley path (`_conley_cluster_override = - None`), so the variance-provenance metadata must reflect that. The - result's `cluster_name` is None and `to_dict()` does not advertise - `cluster_name` — otherwise downstream consumers would be told the - SEs were CR1-clustered when they're actually Conley spatial HAC. + def test_multi_period_did_conley_missing_coords_raises(self): + """MPD + vcov_type='conley' without conley_coords raises a clean + ValueError instead of a raw TypeError on `self.conley_coords[0]`. + Closes Codex P2 #1. """ - from diff_diff import TwoWayFixedEffects - - res = TwoWayFixedEffects( - vcov_type="conley", - conley_coords=("lat", "lon"), - conley_cutoff_km=2000.0, - conley_lag_cutoff=1, - ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") - assert res.cluster_name is None - d = res.to_dict() - assert "cluster_name" not in d + import pandas as pd - def test_twfe_conley_non_numeric_time_fails(self, panel): - """TWFE's `_treatment_post = treated * time` design step requires - numeric `time`. Non-numeric labels (datetime64, pd.Period, strings) - are TWFE-incompatible end-to-end and surface as a clean error before - the Conley path runs. Use MultiPeriodDiD if you need non-numeric - time labels. - """ - from diff_diff import TwoWayFixedEffects + from diff_diff import MultiPeriodDiD - df_str = panel.copy() - df_str["time_str"] = df_str["time"].map({0: "pre", 1: "post"}) - with pytest.raises((TypeError, ValueError)): - TwoWayFixedEffects( - vcov_type="conley", - conley_coords=("lat", "lon"), - conley_cutoff_km=2000.0, - conley_lag_cutoff=1, - ).fit( - df_str, + rng = np.random.default_rng(seed=31) + n_units = 10 + rows = [] + for u in range(n_units): + for t in range(2): + rows.append( + { + "unit": u, + "time": t, + "y": rng.normal(), + "treated": int(u < 5), + } + ) + df_mp = pd.DataFrame(rows) + with pytest.raises(ValueError, match="conley_coords.*conley_cutoff_km"): + MultiPeriodDiD( + vcov_type="conley", + conley_lag_cutoff=1, + ).fit( + df_mp, + outcome="y", + treatment="treated", + time="time", + post_periods=[1], + unit="unit", + reference_period=0, + ) + + +class TestConleyTWFE: + """TwoWayFixedEffects + vcov_type='conley' uses the Phase 2 block-decomposed + panel HAC (matches R conleyreg). The within-transformed scores feed the same + block-decomposed helper that LinearRegression uses; FWL composability + ensures the FE-residualized meat matches the full-dummy-expansion meat. + """ + + @pytest.fixture + def panel(self): + """Build a 2-period panel with geocoords for TWFE tests.""" + rng = np.random.default_rng(seed=17) + n_units = 30 + rows = [] + for u in range(n_units): + lat = rng.uniform(-30, 30) + lon = rng.uniform(-100, 100) + treated = u < 15 + unit_fe = rng.normal(0, 0.3) + for t in range(2): + time_fe = 0.5 if t == 1 else 0.0 + effect = 1.0 if (treated and t == 1) else 0.0 + y = unit_fe + time_fe + effect + rng.normal(0, 0.4) + rows.append( + {"unit": u, "time": t, "y": y, "treated": int(treated), "lat": lat, "lon": lon} + ) + import pandas as pd + + return pd.DataFrame(rows) + + def test_twfe_conley_panel_finite_se(self, panel): + """TWFE + vcov_type='conley' on a balanced panel produces a finite SE.""" + from diff_diff import TwoWayFixedEffects + + res = TwoWayFixedEffects( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + assert np.isfinite(res.att), "ATT must be finite" + assert np.isfinite(res.se) and res.se > 0, "SE must be positive and finite" + + def test_twfe_conley_with_explicit_cluster_combined_kernel(self, panel): + """TWFE + vcov_type='conley' + explicit cluster= applies the + combined spatial + cluster product kernel. The user-supplied cluster + column propagates to ``cluster_name`` (no longer cleared on the + Conley path) and a finite SE is produced. Auto-cluster on the + Conley path remains silently dropped — the user MUST explicitly + opt in to the combined kernel.""" + from diff_diff import TwoWayFixedEffects + + # Add a unit-level region column so the cluster is time-invariant + # within unit (the panel block-decomposed validator's contract). + panel = panel.copy() + panel["region"] = panel["unit"] // 5 + res = TwoWayFixedEffects( + vcov_type="conley", + cluster="region", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + assert np.isfinite(res.att) + assert np.isfinite(res.se) and res.se > 0 + assert res.cluster_name == "region" + d = res.to_dict() + assert d.get("cluster_name") == "region" + + def test_twfe_conley_with_wild_bootstrap_raises(self, panel): + """vcov_type='conley' + inference='wild_bootstrap' raises: wild bootstrap + does not consume the analytical Conley sandwich.""" + from diff_diff import TwoWayFixedEffects + + with pytest.raises(NotImplementedError, match="wild_bootstrap"): + TwoWayFixedEffects( + vcov_type="conley", + inference="wild_bootstrap", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + + def test_twfe_conley_repeated_coords_panel_finite_se(self, panel): + """Phase 2 regression for the Phase-1 silent-bug case: each unit's + coords are time-invariant. The block-decomposed sandwich correctly + sums within-period (period 0 and period 1 separately) plus + within-unit serial (lag=1) so the same-unit cross-time pairs at + d_ij=0 do NOT inflate the meat.""" + from diff_diff import TwoWayFixedEffects + + coord_var = panel.groupby("unit")[["lat", "lon"]].nunique() + assert (coord_var.values == 1).all(), "Fixture coords must be time-invariant" + res = TwoWayFixedEffects( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + assert np.isfinite(res.se) and res.se > 0 + + def test_twfe_conley_missing_lag_cutoff_raises(self, panel): + """conley_lag_cutoff is required; no defensible default per Conley §5.""" + from diff_diff import TwoWayFixedEffects + + with pytest.raises(ValueError, match="conley_lag_cutoff"): + TwoWayFixedEffects( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + + def test_twfe_conley_binary_post_label_normalization(self, panel): + """TWFE with binary `post` (values {0,1}) + `conley_lag_cutoff=1` + produces the same finite vcov as the equivalent dense-period-index + fit. Closes the Codex P1 example — the time-label normalization + means lag is counted in panel periods regardless of how `time` is + encoded (binary post indicator vs. dense period index). + """ + from diff_diff import TwoWayFixedEffects + + # `panel` fixture uses `time` in {0, 1}, identical to a binary post. + # Rename to `post` to make the test scenario explicit. + df_post = panel.rename(columns={"time": "post"}) + res = TwoWayFixedEffects( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(df_post, outcome="y", treatment="treated", time="post", unit="unit") + assert np.isfinite(res.se) and res.se > 0 + + def test_twfe_conley_summary_emits_conley_label(self, panel): + """Panel result summary must label the variance family as Conley + spatial HAC and surface `lag_cutoff` so downstream consumers can tell + which Conley variant produced the SEs. Closes Codex P3 and the + re-review P1 (Maintainability). + """ + from diff_diff import TwoWayFixedEffects + + res = TwoWayFixedEffects( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + summary = res.summary() + assert "Conley spatial HAC" in summary + assert "lag_cutoff=1" in summary + # No explicit cluster on this fit → label must NOT advertise the + # combined kernel. + assert "+ cluster product kernel" not in summary + # The result dataclass also carries the lag for programmatic access. + assert res.conley_lag_cutoff == 1 + + def test_twfe_conley_with_cluster_summary_label_names_kernel_and_cluster(self, panel): + """When an explicit cluster= is combined with Conley, the + summary label must distinguish the combined spatial + cluster + product kernel from bare Conley and name the cluster column. + Codex CI R3 P3 (Maintainability).""" + from diff_diff import TwoWayFixedEffects + + panel = panel.copy() + panel["region"] = panel["unit"] // 5 # time-invariant within unit + res = TwoWayFixedEffects( + vcov_type="conley", + cluster="region", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + summary = res.summary() + assert "Conley spatial HAC" in summary + assert "+ cluster product kernel at region" in summary + assert "lag_cutoff=1" in summary + # Programmatic access via the result dataclass. + assert res.cluster_name == "region" + assert res.conley_lag_cutoff == 1 + + def test_twfe_conley_to_dict_carries_lag_cutoff(self, panel): + """Closes Codex re-review round 4 P1 (Maintainability): the + serialized `to_dict()` must include `vcov_type` and + `conley_lag_cutoff` so downstream programmatic consumers (notebooks, + adapters, pipelines) can tell which Conley variant produced the SEs + without re-deriving from the summary string.""" + from diff_diff import TwoWayFixedEffects + + res = TwoWayFixedEffects( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + d = res.to_dict() + assert d["vcov_type"] == "conley" + assert d["conley_lag_cutoff"] == 1 + + def test_twfe_conley_cluster_name_is_none(self, panel): + """Closes Codex re-review round 5 P1 (Maintainability): TWFE drops + its auto-unit-cluster on the Conley path (`_conley_cluster_override = + None`), so the variance-provenance metadata must reflect that. The + result's `cluster_name` is None and `to_dict()` does not advertise + `cluster_name` — otherwise downstream consumers would be told the + SEs were CR1-clustered when they're actually Conley spatial HAC. + """ + from diff_diff import TwoWayFixedEffects + + res = TwoWayFixedEffects( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(panel, outcome="y", treatment="treated", time="time", unit="unit") + assert res.cluster_name is None + d = res.to_dict() + assert "cluster_name" not in d + + def test_twfe_conley_non_numeric_time_fails(self, panel): + """TWFE's `_treatment_post = treated * time` design step requires + numeric `time`. Non-numeric labels (datetime64, pd.Period, strings) + are TWFE-incompatible end-to-end and surface as a clean error before + the Conley path runs. Use MultiPeriodDiD if you need non-numeric + time labels. + """ + from diff_diff import TwoWayFixedEffects + + df_str = panel.copy() + df_str["time_str"] = df_str["time"].map({0: "pre", 1: "post"}) + with pytest.raises((TypeError, ValueError)): + TwoWayFixedEffects( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit( + df_str, outcome="y", treatment="treated", time="time_str", @@ -1613,563 +2200,1691 @@ def df(self): } ) - def test_did_conley_combinations_all_raise(self, df): - """Every DifferenceInDifferences + vcov_type='conley' combination - rejects unconditionally (DiD is intrinsically a two-period panel; - cross-sectional Conley is unsafe over (unit, time) rows). Asserts - the reject regardless of cluster=, absorb=, or missing coords/cutoff. - Closes CI reviewer P1 #1. - """ + def test_did_conley_combinations(self, df): + """DifferenceInDifferences + vcov_type='conley' validation table: + missing coords/cutoff/lag_cutoff/unit each raise ValueError; + valid full kwarg set succeeds; survey_design + Conley raises + NotImplementedError (Wave A scope: row 121 deferred); + wild_bootstrap + Conley raises NotImplementedError.""" from diff_diff import DifferenceInDifferences - # cluster + conley - with pytest.raises(NotImplementedError, match="conley"): - DifferenceInDifferences( - vcov_type="conley", - cluster="stratum", - conley_coords=("lat", "lon"), - conley_cutoff_km=100.0, - ).fit(df, outcome="y", treatment="treated", time="time") - # missing conley_coords - with pytest.raises(NotImplementedError, match="conley"): - DifferenceInDifferences( - vcov_type="conley", - conley_cutoff_km=100.0, - ).fit(df, outcome="y", treatment="treated", time="time") # missing conley_cutoff_km - with pytest.raises(NotImplementedError, match="conley"): + with pytest.raises(ValueError, match="conley_coords|conley_cutoff_km"): DifferenceInDifferences( vcov_type="conley", conley_coords=("lat", "lon"), - ).fit(df, outcome="y", treatment="treated", time="time") - # unknown coord column (data validation skipped — outer reject fires first) - with pytest.raises(NotImplementedError, match="conley"): + conley_lag_cutoff=1, + ).fit(df, outcome="y", treatment="treated", time="time", unit="unit") + # missing conley_lag_cutoff + with pytest.raises(ValueError, match="conley_lag_cutoff"): DifferenceInDifferences( vcov_type="conley", - conley_coords=("missing_lat", "lon"), + conley_coords=("lat", "lon"), conley_cutoff_km=100.0, - ).fit(df, outcome="y", treatment="treated", time="time") - # absorb + conley - with pytest.raises(NotImplementedError, match="conley"): + ).fit(df, outcome="y", treatment="treated", time="time", unit="unit") + # missing unit + with pytest.raises(ValueError, match=r"`unit=`"): DifferenceInDifferences( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=100.0, - ).fit(df, outcome="y", treatment="treated", time="time", absorb=["unit"]) + conley_lag_cutoff=1, + ).fit(df, outcome="y", treatment="treated", time="time") + # Valid full kwarg set does NOT raise (separate fixture in + # TestConleyEstimatorIntegration covers the finite-SE assertion; + # this fixture's treated/time correlation triggers rank deficiency). + DifferenceInDifferences( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=100.0, + conley_lag_cutoff=1, + rank_deficient_action="silent", + ).fit( + df, + outcome="y", + treatment="treated", + time="time", + unit="unit", + ) + + def test_did_conley_with_survey_design_raises(self, df): + """DiD + Conley + survey_design raises NotImplementedError (deferred + to Wave 2 weighted-Conley).""" + from diff_diff import DifferenceInDifferences, SurveyDesign + + with pytest.raises(NotImplementedError, match="conley.*survey_design"): + DifferenceInDifferences( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=100.0, + conley_lag_cutoff=1, + ).fit( + df, + outcome="y", + treatment="treated", + time="time", + unit="unit", + survey_design=SurveyDesign(strata="stratum"), + ) + + def test_did_conley_with_wild_bootstrap_raises(self, df): + """DiD + Conley + inference='wild_bootstrap' raises.""" + from diff_diff import DifferenceInDifferences + + with pytest.raises(NotImplementedError, match="wild_bootstrap"): + DifferenceInDifferences( + vcov_type="conley", + inference="wild_bootstrap", + cluster="unit", + conley_coords=("lat", "lon"), + conley_cutoff_km=100.0, + conley_lag_cutoff=1, + ).fit(df, outcome="y", treatment="treated", time="time", unit="unit") + + def test_synthetic_did_conley_raises(self): + from diff_diff import SyntheticDiD + + with pytest.raises(TypeError, match="conley"): + SyntheticDiD(vcov_type="conley") # type: ignore[call-arg] + + def test_synthetic_did_conley_kwarg_raises(self): + from diff_diff import SyntheticDiD + + with pytest.raises(TypeError, match="conley"): + SyntheticDiD(conley_cutoff_km=100.0) # type: ignore[call-arg] + + def test_synthetic_did_set_params_conley_raises(self): + """SyntheticDiD.set_params(vcov_type='conley') must raise (mirrors + __init__'s contract — closes the silent-bypass gap CI reviewer flagged + as P1 CQ1).""" + from diff_diff import SyntheticDiD + + est = SyntheticDiD() + # Snapshot pre-call state + before_variance = est.variance_method + before_n_boot = est.n_bootstrap + before_zeta = est.zeta_omega + + with pytest.raises(TypeError, match="conley"): + est.set_params(vcov_type="conley") + # Verify nothing mutated + assert est.variance_method == before_variance + assert est.n_bootstrap == before_n_boot + assert est.zeta_omega == before_zeta + + def test_synthetic_did_set_params_conley_kwarg_raises(self): + from diff_diff import SyntheticDiD + + est = SyntheticDiD() + with pytest.raises(TypeError, match="conley"): + est.set_params(conley_cutoff_km=100.0) + # Verify the conley attr stays None (rejected before mutation) + assert getattr(est, "conley_cutoff_km", None) is None + + def test_synthetic_did_get_params_includes_conley_keys(self): + """get_params() / set_params() round-trip must include the inherited + conley_* keys with None values for sklearn-style API consistency + (CI reviewer P2 CQ3).""" + from diff_diff import SyntheticDiD + + est = SyntheticDiD(variance_method="placebo", n_bootstrap=10) + params = est.get_params() + assert "vcov_type" in params and params["vcov_type"] is None + assert "conley_coords" in params and params["conley_coords"] is None + assert "conley_cutoff_km" in params and params["conley_cutoff_km"] is None + assert "conley_metric" in params and params["conley_metric"] is None + assert "conley_kernel" in params and params["conley_kernel"] is None + # Round-trip: passing None values back into set_params is a no-op + est.set_params(**params) + assert est.variance_method == "placebo" + assert est.n_bootstrap == 10 + + +class TestConleySetParamsAtomicity: + """set_params atomicity for Conley fields. Per + feedback_transactional_set_params: invalid multi-kwarg call must not + leave the estimator in a partial state.""" + + def test_unknown_kwarg_raises_no_mutation(self): + from diff_diff import DifferenceInDifferences + + est = DifferenceInDifferences(conley_coords=("lat", "lon"), conley_cutoff_km=100.0) + # Pre-call snapshot + before_cutoff = est.conley_cutoff_km + before_kernel = est.conley_kernel + # set_params with valid + unknown key → must raise & not mutate + with pytest.raises(ValueError, match="Unknown parameter"): + est.set_params(conley_cutoff_km=200.0, garbage_field="x") + # Verify state did NOT change + assert est.conley_cutoff_km == before_cutoff + assert est.conley_kernel == before_kernel + + def test_valid_kwargs_apply(self): + from diff_diff import DifferenceInDifferences + + est = DifferenceInDifferences(conley_coords=("lat", "lon"), conley_cutoff_km=100.0) + est.set_params(conley_cutoff_km=250.0, conley_kernel="uniform") + assert est.conley_cutoff_km == 250.0 + assert est.conley_kernel == "uniform" + + +class TestConleyParityR: + """R conleyreg parity for the Conley spatial HAC implementation. + + Skips when the golden JSON is absent (CI's isolated-install job copies + only tests/, not benchmarks/). Local regeneration: + cd benchmarks/R && Rscript generate_conley_golden.R + """ + + GOLDEN_PATH = "benchmarks/data/r_conleyreg_conley_golden.json" + PARITY_TOL = 1e-6 # Phase 1 success criterion + + @pytest.fixture(scope="class") + def golden(self): + import json + from pathlib import Path + + repo_root = Path(__file__).resolve().parent.parent + path = repo_root / self.GOLDEN_PATH + if not path.exists(): + pytest.skip( + f"Golden JSON not present at {path}; run " + "`cd benchmarks/R && Rscript generate_conley_golden.R` to generate. " + "Requires conleyreg R package + sf/lwgeom + system libs gdal/proj/geos/udunits." + ) + return json.loads(path.read_text()) + + def _check_fixture(self, golden, name): + entry = golden[name] + X = np.asarray(entry["x"], dtype=np.float64).reshape(entry["x_shape"]) + y = np.asarray(entry["y"], dtype=np.float64) + coords = np.asarray(entry["coords"], dtype=np.float64).reshape(entry["coords_shape"]) + vcov_expected = np.asarray(entry["vcov"], dtype=np.float64).reshape(entry["vcov_shape"]) + + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + vcov_got = compute_robust_vcov( + X, + residuals, + vcov_type="conley", + conley_coords=coords, + conley_cutoff_km=entry["cutoff_km"], + conley_metric=entry["metric"], + conley_kernel=entry["kernel"], + ) + np.testing.assert_allclose( + vcov_got, vcov_expected, atol=self.PARITY_TOL, rtol=self.PARITY_TOL + ) + + def test_parity_small_haversine(self, golden): + self._check_fixture(golden, "small_haversine") + + def test_parity_dense_haversine(self, golden): + self._check_fixture(golden, "dense_haversine") + + def test_parity_lat_lon_realistic(self, golden): + self._check_fixture(golden, "lat_lon_realistic") + + +class TestConleyParitySpacetime: + """R conleyreg parity on the Phase 2 block-decomposed panel form. + + Each fixture has lag_cutoff > 0 and exercises the additive sandwich + (within-period spatial + within-unit Bartlett serial). Earth radius + 6371.01 km. Parity target: atol=1e-6. + """ + + GOLDEN_PATH = "benchmarks/data/r_conleyreg_conley_golden.json" + PARITY_TOL = 1e-6 + + @pytest.fixture(scope="class") + def golden(self): + import json + from pathlib import Path + + repo_root = Path(__file__).resolve().parent.parent + path = repo_root / self.GOLDEN_PATH + if not path.exists(): + pytest.skip( + f"Golden JSON not present at {path}; run " + "`cd benchmarks/R && Rscript generate_conley_golden.R` to generate." + ) + return json.loads(path.read_text()) + + def _check_panel_fixture(self, golden, name): + entry = golden[name] + X = np.asarray(entry["x"], dtype=np.float64).reshape(entry["x_shape"]) + y = np.asarray(entry["y"], dtype=np.float64) + coords = np.asarray(entry["coords"], dtype=np.float64).reshape(entry["coords_shape"]) + vcov_expected = np.asarray(entry["vcov"], dtype=np.float64).reshape(entry["vcov_shape"]) + unit = np.asarray(entry["unit"]) + time = np.asarray(entry["time"]) + + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + vcov_got = compute_robust_vcov( + X, + residuals, + vcov_type="conley", + conley_coords=coords, + conley_cutoff_km=entry["cutoff_km"], + conley_metric=entry["metric"], + conley_kernel=entry["kernel"], + conley_time=time, + conley_unit=unit, + conley_lag_cutoff=int(entry["lag_cutoff"]), + ) + np.testing.assert_allclose( + vcov_got, vcov_expected, atol=self.PARITY_TOL, rtol=self.PARITY_TOL + ) + + def test_parity_panel_haversine_lag1(self, golden): + self._check_panel_fixture(golden, "panel_haversine_lag1") + + def test_parity_panel_haversine_lag2(self, golden): + self._check_panel_fixture(golden, "panel_haversine_lag2") + + def test_parity_panel_lat_lon_realistic_lag1(self, golden): + self._check_panel_fixture(golden, "panel_lat_lon_realistic_lag1") + + +class TestConleyReductionsAddendum: + """Additional reduction tests not covered by the helper-direct class. + + Placeholder: the helper-direct class already covers the essential + reductions (HC0 at tiny cutoff, K=ones at huge cutoff, etc.). + Kept here so future test expansions have a clear class to attach to. + """ + + def test_diagonal_of_meat_equals_HC0_contribution(self): + """For any kernel, K(0/h) = 1 so the diagonal contribution to the + meat is exactly the HC0 term Σ_i X_i ε_i² X_i'.""" + rng = np.random.default_rng(seed=9) + n = 20 + coords = rng.uniform(0, 1000, size=(n, 2)) + X = np.column_stack([np.ones(n), rng.standard_normal(n)]) + eps = rng.standard_normal(n) + # Build kernel with cutoff between min and max pairwise distance + D = _pairwise_distance_matrix(coords, "euclidean") + cutoff = float(D[D > 0].min() * 0.001) # ensure off-diagonal kernel is 0 + # With this cutoff, the Bartlett kernel is 1 on the diagonal and 0 off, + # so meat == HC0. + S = X * eps[:, None] + meat_full = S.T @ _bartlett_kernel(D / cutoff) @ S + meat_hc0 = X.T @ (X * (eps**2)[:, None]) + np.testing.assert_allclose(meat_full, meat_hc0, atol=1e-12) + + +# --------------------------------------------------------------------------- +# TestConleyPanelHelper — _compute_conley_vcov with the block-decomposed +# panel path (R conleyreg lag_cutoff > 0 form). +# --------------------------------------------------------------------------- + + +class TestConleyPanelHelper: + """The Phase 2 panel block-decomposed form (matches R conleyreg).""" + + def _panel_fixture(self, n_units=5, T=3, k=2, seed=42, cutoff_km=5000): + rng = np.random.default_rng(seed) + lat_unit = rng.uniform(-30, 30, size=n_units) + lon_unit = rng.uniform(-100, 100, size=n_units) + unit = np.repeat(np.arange(n_units), T) + time = np.tile(np.arange(1, T + 1), n_units) + lat = lat_unit[unit] + lon = lon_unit[unit] + coords = np.column_stack([lat, lon]) + n = n_units * T + X = np.column_stack([np.ones(n)] + [rng.standard_normal(n) for _ in range(k - 1)]) + beta = np.linspace(0.5, 2.0, k) + y = X @ beta + rng.standard_normal(n) * 0.5 + beta_hat, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ beta_hat + bread = X.T @ X + return X, residuals, coords, time, unit, bread, cutoff_km + + def test_T_eq_1_equals_cross_sectional(self): + """Block-decomposed form with T=1 (single time period) should equal + the Phase 1 cross-sectional form on the same data.""" + X, residuals, coords, _, _, bread, cutoff = self._panel_fixture(n_units=8, T=1, k=2) + unit_single = np.arange(X.shape[0]) + time_single = np.ones(X.shape[0], dtype=int) + # Phase 1 cross-sectional + V_cs = _compute_conley_vcov(X, residuals, coords, cutoff, "haversine", "bartlett", bread) + # Phase 2 panel block-decomposed with T=1 and lag_cutoff > 0 + # (the serial component has nothing to do at T=1 — only one period) + V_panel = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time_single, + unit=unit_single, + lag_cutoff=2, + ) + np.testing.assert_allclose(V_panel, V_cs, atol=1e-12) + + def test_lag_cutoff_zero_drops_serial(self): + """lag_cutoff=0 means the serial component contributes nothing; + only the within-period spatial sandwich applies.""" + X, residuals, coords, time, unit, bread, cutoff = self._panel_fixture() + # lag_cutoff=0 + V0 = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time, + unit=unit, + lag_cutoff=0, + ) + # Manually compute the within-period spatial sandwich + S = X * residuals[:, None] + meat_spatial = np.zeros((X.shape[1], X.shape[1])) + for t_val in np.unique(time): + mask = time == t_val + D_t = _pairwise_distance_matrix(coords[mask], "haversine") + K_t = _bartlett_kernel(D_t / cutoff) + meat_spatial += S[mask].T @ K_t @ S[mask] + V_expected = np.linalg.solve(bread, meat_spatial) + V_expected = np.linalg.solve(bread, V_expected.T).T + np.testing.assert_allclose(V0, V_expected, atol=1e-12) + + def test_lag_cutoff_positive_adds_serial(self): + """lag_cutoff > 0 strictly increases the meat by a positive contribution + (off-diagonals from within-unit cross-time pairs).""" + X, residuals, coords, time, unit, bread, cutoff = self._panel_fixture() + V0 = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time, + unit=unit, + lag_cutoff=0, + ) + V1 = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time, + unit=unit, + lag_cutoff=1, + ) + # The serial sandwich adds non-trivial off-diagonal contributions + # for the panel fixture; V1 should differ from V0 + assert not np.allclose( + V0, V1, atol=1e-8 + ), "lag_cutoff=1 must differ from lag_cutoff=0 with a serial component" + + def test_panel_matches_block_decomposed_reference(self): + """Direct verification that _compute_conley_vcov matches the + hand-coded block decomposition from time_dist.cpp at machine precision.""" + X, residuals, coords, time, unit, bread, cutoff = self._panel_fixture(seed=314) + bread_inv = np.linalg.inv(bread) + S = X * residuals[:, None] + # Hand-coded reference (matches R conleyreg per the spike) + for L in (0, 1, 2): + meat = np.zeros((X.shape[1], X.shape[1])) + for t_val in np.unique(time): + mask = time == t_val + D_t = _pairwise_distance_matrix(coords[mask], "haversine") + K_t = _bartlett_kernel(D_t / cutoff) + meat += S[mask].T @ K_t @ S[mask] + if L > 0: + for u_val in np.unique(unit): + mask = unit == u_val + S_u = S[mask] + t_u = time[mask].astype(np.float64) + lag = np.abs(t_u[:, None] - t_u[None, :]) + K_u = ((lag <= L) & (lag != 0)).astype(np.float64) * (1.0 - lag / (L + 1.0)) + meat += S_u.T @ K_u @ S_u + V_ref = bread_inv @ meat @ bread_inv + + V_helper = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time, + unit=unit, + lag_cutoff=L, + ) + np.testing.assert_allclose(V_helper, V_ref, atol=1e-12) + + def test_time_label_normalization_non_unit_spaced_int(self): + """Year-like int labels (2020, 2021, 2022) and YYYYMM labels + (202011, 202012, 202101) produce the same vcov as the equivalent + dense codes (0, 1, 2). Closes Codex P1: `conley_lag_cutoff` is a + count of panel periods, not raw label difference.""" + X, residuals, coords, _, unit, bread, cutoff = self._panel_fixture(n_units=8, T=3, k=2) + time_dense = np.tile([1, 2, 3], 8) + time_years = np.tile([2020, 2021, 2022], 8) + time_yyyymm = np.tile([202011, 202012, 202101], 8) + V_dense = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time_dense, + unit=unit, + lag_cutoff=1, + ) + for time_alt in (time_years, time_yyyymm): + V_alt = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time_alt, + unit=unit, + lag_cutoff=1, + ) + np.testing.assert_allclose(V_alt, V_dense, atol=1e-12) + + def test_time_label_normalization_datetime64(self): + """datetime64 time labels normalize to dense codes via np.unique.""" + X, residuals, coords, _, unit, bread, cutoff = self._panel_fixture(n_units=6, T=3, k=2) + time_dense = np.tile([0, 1, 2], 6) + time_dt = np.tile( + np.array(["2024-01-01", "2024-04-01", "2024-08-01"], dtype="datetime64[D]"), + 6, + ) + V_dense = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time_dense, + unit=unit, + lag_cutoff=1, + ) + V_dt = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time_dt, + unit=unit, + lag_cutoff=1, + ) + np.testing.assert_allclose(V_dt, V_dense, atol=1e-12) + + def test_serial_kernel_bartlett_hardcoded_even_when_kernel_uniform(self): + """conleyreg::time_dist hardcodes Bartlett-style temporal kernel + regardless of the user's `kernel` choice. We mirror that asymmetry.""" + # Two panels: same data, one bartlett spatial, one uniform spatial. + # The serial contribution should be IDENTICAL because the temporal + # kernel is Bartlett-hardcoded. + X, residuals, coords, time, unit, bread, cutoff = self._panel_fixture() + V_bartlett_L0 = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time, + unit=unit, + lag_cutoff=0, + ) + V_bartlett_L2 = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + time=time, + unit=unit, + lag_cutoff=2, + ) + V_uniform_L0 = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "uniform", + bread, + time=time, + unit=unit, + lag_cutoff=0, + ) + V_uniform_L2 = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "uniform", + bread, + time=time, + unit=unit, + lag_cutoff=2, + ) + # The serial delta should be the same regardless of spatial kernel. + # Convert vcov back to meat: meat = bread @ V @ bread + delta_bartlett = bread @ (V_bartlett_L2 - V_bartlett_L0) @ bread + delta_uniform = bread @ (V_uniform_L2 - V_uniform_L0) @ bread + np.testing.assert_allclose(delta_bartlett, delta_uniform, atol=1e-10) - def test_synthetic_did_conley_raises(self): - from diff_diff import SyntheticDiD - with pytest.raises(TypeError, match="conley"): - SyntheticDiD(vcov_type="conley") # type: ignore[call-arg] +# --------------------------------------------------------------------------- +# TestConleySparse — sparse k-d-tree fast path (Wave A item #120). +# --------------------------------------------------------------------------- - def test_synthetic_did_conley_kwarg_raises(self): - from diff_diff import SyntheticDiD - with pytest.raises(TypeError, match="conley"): - SyntheticDiD(conley_cutoff_km=100.0) # type: ignore[call-arg] +class TestConleySparse: + """Sparse k-d-tree fast path for the spatial Bartlett meat. - def test_synthetic_did_set_params_conley_raises(self): - """SyntheticDiD.set_params(vcov_type='conley') must raise (mirrors - __init__'s contract — closes the silent-bypass gap CI reviewer flagged - as P1 CQ1).""" - from diff_diff import SyntheticDiD + The sparse path is gated by three conditions: total n above the + threshold, metric in {"haversine", "euclidean"} (no callable), and + kernel == "bartlett". Each of these tests exercises one of the + gates plus the bit-identity parity claim vs the dense path. + """ - est = SyntheticDiD() - # Snapshot pre-call state - before_variance = est.variance_method - before_n_boot = est.n_bootstrap - before_zeta = est.zeta_omega + def _euclidean_fixture(self, n=1000, k=3, cutoff=15.0, seed=11): + rng = np.random.default_rng(seed) + coords = rng.uniform(0.0, 100.0, size=(n, 2)) + X = np.column_stack([np.ones(n)] + [rng.standard_normal(n) for _ in range(k - 1)]) + beta = np.linspace(0.5, 2.0, k) + y = X @ beta + rng.standard_normal(n) * 0.5 + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + bread = X.T @ X + return X, residuals, coords, bread, cutoff - with pytest.raises(TypeError, match="conley"): - est.set_params(vcov_type="conley") - # Verify nothing mutated - assert est.variance_method == before_variance - assert est.n_bootstrap == before_n_boot - assert est.zeta_omega == before_zeta + def _haversine_fixture(self, n=1000, k=3, cutoff_km=500.0, seed=13): + rng = np.random.default_rng(seed) + lats = rng.uniform(-30.0, 30.0, size=n) + lons = rng.uniform(-100.0, 100.0, size=n) + coords = np.column_stack([lats, lons]) + X = np.column_stack([np.ones(n)] + [rng.standard_normal(n) for _ in range(k - 1)]) + beta = np.linspace(0.5, 2.0, k) + y = X @ beta + rng.standard_normal(n) * 0.5 + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + bread = X.T @ X + return X, residuals, coords, bread, cutoff_km - def test_synthetic_did_set_params_conley_kwarg_raises(self): - from diff_diff import SyntheticDiD + def test_sparse_vs_dense_bit_identity_euclidean_cross_sectional(self): + """Sparse and dense paths produce the same meat on a 1000-row + euclidean+bartlett fixture (atol=1e-10). Headroom over the ~1e-14 + roundoff in chord-projection (haversine) and matmul ordering.""" + X, residuals, coords, bread, cutoff = self._euclidean_fixture(n=1000) + V_dense = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "euclidean", + "bartlett", + bread, + _conley_sparse=False, + ) + V_sparse = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "euclidean", + "bartlett", + bread, + _conley_sparse=True, + ) + np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10) - est = SyntheticDiD() - with pytest.raises(TypeError, match="conley"): - est.set_params(conley_cutoff_km=100.0) - # Verify the conley attr stays None (rejected before mutation) - assert getattr(est, "conley_cutoff_km", None) is None + def test_sparse_vs_dense_bit_identity_haversine_cross_sectional(self): + """Sparse and dense paths produce the same meat on a 1000-row + haversine+bartlett fixture (atol=1e-10). Haversine adds the + chord-projection roundoff that the tolerance must absorb.""" + X, residuals, coords, bread, cutoff = self._haversine_fixture(n=1000) + V_dense = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + _conley_sparse=False, + ) + V_sparse = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "haversine", + "bartlett", + bread, + _conley_sparse=True, + ) + np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10) - def test_synthetic_did_get_params_includes_conley_keys(self): - """get_params() / set_params() round-trip must include the inherited - conley_* keys with None values for sklearn-style API consistency - (CI reviewer P2 CQ3).""" - from diff_diff import SyntheticDiD + def test_auto_toggle_above_threshold_uses_sparse(self, monkeypatch): + """n > _CONLEY_SPARSE_N_THRESHOLD with bartlett + euclidean must + auto-route through the sparse helper. Verified by spying on the + sparse helper call count.""" + import diff_diff.conley as conley_module - est = SyntheticDiD(variance_method="placebo", n_bootstrap=10) - params = est.get_params() - assert "vcov_type" in params and params["vcov_type"] is None - assert "conley_coords" in params and params["conley_coords"] is None - assert "conley_cutoff_km" in params and params["conley_cutoff_km"] is None - assert "conley_metric" in params and params["conley_metric"] is None - assert "conley_kernel" in params and params["conley_kernel"] is None - # Round-trip: passing None values back into set_params is a no-op - est.set_params(**params) - assert est.variance_method == "placebo" - assert est.n_bootstrap == 10 + X, residuals, coords, bread, cutoff = self._euclidean_fixture( + n=_CONLEY_SPARSE_N_THRESHOLD + 1 + ) + calls = {"n": 0} + orig = conley_module._compute_spatial_bartlett_meat_sparse + + def _spy(*args, **kwargs): + calls["n"] += 1 + return orig(*args, **kwargs) + + monkeypatch.setattr(conley_module, "_compute_spatial_bartlett_meat_sparse", _spy) + _compute_conley_vcov(X, residuals, coords, cutoff, "euclidean", "bartlett", bread) + assert calls["n"] >= 1, "Sparse helper not called when n > threshold." + + def test_auto_toggle_below_threshold_stays_dense(self, monkeypatch): + """n <= _CONLEY_SPARSE_N_THRESHOLD must use the dense path even + when other sparse conditions (bartlett + euclidean) are met.""" + import diff_diff.conley as conley_module + + X, residuals, coords, bread, cutoff = self._euclidean_fixture(n=_CONLEY_SPARSE_N_THRESHOLD) + calls = {"n": 0} + orig = conley_module._compute_spatial_bartlett_meat_sparse + + def _spy(*args, **kwargs): + calls["n"] += 1 + return orig(*args, **kwargs) + + monkeypatch.setattr(conley_module, "_compute_spatial_bartlett_meat_sparse", _spy) + _compute_conley_vcov(X, residuals, coords, cutoff, "euclidean", "bartlett", bread) + assert calls["n"] == 0, "Sparse helper called below threshold." + + def test_auto_toggle_callable_metric_stays_dense(self, monkeypatch): + """A callable conley_metric forces the dense path even at large n — + the kd-tree query needs a vectorizable metric, and callables are + not supported via projection.""" + import diff_diff.conley as conley_module + + X, residuals, coords, bread, cutoff = self._euclidean_fixture( + n=_CONLEY_SPARSE_N_THRESHOLD + 100 + ) + def callable_metric(c1, c2): + diff = c1[:, None, :] - c2[None, :, :] + return np.sqrt(np.sum(diff * diff, axis=-1)) -class TestConleySetParamsAtomicity: - """set_params atomicity for Conley fields. Per - feedback_transactional_set_params: invalid multi-kwarg call must not - leave the estimator in a partial state.""" + calls = {"n": 0} + orig = conley_module._compute_spatial_bartlett_meat_sparse - def test_unknown_kwarg_raises_no_mutation(self): - from diff_diff import DifferenceInDifferences + def _spy(*args, **kwargs): + calls["n"] += 1 + return orig(*args, **kwargs) - est = DifferenceInDifferences(conley_coords=("lat", "lon"), conley_cutoff_km=100.0) - # Pre-call snapshot - before_cutoff = est.conley_cutoff_km - before_kernel = est.conley_kernel - # set_params with valid + unknown key → must raise & not mutate - with pytest.raises(ValueError, match="Unknown parameter"): - est.set_params(conley_cutoff_km=200.0, garbage_field="x") - # Verify state did NOT change - assert est.conley_cutoff_km == before_cutoff - assert est.conley_kernel == before_kernel + monkeypatch.setattr(conley_module, "_compute_spatial_bartlett_meat_sparse", _spy) + _compute_conley_vcov(X, residuals, coords, cutoff, callable_metric, "bartlett", bread) + assert calls["n"] == 0, "Sparse helper called for callable metric." - def test_valid_kwargs_apply(self): - from diff_diff import DifferenceInDifferences + def test_auto_toggle_uniform_kernel_stays_dense(self, monkeypatch): + """uniform kernel forces dense path — bartlett has K(u=1) == 0 which + the sparse path relies on; uniform has K(u=1) == 1 which would + require a closed-interval query semantic the chord projection + cannot reliably preserve.""" + import diff_diff.conley as conley_module - est = DifferenceInDifferences(conley_coords=("lat", "lon"), conley_cutoff_km=100.0) - est.set_params(conley_cutoff_km=250.0, conley_kernel="uniform") - assert est.conley_cutoff_km == 250.0 - assert est.conley_kernel == "uniform" + X, residuals, coords, bread, cutoff = self._euclidean_fixture( + n=_CONLEY_SPARSE_N_THRESHOLD + 100 + ) + calls = {"n": 0} + orig = conley_module._compute_spatial_bartlett_meat_sparse + + def _spy(*args, **kwargs): + calls["n"] += 1 + return orig(*args, **kwargs) + + monkeypatch.setattr(conley_module, "_compute_spatial_bartlett_meat_sparse", _spy) + _compute_conley_vcov(X, residuals, coords, cutoff, "euclidean", "uniform", bread) + assert calls["n"] == 0, "Sparse helper called for uniform kernel." + + def test_force_sparse_with_uniform_raises(self): + """Explicit _conley_sparse=True with uniform kernel raises rather + than silently falling back, so callers see the mismatch.""" + X, residuals, coords, bread, cutoff = self._euclidean_fixture(n=100) + with pytest.raises(ValueError, match="_conley_sparse=True requires"): + _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "euclidean", + "uniform", + bread, + _conley_sparse=True, + ) + def test_force_sparse_with_callable_metric_raises(self): + """Explicit _conley_sparse=True with a callable metric raises.""" + X, residuals, coords, bread, cutoff = self._euclidean_fixture(n=100) -class TestConleyParityR: - """R conleyreg parity for the Conley spatial HAC implementation. + def callable_metric(c1, c2): + diff = c1[:, None, :] - c2[None, :, :] + return np.sqrt(np.sum(diff * diff, axis=-1)) - Skips when the golden JSON is absent (CI's isolated-install job copies - only tests/, not benchmarks/). Local regeneration: - cd benchmarks/R && Rscript generate_conley_golden.R - """ + with pytest.raises(ValueError, match="_conley_sparse=True requires"): + _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + callable_metric, + "bartlett", + bread, + _conley_sparse=True, + ) - GOLDEN_PATH = "benchmarks/data/r_conleyreg_conley_golden.json" - PARITY_TOL = 1e-6 # Phase 1 success criterion + def test_force_dense_with_sparse_eligible_inputs(self): + """_conley_sparse=False overrides the auto-toggle and stays dense + even when n is above the threshold.""" + import diff_diff.conley as conley_module - @pytest.fixture(scope="class") - def golden(self): - import json - from pathlib import Path + X, residuals, coords, bread, cutoff = self._euclidean_fixture( + n=_CONLEY_SPARSE_N_THRESHOLD + 100 + ) + calls = {"n": 0} + orig = conley_module._compute_spatial_bartlett_meat_sparse - repo_root = Path(__file__).resolve().parent.parent - path = repo_root / self.GOLDEN_PATH - if not path.exists(): - pytest.skip( - f"Golden JSON not present at {path}; run " - "`cd benchmarks/R && Rscript generate_conley_golden.R` to generate. " - "Requires conleyreg R package + sf/lwgeom + system libs gdal/proj/geos/udunits." - ) - return json.loads(path.read_text()) + def _spy(*args, **kwargs): + calls["n"] += 1 + return orig(*args, **kwargs) - def _check_fixture(self, golden, name): - entry = golden[name] - X = np.asarray(entry["x"], dtype=np.float64).reshape(entry["x_shape"]) - y = np.asarray(entry["y"], dtype=np.float64) - coords = np.asarray(entry["coords"], dtype=np.float64).reshape(entry["coords_shape"]) - vcov_expected = np.asarray(entry["vcov"], dtype=np.float64).reshape(entry["vcov_shape"]) + # The monkeypatch fixture isn't available here; use plain attribute swap. + conley_module._compute_spatial_bartlett_meat_sparse = _spy + try: + _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "euclidean", + "bartlett", + bread, + _conley_sparse=False, + ) + finally: + conley_module._compute_spatial_bartlett_meat_sparse = orig + assert calls["n"] == 0, "Sparse helper called when _conley_sparse=False." + def test_helper_direct_matches_dense_meat_euclidean(self): + """Call _compute_spatial_bartlett_meat_sparse directly and compare + to the dense matmul S' K S on the same data.""" + X, residuals, coords, _, cutoff = self._euclidean_fixture(n=500) + S = X * residuals[:, None] + D = _pairwise_distance_matrix(coords, "euclidean") + K = _bartlett_kernel(D / cutoff) + meat_dense = S.T @ K @ S + meat_sparse = _compute_spatial_bartlett_meat_sparse(S, coords, cutoff, "euclidean") + np.testing.assert_allclose(meat_sparse, meat_dense, atol=1e-10, rtol=1e-10) + + def test_helper_direct_matches_dense_meat_haversine(self): + """Helper direct match for haversine — exercises the chord-projection + + exact-distance refinement path.""" + X, residuals, coords, _, cutoff = self._haversine_fixture(n=500) + S = X * residuals[:, None] + D = _pairwise_distance_matrix(coords, "haversine") + K = _bartlett_kernel(D / cutoff) + meat_dense = S.T @ K @ S + meat_sparse = _compute_spatial_bartlett_meat_sparse(S, coords, cutoff, "haversine") + np.testing.assert_allclose(meat_sparse, meat_dense, atol=1e-10, rtol=1e-10) + + def test_sparse_haversine_cutoff_above_half_earth_circumference(self): + """Sparse haversine path with conley_cutoff_km > π·R_earth (~20,015 km) + must include all geometrically eligible pairs. Without the arc-radians + clamp, the chord-radius formula 2·sin(arc/2) shrinks for arc > π and + the kd-tree silently drops pairs that still have positive Bartlett + weight. The dense path saturates at π·R via _haversine_km's clip; + the sparse path matches via the clamp. Codex Wave A R1 P0 #1. + """ + rng = np.random.default_rng(seed=101) + n = 200 + lats = rng.uniform(-90.0, 90.0, size=n) + lons = rng.uniform(-180.0, 180.0, size=n) + coords = np.column_stack([lats, lons]) + X = np.column_stack([np.ones(n), rng.standard_normal(n), rng.standard_normal(n)]) + y = X @ np.array([1.0, 2.0, -0.5]) + rng.standard_normal(n) * 0.4 coefs, *_ = np.linalg.lstsq(X, y, rcond=None) residuals = y - X @ coefs - vcov_got = compute_robust_vcov( + bread = X.T @ X + # Cutoff well above half-Earth circumference (~20,015 km). Without + # the clamp, the sparse path drops antipodal pairs and the meat + # diverges from the dense path. + cutoff_km = 25_000.0 # > π·R_earth ≈ 20015 km + V_dense = _compute_conley_vcov( X, residuals, - vcov_type="conley", - conley_coords=coords, - conley_cutoff_km=entry["cutoff_km"], - conley_metric=entry["metric"], - conley_kernel=entry["kernel"], + coords, + cutoff_km, + "haversine", + "bartlett", + bread, + _conley_sparse=False, ) - np.testing.assert_allclose( - vcov_got, vcov_expected, atol=self.PARITY_TOL, rtol=self.PARITY_TOL + V_sparse = _compute_conley_vcov( + X, + residuals, + coords, + cutoff_km, + "haversine", + "bartlett", + bread, + _conley_sparse=True, ) - - def test_parity_small_haversine(self, golden): - self._check_fixture(golden, "small_haversine") - - def test_parity_dense_haversine(self, golden): - self._check_fixture(golden, "dense_haversine") - - def test_parity_lat_lon_realistic(self, golden): - self._check_fixture(golden, "lat_lon_realistic") - - -class TestConleyParitySpacetime: - """R conleyreg parity on the Phase 2 block-decomposed panel form. - - Each fixture has lag_cutoff > 0 and exercises the additive sandwich - (within-period spatial + within-unit Bartlett serial). Earth radius - 6371.01 km. Parity target: atol=1e-6. - """ - - GOLDEN_PATH = "benchmarks/data/r_conleyreg_conley_golden.json" - PARITY_TOL = 1e-6 - - @pytest.fixture(scope="class") - def golden(self): - import json - from pathlib import Path - - repo_root = Path(__file__).resolve().parent.parent - path = repo_root / self.GOLDEN_PATH - if not path.exists(): - pytest.skip( - f"Golden JSON not present at {path}; run " - "`cd benchmarks/R && Rscript generate_conley_golden.R` to generate." + np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10) + + def test_sparse_density_gate_falls_back_to_dense_and_warns(self): + """When neighbor density exceeds the threshold (default 30%), the + sparse helper returns None and the dispatcher falls back to dense. + A UserWarning surfaces the reason so users with large cutoffs aren't + surprised by the "sparse" path materializing a near-dense matrix + and using more memory than dense float64. Codex CI R6 P2. + """ + # Tight cluster of points + large cutoff → 100% density. + rng = np.random.default_rng(seed=151) + n = 100 + coords = rng.uniform(0.0, 10.0, size=(n, 2)) + X = np.column_stack([np.ones(n), rng.standard_normal(n)]) + y = X @ np.array([1.0, 1.5]) + rng.standard_normal(n) * 0.4 + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + bread = X.T @ X + # Cutoff = 1e6 → every pair is within range (density = 100%) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + V_forced_sparse = _compute_conley_vcov( + X, + residuals, + coords, + 1e6, + "euclidean", + "bartlett", + bread, + _conley_sparse=True, + ) + density_warnings = [msg for msg in w if "exceeds threshold" in str(msg.message)] + assert len(density_warnings) == 1, "Expected exactly one density-gate UserWarning" + # Result must equal the dense path (the dispatcher fell back). + V_dense = _compute_conley_vcov( + X, + residuals, + coords, + 1e6, + "euclidean", + "bartlett", + bread, + _conley_sparse=False, + ) + np.testing.assert_allclose(V_forced_sparse, V_dense, atol=1e-10, rtol=1e-10) + + def test_sparse_density_gate_does_not_trigger_below_threshold(self): + """At realistic Conley cutoffs (neighbor density well below 30%), + the sparse path runs normally without the density warning.""" + # 1000 points spread over a wide area; cutoff small relative to span. + rng = np.random.default_rng(seed=157) + n = 1000 + coords = rng.uniform(0.0, 100.0, size=(n, 2)) + X = np.column_stack([np.ones(n), rng.standard_normal(n)]) + y = X @ np.array([1.0, 1.5]) + rng.standard_normal(n) * 0.4 + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + bread = X.T @ X + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _compute_conley_vcov( + X, + residuals, + coords, + 10.0, # small cutoff → sparse density << 30% + "euclidean", + "bartlett", + bread, + _conley_sparse=True, ) - return json.loads(path.read_text()) - - def _check_panel_fixture(self, golden, name): - entry = golden[name] - X = np.asarray(entry["x"], dtype=np.float64).reshape(entry["x_shape"]) - y = np.asarray(entry["y"], dtype=np.float64) - coords = np.asarray(entry["coords"], dtype=np.float64).reshape(entry["coords_shape"]) - vcov_expected = np.asarray(entry["vcov"], dtype=np.float64).reshape(entry["vcov_shape"]) - unit = np.asarray(entry["unit"]) - time = np.asarray(entry["time"]) - + density_warnings = [msg for msg in w if "exceeds threshold" in str(msg.message)] + assert ( + len(density_warnings) == 0 + ), "Density gate should not have triggered at low density" + + def test_sparse_with_cluster_matches_dense(self): + """Sparse + combined cluster kernel matches the dense path bit-for-bit + at atol=1e-10 (cross-sectional). Codex CI R8 P1: load-bearing + sparse+cluster regression that was missing prior.""" + rng = np.random.default_rng(seed=181) + n = 600 + coords = rng.uniform(0.0, 100.0, size=(n, 2)) + cluster_ids = rng.integers(0, 8, size=n) # 8 clusters + X = np.column_stack([np.ones(n), rng.standard_normal(n), rng.standard_normal(n)]) + y = X @ np.array([1.0, 1.5, -0.5]) + rng.standard_normal(n) * 0.4 coefs, *_ = np.linalg.lstsq(X, y, rcond=None) residuals = y - X @ coefs - vcov_got = compute_robust_vcov( + bread = X.T @ X + V_dense = _compute_conley_vcov( X, residuals, - vcov_type="conley", - conley_coords=coords, - conley_cutoff_km=entry["cutoff_km"], - conley_metric=entry["metric"], - conley_kernel=entry["kernel"], - conley_time=time, - conley_unit=unit, - conley_lag_cutoff=int(entry["lag_cutoff"]), + coords, + 15.0, + "euclidean", + "bartlett", + bread, + cluster_ids=cluster_ids, + _conley_sparse=False, ) - np.testing.assert_allclose( - vcov_got, vcov_expected, atol=self.PARITY_TOL, rtol=self.PARITY_TOL + V_sparse = _compute_conley_vcov( + X, + residuals, + coords, + 15.0, + "euclidean", + "bartlett", + bread, + cluster_ids=cluster_ids, + _conley_sparse=True, ) - - def test_parity_panel_haversine_lag1(self, golden): - self._check_panel_fixture(golden, "panel_haversine_lag1") - - def test_parity_panel_haversine_lag2(self, golden): - self._check_panel_fixture(golden, "panel_haversine_lag2") - - def test_parity_panel_lat_lon_realistic_lag1(self, golden): - self._check_panel_fixture(golden, "panel_lat_lon_realistic_lag1") - - -class TestConleyReductionsAddendum: - """Additional reduction tests not covered by the helper-direct class. - - Placeholder: the helper-direct class already covers the essential - reductions (HC0 at tiny cutoff, K=ones at huge cutoff, etc.). - Kept here so future test expansions have a clear class to attach to. - """ - - def test_diagonal_of_meat_equals_HC0_contribution(self): - """For any kernel, K(0/h) = 1 so the diagonal contribution to the - meat is exactly the HC0 term Σ_i X_i ε_i² X_i'.""" - rng = np.random.default_rng(seed=9) - n = 20 - coords = rng.uniform(0, 1000, size=(n, 2)) - X = np.column_stack([np.ones(n), rng.standard_normal(n)]) - eps = rng.standard_normal(n) - # Build kernel with cutoff between min and max pairwise distance - D = _pairwise_distance_matrix(coords, "euclidean") - cutoff = float(D[D > 0].min() * 0.001) # ensure off-diagonal kernel is 0 - # With this cutoff, the Bartlett kernel is 1 on the diagonal and 0 off, - # so meat == HC0. - S = X * eps[:, None] - meat_full = S.T @ _bartlett_kernel(D / cutoff) @ S - meat_hc0 = X.T @ (X * (eps**2)[:, None]) - np.testing.assert_allclose(meat_full, meat_hc0, atol=1e-12) - - -# --------------------------------------------------------------------------- -# TestConleyPanelHelper — _compute_conley_vcov with the block-decomposed -# panel path (R conleyreg lag_cutoff > 0 form). -# --------------------------------------------------------------------------- - - -class TestConleyPanelHelper: - """The Phase 2 panel block-decomposed form (matches R conleyreg).""" - - def _panel_fixture(self, n_units=5, T=3, k=2, seed=42, cutoff_km=5000): - rng = np.random.default_rng(seed) - lat_unit = rng.uniform(-30, 30, size=n_units) - lon_unit = rng.uniform(-100, 100, size=n_units) + np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10) + + def test_sparse_with_cluster_panel_matches_dense(self): + """Sparse + cluster + panel block-decomposed sandwich matches the + dense path on a 3-period panel with time-invariant cluster. + Codex CI R8 P1.""" + rng = np.random.default_rng(seed=191) + n_units = 200 + T = 3 unit = np.repeat(np.arange(n_units), T) - time = np.tile(np.arange(1, T + 1), n_units) - lat = lat_unit[unit] - lon = lon_unit[unit] - coords = np.column_stack([lat, lon]) + time = np.tile(np.arange(T), n_units) n = n_units * T - X = np.column_stack([np.ones(n)] + [rng.standard_normal(n) for _ in range(k - 1)]) - beta = np.linspace(0.5, 2.0, k) - y = X @ beta + rng.standard_normal(n) * 0.5 - beta_hat, *_ = np.linalg.lstsq(X, y, rcond=None) - residuals = y - X @ beta_hat + # Time-invariant coords + cluster (per-unit) + unit_coords = rng.uniform(0.0, 100.0, size=(n_units, 2)) + coords = unit_coords[unit] + cluster_per_unit = rng.integers(0, 5, size=n_units) + cluster_ids = cluster_per_unit[unit] + X = np.column_stack([np.ones(n), rng.standard_normal(n), rng.standard_normal(n)]) + y = X @ np.array([1.0, 1.5, -0.3]) + rng.standard_normal(n) * 0.4 + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs bread = X.T @ X - return X, residuals, coords, time, unit, bread, cutoff_km - - def test_T_eq_1_equals_cross_sectional(self): - """Block-decomposed form with T=1 (single time period) should equal - the Phase 1 cross-sectional form on the same data.""" - X, residuals, coords, _, _, bread, cutoff = self._panel_fixture(n_units=8, T=1, k=2) - unit_single = np.arange(X.shape[0]) - time_single = np.ones(X.shape[0], dtype=int) - # Phase 1 cross-sectional - V_cs = _compute_conley_vcov(X, residuals, coords, cutoff, "haversine", "bartlett", bread) - # Phase 2 panel block-decomposed with T=1 and lag_cutoff > 0 - # (the serial component has nothing to do at T=1 — only one period) - V_panel = _compute_conley_vcov( + V_dense = _compute_conley_vcov( X, residuals, coords, - cutoff, - "haversine", + 15.0, + "euclidean", "bartlett", bread, - time=time_single, - unit=unit_single, - lag_cutoff=2, + time=time, + unit=unit, + lag_cutoff=1, + cluster_ids=cluster_ids, + _conley_sparse=False, ) - np.testing.assert_allclose(V_panel, V_cs, atol=1e-12) - - def test_lag_cutoff_zero_drops_serial(self): - """lag_cutoff=0 means the serial component contributes nothing; - only the within-period spatial sandwich applies.""" - X, residuals, coords, time, unit, bread, cutoff = self._panel_fixture() - # lag_cutoff=0 - V0 = _compute_conley_vcov( + V_sparse = _compute_conley_vcov( X, residuals, coords, - cutoff, - "haversine", + 15.0, + "euclidean", "bartlett", bread, time=time, unit=unit, - lag_cutoff=0, + lag_cutoff=1, + cluster_ids=cluster_ids, + _conley_sparse=True, ) - # Manually compute the within-period spatial sandwich - S = X * residuals[:, None] - meat_spatial = np.zeros((X.shape[1], X.shape[1])) - for t_val in np.unique(time): - mask = time == t_val - D_t = _pairwise_distance_matrix(coords[mask], "haversine") - K_t = _bartlett_kernel(D_t / cutoff) - meat_spatial += S[mask].T @ K_t @ S[mask] - V_expected = np.linalg.solve(bread, meat_spatial) - V_expected = np.linalg.solve(bread, V_expected.T).T - np.testing.assert_allclose(V0, V_expected, atol=1e-12) - - def test_lag_cutoff_positive_adds_serial(self): - """lag_cutoff > 0 strictly increases the meat by a positive contribution - (off-diagonals from within-unit cross-time pairs).""" - X, residuals, coords, time, unit, bread, cutoff = self._panel_fixture() - V0 = _compute_conley_vcov( + np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10) + + def test_sparse_density_gate_cluster_aware(self): + """High global spatial density + many small clusters: within-cluster + density is low, so the sparse path should NOT spuriously fall back + to dense. Tests the cluster-aware refinement of the density gate. + Codex CI R8 P1.""" + # Tight spatial cluster of points → 100% global spatial density, + # but split into many small disjoint clusters so post-mask density + # is well below 30%. + rng = np.random.default_rng(seed=199) + n = 200 + coords = rng.uniform(0.0, 5.0, size=(n, 2)) # tight cluster + # 50 clusters of 4 points each → within-cluster nnz = 4*4 = 16 per + # cluster, total = 50*16 = 800 = 800/(200*200) = 2% density << 30%. + cluster_ids = np.repeat(np.arange(50), 4) + X = np.column_stack([np.ones(n), rng.standard_normal(n)]) + y = X @ np.array([1.0, 1.5]) + rng.standard_normal(n) * 0.4 + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + bread = X.T @ X + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + V_sparse = _compute_conley_vcov( + X, + residuals, + coords, + 1e6, # cutoff > data span → 100% global density + "euclidean", + "bartlett", + bread, + cluster_ids=cluster_ids, + _conley_sparse=True, + ) + density_warnings = [msg for msg in w if "exceeds threshold" in str(msg.message)] + assert len(density_warnings) == 0, ( + "Density gate spuriously fell back to dense even though " + "within-cluster density is low — cluster-aware refinement " + "is not working." + ) + # Result must equal the dense path (sparse path executed correctly) + V_dense = _compute_conley_vcov( X, residuals, coords, - cutoff, + 1e6, + "euclidean", + "bartlett", + bread, + cluster_ids=cluster_ids, + _conley_sparse=False, + ) + np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10) + + def test_sparse_haversine_cutoff_at_exactly_half_earth_circumference(self): + """Cutoff = π·R_earth: chord radius = 2 (sphere diameter); all + pairs are included. Bartlett at u=1 returns 0, so the antipodal + pair contributes zero — but pairs at all other distances + contribute. Sparse and dense paths must agree.""" + rng = np.random.default_rng(seed=103) + n = 150 + lats = rng.uniform(-90.0, 90.0, size=n) + lons = rng.uniform(-180.0, 180.0, size=n) + coords = np.column_stack([lats, lons]) + X = np.column_stack([np.ones(n), rng.standard_normal(n)]) + y = X @ np.array([1.0, 1.5]) + rng.standard_normal(n) * 0.5 + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + bread = X.T @ X + cutoff_km = float(np.pi * _CONLEY_EARTH_RADIUS_KM) # ≈ 20015.16 km + V_dense = _compute_conley_vcov( + X, + residuals, + coords, + cutoff_km, "haversine", "bartlett", bread, + _conley_sparse=False, + ) + V_sparse = _compute_conley_vcov( + X, + residuals, + coords, + cutoff_km, + "haversine", + "bartlett", + bread, + _conley_sparse=True, + ) + np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10) + + def test_panel_block_decomposed_sparse_matches_dense(self): + """Panel block-decomposed sandwich produces the same vcov whether + the spatial component is computed dense or sparse. The serial + component is always dense regardless of the flag.""" + X, residuals, coords, _, cutoff = self._euclidean_fixture(n=900, seed=21) + # Synthetic 3-period panel with 300 units per period + time = np.repeat(np.arange(3), 300) + unit = np.tile(np.arange(300), 3) + bread = X.T @ X + V_dense = _compute_conley_vcov( + X, + residuals, + coords, + cutoff, + "euclidean", + "bartlett", + bread, time=time, unit=unit, - lag_cutoff=0, + lag_cutoff=1, + _conley_sparse=False, ) - V1 = _compute_conley_vcov( + V_sparse = _compute_conley_vcov( X, residuals, coords, cutoff, - "haversine", + "euclidean", "bartlett", bread, time=time, unit=unit, lag_cutoff=1, + _conley_sparse=True, ) - # The serial sandwich adds non-trivial off-diagonal contributions - # for the panel fixture; V1 should differ from V0 - assert not np.allclose( - V0, V1, atol=1e-8 - ), "lag_cutoff=1 must differ from lag_cutoff=0 with a serial component" + np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10) - def test_panel_matches_block_decomposed_reference(self): - """Direct verification that _compute_conley_vcov matches the - hand-coded block decomposition from time_dist.cpp at machine precision.""" - X, residuals, coords, time, unit, bread, cutoff = self._panel_fixture(seed=314) - bread_inv = np.linalg.inv(bread) - S = X * residuals[:, None] - # Hand-coded reference (matches R conleyreg per the spike) - for L in (0, 1, 2): - meat = np.zeros((X.shape[1], X.shape[1])) - for t_val in np.unique(time): - mask = time == t_val - D_t = _pairwise_distance_matrix(coords[mask], "haversine") - K_t = _bartlett_kernel(D_t / cutoff) - meat += S[mask].T @ K_t @ S[mask] - if L > 0: - for u_val in np.unique(unit): - mask = unit == u_val - S_u = S[mask] - t_u = time[mask].astype(np.float64) - lag = np.abs(t_u[:, None] - t_u[None, :]) - K_u = ((lag <= L) & (lag != 0)).astype(np.float64) * (1.0 - lag / (L + 1.0)) - meat += S_u.T @ K_u @ S_u - V_ref = bread_inv @ meat @ bread_inv - V_helper = _compute_conley_vcov( - X, - residuals, - coords, - cutoff, - "haversine", - "bartlett", - bread, - time=time, - unit=unit, - lag_cutoff=L, +class TestConleySparseRParityForced: + """R conleyreg parity at atol=1e-6 with the sparse path FORCED on the + three panel R fixtures (bartlett kernel, haversine metric).""" + + GOLDEN_PATH = "benchmarks/data/r_conleyreg_conley_golden.json" + PARITY_TOL = 1e-6 + + @pytest.fixture(scope="class") + def golden(self): + import json + from pathlib import Path + + repo_root = Path(__file__).resolve().parent.parent + path = repo_root / self.GOLDEN_PATH + if not path.exists(): + pytest.skip( + f"Golden JSON not present at {path}; run " + "`cd benchmarks/R && Rscript generate_conley_golden.R` to generate." ) - np.testing.assert_allclose(V_helper, V_ref, atol=1e-12) + return json.loads(path.read_text()) - def test_time_label_normalization_non_unit_spaced_int(self): - """Year-like int labels (2020, 2021, 2022) and YYYYMM labels - (202011, 202012, 202101) produce the same vcov as the equivalent - dense codes (0, 1, 2). Closes Codex P1: `conley_lag_cutoff` is a - count of panel periods, not raw label difference.""" - X, residuals, coords, _, unit, bread, cutoff = self._panel_fixture(n_units=8, T=3, k=2) - time_dense = np.tile([1, 2, 3], 8) - time_years = np.tile([2020, 2021, 2022], 8) - time_yyyymm = np.tile([202011, 202012, 202101], 8) - V_dense = _compute_conley_vcov( + def _check_panel_forced_sparse(self, golden, name): + entry = golden[name] + # Sparse path requires bartlett kernel; skip if fixture is uniform. + if entry["kernel"] != "bartlett": + pytest.skip(f"Fixture {name!r} is not bartlett; skipped for sparse parity.") + X = np.asarray(entry["x"], dtype=np.float64).reshape(entry["x_shape"]) + y = np.asarray(entry["y"], dtype=np.float64) + coords = np.asarray(entry["coords"], dtype=np.float64).reshape(entry["coords_shape"]) + vcov_expected = np.asarray(entry["vcov"], dtype=np.float64).reshape(entry["vcov_shape"]) + unit = np.asarray(entry["unit"]) + time = np.asarray(entry["time"]) + + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + bread = X.T @ X + vcov_got = _compute_conley_vcov( X, residuals, coords, - cutoff, - "haversine", - "bartlett", + entry["cutoff_km"], + entry["metric"], + entry["kernel"], bread, - time=time_dense, + time=time, unit=unit, - lag_cutoff=1, + lag_cutoff=int(entry["lag_cutoff"]), + _conley_sparse=True, + ) + np.testing.assert_allclose( + vcov_got, vcov_expected, atol=self.PARITY_TOL, rtol=self.PARITY_TOL ) - for time_alt in (time_years, time_yyyymm): - V_alt = _compute_conley_vcov( - X, - residuals, - coords, - cutoff, - "haversine", - "bartlett", - bread, - time=time_alt, - unit=unit, - lag_cutoff=1, - ) - np.testing.assert_allclose(V_alt, V_dense, atol=1e-12) - def test_time_label_normalization_datetime64(self): - """datetime64 time labels normalize to dense codes via np.unique.""" - X, residuals, coords, _, unit, bread, cutoff = self._panel_fixture(n_units=6, T=3, k=2) - time_dense = np.tile([0, 1, 2], 6) - time_dt = np.tile( - np.array(["2024-01-01", "2024-04-01", "2024-08-01"], dtype="datetime64[D]"), - 6, + def test_sparse_parity_panel_haversine_lag1(self, golden): + self._check_panel_forced_sparse(golden, "panel_haversine_lag1") + + def test_sparse_parity_panel_haversine_lag2(self, golden): + self._check_panel_forced_sparse(golden, "panel_haversine_lag2") + + def test_sparse_parity_panel_lat_lon_realistic_lag1(self, golden): + self._check_panel_forced_sparse(golden, "panel_lat_lon_realistic_lag1") + + +# --------------------------------------------------------------------------- +# TestConleyCluster — combined spatial + cluster product kernel (Wave A #119). +# --------------------------------------------------------------------------- + + +class TestConleyCluster: + """Combined spatial + cluster product kernel: K(d_ij/h) * 1{c_i = c_j}. + + Wave A item #119. Lifts the prior linalg-level and TWFE-level rejects of + ``vcov_type='conley' + cluster_ids``. The cluster mask multiplies the + spatial kernel on both cross-sectional and panel block-decomposed paths. + On the panel path the validator enforces that cluster membership is + constant within each unit across periods (so the within-unit serial + sandwich's mask is trivially all-ones — no per-unit-time mask needed). + """ + + def _cross_sectional(self, n=24, k=2, seed=11): + rng = np.random.default_rng(seed) + coords = rng.uniform(0.0, 50.0, size=(n, 2)) + X = np.column_stack([np.ones(n)] + [rng.standard_normal(n) for _ in range(k - 1)]) + y = X @ np.array([1.0, 2.0])[:k] + rng.standard_normal(n) * 0.4 + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + bread = X.T @ X + return X, residuals, coords, bread + + def test_cross_sectional_cluster_no_longer_raises(self): + """compute_robust_vcov + vcov_type='conley' + cluster_ids no longer + raises (was the linalg validator's NotImplementedError).""" + X, residuals, coords, _ = self._cross_sectional() + cluster_ids = np.arange(X.shape[0]) // 4 + V = compute_robust_vcov( + X, + residuals, + cluster_ids=cluster_ids, + vcov_type="conley", + conley_coords=coords, + conley_cutoff_km=20.0, ) - V_dense = _compute_conley_vcov( + assert V.shape == (X.shape[1], X.shape[1]) + assert np.all(np.isfinite(V)) + + def test_combined_kernel_matches_hadamard_dense(self): + """The combined kernel matches the explicit Hadamard + ``K_space * cluster_mask`` on the same data.""" + X, residuals, coords, bread = self._cross_sectional(n=30, seed=7) + cluster_ids = np.array([i % 4 for i in range(X.shape[0])]) + cutoff = 15.0 + V_helper = _compute_conley_vcov( X, residuals, coords, cutoff, - "haversine", + "euclidean", "bartlett", bread, - time=time_dense, - unit=unit, - lag_cutoff=1, + cluster_ids=cluster_ids, ) - V_dt = _compute_conley_vcov( + S = X * residuals[:, None] + D = _pairwise_distance_matrix(coords, "euclidean") + K = _bartlett_kernel(D / cutoff) * (cluster_ids[:, None] == cluster_ids[None, :]) + meat = S.T @ K @ S + bread_inv = np.linalg.inv(bread) + V_manual = bread_inv @ meat @ bread_inv + np.testing.assert_allclose(V_helper, V_manual, atol=1e-12) + + def test_combined_kernel_reduces_to_hc0_when_all_unique_clusters(self): + """Every observation in its own cluster → cluster_mask is the identity, + so the meat reduces to the diagonal HC0 contribution.""" + X, residuals, coords, bread = self._cross_sectional(n=20, seed=13) + cluster_ids = np.arange(X.shape[0]) # all unique → cluster_mask = I + V_combined = _compute_conley_vcov( X, residuals, coords, - cutoff, - "haversine", + 10.0, + "euclidean", "bartlett", bread, - time=time_dt, - unit=unit, - lag_cutoff=1, + cluster_ids=cluster_ids, ) - np.testing.assert_allclose(V_dt, V_dense, atol=1e-12) - - def test_serial_kernel_bartlett_hardcoded_even_when_kernel_uniform(self): - """conleyreg::time_dist hardcodes Bartlett-style temporal kernel - regardless of the user's `kernel` choice. We mirror that asymmetry.""" - # Two panels: same data, one bartlett spatial, one uniform spatial. - # The serial contribution should be IDENTICAL because the temporal - # kernel is Bartlett-hardcoded. - X, residuals, coords, time, unit, bread, cutoff = self._panel_fixture() - V_bartlett_L0 = _compute_conley_vcov( + # Manual HC0 + S = X * residuals[:, None] + meat_hc0 = X.T @ (X * (residuals**2)[:, None]) + bread_inv = np.linalg.inv(bread) + V_hc0 = bread_inv @ meat_hc0 @ bread_inv + np.testing.assert_allclose(V_combined, V_hc0, atol=1e-12) + del S + + def test_combined_kernel_reduces_to_pure_cluster_at_huge_cutoff(self): + """Cutoff so large that K_space is identically 1 → combined kernel + reduces to the pure within-cluster sum (cluster mask alone). Uses + the UNIFORM kernel: K_uniform(u) = 1 for |u| <= 1, so at huge_cutoff + the kernel is exactly 1 on every in-range pair (i.e. all pairs). + Bartlett would give `1 - d_ij/cutoff < 1` for all pairs with d > 0, + so the reduction is only asymptotic, not exact (codex CI R6 P3). + """ + X, residuals, coords, bread = self._cross_sectional(n=24, seed=19) + cluster_ids = np.array([i // 3 for i in range(X.shape[0])]) + huge_cutoff = 1e9 # K_space (uniform) = 1 on every pair + V_combined = _compute_conley_vcov( + X, + residuals, + coords, + huge_cutoff, + "euclidean", + "uniform", + bread, + cluster_ids=cluster_ids, + ) + # Manual pure-cluster meat + S = X * residuals[:, None] + K_cluster = (cluster_ids[:, None] == cluster_ids[None, :]).astype(np.float64) + meat = S.T @ K_cluster @ S + bread_inv = np.linalg.inv(bread) + V_expected = bread_inv @ meat @ bread_inv + np.testing.assert_allclose(V_combined, V_expected, atol=1e-12) + + def test_combined_kernel_panel_serial_unchanged_when_cluster_per_unit(self): + """When cluster is constant within unit, the SERIAL component of the + panel sandwich is identical to the no-cluster case (the within-unit + cluster mask is trivially all-ones). Only the spatial component + differs.""" + rng = np.random.default_rng(seed=23) + n_units = 6 + T = 3 + unit = np.repeat(np.arange(n_units), T) + time = np.tile(np.arange(T), n_units) + n = n_units * T + coords = np.column_stack([rng.uniform(-10, 10, size=n), rng.uniform(-10, 10, size=n)]) + X = np.column_stack([np.ones(n), rng.standard_normal(n)]) + y = X @ np.array([1.0, 1.5]) + rng.standard_normal(n) * 0.3 + coefs, *_ = np.linalg.lstsq(X, y, rcond=None) + residuals = y - X @ coefs + bread = X.T @ X + # Time-invariant cluster: one cluster per unit (cluster_per_unit) + cluster_per_unit = np.repeat(rng.integers(0, 3, size=n_units), T) + cutoff = 8.0 + # Two variants: lag=1 with and without cluster + V_no_cluster_l0 = _compute_conley_vcov( X, residuals, coords, cutoff, - "haversine", + "euclidean", "bartlett", bread, time=time, unit=unit, lag_cutoff=0, ) - V_bartlett_L2 = _compute_conley_vcov( + V_no_cluster_l1 = _compute_conley_vcov( X, residuals, coords, cutoff, - "haversine", + "euclidean", "bartlett", bread, time=time, unit=unit, - lag_cutoff=2, + lag_cutoff=1, ) - V_uniform_L0 = _compute_conley_vcov( + V_cluster_l0 = _compute_conley_vcov( X, residuals, coords, cutoff, - "haversine", - "uniform", + "euclidean", + "bartlett", bread, time=time, unit=unit, lag_cutoff=0, + cluster_ids=cluster_per_unit, ) - V_uniform_L2 = _compute_conley_vcov( + V_cluster_l1 = _compute_conley_vcov( X, residuals, coords, cutoff, - "haversine", - "uniform", + "euclidean", + "bartlett", bread, time=time, unit=unit, - lag_cutoff=2, + lag_cutoff=1, + cluster_ids=cluster_per_unit, ) - # The serial delta should be the same regardless of spatial kernel. - # Convert vcov back to meat: meat = bread @ V @ bread - delta_bartlett = bread @ (V_bartlett_L2 - V_bartlett_L0) @ bread - delta_uniform = bread @ (V_uniform_L2 - V_uniform_L0) @ bread - np.testing.assert_allclose(delta_bartlett, delta_uniform, atol=1e-10) + # Serial delta should be identical under cluster vs no-cluster — the + # within-unit mask is all-ones when cluster is constant within unit. + delta_no_cluster = bread @ (V_no_cluster_l1 - V_no_cluster_l0) @ bread + delta_cluster = bread @ (V_cluster_l1 - V_cluster_l0) @ bread + np.testing.assert_allclose(delta_cluster, delta_no_cluster, atol=1e-10) + + def test_panel_time_varying_cluster_raises(self): + """Panel block-decomposed path with a cluster that varies across + periods within a unit raises ValueError naming the violating units.""" + rng = np.random.default_rng(seed=29) + n_units = 4 + T = 3 + unit = np.repeat(np.arange(n_units), T) + time = np.tile(np.arange(T), n_units) + n = n_units * T + coords = np.column_stack([rng.uniform(-10, 10, size=n), rng.uniform(-10, 10, size=n)]) + # Unit 1 changes cluster from 0 -> 1 -> 1 across periods + cluster_ids = np.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2]) + with pytest.raises(ValueError, match="constant within each unit"): + _validate_conley_kwargs( + coords=coords, + cutoff=10.0, + metric="euclidean", + kernel="bartlett", + n=n, + time=time, + unit=unit, + lag_cutoff=1, + cluster_ids=cluster_ids, + ) + + def test_cross_sectional_time_varying_cluster_ok(self): + """Cross-sectional path (no time/unit/lag_cutoff) has NO time- + invariance constraint — the validator should accept any cluster.""" + X, _, coords, _ = self._cross_sectional(n=20, seed=31) + cluster_ids = np.arange(X.shape[0]) % 3 + # Should not raise + _validate_conley_kwargs( + coords=coords, + cutoff=10.0, + metric="euclidean", + kernel="bartlett", + n=X.shape[0], + cluster_ids=cluster_ids, + ) + + def test_cluster_wrong_shape_raises(self): + X, _, coords, _ = self._cross_sectional(n=15) + with pytest.raises(ValueError, match="cluster_ids must be a 1-D array"): + _validate_conley_kwargs( + coords=coords, + cutoff=10.0, + metric="euclidean", + kernel="bartlett", + n=15, + cluster_ids=np.zeros((10,)), + ) + + def test_cluster_nan_raises(self): + X, _, coords, _ = self._cross_sectional(n=10) + cluster_ids = np.array([0, 0, 1, 1, np.nan, 2, 2, 2, 0, 1], dtype=object) + with pytest.raises(ValueError, match="cluster_ids contains NaN"): + _validate_conley_kwargs( + coords=coords, + cutoff=10.0, + metric="euclidean", + kernel="bartlett", + n=10, + cluster_ids=cluster_ids, + ) + + def test_twfe_explicit_cluster_propagates_to_cluster_name(self): + """TWFE + Conley + explicit cluster= sets res.cluster_name to + the user's column AND to_dict()['cluster_name'] reflects it.""" + from diff_diff import TwoWayFixedEffects + + rng = np.random.default_rng(seed=37) + rows = [] + n_units = 10 + for u in range(n_units): + treated = u >= 5 + lat = rng.uniform(-5, 5) + lon = rng.uniform(-5, 5) + region = u // 5 # time-invariant within unit + for t in range(2): + effect = 1.0 if (treated and t == 1) else 0.0 + yv = effect + rng.normal(0, 0.5) + rows.append( + { + "unit": u, + "time": t, + "y": yv, + "treated": int(treated), + "lat": lat, + "lon": lon, + "region": region, + } + ) + import pandas as _pd + + df = _pd.DataFrame(rows) + res = TwoWayFixedEffects( + vcov_type="conley", + cluster="region", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(df, outcome="y", treatment="treated", time="time", unit="unit") + assert res.cluster_name == "region" + d = res.to_dict() + assert d.get("cluster_name") == "region" + + def _multi_period_panel_with_region(self, n_units=12, T=4, seed=41): + """Multi-period panel with a time-invariant `region` column for + combined-kernel estimator tests.""" + import pandas as _pd + + rng = np.random.default_rng(seed=seed) + rows = [] + for u in range(n_units): + treated = u >= n_units // 2 + lat = rng.uniform(-30, 30) + lon = rng.uniform(-100, 100) + region = u // 3 # time-invariant within unit; spans multiple units + for t in range(T): + effect = 1.0 if (treated and t >= T // 2) else 0.0 + yv = 0.2 * t + effect + rng.normal(0, 0.3) + rows.append( + { + "unit": u, + "time": t, + "y": yv, + "treated": int(treated), + "lat": lat, + "lon": lon, + "region": region, + } + ) + return _pd.DataFrame(rows) + + def test_did_combined_kernel_finite_se_and_cluster_name(self): + """DifferenceInDifferences(vcov_type='conley', cluster='region') on + a 2-period panel produces a finite SE, propagates `region` to + res.cluster_name and to_dict(), and differs from the no-cluster + baseline (combined kernel zeros out cross-cluster off-diagonals).""" + from diff_diff import DifferenceInDifferences + + df = self._multi_period_panel_with_region(n_units=12, T=2, seed=43) + kwargs = dict( + vcov_type="conley", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ) + res_combined = DifferenceInDifferences(cluster="region", **kwargs).fit( + df, outcome="y", treatment="treated", time="time", unit="unit" + ) + res_bare = DifferenceInDifferences(**kwargs).fit( + df, outcome="y", treatment="treated", time="time", unit="unit" + ) + assert np.isfinite(res_combined.att) + assert np.isfinite(res_combined.se) and res_combined.se > 0 + assert res_combined.cluster_name == "region" + d = res_combined.to_dict() + assert d.get("cluster_name") == "region" + # Combined kernel zeros out off-cluster pairs → SE differs from bare + assert not np.isclose(res_combined.se, res_bare.se, atol=1e-8) + + def test_did_combined_kernel_time_varying_cluster_raises(self): + """DiD + Conley + cluster= on the panel block-decomposed path + must raise when the cluster column varies across periods within a + unit (time-invariance contract). Codex CI R1 P1 #2.""" + from diff_diff import DifferenceInDifferences + + df = self._multi_period_panel_with_region(n_units=10, T=2, seed=47) + # Make region time-varying for unit 0 (different region in t=1) + mask_u0_t1 = (df["unit"] == 0) & (df["time"] == 1) + df.loc[mask_u0_t1, "region"] = 99 + with pytest.raises(ValueError, match="constant within each unit"): + DifferenceInDifferences( + vcov_type="conley", + cluster="region", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit(df, outcome="y", treatment="treated", time="time", unit="unit") + + def test_mpd_combined_kernel_finite_se_and_cluster_name(self): + """MultiPeriodDiD(vcov_type='conley', cluster='region') on a 4-period + panel produces a finite SE and propagates `region` to cluster_name + on the result + to_dict().""" + from diff_diff import MultiPeriodDiD + + df = self._multi_period_panel_with_region(n_units=12, T=4, seed=53) + res = MultiPeriodDiD( + vcov_type="conley", + cluster="region", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit( + df, + outcome="y", + treatment="treated", + time="time", + unit="unit", + post_periods=[2, 3], + reference_period=0, + ) + assert np.isfinite(res.avg_att) + assert np.isfinite(res.avg_se) and res.avg_se > 0 + assert res.cluster_name == "region" + d = res.to_dict() + assert d.get("cluster_name") == "region" + + def test_mpd_combined_kernel_time_varying_cluster_raises(self): + """MultiPeriodDiD + Conley + cluster= with a cluster that + varies across periods within a unit raises ValueError (same time- + invariance contract as the linalg validator). Codex CI R1 P1 #2.""" + from diff_diff import MultiPeriodDiD + + df = self._multi_period_panel_with_region(n_units=10, T=3, seed=59) + mask_violator = (df["unit"] == 2) & (df["time"] == 2) + df.loc[mask_violator, "region"] = 77 + with pytest.raises(ValueError, match="constant within each unit"): + MultiPeriodDiD( + vcov_type="conley", + cluster="region", + conley_coords=("lat", "lon"), + conley_cutoff_km=2000.0, + conley_lag_cutoff=1, + ).fit( + df, + outcome="y", + treatment="treated", + time="time", + unit="unit", + post_periods=[1, 2], + reference_period=0, + )