From 0ca36921df7bd43eac42c481c7961915cb00f586 Mon Sep 17 00:00:00 2001 From: igerber Date: Tue, 7 Jul 2026 17:53:44 -0400 Subject: [PATCH 1/2] perf(spillover): sparse cKDTree branch for staggered nearest-treated distances MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The staggered cohort loop always built a dense (n_units, n_treated_by_onset) distance matrix per cohort. It now dispatches per cohort to the same cKDTree helper the static path uses (auto-activated past _CONLEY_SPARSE_N_THRESHOLD, built-in metrics only, cutoff_km = the outermost ring edge / _effective_d_bar). Within-cutoff distances are exact (the helper recomputes the true metric for in-range matches); beyond-cutoff units get inf — semantics-preserving because every staggered d_it consumer (ring membership, S_it, the far-away check, the event-study d_bar trigger) compares against thresholds at or below that cutoff. Helper- and fit-level equality tests pin the sparse arm against dense (atol 1e-12 end-to-end, trigger array identical). Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 10 ++++ TODO.md | 1 - diff_diff/spillover.py | 46 +++++++++++++++--- tests/test_spillover.py | 105 +++++++++++++++++++++++++++++++++++++--- 4 files changed, 149 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 09b45f9c..a3e0ff09 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -489,6 +489,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 h=1 clamp-parity test). The symbol is imported independently (mixed-version safe — a stale extension degrades HC2 to NumPy without disabling older Rust accelerations). `return_dof` / weighted / CR2-BM requests stay on NumPy (CR2-BM tracked in TODO). +- **`SpilloverDiD` staggered nearest-treated distances gain the sparse cKDTree branch.** + The staggered cohort loop always built a dense `(n_units, n_treated_by_onset)` distance + matrix per cohort; it now dispatches per cohort to the same cKDTree helper the static + path uses (auto-activated when `n_units` exceeds the sparse threshold, built-in metrics + only, `cutoff_km = ` the outermost ring edge). Within-cutoff distances are exact (the + helper recomputes the true great-circle/planar metric for in-range matches) and + beyond-cutoff units get `inf` — semantics-preserving because every staggered `d_it` + consumer (ring membership, `S_it`, the far-away check, the event-study `d_bar` trigger) + compares against thresholds at or below that cutoff. Helper- and fit-level equality + tests pin the sparse arm against the dense path (atol 1e-12 end-to-end). - **`CallawaySantAnna` per-(g,t) IF scatters converted from `np.add.at` to fancy `+=`** (`staggered.py::_cluster_robust_se_from_per_gt_if` — runs once per (g,t) cell when `cluster=` is set — and the general combined-IF assembly path in diff --git a/TODO.md b/TODO.md index 97444378..111702c6 100644 --- a/TODO.md +++ b/TODO.md @@ -49,7 +49,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low | | Rust-backend CR2 Bell-McCaffrey: falls through to NumPy (the leverage/Satterthwaite-DOF path needs `return_dof` support, which the Rust vcov dispatch excludes). The one-way HC2 kernel landed 2026-07-07 (`compute_robust_vcov_hc2`, mirrors the NumPy hc2 branch at ~1e-15; near-singular hat-diagonal sentinel + Python-side warn-and-HC1-fallback). | `rust/src/linalg.rs` | Phase 1a | Mid | Low | | Wild cluster bootstrap CI inversion calls `_t_star(r)` ~O(100) times, each materializing a fresh `(B×n)` `y_star` + `(k×B)` refit + `(n×B)` residual arrays. Acceptable for the few-cluster regime; for large-`n`/large-`B`, chunk `_t_star` over draws or precompute the `r`-independent cluster-level pieces (restricted residuals are linear in `r`). | `utils.py::wild_bootstrap_se._t_star` | #543 | Mid | Low | -| `SpilloverDiD` sparse cKDTree path for the staggered nearest-treated-distance helper (mirrors the static helper's sparse branch). `_compute_nearest_treated_distance_staggered` always builds dense `(n_units, n_treated_by_onset)` matrices per cohort; add a sparse branch gated on `n > _CONLEY_SPARSE_N_THRESHOLD`. | `spillover.py` | Wave B | Mid | Low | | `HeterogeneousAdoptionDiD` Phase 3 Stute: Appendix-D vectorized form replaces the per-iteration OLS refit with a single precomputed `M = I - X(X'X)^{-1}X'` applied to `eps*eta` (~2× faster, functionally identical). Shipped the literal-refit form to match paper text. | `had_pretests.py::stute_test` | Phase 3 | Mid | Low | | Migrate `spillover._iterative_fe_subset` onto the shared `diff_diff.utils._iterative_fe_solve` (same Gauss-Seidel-on-codes recursion, specialized to a masked Butts subsample; ImputationDiD/TwoStageDiD already route through the shared helper — this takes the FE-solver copy count from 2 to 1). Preserve the SpilloverDiD positive-weight/NaN-FE REGISTRY contract and `_FE_ITER_MAX=100` budget (or align it to 10k with a REGISTRY note). | `spillover.py` | demean-modernization | Mid | Low | | Adopt `snap_absorbed_regressors` (FE-spanned regressor two-stage snap + LSMR confirmation) on the ImputationDiD lead-indicator path — lead columns are the most plausible FE-spanned regressors, and the truncated-MAP-iterate exposure documented at `utils.py` applies; today only `solve_ols` rank detection guards it. Behavior change beyond the demean-modernization refactor, so deferred from that PR. | `imputation.py::_compute_lead_coefficients` | demean-modernization | Mid | Low | diff --git a/diff_diff/spillover.py b/diff_diff/spillover.py index cb343f81..f22423e4 100644 --- a/diff_diff/spillover.py +++ b/diff_diff/spillover.py @@ -341,6 +341,7 @@ def _compute_nearest_treated_distance_staggered( metric: SpilloverMetric, first_treat_by_unit: Dict[Any, Any], d_bar: Optional[float] = None, + cutoff_km: Optional[float] = None, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]: """Return per-row nearest-treated distance for the staggered case. @@ -372,11 +373,20 @@ def _compute_nearest_treated_distance_staggered( duplicate cohort pass on the event-study path (PR #456 R6 performance fix). - Notes - ----- - The staggered helper currently always uses dense pairwise distance per - cohort. A sparse cKDTree branch (mirroring the static helper) is queued - as a follow-up — see TODO.md. + cutoff_km : float, optional + When set, ``n_units > _CONLEY_SPARSE_N_THRESHOLD``, and the metric + is a built-in string, each cohort's nearest-treated distances are + computed via the sparse cKDTree helper + (:func:`_compute_nearest_treated_distance_sparse`) instead of the + dense (n_units × n_treated_by_onset) matrix. Units with no treated + neighbor within ``cutoff_km`` for a cohort get ``inf`` for that + cohort — identical downstream semantics to the dense path because + every consumer of the staggered ``d_it`` compares against + thresholds ≤ the outermost ring edge (ring membership, ``S_it``, + the far-away check, and the ``d_bar`` trigger), so the caller + passes ``cutoff_km = _effective_d_bar``. Within-cutoff distances + are exact (the sparse helper recomputes the true metric for + in-range matches). Returns ------- @@ -432,7 +442,25 @@ def _compute_nearest_treated_distance_staggered( continue treated_coords = all_coords[treated_positions] # Compute per-unit nearest distance to this cohort's treated set. - dists_to_cohort = _pairwise_ring_distances(all_coords, treated_coords, metric).min(axis=1) + # Sparse cKDTree branch (mirrors the static helper's dispatch at + # _compute_nearest_treated_distance_static): per-cohort tree on the + # treated-by-onset subset, exact metric recomputed for in-range + # matches; beyond-cutoff units get inf (see the cutoff_km doc above). + if ( + cutoff_km is not None + and len(unit_index) > _CONLEY_SPARSE_N_THRESHOLD + and metric in ("haversine", "euclidean") + ): + dists_to_cohort = _compute_nearest_treated_distance_sparse( + all_coords=all_coords, + treated_coords=treated_coords, + metric=metric, # type: ignore[arg-type] + cutoff_km=float(cutoff_km), + ) + else: + dists_to_cohort = _pairwise_ring_distances(all_coords, treated_coords, metric).min( + axis=1 + ) # Update rows whose period t >= onset: take min of current d_it and the # newly-available cohort distance. affected_rows = row_time >= onset @@ -2483,6 +2511,12 @@ def fit( metric=self.conley_metric, first_treat_by_unit=effective_onsets, d_bar=self._effective_d_bar if self.event_study else None, + # Sparse cKDTree auto-activates past the threshold for + # built-in metrics; every staggered d_it consumer + # compares against thresholds <= _effective_d_bar, so + # beyond-cutoff inf is semantics-preserving (mirrors + # the static call below). + cutoff_km=self._effective_d_bar, ) ) else: diff --git a/tests/test_spillover.py b/tests/test_spillover.py index 7f8ee611..eab32793 100644 --- a/tests/test_spillover.py +++ b/tests/test_spillover.py @@ -8856,9 +8856,7 @@ def force_drop(A, **kwargs): warnings.simplefilter("always") result = est.fit(df, outcome="y", unit="unit", time="time", treatment="D") msgs = [str(w.message) for w in caught] - assert any( - "SpilloverDiD Wave D bread" in m and "rank-deficient" in m for m in msgs - ), msgs + assert any("SpilloverDiD Wave D bread" in m and "rank-deficient" in m for m in msgs), msgs assert est.is_fitted_ and np.isfinite(result.att) def test_dropped_ring_coefficient_propagates_nan_inference(self): @@ -8895,11 +8893,106 @@ def drop_last(A, **kwargs): eff = res.spillover_effects nan_rows = eff[np.isnan(eff["se"])] fin_rows = eff[np.isfinite(eff["se"]) & (eff["se"] > 0)] - assert len(nan_rows) >= 1, ( - f"a dropped Wave D coord should NaN a ring SE; got {eff['se'].tolist()}" - ) + assert ( + len(nan_rows) >= 1 + ), f"a dropped Wave D coord should NaN a ring SE; got {eff['se'].tolist()}" assert len(fin_rows) >= 1, "identified rings should keep finite SE" # The NaN-SE ring's FULL inference must be NaN, not just se. for _, r in nan_rows.iterrows(): assert np.isnan(r["t_stat"]) and np.isnan(r["p_value"]) assert np.isnan(r["ci_low"]) and np.isnan(r["ci_high"]) + + +class TestStaggeredSparseKDTreeBranch: + """The staggered cohort loop's sparse cKDTree branch (activated when + ``cutoff_km`` is set, ``n_units > _CONLEY_SPARSE_N_THRESHOLD``, and the + metric is a built-in string) must reproduce the dense path exactly for + every within-cutoff distance, and produce identical fits end-to-end — + every staggered ``d_it`` consumer compares against thresholds <= + ``_effective_d_bar``, so beyond-cutoff ``inf`` is semantics-preserving. + Mirrors the static helper's sparse branch (auto-activated via the same + threshold) and its tests above.""" + + def test_helper_sparse_matches_dense_within_cutoff(self, staggered_panel, monkeypatch): + import diff_diff.spillover as sp + + df, ft = staggered_panel + kwargs = dict( + unit="unit", + time="time", + coords=("lat", "lon"), + metric="haversine", + first_treat_by_unit=ft, + d_bar=1200.0, + ) + d_dense, ru, rt, trig_dense = sp._compute_nearest_treated_distance_staggered(df, **kwargs) + monkeypatch.setattr(sp, "_CONLEY_SPARSE_N_THRESHOLD", 0) + d_sparse, _, _, trig_sparse = sp._compute_nearest_treated_distance_staggered( + df, cutoff_km=1200.0, **kwargs + ) + in_range = d_dense <= 1200.0 * (1 + 1e-6) + np.testing.assert_allclose(d_sparse[in_range], d_dense[in_range], atol=1e-8) + # Beyond-cutoff entries are inf on the sparse path. + assert np.isinf(d_sparse[~in_range]).all() + # The d_bar trigger consumes distances <= d_bar (== cutoff), so it + # must be identical between the paths (NaN pattern included). + np.testing.assert_array_equal(np.isnan(trig_dense), np.isnan(trig_sparse)) + both = ~np.isnan(trig_dense) + np.testing.assert_array_equal(trig_dense[both], trig_sparse[both]) + + def test_fit_sparse_matches_dense_end_to_end(self, monkeypatch): + """Force the sparse branch on a small staggered fit: att, ring + coefficients and SEs must match the dense fit (the within-cutoff + distances are exact; beyond-cutoff rows land in the far-away + control group on both paths).""" + import diff_diff.spillover as sp + + rng = np.random.default_rng(42) + rows = [] + # 3 cohorts of treated units near the origin + controls in/beyond rings. + units = {} + uid = 0 + for k in range(6): # treated: onset staggered 1/2 + units[f"T{uid}"] = (0.05 * k, 0.02 * k, 1 + (k % 2)) + uid += 1 + for k in range(10): # near controls within ~40 km + units[f"C{uid}"] = (0.1 + 0.02 * k, 0.15 + 0.02 * k, np.inf) + uid += 1 + for k in range(8): # far controls (>5 deg away, far outside rings) + units[f"F{uid}"] = (6.0 + 0.1 * k, 6.0, np.inf) + uid += 1 + for u, (lat, lon, ft) in units.items(): + for t in range(4): + treated_now = np.isfinite(ft) and t >= ft + rows.append( + { + "unit": u, + "time": t, + "lat": lat, + "lon": lon, + "first_treat": ft if np.isfinite(ft) else 0, + "y": 1.0 + 0.1 * t + (0.5 if treated_now else 0.0) + rng.normal(0, 0.05), + } + ) + df = pd.DataFrame(rows) + fit_kwargs = dict(outcome="y", unit="unit", time="time", first_treat="first_treat") + + dense = SpilloverDiD(rings=[0.0, 50.0], conley_coords=("lat", "lon")).fit(df, **fit_kwargs) + monkeypatch.setattr(sp, "_CONLEY_SPARSE_N_THRESHOLD", 0) + sparse_res = SpilloverDiD(rings=[0.0, 50.0], conley_coords=("lat", "lon")).fit( + df, **fit_kwargs + ) + + assert sparse_res.is_staggered is True and dense.is_staggered is True + np.testing.assert_allclose(sparse_res.att, dense.att, rtol=0, atol=1e-12) + np.testing.assert_allclose(sparse_res.se, dense.se, rtol=0, atol=1e-12) + # spillover_effects is a per-ring DataFrame; compare its numeric columns. + sp_num = sparse_res.spillover_effects.select_dtypes("number") + de_num = dense.spillover_effects.select_dtypes("number") + np.testing.assert_allclose( + sp_num.to_numpy(dtype=float), + de_num.to_numpy(dtype=float), + rtol=0, + atol=1e-12, + equal_nan=True, + ) From 8e8050249b10da11d705e23d204e7623b026232b Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 8 Jul 2026 07:13:29 -0400 Subject: [PATCH 2/2] docs(changelog): fix cutoff_km backtick typo (review P3) Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a3e0ff09..1a8acc4e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -493,7 +493,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 The staggered cohort loop always built a dense `(n_units, n_treated_by_onset)` distance matrix per cohort; it now dispatches per cohort to the same cKDTree helper the static path uses (auto-activated when `n_units` exceeds the sparse threshold, built-in metrics - only, `cutoff_km = ` the outermost ring edge). Within-cutoff distances are exact (the + only, `cutoff_km` set to the outermost ring edge). Within-cutoff distances are exact (the helper recomputes the true great-circle/planar metric for in-range matches) and beyond-cutoff units get `inf` — semantics-preserving because every staggered `d_it` consumer (ring membership, `S_it`, the far-away check, the event-study `d_bar` trigger)