Skip to content

Commit dedc2a6

Browse files
igerberclaude
andcommitted
Gate auto-inject on varying-strata + gate per-period tensors on survey
Public contract fix: when `survey_design.psu` is omitted, fit() auto-injects psu=<group_col> with nest=False. If strata vary within group, the synthesized PSU column reuses group labels across strata and downstream survey resolution fails on the cross-stratum PSU uniqueness check. fit() now detects that combination before auto- inject runs (via the existing _strata_psu_vary_within_group helper) and raises a targeted ValueError pointing users at the explicit `psu=<col>, nest=True` path. Performance fix: the cell-period allocator's per-(g,t) tensor (U_per_period) was being allocated unconditionally in _compute_full_per_group_contributions, _compute_per_group_if_multi_horizon, _compute_per_group_if_placebo_horizon, and their _compute_cohort_recentered_inputs aggregator. That added O(G*T) memory to the main ATT path and O(G*T*L_max) to the event-study / placebo paths on every fit — including non-survey fits that never consume the tensor. Adds a `compute_per_period: bool = True` param through each helper, gates allocation + accumulation behind the flag, and threads `compute_per_period=(_obs_survey_info is not None)` from fit(). Downstream consumers now None-guard the per-period recentering / _compute_se calls so the plug-in SE path runs with no cell-level work. Adds `test_auto_inject_with_varying_strata_raises` under `TestSurveyWithinGroupValidation` covering the targeted error path. Narrows the REGISTRY.md survey IF expansion Note to state the nest=True requirement under varying strata. Softens two stale comments (`_cohort_recenter_per_period` docstring, `_compute_cohort_recentered_inputs` inline comment) so they match the convention framing already in REGISTRY.md and TODO.md rather than claiming methodological correctness. All 352 tests pass (slow MC coverage sim included). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent b14aa22 commit dedc2a6

3 files changed

Lines changed: 165 additions & 54 deletions

File tree

0 commit comments

Comments
 (0)