Skip to content

perf(callaway): stop materializing per-cell unit-label arrays in influence_func_info#640

Open
igerber wants to merge 1 commit into
mainfrom
perf/cs-lazy-unit-label-arrays
Open

perf(callaway): stop materializing per-cell unit-label arrays in influence_func_info#640
igerber wants to merge 1 commit into
mainfrom
perf/cs-lazy-unit-label-arrays

Conversation

@igerber

@igerber igerber commented Jul 7, 2026

Copy link
Copy Markdown
Owner

Summary

  • Every (g,t) cell's internal influence-function record carried treated_units / control_units label arrays (all_units[positions] — two O(n_control) fancy-index allocations per cell) that no in-package code path ever read. Producers (CallawaySantAnna panel + vectorized + RCS + reference cells, StaggeredTripleDifference) stop materializing them; labels remain recoverable as all_units[treated_idx].
  • The only consumers were the precomputed-less fallbacks of the combined-IF assembly (staggered_aggregation.py general-path local branch) and the multiplier bootstrap (staggered_bootstrap.py) — unreachable from any in-package fit (all callers thread precomputed / global unit ids; verified across staggered.py, staggered_triple_diff.py, efficient_did.py). Those fallbacks now raise an actionable ValueError when reached without label arrays instead of a bare KeyError (no-silent-failures posture).
  • Closes the CS-scaling Mid TODO row (per-cell label arrays). Internal dict only — the _influence_func_info whitebox surface keeps treated_idx/control_idx/treated_inf/control_inf; no public API or numerical change.

Methodology references (required if estimator / math changes)

  • Method name(s): Callaway & Sant'Anna (2021) group-time ATT — no formula change (internal record slimming; the combined-IF fast-path contract and duplicate-free index invariant are unchanged)
  • Paper / source link(s): REGISTRY.md § CallawaySantAnna (combined-IF assembly Note)
  • Any intentional deviations from the source (and why): None

Validation

  • Tests added/updated: none needed — the frozen-parity oracle (TestCombinedIFFrozenParity, which byte-copies the pre-rewrite path) and all direct-call tests thread global unit ids, so they exercise the unchanged branches; tests/test_staggered_aggregation.py, tests/test_staggered.py, tests/test_se_accuracy.py, tests/test_methodology_staggered_triple_diff.py, tests/test_survey_phase5.py (399 tests) plus tests/test_methodology_callaway.py + tests/test_csdid_ported.py (144) all pass locally.
  • Backtest / simulation / notebook evidence (if applicable): N/A (allocation-only change; outputs unchanged by construction — the removed arrays were write-only in-package).

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. No unmitigated P0/P1 findings.

Executive Summary

  • Affected paths: CallawaySantAnna, StaggeredTripleDifference, combined-IF aggregation, and multiplier bootstrap internals.
  • Methodology check passes: removed unit-label arrays are not part of the ATT, WIF, SE, or bootstrap formulas.
  • In-package callers still pass precomputed / global unit IDs, so numerical paths remain unchanged.
  • New fallback behavior fails closed with actionable ValueError, not silent wrong output.
  • No new inline inference or partial NaN guard anti-patterns found.
  • One P3 documentation/test suggestion only.

Methodology

  • Severity: None
  • Impact: No estimator equation, weighting, variance/SE, identification assumption, or public default behavior appears changed. The retained treated_idx / control_idx and IF arrays still satisfy the REGISTRY combined-IF/WIF contract.
  • Concrete fix: N/A.

Code Quality

  • Severity: None
  • Impact: The new guards in diff_diff/staggered_aggregation.py:L508-L530 and diff_diff/staggered_bootstrap.py:L212-L231 are explicit and preferable to the prior implicit KeyError.
  • Concrete fix: N/A.

Performance

  • Severity: None
  • Impact: The PR removes two per-cell label-array allocations while preserving the index-based fast path. This is consistent with the documented CS-scaling TODO removal.
  • Concrete fix: N/A.

Maintainability

  • Severity: P3
  • Impact: docs/methodology/REGISTRY.md:L599-L614 still describes the general fallback as preserved for absent precomputed structures; the changelog documents the new “labels required or pass global IDs” behavior, but the registry note could be clearer for future reviewers.
  • Concrete fix: Add one sentence to the combined-IF note: fallback without precomputed / global IDs requires per-cell treated_units / control_units label arrays.

Tech Debt

  • Severity: None
  • Impact: The TODO row for per-cell unit-label arrays was removed, matching the implemented allocation reduction.
  • Concrete fix: N/A.

Security

  • Severity: None
  • Impact: No secrets, external input handling, filesystem, network, or serialization risk introduced.
  • Concrete fix: N/A.

Documentation/Tests

  • Severity: P3
  • Impact: Existing fast-path tests cover the in-package path, but the new direct-caller ValueError guards are not directly regression-tested.
  • Concrete fix: Add small unit tests asserting the no-precomputed / no-label cases raise ValueError for _compute_combined_influence_function() and _run_multiplier_bootstrap().

Tests not run; this was a static review against the diff and surrounding source.

@igerber igerber added the ready-for-ci Triggers CI test workflows label Jul 7, 2026
@igerber igerber force-pushed the perf/cs-lazy-unit-label-arrays branch from 65cf161 to 9186e3d Compare July 7, 2026 21:06
@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown

🔁 AI review rerun (requested by @igerber)

Head SHA: 9186e3db0f64b5dae2b2dbf9df2c694e73a08319


Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

  • Affected methods: CallawaySantAnna, StaggeredTripleDifference, combined-IF aggregation, and multiplier bootstrap internals.
  • Methodology check passes: removed unit-label arrays are not part of ATT, WIF, SE, or bootstrap formulas.
  • In-package callers still pass precomputed / global unit IDs, so numerical paths remain unchanged.
  • New fallback behavior fails explicitly with ValueError, not silent wrong output.
  • Prior P3 documentation/test suggestions remain non-blocking.

Methodology

  • Severity: None
  • Impact: No estimator equation, weighting rule, variance/SE formula, assumption check, or default behavior changes. The registry’s combined-IF/WIF and bootstrap contracts remain satisfied (docs/methodology/REGISTRY.md:L596-L616, docs/methodology/REGISTRY.md:L2695-L2738). Producers still retain treated_idx / control_idx and IF arrays (diff_diff/staggered.py:L1101-L1121, diff_diff/staggered.py:L3992-L4001, diff_diff/staggered_triple_diff.py:L504-L509).
  • Concrete fix: N/A.

Code Quality

  • Severity: None
  • Impact: The direct-caller fallback branches now raise clear ValueErrors when label arrays are absent (diff_diff/staggered_aggregation.py:L456-L466, diff_diff/staggered_bootstrap.py:L221-L230). No new inline inference or partial NaN guard anti-pattern was introduced.
  • Concrete fix: N/A.

Performance

  • Severity: None
  • Impact: Removes two per-cell unit-label fancy-index allocations while preserving index-based IF assembly.
  • Concrete fix: N/A.

Maintainability

  • Severity: P3
  • Impact: Prior P3 remains: REGISTRY.md still says the general path remains the fallback for absent precomputed structures, but does not mention that absent-precomputed fallback now requires treated_units / control_units label arrays (docs/methodology/REGISTRY.md:L612-L616). The changelog documents it, so this is informational only.
  • Concrete fix: Add one registry sentence: fallback without precomputed / global IDs requires per-cell label arrays.

Tech Debt

  • Severity: None
  • Impact: The TODO row for per-cell label-array allocation was removed consistently with the implemented change (TODO.md:L48-L60).
  • Concrete fix: N/A.

Security

  • Severity: None
  • Impact: No secrets, filesystem/network behavior, serialization, or external input handling risk introduced.
  • Concrete fix: N/A.

Documentation/Tests

  • Severity: P3
  • Impact: Prior P3 remains: the new ValueError guard paths are not directly regression-tested. This is a direct-caller/internal fallback surface, not a numerical correctness blocker.
  • Concrete fix: Add small tests asserting missing label arrays raise for _compute_combined_influence_function() and _run_multiplier_bootstrap() when called without precomputed / global IDs.

Tests not run; this was a static review of the diff and surrounding source.

…uence_func_info

Every (g,t) cell's internal IF record carried treated_units /
control_units label arrays (all_units[positions], two O(n_control)
allocations per cell) that no in-package code path ever read — the only
consumers were the precomputed-less fallbacks of the combined-IF
assembly and the multiplier bootstrap, which every in-package caller
bypasses by threading precomputed structures. Labels stay recoverable
as all_units[treated_idx]. The two fallbacks now raise an actionable
ValueError when reached without label arrays instead of silently
KeyErroring. Applies to CallawaySantAnna (panel, vectorized, RCS,
reference cells) and StaggeredTripleDifference producers.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@igerber igerber force-pushed the perf/cs-lazy-unit-label-arrays branch from 9186e3d to a8ca356 Compare July 7, 2026 23:20
@github-actions

github-actions Bot commented Jul 8, 2026

Copy link
Copy Markdown

🔁 AI review rerun (requested by @igerber)

Head SHA: a8ca356583f851b187aecccf1c2fedbdcb601b65


Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

  • Re-review scope: prior review had no P0/P1 findings; the prior P3 documentation/test suggestions remain informational.
  • Affected methods: CallawaySantAnna, StaggeredTripleDifference, combined-IF aggregation, and multiplier bootstrap internals.
  • Methodology check passes: removed unit-label arrays are not part of ATT, WIF, aggregation, SE, or bootstrap formulas.
  • In-package callers still pass precomputed / global indexing, so the numerical paths continue to use treated_idx / control_idx plus IF values.
  • New fallback behavior fails explicitly with ValueError instead of silently producing wrong output or raising a bare KeyError.

Methodology

  • Severity: None
  • Impact: No estimator equation, weighting rule, identification assumption, variance/SE formula, or default behavior changed. The registry’s combined-IF/WIF and bootstrap contracts remain the relevant source: docs/methodology/REGISTRY.md:L602-L622, docs/methodology/REGISTRY.md:L2701-L2744. Producers still retain index arrays and IF values: diff_diff/staggered.py:L1101-L1121, diff_diff/staggered.py:L1297-L1305, diff_diff/staggered.py:L1723-L1734, diff_diff/staggered.py:L2596-L2608, diff_diff/staggered.py:L3995-L4001, diff_diff/staggered_triple_diff.py:L504-L509.
  • Concrete fix: N/A.

Code Quality

  • Severity: None
  • Impact: The changed fallback branches now raise clear ValueErrors when direct callers omit both precomputed / global IDs and per-cell label arrays: diff_diff/staggered_aggregation.py:L445-L466, diff_diff/staggered_bootstrap.py:L213-L231. No new inline inference or partial NaN-guard anti-pattern was introduced.
  • Concrete fix: N/A.

Performance

  • Severity: None
  • Impact: The PR removes two per-cell label-array fancy-index allocations while preserving index-based IF assembly. This is a memory-allocation cleanup with no numerical change.
  • Concrete fix: N/A.

Maintainability

  • Severity: P3-informational
  • Impact: REGISTRY.md still says the general combined-IF path remains the fallback for absent precomputed structures, but does not mention that this local fallback now requires per-cell treated_units / control_units label arrays unless global IDs are supplied: docs/methodology/REGISTRY.md:L618-L622. The changelog documents the behavior at CHANGELOG.md:L457-L467, so this is not blocking.
  • Concrete fix: Add one registry sentence near docs/methodology/REGISTRY.md:L618-L622 stating that absent-precomputed local fallback requires label arrays, or callers must pass global indexing.

Tech Debt

  • Severity: None
  • Impact: The completed TODO row for per-cell treated_units / control_units allocations was removed consistently: TODO.md:L44-L60.
  • Concrete fix: N/A.

Security

  • Severity: None
  • Impact: No secrets, external input handling, filesystem/network behavior, serialization, or security-sensitive surface changed.
  • Concrete fix: N/A.

Documentation/Tests

  • Severity: P3-informational
  • Impact: The new ValueError guard paths are not directly regression-tested. This is an internal/direct-caller fallback surface, not a numerical correctness blocker.
  • Concrete fix: Add small tests asserting missing label arrays raise for _compute_combined_influence_function() without precomputed / global IDs and _run_multiplier_bootstrap() without precomputed.

Tests not run; this was a static re-review of the diff and surrounding source.

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