Skip to content

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

Open
igerber wants to merge 2 commits into
mainfrom
perf/spillover-staggered-sparse-kdtree
Open

perf(spillover): sparse cKDTree branch for staggered nearest-treated distances#645
igerber wants to merge 2 commits into
mainfrom
perf/spillover-staggered-sparse-kdtree

Conversation

@igerber

@igerber igerber commented Jul 7, 2026

Copy link
Copy Markdown
Owner

Summary

  • The staggered cohort loop in _compute_nearest_treated_distance_staggered 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 already uses (_compute_nearest_treated_distance_sparse) — auto-activated when n_units > _CONLEY_SPARSE_N_THRESHOLD, built-in metrics only, with cutoff_km = _effective_d_bar (the outermost ring edge), mirroring the static call site exactly. Closes the Wave-B TODO row.
  • Semantics-preserving by construction: within-cutoff distances are exact (the sparse helper recomputes the true great-circle/planar metric for in-range matches); beyond-cutoff units get inf, and every staggered d_it consumer — ring membership, S_it, the far-away-control check, and the event-study d_bar trigger — compares against thresholds ≤ that cutoff, so inf vs large-finite is indistinguishable downstream (same argument as the documented static-path behavior).
  • New TestStaggeredSparseKDTreeBranch: helper-level equality (within-cutoff distances match dense at 1e-8; beyond-cutoff all inf; the d_bar trigger array identical incl. NaN pattern) and end-to-end fit identity (att/SE/per-ring effects at atol 1e-12) with the threshold monkeypatched to force the sparse arm.

Methodology references (required if estimator / math changes)

  • Method name(s): Butts (2021) ring-exposure construction (distance-computation routing only; ring/exposure/trigger semantics unchanged)
  • Paper / source link(s): REGISTRY.md § SpilloverDiD (ring membership + far-away contract; the static sparse-path convention this mirrors is already documented at the static call site)
  • Any intentional deviations from the source (and why): None

Validation

  • Tests added/updated: tests/test_spillover.py::TestStaggeredSparseKDTreeBranch (2 tests); full tests/test_spillover.py passes (321).
  • Backtest / simulation / notebook evidence (if applicable): N/A

Security / privacy

  • Confirm no secrets/PII in this PR: Yes

Generated with Claude Code

@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown

Overall Assessment: ✅ Looks Good

Executive Summary

  • No P0/P1 methodology or correctness findings.
  • Affected method: SpilloverDiD staggered nearest-treated distance routing in diff_diff/spillover.py.
  • The sparse branch preserves the registry contract: ring membership, S_it, far-away controls, and event-study trigger all compare only against d_bar, which is validated equal to the outer ring edge.
  • This is an implementation choice for distance computation, not a methodology deviation from Butts ring-exposure construction. Butts supports local spillover exposure, distance bins/rings, and staggered settings. (arxiv.org)
  • I could not run the tests locally because pytest is not installed in this environment.

Methodology

  • Severity: P3 informational
  • Impact: The PR affects SpilloverDiD ring-exposure distance routing only: _compute_nearest_treated_distance_staggered now uses the existing sparse helper when cutoff_km=self._effective_d_bar, n_units exceeds the sparse threshold, and the metric is built-in. The estimator’s exposure definition remains aligned with docs/methodology/REGISTRY.md: nearest currently-treated unit, time-varying staggered rings, Omega_0 = {D_it = 0 AND S_it = 0}, and far-away controls defined by d_it > d_bar.
  • Concrete fix: None required. The inf substitution for beyond-cutoff distances is not user-visible under the current contract because all affected downstream checks are threshold comparisons at <= d_bar or > d_bar: diff_diff/spillover.py:L484-L486, diff_diff/spillover.py:L1073-L1130, diff_diff/spillover.py:L1891-L1908, diff_diff/spillover.py:L2604-L2624.

Code Quality

  • No findings.

Performance

  • No findings. The change removes the documented TODO.md performance backlog item for staggered sparse nearest-treated distances and scopes the sparse path to built-in metrics only: diff_diff/spillover.py:L449-L459.

Maintainability

  • No findings. The staggered path reuses the existing static sparse helper rather than adding a second implementation: diff_diff/spillover.py:L234-L332.

Tech Debt

  • No findings. The corresponding backlog row was removed from TODO.md after implementation.

Security

  • No findings. No secrets, external execution, or unsafe input handling changes observed.

Documentation/Tests

  • Severity: P3

  • Impact: Minor changelog typo: cutoff_km = has a stray space/backtick gap, which slightly weakens the release note clarity.

  • Concrete fix: Change cutoff_km = the outermost ring edge to ``cutoff_km = <outermost ring edge>`` or ``cutoff_km set to the outermost ring edgeinCHANGELOG.md:L437-L446`.

  • Severity: P3 informational

  • Impact: New helper-level and end-to-end tests cover sparse vs dense parity for staggered distances and fit outputs: tests/test_spillover.py:L8906-L8998.

  • Concrete fix: None required. I was unable to execute them here because pytest / python -m pytest is unavailable.

@igerber igerber added the ready-for-ci Triggers CI test workflows label Jul 7, 2026
@igerber igerber force-pushed the perf/spillover-staggered-sparse-kdtree branch from 18393e1 to 1fbea09 Compare July 8, 2026 09:58
igerber and others added 2 commits July 8, 2026 07:13
…distances

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 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@igerber igerber force-pushed the perf/spillover-staggered-sparse-kdtree branch from 1fbea09 to 8e80502 Compare July 8, 2026 11:13
@github-actions

github-actions Bot commented Jul 8, 2026

Copy link
Copy Markdown

🔁 AI review rerun (requested by @igerber)

Head SHA: 8e8050249b10da11d705e23d204e7623b026232b


Overall Assessment: ✅ Looks Good

Executive Summary

  • No unmitigated P0/P1 findings.
  • Affected method: SpilloverDiD staggered nearest-treated distance routing.
  • The sparse branch is methodology-neutral: it preserves all threshold comparisons at rings[-1] == d_bar.
  • Prior review changelog wording issue is resolved at CHANGELOG.md:L492-L501.
  • I could not run tests here because pytest and numpy are not installed.

Methodology

  • Severity: P3 informational
    Impact: The PR changes distance computation routing only, not the Butts ring-exposure estimand. The registry requires nearest currently-treated-unit rings, row-level Omega_0 = {D_it = 0 AND S_it = 0}, and event-study spillover triggers based on d(i,j) <= d_bar; the new sparse branch returns exact within-cutoff distances and inf beyond the same cutoff used by those comparisons. See docs/methodology/REGISTRY.md:L3986-L4018, diff_diff/spillover.py:L335-L389, diff_diff/spillover.py:L449-L459, diff_diff/spillover.py:L2505-L2520.
    Concrete fix: None required.

Code Quality

  • No findings. The staggered path reuses the existing sparse helper instead of adding a second implementation.

Performance

  • No findings. This directly addresses the removed TODO backlog item for staggered sparse nearest-treated distances; current surrounding backlog is at TODO.md:L43-L54.

Maintainability

  • No findings. The cutoff invariant is enforced by validation: d_bar defaults to rings[-1], and explicit d_bar must equal rings[-1], so the sparse cutoff cannot drift from the ring boundary in public fit() usage. See diff_diff/spillover.py:L1891-L1908.

Tech Debt

  • No findings.

Security

  • No findings. No secrets, unsafe input handling, networking, or external execution changes observed.

Documentation/Tests

  • Severity: P3 informational
    Impact: Added tests cover helper-level sparse/dense parity, beyond-cutoff inf, trigger equality, and aggregate fit parity at tests/test_spillover.py:L8906-L8998. I could not execute them locally: python -m pytest failed with No module named pytest, and a direct import smoke test failed with No module named numpy.
    Concrete fix: None for the PR; rely on CI or run the targeted class in an environment with test dependencies installed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready-for-ci Triggers CI test workflows

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant