perf(spillover): sparse cKDTree branch for staggered nearest-treated distances

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` 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)
+  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

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_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:

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,
+        )