Address PR #363 R6 review (1 P2 + 1 P3)

R6 P2 (event-study opt-out overhead): cband=False on the weighted
event-study path no longer allocates the stacked (G, H) IF matrix or
forces per-horizon IF return on the weights= shortcut. Split two
flags internally:

  needs_per_horizon_if = survey= path OR (weights= AND cband=True)
  needs_stacked_if_matrix = weights= AND cband=True  (alias for
                           weighted_es AND cband)

- Psi allocation gated on needs_stacked_if_matrix.
- _fit_continuous force_return_influence gated on
  (needs_stacked_if_matrix AND resolved_survey_unit_full is None) —
  under survey= path, _fit_continuous returns the IF anyway via its
  resolved_survey_unit gate, so no extra cost.
- _fit_mass_point_2sls return_influence gated on
  needs_per_horizon_if — survey= path needs the per-horizon IF for
  the Binder-TSL override regardless of cband.

Net effect: cband=False + weights= shortcut + weighted_es skips the
O(GH) Psi allocation and the per-horizon IF work entirely. cband=True
paths and survey= paths unchanged.

R6 P3 (event-study survey= integration coverage): added two end-to-end
integration tests for the previously-unguarded positive-path
estimator-level survey= + aggregate='event_study' dispatch:

- test_survey_event_study_continuous_end_to_end: continuous_at_zero +
  SurveyDesign(strata='stratum') — asserts variance_formula=
  'survey_binder_tsl', survey_metadata.df_survey=G-n_strata, cband_*
  populated, PSU dispatch through _aggregate_unit_resolved_survey.
- test_survey_event_study_mass_point_end_to_end: mass_point +
  SurveyDesign(strata=...) — asserts variance_formula=
  'survey_binder_tsl_2sls' and that the 2SLS IF flows through
  per-horizon Binder-TSL + sup-t bootstrap.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: igerber

diff_diff/had.py

@@ -4127,11 +4127,18 @@ def _fit_event_study(
         # carries the full-design effective_n / n_psu / etc.).
         n_obs_arr = np.full(n_horizons, n_units, dtype=np.int64)
 
-        # Per-horizon IF matrix on the weighted path (shape (G, H)); drives
-        # both per-horizon Binder-TSL variance (already composed inside
-        # ``_fit_continuous`` for continuous, or explicitly below for
-        # mass-point) AND the shared-PSU multiplier bootstrap for sup-t.
-        if weighted_es:
+        # Two IF-consumption flags (review R6 P2): the PER-HORIZON IF is
+        # needed when the survey= path composes Binder-TSL variance (via
+        # compute_survey_if_variance inside _fit_continuous or the
+        # mass-point override below); the STACKED (G, H) IF matrix is
+        # needed only when the sup-t multiplier bootstrap runs
+        # (``cband=True`` on the weighted path). Splitting them avoids
+        # allocating / filling Psi on the common opt-out path
+        # ``cband=False`` + weights= shortcut, where no IF consumer
+        # exists.
+        needs_per_horizon_if = resolved_survey_unit_full is not None or (weighted_es and cband)
+        needs_stacked_if_matrix = weighted_es and cband
+        if needs_stacked_if_matrix:
             Psi = np.full((G_full, n_horizons), np.nan, dtype=np.float64)
         else:
             Psi = np.zeros((0, 0), dtype=np.float64)  # sentinel, not used
@@ -4173,23 +4180,28 @@ def _fit_event_study(
                     d_lower_val,
                     weights_arr=weights_unit_full,
                     resolved_survey_unit=resolved_survey_unit_full,
-                    # Force IF return on the weighted event-study path
-                    # (needed for the sup-t bootstrap). Does NOT change
-                    # the per-horizon SE formula — that still follows
-                    # the static-path convention (Binder-TSL under
-                    # survey=, bc_fit.se_robust under weights= shortcut).
-                    force_return_influence=weighted_es,
+                    # Force IF return only when the sup-t bootstrap
+                    # needs the stacked matrix AND the survey= gate
+                    # won't already produce it. Under survey= path,
+                    # _fit_continuous returns the IF automatically
+                    # (resolved_survey_unit_full != None); under the
+                    # weights= shortcut + cband=True, force it here;
+                    # otherwise skip the O(G) IF work (review R6 P2).
+                    force_return_influence=(
+                        needs_stacked_if_matrix and resolved_survey_unit_full is None
+                    ),
                 )
                 if bc_fits is not None:
                     bc_fits.append(bc_fit_e)
                 if bw_diags is not None:
                     bw_diags.append(bw_diag_e)
-                # Collect per-unit IF on β̂-scale (psi_bc / den) so the
-                # sup-t bootstrap operates on the same θ̂-scale IF that
-                # the analytical variance sees. Per continuous-path
-                # construction in _fit_continuous, bc_fit.influence_function
-                # is the numerator IF; dividing by |den| yields the β̂ IF.
-                if weighted_es and bc_fit_e is not None and bc_fit_e.influence_function is not None:
+                # Collect per-unit IF on β̂-scale (psi_bc / den) into
+                # Psi ONLY when the sup-t bootstrap will consume it.
+                if (
+                    needs_stacked_if_matrix
+                    and bc_fit_e is not None
+                    and bc_fit_e.influence_function is not None
+                ):
                     if resolved_design == "continuous_at_zero":
                         den_e = float(np.average(d_arr_full, weights=weights_unit_full))
                     else:
@@ -4209,7 +4221,12 @@ def _fit_event_study(
                     cluster_arr,
                     vcov_requested,
                     weights=weights_unit_full,
-                    return_influence=resolved_survey_unit_full is not None or weighted_es,
+                    # Return IF only when a consumer exists: survey=
+                    # path needs it for per-horizon Binder-TSL override;
+                    # weights= shortcut + cband=True needs it for the
+                    # bootstrap. weights= shortcut + cband=False skips
+                    # IF computation entirely (review R6 P2).
+                    return_influence=needs_per_horizon_if,
                 )
                 # Survey path: override analytical sandwich SE with
                 # Binder-TSL via compute_survey_if_variance (matches
@@ -4222,7 +4239,7 @@ def _fit_event_study(
                         se_e = float(np.sqrt(v_survey))
                     else:
                         se_e = float("nan")
-                if weighted_es and psi_e is not None:
+                if needs_stacked_if_matrix and psi_e is not None:
                     Psi[:, i] = psi_e
             else:
                 raise ValueError(f"Internal error: unhandled design={resolved_design!r}.")

tests/test_had.py

@@ -5351,6 +5351,108 @@ def test_mass_point_default_vcov_event_study_cband_rejected(self):
                     cband=True,
                 )
 
+    def test_survey_event_study_continuous_end_to_end(self):
+        """Review R6 P3: estimator-level
+        ``fit(aggregate='event_study', survey=SurveyDesign(...))``
+        integration lock for the continuous path. Verifies
+        variance_formula, survey_metadata.df_survey (t-inference path),
+        cband_* population, and stratified PSU dispatch through
+        _aggregate_unit_resolved_survey."""
+        from diff_diff.survey import SurveyDesign
+
+        rng = np.random.default_rng(70)
+        G, T, n_strata = 200, 4, 4
+        d_post = rng.uniform(0.0, 1.0, G)
+        strata_per_unit = np.repeat(np.arange(n_strata), G // n_strata)
+        rng.shuffle(strata_per_unit)
+        rows = []
+        for t in range(T):
+            for g in range(G):
+                dose = d_post[g] if t == T - 1 else 0.0
+                y = 0.2 * t + (2.0 * dose if t == T - 1 else 0.0) + 0.5 * rng.standard_normal()
+                rows.append((g, t, dose, y, strata_per_unit[g]))
+        panel = pd.DataFrame(
+            rows,
+            columns=["unit", "period", "dose", "outcome", "stratum"],
+        )
+        w_unit = 1.0 + 0.3 * np.abs(rng.standard_normal(G))
+        panel["w"] = panel["unit"].map(lambda g: w_unit[g])
+        sd = SurveyDesign(weights="w", strata="stratum")
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", UserWarning)
+            est = HeterogeneousAdoptionDiD(design="continuous_at_zero", seed=0, n_bootstrap=200)
+            r = est.fit(
+                panel,
+                "outcome",
+                "dose",
+                "period",
+                "unit",
+                aggregate="event_study",
+                survey=sd,
+            )
+        assert r.variance_formula == "survey_binder_tsl"
+        assert r.survey_metadata is not None
+        assert r.survey_metadata.n_strata == n_strata
+        assert r.survey_metadata.n_psu == G
+        assert r.survey_metadata.df_survey == G - n_strata
+        assert r.cband_crit_value is not None and np.isfinite(r.cband_crit_value)
+        assert r.cband_method == "multiplier_bootstrap"
+        assert r.cband_n_bootstrap == 200
+        assert r.cband_low is not None and r.cband_high is not None
+        assert np.all(np.isfinite(r.se))
+
+    def test_survey_event_study_mass_point_end_to_end(self):
+        """Review R6 P3: estimator-level
+        ``fit(design='mass_point', aggregate='event_study',
+        survey=...)`` integration lock. Verifies
+        variance_formula='survey_binder_tsl_2sls' and that the
+        weighted 2SLS IF flows correctly through per-horizon
+        Binder-TSL + sup-t bootstrap."""
+        from diff_diff.survey import SurveyDesign
+
+        rng = np.random.default_rng(71)
+        G, T = 200, 4
+        d_mp = np.concatenate([np.full(40, 0.3), rng.uniform(0.3, 1.0, G - 40)])
+        rng.shuffle(d_mp)
+        strata_per_unit = np.repeat(np.arange(4), G // 4)
+        rng.shuffle(strata_per_unit)
+        rows = []
+        for t in range(T):
+            for g in range(G):
+                dose = d_mp[g] if t == T - 1 else 0.0
+                y = 0.2 * t + (2.0 * dose if t == T - 1 else 0.0) + 0.5 * rng.standard_normal()
+                rows.append((g, t, dose, y, strata_per_unit[g]))
+        panel = pd.DataFrame(
+            rows,
+            columns=["unit", "period", "dose", "outcome", "stratum"],
+        )
+        w_unit = 1.0 + 0.3 * np.abs(rng.standard_normal(G))
+        panel["w"] = panel["unit"].map(lambda g: w_unit[g])
+        sd = SurveyDesign(weights="w", strata="stratum")
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", UserWarning)
+            est = HeterogeneousAdoptionDiD(
+                design="mass_point",
+                vcov_type="hc1",
+                seed=0,
+                n_bootstrap=200,
+            )
+            r = est.fit(
+                panel,
+                "outcome",
+                "dose",
+                "period",
+                "unit",
+                aggregate="event_study",
+                survey=sd,
+            )
+        assert r.variance_formula == "survey_binder_tsl_2sls"
+        assert r.survey_metadata is not None
+        assert r.survey_metadata.n_strata == 4
+        assert r.cband_crit_value is not None and np.isfinite(r.cband_crit_value)
+        assert r.cband_method == "multiplier_bootstrap"
+        assert np.all(np.isfinite(r.se))
+
     def test_mass_point_default_vcov_robust_true_survey_allowed(self):
         """Complement: robust=True on the default path resolves to
         hc1, so the survey= mass-point fit is allowed with no explicit