Merge pull request #420 from igerber/fix-audit-402

Address residual P2s+P3 from re-audit of PR #402

Author: igerber

diff_diff/guides/llms-full.txt

@@ -636,21 +636,26 @@ had.fit(
 ```python
 from diff_diff import HeterogeneousAdoptionDiD, did_had_pretest_workflow
 
-# Vet the testable identifying assumptions first:
+# `aggregate="overall"` (the default) is two-period-only: did_had_pretest_workflow
+# reduces to a single first-difference and HeterogeneousAdoptionDiD.fit hard-rejects
+# panels with more than two periods. `aggregate="event_study"` requires a
+# multi-period panel. Use distinct data objects for the two regimes.
+
+# Vet the testable identifying assumptions on the two-period panel first:
 report = did_had_pretest_workflow(
-    data, outcome_col='y', unit_col='unit', time_col='t',
+    data_2p, outcome_col='y', unit_col='unit', time_col='t',
     dose_col='d', first_treat_col='first_treat')
 print(report.summary())
 
-# Single-period scalar WAS (aggregate="overall" default):
+# Single-period scalar WAS (aggregate="overall" default) on the two-period panel:
 est = HeterogeneousAdoptionDiD()
-results = est.fit(data, outcome_col='y', unit_col='unit',
+results = est.fit(data_2p, outcome_col='y', unit_col='unit',
                   time_col='t', dose_col='d',
                   first_treat_col='first_treat')
 print(results.summary())
 
-# Multi-period per-horizon WAS:
-es = est.fit(data, outcome_col='y', unit_col='unit',
+# Multi-period per-horizon WAS on the multi-period panel:
+es = est.fit(data_mp, outcome_col='y', unit_col='unit',
              time_col='t', dose_col='d',
              first_treat_col='first_treat',
              aggregate='event_study')

diff_diff/practitioner.py

@@ -871,14 +871,18 @@ def _handle_had(results: Any):
                 "workflow still sets pretrends_joint=None, all_pass=False, "
                 "and a 'joint pre-trends skipped (no earlier pre-period)' "
                 "verdict suffix - in that case step 2 stays uncovered "
-                "even on the event-study path. On survey-weighted "
-                "fits (survey_design= / survey= / weights=) the workflow "
-                "skips QUG with a UserWarning (permanent Phase 4.5 C0 "
-                "deferral - extreme order statistics are not smooth "
-                "functionals of the empirical CDF) and returns a "
-                "linearity-conditional verdict only - so step 1 coverage "
-                "is unweighted-only and the reported verdict on weighted "
-                "fits is conditional on QUG holding by assumption. "
+                "even on the event-study path. On supported survey-weighted "
+                "fits (pweight + PSU/FPC under survey_design= / survey= / "
+                "weights=) the workflow skips QUG with a UserWarning "
+                "(permanent Phase 4.5 C0 deferral - extreme order statistics "
+                "are not smooth functionals of the empirical CDF) and returns "
+                "a linearity-conditional verdict only - so step 1 coverage "
+                "is unweighted-only and the reported verdict on supported "
+                "weighted fits is conditional on QUG holding by assumption. "
+                "Stratified (SurveyDesign(strata=...)) and replicate-weight "
+                "(BRR/Fay/JK1/JKn/SDR) designs raise NotImplementedError on "
+                "the linearity kernels and have no pretest workflow path "
+                "yet - deferred to a follow-up. "
                 "Assumptions 3 / 5 / 6 (uniform continuity at the "
                 "boundary, Design 1 sign / WAS_d_lower identification) "
                 "are NOT testable via pre-trends - the workflow vets only "
@@ -1024,16 +1028,22 @@ def _handle_had_event_study(results: Any):
                 "trends_lin=True where the consumed F-2 placebo gets "
                 "dropped), pretrends_joint=None, all_pass=False, and the "
                 "verdict carries 'joint pre-trends skipped (no earlier "
-                "pre-period)' - step 2 stays uncovered. On survey-weighted fits (survey_design= / survey= / "
-                "weights=) the workflow skips QUG with a UserWarning "
-                "(permanent Phase 4.5 C0 deferral) and returns a "
-                "linearity-conditional verdict only - so step 1 coverage "
-                "is unweighted-only on the event-study path too, and the "
-                "weighted verdict is conditional on QUG holding by "
-                "assumption. The joint Stute pre-trends and joint "
+                "pre-period)' - step 2 stays uncovered. On supported "
+                "survey-weighted fits (pweight + PSU/FPC under "
+                "survey_design= / survey= / weights=) the workflow skips "
+                "QUG with a UserWarning (permanent Phase 4.5 C0 deferral) "
+                "and returns a linearity-conditional verdict only - so "
+                "step 1 coverage is unweighted-only on the event-study "
+                "path too, and the weighted verdict is conditional on QUG "
+                "holding by assumption. Stratified (SurveyDesign("
+                "strata=...)) and replicate-weight (BRR/Fay/JK1/JKn/SDR) "
+                "designs raise NotImplementedError on the linearity "
+                "kernels and have no pretest workflow path on the "
+                "event-study path yet - deferred to a follow-up. "
+                "The joint Stute pre-trends and joint "
                 "homogeneity-linearity tests themselves remain available "
-                "under survey weighting via PSU-level Mammen multiplier "
-                "bootstrap."
+                "under supported survey weighting via PSU-level Mammen "
+                "multiplier bootstrap."
             ),
             code=(
                 "from diff_diff import did_had_pretest_workflow\n"

docs/methodology/REGISTRY.md

@@ -2552,8 +2552,8 @@ Shipped in `diff_diff/had_pretests.py` as `stute_joint_pretest()` (residuals-in
 - [x] Phase 3: `did_had_pretest_workflow()` composite helper. Two-period `data`-only entry point (Phase 2a overall-path dispatch); reduces panel via `_aggregate_first_difference` and runs all three IMPLEMENTED tests at a shared `alpha`. `seed` forwards to `stute_test` only (QUG and Yatchew are deterministic). Returns `HADPretestReport` with priority-ordered verdict string. Because Phase 3 ships steps 1 + 3 of the paper's four-step workflow but **not** step 2 (Assumption 7 pre-trends test via Equation 18), the fail-to-reject verdict explicitly flags the Assumption 7 gap rather than claiming unconditional TWFE safety: `"QUG and linearity diagnostics fail-to-reject; Assumption 7 pre-trends test NOT run (paper step 2 deferred to Phase 3 follow-up)"`. Verdict priority follows the paper's one-way rule (TWFE admissible only if NO test rejects): **conclusive rejections are the primary verdict and are NEVER hidden by inconclusive status** — any unresolved-step note is appended via `"; additional steps unresolved: ..."` rather than replacing the rejection. The pure `"inconclusive - QUG NaN"` / `"inconclusive - both Stute and Yatchew linearity tests NaN"` forms only fire when NO conclusive test rejects AND a required step is unresolved. The partial-workflow fail-to-reject verdict may carry a `"(Yatchew NaN - skipped)"` (or Stute) suffix when one linearity test is NaN but the other is conclusive (step 3 resolved via the paper's "Stute OR Yatchew" wording). Bundled rejection-reason strings name each failed assumption in the conclusive-rejection case. `all_pass` is `True` iff QUG is conclusive AND at least one of Stute/Yatchew is conclusive AND no conclusive test rejects. **Non-negative-dose contract**: all three raw linearity helpers (`qug_test`, `stute_test`, `yatchew_hr_test`) raise a front-door `ValueError` on any `d < 0`, mirroring the `_validate_had_panel` guard (paper Section 2 HAD support restriction). Multi-period panels pre-slice to `(F-1, F)` before calling; joint-horizon dispatch deferred to Phase 3 follow-up.
 - [ ] Phase 4: Pierce-Schott (2016) replication harness reproduces Figure 2 values.
 - [ ] Phase 4: Full DGP 1/2/3 coverage-rate reproduction from Table 1.
-- [x] Phase 5 (wave 1, PR #402): `practitioner_next_steps()` integration for HAD results - `_handle_had` and `_handle_had_event_study` route both result classes through HAD-specific Baker et al. (2025) step guidance with bidirectional HAD ↔ ContinuousDiD Step-4 routing closure. The `_check_nan_att` helper extends to ndarray `att` (HAD event-study) via `np.all(np.isnan(arr))` semantics; scalar path bit-exact preserved.
-- [x] Phase 5 (wave 1, PR #402): `llms-full.txt` HeterogeneousAdoptionDiD section + result-class blocks + `## HAD Pretests` index + Choosing-an-Estimator row landed; constructor / fit() signatures match the real API (regression-tested via `inspect.signature`); result-class field tables enumerate every public dataclass field (regression-tested via `dataclasses.fields()`); `llms-practitioner.txt` Step 4 decision tree distinguishes ContinuousDiD (per-dose ATT(d), needs never-treated) from HeterogeneousAdoptionDiD (WAS, universal-rollout-compatible).
+- [x] Phase 5 (wave 1, PR #402): `practitioner_next_steps()` integration for HAD results - `_handle_had` and `_handle_had_event_study` route both result classes through HAD-specific Baker et al. (2025) step guidance with bidirectional HAD ↔ ContinuousDiD Step-4 routing closure. The `_check_nan_att` helper extends to ndarray `att` (HAD event-study) via `np.all(np.isnan(arr))` semantics; scalar path bit-exact preserved. The `llms-full.txt` HAD section's documented constructor and `fit()` parameter lists are regression-locked against `inspect.signature(HeterogeneousAdoptionDiD.__init__)` and `HeterogeneousAdoptionDiD.fit` for parameter-name presence (defaults, type annotations, and return-type unions are not pinned by the current test).
+- [x] Phase 5 (wave 1, PR #402): `llms-full.txt` HeterogeneousAdoptionDiD section + result-class blocks + `## HAD Pretests` index + Choosing-an-Estimator row landed; constructor / fit() parameter names are regression-locked against `inspect.signature(HeterogeneousAdoptionDiD.__init__)` and `HeterogeneousAdoptionDiD.fit` (parameter-name presence only - defaults, type annotations, and return-type unions are not pinned); result-class field tables enumerate every public dataclass field (regression-tested via `dataclasses.fields()`); `llms-practitioner.txt` Step 4 decision tree distinguishes ContinuousDiD (per-dose ATT(d), needs never-treated) from HeterogeneousAdoptionDiD (WAS, universal-rollout-compatible).
 - [x] Phase 5 (partial): README catalog one-liner, bundled `llms.txt` `## Estimators` entry, `docs/api/had.rst` (autoclass for the three classes), and `docs/references.rst` citation landed in PR #372 docs refresh.
 - [x] Phase 5 (wave 2 first slice, PR #409): T21 HAD pretest workflow tutorial (`docs/tutorials/21_had_pretest_workflow.ipynb`) — composite pre-test walkthrough for `did_had_pretest_workflow`. Uses a `Uniform[$0.01K, $50K]` dose-distribution variant of T20's brand-campaign panel (true support strictly positive but near-zero, chosen so QUG fails-to-reject `H0: d_lower = 0` in finite sample). Walks through `aggregate="overall"` (Steps 1 + 3 only, verdict explicitly flags Step 2 deferral) and upgrades to `aggregate="event_study"` (joint pre-trends Stute + joint homogeneity Stute close the gap). Side panel exercises both `yatchew_hr_test` null modes (`linearity` vs `mean_independence`). Companion drift-test file `tests/test_t21_had_pretest_workflow_drift.py` (16 tests pinning panel composition, both verdict pivots, structural anchors, deterministic stats, bootstrap p-value tolerance bands per backend, and `HAD(design="auto")` resolution to `continuous_at_zero` on this panel).
 - [ ] Phase 5 (remaining): T22 weighted/survey HAD tutorial - tracked in `TODO.md`.

tests/test_practitioner.py

@@ -894,6 +894,47 @@ def test_had_step_3_flags_qug_under_survey_deferral(
                 "assumption."
             )
 
+    def test_had_step_3_qualifies_supported_survey_scope(
+        self, mock_had_results, mock_had_event_study_results
+    ):
+        # Per diff_diff/had_pretests.py:1725-1740 + :1927-1940, only
+        # pweight + PSU/FPC survey designs are supported on HAD
+        # pretests. Stratified (SurveyDesign(strata=...)) and
+        # replicate-weight (BRR/Fay/JK1/JKn/SDR) designs raise
+        # NotImplementedError on the linearity kernels. Both HAD
+        # handlers' Step-3 text must call out the supported subset
+        # and the deferred regimes so agents don't generate
+        # `practitioner_next_steps` outputs that overstate what the
+        # workflow will run on a given survey design.
+        for fixture in (mock_had_results, mock_had_event_study_results):
+            output = practitioner_next_steps(fixture, verbose=False)
+            step_3_steps = [s for s in output["next_steps"] if s["baker_step"] == 3]
+            assert len(step_3_steps) == 1
+            text = step_3_steps[0].get("why", "").lower()
+            # Supported subset must be named explicitly.
+            assert "pweight" in text and "psu" in text and "fpc" in text, (
+                "Step-3 text must name the supported survey-pretest scope "
+                "(pweight + PSU/FPC) so agents do not assume any "
+                "survey_design= path is supported."
+            )
+            # Deferred regimes must be flagged explicitly so agents
+            # know not to attempt them.
+            assert "stratif" in text, (
+                "Step-3 text must explicitly note that stratified "
+                "(SurveyDesign(strata=...)) survey designs are not yet "
+                "supported on HAD pretests."
+            )
+            assert "replicate" in text, (
+                "Step-3 text must explicitly note that replicate-weight "
+                "(BRR/Fay/JK1/JKn/SDR) survey designs are not yet "
+                "supported on HAD pretests."
+            )
+            assert "notimplementederror" in text, (
+                "Step-3 text must name the actual exception raised "
+                "(NotImplementedError) so agents can match it in "
+                "error-handling paths."
+            )
+
     def test_had_step_3_pretest_assumption_labels_correct(self, mock_had_results):
         # Per docs/methodology/REGISTRY.md and diff_diff/had_pretests.py
         # docstrings: