two-stage-did: explicit vcov_type on convenience fn + full NaN-contract tests (codex P3)

- two_stage_did(): expose vcov_type="hc1" explicitly (was hidden behind **kwargs)
  and forward it, matching the imputation_did/efficient_did sibling wrappers — the
  convenience API surface, generated signature, and IDE help now show the param.
- Degenerate-bootstrap tests now assert the FULL public NaN-propagation contract
  (overall t_stat/p_value/conf_int + every event-study/group inference field) via
  a shared _assert_full_bootstrap_nan helper, not just overall_se, so a partial
  regression in _build_nan_bootstrap_results can't slip through.

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

Author: igerber

diff_diff/two_stage.py

@@ -3361,6 +3361,7 @@ def two_stage_did(
     aggregate: Optional[str] = None,
     balance_e: Optional[int] = None,
     survey_design: object = None,
+    vcov_type: str = "hc1",
     **kwargs,
 ) -> TwoStageDiDResults:
     """
@@ -3392,6 +3393,11 @@ def two_stage_did(
         PSU, and FPC for design-based GMM sandwich variance. Strata enters
         survey df for t-distribution inference.
         Both analytical (n_bootstrap=0) and bootstrap inference are supported.
+    vcov_type : str, default="hc1"
+        Variance estimator family; permanently narrow to ``{"hc1"}`` (the Gardner
+        2022 two-stage GMM cluster-sandwich). Analytical-sandwich families
+        ``{"classical", "hc2", "hc2_bm"}`` and ``"conley"`` are rejected. See
+        :class:`TwoStageDiD`.
     **kwargs
         Additional keyword arguments passed to TwoStageDiD constructor.
 
@@ -3408,7 +3414,7 @@ def two_stage_did(
     ...                         'first_treat', aggregate='event_study')
     >>> results.print_summary()
     """
-    est = TwoStageDiD(**kwargs)
+    est = TwoStageDiD(vcov_type=vcov_type, **kwargs)
     return est.fit(
         data,
         outcome=outcome,

tests/test_two_stage.py

@@ -1979,6 +1979,43 @@ def _eq(a, b):
                 assert _eq(e0[k][f], e1[k][f]), f"{attr}[{k}][{f}] differs"
 
 
+def _assert_full_bootstrap_nan(r):
+    """Assert the FULL public NaN-propagation contract under a degenerate
+    bootstrap (n_clusters<2 / n_psu<2): every overall + per-horizon + per-group
+    inference field is NaN, not just the SE (REGISTRY NaN-inference contract).
+    """
+    # Overall inference fields
+    assert np.isnan(r.overall_se)
+    assert np.isnan(r.overall_t_stat)
+    assert np.isnan(r.overall_p_value)
+    assert all(np.isnan(x) for x in r.overall_conf_int)
+    assert np.isnan(r.coef_var)
+    # Bootstrap payload
+    b = r.bootstrap_results
+    assert np.isnan(b.overall_att_se)
+    assert np.isnan(b.overall_att_p_value)
+    assert all(np.isnan(x) for x in b.overall_att_ci)
+    if b.event_study_ses:
+        assert all(np.isnan(v) for v in b.event_study_ses.values())
+    if b.group_ses:
+        assert all(np.isnan(v) for v in b.group_ses.values())
+    # Per-horizon event-study inference fields (skip reference-period markers,
+    # which carry n_obs == 0 and are not real effects).
+    for eff in (r.event_study_effects or {}).values():
+        if eff.get("n_obs", 1) == 0:
+            continue
+        assert np.isnan(eff["se"])
+        assert np.isnan(eff["t_stat"])
+        assert np.isnan(eff["p_value"])
+        assert all(np.isnan(x) for x in eff["conf_int"])
+    # Per-group inference fields
+    for eff in (r.group_effects or {}).values():
+        assert np.isnan(eff["se"])
+        assert np.isnan(eff["t_stat"])
+        assert np.isnan(eff["p_value"])
+        assert all(np.isnan(x) for x in eff["conf_int"])
+
+
 class TestTwoStageDiDVcovType:
     """Phase 1b interstitial #5 (final): vcov_type input contract on TwoStageDiD.
 
@@ -2164,14 +2201,7 @@ def test_bootstrap_single_cluster_returns_nan(self):
                 first_treat="first_treat",
                 aggregate="all",
             )
-        assert np.isnan(r.overall_se)
-        assert np.isnan(r.coef_var)
-        assert np.isnan(r.bootstrap_results.overall_att_se)
-        assert all(np.isnan(x) for x in r.bootstrap_results.overall_att_ci)
-        if r.bootstrap_results.event_study_ses:
-            assert all(np.isnan(v) for v in r.bootstrap_results.event_study_ses.values())
-        if r.bootstrap_results.group_ses:
-            assert all(np.isnan(v) for v in r.bootstrap_results.group_ses.values())
+        _assert_full_bootstrap_nan(r)
 
     def test_bootstrap_single_psu_survey_returns_nan(self):
         data, _ = _add_survey_cols(generate_test_data(n_units=80, seed=17))
@@ -2188,8 +2218,7 @@ def test_bootstrap_single_psu_survey_returns_nan(self):
                 survey_design=design,
                 aggregate="event_study",
             )
-        assert np.isnan(r.overall_se)
-        assert np.isnan(r.bootstrap_results.overall_att_se)
+        _assert_full_bootstrap_nan(r)
 
     # ---- metadata reflects the post-drop fit sample (codex P2) ----