diff --git a/CHANGELOG.md b/CHANGELOG.md
index d0d15eaa..b162041c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,94 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]
### Added
+- **TripleDifference serves both DDD designs** (v4 program Phase 3(b); ledger
+ rows [M-013] shimmed, [M-064]): `TripleDifference().fit(..., unit=, time=,
+ first_treat=, partition=)` estimates the staggered-adoption DDD design that
+ `StaggeredTripleDifference` used to own, while the existing
+ `fit(df, outcome, group, partition, post)` call keeps serving the 2x2x2
+ design unchanged. `first_treat=` selects the engine; mixing the two
+ parameter sets raises rather than guessing. This mirrors the reference
+ implementation, whose `triplediff::ddd()` also serves both designs from one
+ signature.
+ - The estimation cores are UNCHANGED - both classes now share one relocated
+ engine, so the staggered numbers are identical by construction (pinned
+ bit-exactly, including per-`(g,t)` tables and seeded bootstrap draws).
+ - The staggered-only fit params (`unit`, `first_treat`, `aggregate`,
+ `balance_e`) are keyword-only; every existing positional slot, `time=`
+ included, is unchanged.
+ - The constructor gains `control_group`, `anticipation`, `base_period`,
+ `n_bootstrap`, `bootstrap_weights`, `seed` and `cband`. `control_group`
+ takes the underscored `"not_yet_treated"`/`"never_treated"`.
+ - `time=` keeps both meanings without ambiguity: the calendar column in
+ staggered mode (no warning), the deprecated alias for `post=` in 2x2x2
+ mode (row [M-031], warns).
+ - 2x2x2 mode returns `TripleDifferenceResults` and staggered mode returns
+ `StaggeredTripleDiffResults`; the containers unify at 4.0 (row [M-014]),
+ so downstream consumers are unaffected.
+
+### Changed
+- **`TripleDifference(pscore_trim=)` is now validated** (ledger row [M-142]):
+ values outside `(0, 0.5)` raise instead of being accepted. The value feeds
+ `np.clip(pscore, trim, 1 - trim)`, so `pscore_trim=0` silently disabled the
+ overlap guard that keeps the `1/(1-p)` IPW/DR weights finite, and
+ `>= 0.5` inverted the clip bounds.
+- **The single-PSU bootstrap warning now names the estimator that was fit.**
+ It lives in the mixin shared by `CallawaySantAnna` and both DDD classes and
+ was hard-coded to `"CallawaySantAnna bootstrap ..."`, so a DDD fit failing
+ closed on a degenerate design pointed diagnosis at the wrong estimator.
+ `CallawaySantAnna`'s own message is byte-identical to before.
+- **Degenerate enabling cohorts are now reported** (both DDD staggered
+ surfaces). A positive `first_treat` cohort whose units are all
+ `partition == 0` identifies no `ATT(g,t)` and contributes to no aggregate,
+ but was still counted in `results.groups`/`n_groups` with no warning naming
+ it — so the estimate silently covered fewer cohorts than the metadata
+ claimed. A `UserWarning` now names the cohort, and `groups` reflects the
+ cohorts that actually produced a `(g, t)` cell. **Estimates are unchanged**;
+ this is a reporting fix.
+- **Negative `first_treat` cohort values now raise** (both `TripleDifference`
+ staggered mode and the deprecated `StaggeredTripleDifference`, which share the
+ engine). Never-treated units must be `0`, or `+inf` (recoded to `0` with a
+ warning, unchanged); treated cohorts must be positive period labels. A unit
+ encoded with the common `-1` never-treated convention previously belonged to
+ neither the treated nor the comparison population: it still counted toward
+ `n_obs` but entered no ATT comparison, so the fit returned a plausible finite
+ estimate for a silently different population. **This is a bug fix that changes
+ results:** such a fit now raises instead of returning a number. Fits using
+ `0`/`+inf` are byte-for-byte unchanged.
+- **`anticipation` is validated** as a non-negative integer (`bool` rejected)
+ via the shared `utils.validate_anticipation`, on `TripleDifference` at
+ construction and in the staggered engine for both surfaces. The window feeds
+ both the base-period rule and the not-yet-treated threshold, so a negative
+ value would have made the universal base period an already-treated period and
+ admitted cohorts treated at the evaluation period as clean controls - neither
+ visible in the output. On the deprecated class, construction still succeeds
+ and `fit()` raises: an identification guard, not a signature change.
+- **`cluster=` raises in staggered DDD mode** on the merged surface, steering
+ to `n_bootstrap > 0` (unit-level clustering via the multiplier bootstrap).
+ Cluster-robust analytical SEs are not implemented for the staggered engine;
+ the deprecated `StaggeredTripleDifference` keeps its 3.x behavior of
+ accepting, warning and ignoring. 2x2x2 mode is unaffected (Liang-Zeger CR1).
+- **Power analysis rejects a staggered-configured `TripleDifference`** at
+ `simulate_power`, `simulate_mde` and `simulate_sample_size`. Both registered
+ DDD data generators produce 2x2x2 data, so a staggered configuration would
+ have been simulated under the wrong design. Staggered-DDD power is not
+ supported yet. "Staggered-configured" means a non-default `control_group`,
+ `anticipation`, `base_period` or `n_bootstrap`; `bootstrap_weights`, `seed`
+ and `cband` are accepted, matching `fit()`, since they are inert without
+ `n_bootstrap > 0`.
+
+### Deprecated
+- **`StaggeredTripleDifference` and its `SDDD` alias** (ledger rows [M-013],
+ [M-064]): deprecated in 3.9, removed in 4.0. Migration:
+ `StaggeredTripleDifference(...).fit(df, outcome, unit, time, first_treat,
+ eligibility)` becomes `TripleDifference(...).fit(df, outcome, unit=...,
+ time=..., first_treat=..., partition=...)`. Two vocabulary changes:
+ `eligibility=` is named `partition=`, and `control_group` takes
+ `"not_yet_treated"`/`"never_treated"` instead of R's compact
+ `"notyettreated"`/`"nevertreated"`. The deprecated class keeps its 3.x
+ parameter names and values until removal, so existing code keeps working
+ (with a `FutureWarning`) until 4.0.
+
- **TwoWayFixedEffects event-study mode** (v4 program Phase 3(a); ledger
row [M-010] shimmed): `TWFE().fit(..., event_study=True, time="period",
spec="within"|"pooled", reference_period=None, post_periods=[...])`
diff --git a/DEFERRED.md b/DEFERRED.md
index 316f4b1b..8223aa15 100644
--- a/DEFERRED.md
+++ b/DEFERRED.md
@@ -45,9 +45,9 @@ exists but parity can't be verified without a local toolchain.
| Issue | Location | PR | Priority |
|-------|----------|----|----------|
-| `StaggeredTripleDifference` R cross-validation: CSV fixtures not committed (gitignored); tests skip without local R + `triplediff`. Commit fixtures or generate deterministically. | `tests/test_methodology_staggered_triple_diff.py` | #245 | Medium |
-| `StaggeredTripleDifference` R parity: benchmark only tests the no-covariate path (`xformla=~1`). Add covariate-adjusted scenarios + aggregation-SE parity assertions. | `benchmarks/R/benchmark_staggered_triplediff.R` | #245 | Medium |
-| `StaggeredTripleDifference` per-cohort group-effect SEs include WIF (conservative vs R's `wif=NULL`); documented in REGISTRY. Could override the mixin for an exact R match (verification needs R `triplediff`). | `staggered_triple_diff.py` | #245 | Low |
+| `StaggeredTripleDifference` R cross-validation (the engine is shared with `TripleDifference`'s staggered mode since 3(b), so this covers both surfaces): CSV fixtures not committed (gitignored); tests skip without local R + `triplediff`. Commit fixtures or generate deterministically. | `tests/test_methodology_staggered_triple_diff.py` | #245 | Medium |
+| Staggered DDD R parity (both surfaces - one shared engine since 3(b)): benchmark only tests the no-covariate path (`xformla=~1`). Add covariate-adjusted scenarios + aggregation-SE parity assertions. | `benchmarks/R/benchmark_staggered_triplediff.R` | #245 | Medium |
+| Staggered DDD per-cohort group-effect SEs include WIF (both surfaces - one shared engine since 3(b)) (conservative vs R's `wif=NULL`); documented in REGISTRY. Could override the mixin for an exact R match (verification needs R `triplediff`). | `_staggered_triple_diff_engine.py` | #245 | Low |
| **WooldridgeDiD follow-up cluster** (PR-B Stage D/E fail-closed surfaces; re-enable after R/Stata validation):
• QMLE sandwich uses `aweight` cluster adjustment `(G/(G-1))·(n-1)/(n-k)` vs Stata's `G/(G-1)` (conservative); add a `qmle` weight type if Stata goldens confirm a material difference (`wooldridge.py`, `linalg.py`).
• response-scale APE / log-link coefficient bridge for R `etwfe(family=poisson|logit)` cell-level parity — needs `emfx()` APE extraction or link-inversion with baseline-mean adjustment (`generate_wooldridge_golden.R`, `test_methodology_wooldridge.py`).
• `aggregate(weights="cohort_share")` on survey-weighted fits: `_n_g_per_cohort` uses raw `unit.nunique()`; implement design-weighted unit totals per cohort (paper W2025 §7) and lift the `ValueError` gate (`wooldridge.py`, `wooldridge_results.py`).
• unconditional inference for `cohort_share` accounting for ω̂_g sampling uncertainty (W2025 §7.5); currently NaN-closed (`wooldridge_results.py`).
• `cohort_trends=True × survey_design` and `× control_group="never_treated"` raise `NotImplementedError` (unvalidated TSL variance / trend columns spanned jointly by the placebo cells and the unit FE, which absorb the cohort indicator and recover the omitted reference) (`wooldridge.py`).
• ~~Stata `jwdid` golden-value `TestReferenceValues`~~ RESOLVED: the golden ships four arms pinned by `tests/test_etwfe_cs_stata_parity.py` (no `TestReferenceValues` symbol was ever added). The QMLE bullet above REMAINS OPEN -- every golden arm is linear `jwdid`, so no QMLE cluster-SE reference exists, and SEs are pinned only as a ratio. | `wooldridge.py`, `wooldridge_results.py`, `linalg.py`, benchmarks | #216 · PR-B | Med-Low |
| Extend `WooldridgeDiD` `method ∈ {logit, poisson}` with `vcov_type ∈ {classical, hc2, hc2_bm}`: composing HC2 leverage + Bell-McCaffrey DOF with the QMLE pseudo-residual sandwich needs derivation + R parity vs `clubSandwich::vcovCR(glm, type="CR2")`. Rejected at `__init__`. | `wooldridge.py` | follow-up | Medium |
| Multi-constraint CR2 parallel-trends test (AHT/HTZ) for `hc2_bm` fits: DiagnosticReport's PT check routes `vcov_type="hc2_bm"` sources to Bonferroni over the BM-adjusted per-row p-values because the generic chi-square joint Wald would discard the CR2 small-sample correction (see REPORTING.md "hc2_bm parallel-trends policy"). The proper joint test is the AHT/HTZ Wald with a Satterthwaite-style denominator df over the pre-period contrast block; needs derivation for the stacked/pooled WLS-CR2 layout + parity vs `clubSandwich::Wald_test(..., test="HTZ")`. | `diagnostic_report.py`, `linalg.py` | vcov/df round-trip PR | Low |
diff --git a/METHODOLOGY_REVIEW.md b/METHODOLOGY_REVIEW.md
index 5dde4997..7ef3e7dd 100644
--- a/METHODOLOGY_REVIEW.md
+++ b/METHODOLOGY_REVIEW.md
@@ -988,7 +988,7 @@ These three are feature deferrals (paper-supported extensions that the library h
| Field | Value |
|-------|-------|
-| Module | `staggered_triple_diff.py`, `staggered_triple_diff_results.py` |
+| Module | `_staggered_triple_diff_engine.py` (the shared engine), `staggered_triple_diff.py` (deprecated class surface), `staggered_triple_diff_results.py` |
| Primary Reference | Ortiz-Villavicencio & Sant'Anna (2025) — same paper as TripleDifference, staggered case |
| R Reference | `triplediff::ddd(panel=TRUE)` + `agg_ddd()` (per `benchmarks/R/benchmark_staggered_triplediff.R`) |
| Status | **Complete** |
diff --git a/README.md b/README.md
index 818f0875..6af027de 100644
--- a/README.md
+++ b/README.md
@@ -109,14 +109,14 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`.
- [SpilloverDiD](https://diff-diff.readthedocs.io/en/stable/api/spillover.html) - Butts (2021) ring-indicator spillover-aware DiD identifying direct effect on treated + per-ring spillover on near-control units; handles non-staggered and staggered timing; supports survey-design variance under `survey_design=` for HC1 / CR1 (Wave E.1 Binder TSL) and Conley (Wave E.2 panel-aware stratified-Conley sandwich on per-period PSU totals; extended in Wave E.2 follow-up to `conley_lag_cutoff > 0` via panel-block composition with within-PSU serial Bartlett HAC — `lag>0` requires an effective PSU via explicit `survey_design.psu` or injected `cluster=`); `SurveyDesign.subpopulation()` preserves full-design `n_psu` / `df_survey` via zero-padded scores (Wave E.3, R `svyrecvar(subset())` form)
- [SyntheticDiD](https://diff-diff.readthedocs.io/en/stable/api/estimators.html) - Synthetic DiD combining standard DiD and synthetic control for few treated units
- [SyntheticControl](https://diff-diff.readthedocs.io/en/stable/api/synthetic_control.html) - Abadie, Diamond & Hainmueller (2010) classic synthetic control for a single treated unit (donor-weight counterfactual, nested/cv/inverse-variance/custom V; in-space placebo permutation inference via `in_space_placebo()`, plus ADH-2015 `leave_one_out()` + `in_time_placebo()` robustness, Firpo-Possebom (2018) test-inversion confidence sets, and Chernozhukov-Wüthrich-Zhu (2021) conformal inference)
-- [TripleDifference](https://diff-diff.readthedocs.io/en/stable/api/triple_diff.html) - triple difference (DDD) estimator for designs requiring two criteria for treatment eligibility
+- [TripleDifference](https://diff-diff.readthedocs.io/en/stable/api/triple_diff.html) - triple difference (DDD) estimator for designs requiring two criteria for treatment eligibility; serves both the 2x2x2 and the staggered-adoption design from one signature (`fit(..., first_treat=)` selects the staggered engine)
- [ContinuousDiD](https://diff-diff.readthedocs.io/en/stable/api/continuous_did.html) - Callaway, Goodman-Bacon & Sant'Anna (2024) continuous treatment DiD with dose-response curves
- [HeterogeneousAdoptionDiD](https://diff-diff.readthedocs.io/en/stable/api/had.html) - de Chaisemartin, Ciccia, D'Haultfœuille & Knau (2026) for designs where **no unit remains untreated**; local-linear estimator at the dose support boundary returning Weighted Average Slope (WAS) on Design 1' (`d̲ = 0` / QUG) or `WAS_{d̲}` on Design 1 (`d̲ > 0`, continuous-near-d̲ or mass-point), with a multi-period event-study extension (last-treatment cohort, pointwise CIs). **Panel-only** in this release - repeated cross-sections rejected by the validator. Alias `HAD`.
- [RegressionDiscontinuity](https://diff-diff.readthedocs.io/en/stable/api/regression_discontinuity.html) - Calonico, Cattaneo & Titiunik (2014) sharp, fuzzy, AND covariate-adjusted regression discontinuity with robust bias-corrected inference and rdrobust-parity bandwidth selection (all 10 selectors, mass-point handling; fuzzy via `takeup=` with a first-stage block and weak-identification warning; covariates via `covariates=` - CCFT 2019, same estimand, covariate-aware bandwidths). Canonical `att` is the bias-corrected estimate with a coherent robust CI (rdrobust's printed headline is `att_conventional`). Alias `RDD`.
- [StackedDiD](https://diff-diff.readthedocs.io/en/stable/api/stacked_did.html) - Wing, Freedman & Hollingsworth (2024) stacked DiD with Q-weights and sub-experiments; optional covariate balancing (Ustyuzhanin 2026)
- [EfficientDiD](https://diff-diff.readthedocs.io/en/stable/api/efficient_did.html) - Chen, Sant'Anna & Xie (2025) efficient DiD with optimal weighting for tighter SEs
- [TROP](https://diff-diff.readthedocs.io/en/stable/api/trop.html) - Triply Robust Panel estimator (Athey et al. 2025) with nuclear norm factor adjustment
-- [StaggeredTripleDifference](https://diff-diff.readthedocs.io/en/stable/api/staggered.html#staggeredtripledifference) - Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD with group-time ATT
+- [StaggeredTripleDifference](https://diff-diff.readthedocs.io/en/stable/api/staggered.html#staggeredtripledifference) - Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD with group-time ATT (deprecated 3.9 - use `TripleDifference` with `first_treat=`)
- [WooldridgeDiD](https://diff-diff.readthedocs.io/en/stable/api/wooldridge_etwfe.html) - Wooldridge (2023, 2025) ETWFE: saturated OLS, logit/Poisson QMLE (ASF-based ATT). Alias `ETWFE`.
- [LPDiD](https://diff-diff.readthedocs.io/en/stable/api/lpdid.html) - Dube, Girardi, Jorda & Taylor (2025) Local Projections DiD: per-horizon long-difference event study on clean controls (no negative weighting), variance- or equally-weighted ATT, for absorbing or non-absorbing (reversible) treatment
- [ChangesInChanges](https://diff-diff.readthedocs.io/en/stable/api/changes_in_changes.html) - Athey & Imbens (2006) nonlinear/distributional DiD for the 2x2 design: full counterfactual distribution and quantile treatment effects via CDF transformation, plus the QDiD comparison estimator; bootstrap inference; R qte parity. Alias `CiC`
diff --git a/TODO.md b/TODO.md
index f7c9dd07..1232dcd8 100644
--- a/TODO.md
+++ b/TODO.md
@@ -21,6 +21,12 @@ Related tracking surfaces:
| Issue | Location | Origin | Effort | Priority |
|-------|----------|--------|--------|----------|
+| Post-fit `aggregate()` for the staggered DDD container: `StaggeredTripleDiffResults` carries no `AggregationMixin`, which is why the phase-3(b) merge had to carry fit-time `aggregate=`/`balance_e=` onto the surviving `TripleDifference` (rows M-140/M-141) as the ONE documented exception to the section-6 aggregate-postfit program. Porting the container onto the M-122 aggregation contract retires both rows; note the bootstrapped-fit recompute levels will need draw retention or a fail-closed relay, the same problem tracked for CS/EfficientDiD/ImputationDiD | `diff_diff/staggered_triple_diff_results.py`, `diff_diff/aggregation.py` | 3(b) | Heavy | Medium |
+| Staggered-DDD power support: `simulate_power`/`simulate_mde`/`simulate_sample_size` now REJECT a staggered-configured `TripleDifference` (both registered DDD generators emit 2x2x2 data and fit with `(group, partition, post)`, so a staggered config would be simulated under the wrong design). Support needs a staggered DDD DGP profile plus fit-kwargs builder, and a decision on whether the mode is selected by profile or by the estimator's own config | `diff_diff/power.py` | 3(b) | Mid | Low |
+| Bootstrap-`seed` provenance on multiplier-bootstrap results containers: neither `StaggeredTripleDiffResults` nor `CallawaySantAnnaResults` carries the `seed` that generated its bootstrap SEs / p-values / sup-t bands, so a serialized result cannot report the random configuration behind its inference. NOT a 3(b) regression - `seed` reaches the engine and `get_params()` correctly (same seed reproduces the SE bit-exactly, a different seed moves it), the gap is results-object observability only, it predates the merge, and both containers inherit it from the shared `CallawaySantAnnaBootstrapMixin`. Add `seed` (and consider `n_bootstrap`/`bootstrap_weights`/`cband`) to BOTH containers plus `to_dict()`, with seeded and unseeded pins; sequence it with the M-014 container unification rather than schema-changing one container mid-merge. Precedent for exposing it: `ContinuousDiDResults`, `EfficientDiDResults`, `SyntheticDiDResults` already do | `diff_diff/staggered_triple_diff_results.py`, `diff_diff/staggered_results.py` | 3(b) | Quick | Low |
+| Library-wide `anticipation` domain validation: `TripleDifference` now rejects non-integral / negative / `bool` windows at construction (phase 3(b)) because the value feeds BOTH the base-period rule and the not-yet-treated threshold, so `anticipation=-1` silently makes the universal base period `g` (already treated) and admits cohorts treated at the evaluation period as clean controls. Only `spillover.py` and `wooldridge.py` validate it today (and neither rejects `bool`, which coerces to a silent one-period window); `CallawaySantAnna`, `SunAbraham`, `ImputationDiD`, `TwoStageDiD`, `StackedDiD`, `ContinuousDiD`, `EfficientDiD` and the deprecated `StaggeredTripleDifference` do not. The shared validator now EXISTS - `utils.validate_anticipation`, adopted by `TripleDifference.__init__` and by the staggered engine (so `StaggeredTripleDifference` fails closed at fit too); aligning the remaining seven estimators is a matter of calling it from each constructor | `diff_diff/staggered.py`, `diff_diff/sun_abraham.py`, `diff_diff/imputation.py`, `diff_diff/two_stage.py`, `diff_diff/stacked_did.py`, `diff_diff/continuous_did.py`, `diff_diff/efficient_did.py`, `diff_diff/spillover.py`, `diff_diff/wooldridge.py` | 3(b) | Mid | Medium |
+| `ContinuousDiD.pscore_trim` still validates `0.0 <= x < 0.5`, i.e. it admits `0`, while `TripleDifference` tightened to `0 < x < 0.5` in phase 3(b) (row M-142) on the grounds that `trim=0` disables the `np.clip(pscore, trim, 1-trim)` overlap guard keeping the `1/(1-p)` weights finite. The same argument applies to ContinuousDiD; aligning it was out of scope for a DDD merge and is recorded in the REGISTRY staggered-mode Note rather than left as silent drift. `TripleDifference` additionally gained a TYPE guard in 3(b) (reject bool/non-real-scalar/non-finite BEFORE the range comparison) because a bare `0 < x < 0.5` raises an incidental `TypeError` on `None`/str/complex/list, an ambiguous-truth error on a multi-element array, and silently ACCEPTS a 1-element array as the parameter; `ContinuousDiD`'s `np.isfinite(self.pscore_trim) and ...` has the same hole. Aligning both is one change - promote the guard to a shared `utils.validate_pscore_trim(value, *, allow_zero)` alongside `validate_n_bootstrap` rather than copying it | `diff_diff/continuous_did.py`, `diff_diff/utils.py` | 3(b) | Quick | Low |
+| Staggered-mode cluster-robust ANALYTICAL SEs: `cluster=` raises in `TripleDifference`'s staggered mode (and is accepted-then-ignored on the deprecated class), so clustered inference there is bootstrap-only. Implementing a clustered analytical path for the GMM-combined influence function would let the raise become a real lane | `diff_diff/_staggered_triple_diff_engine.py` | 3(b) | Heavy | Low |
| diagnostic_report admission for `EventStudyResults` surfaces (the TWFE event-study mode + `aggregate('event_study')` containers): DiagnosticReport/BusinessReport now REJECT the surface explicitly (Phase 3(a); previously a silent zero-check report / all-null headline) and practitioner_next_steps serves the generic fall-through - admission needs source-aware routing (the type-name-keyed `_APPLICABILITY`/`_HANDLERS` registries cannot discriminate the unified container's producers) and a scalar-vs-per-period headline design; MPD-native results received {parallel_trends, pretrends_power, sensitivity, bacon, design_effect} | `diff_diff/diagnostic_report.py`, `diff_diff/business_report.py`, `diff_diff/practitioner.py` | 3(a) | Mid | Medium |
| Align the NATIVE pretrends pre_periods= contract with the container routes' fail-closed validation: `_extract_pre_period_params`'s MPD branch silently filters an explicit `pre_periods=` selection (unknown labels, the reference, unusable-inference rows dropped without error, caller order preserved) while both container routes validate every requested label and enforce calendar chronology (the relative route since M-024; the calendar route since 3(a) R8) - the relative route also silently collapses DUPLICATE requested labels (the calendar route rejects them), and both filter on SE only while the calendar route additionally requires a finite EFFECT (3(a) R9); one contract across all three routes, with pinned rejection messages | `diff_diff/pretrends.py` | 3(a) R8 | Quick | Low |
| `EventStudyResults` inference-provenance fields: the container records no `vcov_type`/`cluster_name`/`n_clusters`/`df_convention`/Conley metadata, so a serialized surface cannot distinguish unit auto-clustering from explicit clustering, survey, Conley, or the one-way carve-out (3(a) R9 review). Adding them is a cross-producer M-092 schema amendment (six builders, to_dict/summary rendering, surface-suite pins) - follow the pre-cut amendment convention (optional fields appended last, ledger note same-diff) rather than bolting onto one producer | `diff_diff/results_base.py` | 3(a) R9 | Mid | Low |
diff --git a/diff_diff/_staggered_triple_diff_engine.py b/diff_diff/_staggered_triple_diff_engine.py
new file mode 100644
index 00000000..f8707a46
--- /dev/null
+++ b/diff_diff/_staggered_triple_diff_engine.py
@@ -0,0 +1,1699 @@
+"""Shared staggered-DDD engine (private).
+
+Holds the estimation core both DDD surfaces run: `TripleDifference` in its
+staggered mode (ledger row M-013) and the deprecated `StaggeredTripleDifference`
+until its 4.0 removal. Relocated VERBATIM in phase 3(b) - the only changes are
+the three parameterization knobs described below, so the merge is an API merge
+and the numbers cannot move.
+
+Knobs, all passed explicitly by the two callers:
+
+ estimator_name interpolated into every message that names a class, so the
+ merged surface never steers users at the deprecated one
+ partition_label "partition" on the merged surface, "eligibility" on the
+ dying one - the third-dimension vocabulary
+ _frame_offset added to every user-attributed `stacklevel`, so warning
+ attribution is identical through both call depths
+ (user -> fit -> core on the dying class = 1;
+ user -> fit -> _fit_staggered -> core on the merged one = 2)
+
+The offset is ALSO mirrored onto `self._warn_frame_offset` for the duration of a
+fit, because four helper methods and the shared bootstrap mixin warn from depths
+this function cannot reach by argument.
+
+Private module: not exported, and mapped as a member of the
+`staggered_triple_diff` group in docs/doc-deps.yaml so /docs-impact still
+resolves this engine's methodology docs.
+"""
+
+import warnings
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+
+from diff_diff.linalg import (
+ _check_propensity_diagnostics,
+ _rank_guarded_inv,
+ solve_logit,
+)
+from diff_diff.staggered_triple_diff_results import StaggeredTripleDiffResults
+from diff_diff.utils import safe_inference, validate_anticipation
+
+if TYPE_CHECKING:
+ from diff_diff.survey import SurveyDesign
+
+# Type alias for pre-computed structures
+PrecomputedData = Dict[str, Any]
+
+# The dying class accepts R's compact spellings; the merged class uses the
+# underscored library vocabulary from birth (M-013). The engine reads both so a
+# single core can serve both surfaces without a value shim on either.
+_NEVER_TREATED = frozenset({"nevertreated", "never_treated"})
+
+
+def _is_never_treated(control_group: str) -> bool:
+ """True for either spelling of the never-treated comparison group."""
+ return control_group in _NEVER_TREATED
+
+
+class _StaggeredTripleDiffEngineMixin:
+ """Estimation core for staggered DDD. Not usable standalone.
+
+ Mixed into both `TripleDifference` and `StaggeredTripleDifference`, which
+ supply the constructor attributes and the CS aggregation/bootstrap mixins
+ this core calls. The annotations below exist because mypy type-checks this
+ class independently of its hosts (`attr-defined` is not disabled) - they are
+ declarations, never assignments.
+ """
+
+ # Constructor attributes read from the host class.
+ estimation_method: str
+ control_group: str
+ alpha: float
+ anticipation: int
+ base_period: str
+ n_bootstrap: int
+ cband: bool
+ pscore_trim: float
+ cluster: Optional[str]
+ rank_deficient_action: str
+ epv_threshold: float
+ pscore_fallback: str
+
+ # Fit-time state this core assigns.
+ is_fitted_: bool
+ _lstsq_fallback_tracker: Any
+ _ps_lstsq_fallback_tracker: Any
+
+ # Warning-attribution offset, mirrored from _fit_staggered_core's
+ # _frame_offset for the duration of a fit. A PLAIN class attribute, not a
+ # ClassVar: `self._warn_frame_offset = ...` on a ClassVar is a mypy [misc]
+ # error ("cannot assign to class variable via instance").
+ _warn_frame_offset = 0
+
+ if TYPE_CHECKING:
+ # Supplied by the CS aggregation/bootstrap mixins on every host class.
+ # Same house pattern as CallawaySantAnnaBootstrapMixin's own stub block,
+ # and the same caveat: these signatures must MATCH the implementations
+ # they shadow, or the two base classes become incompatible wherever both
+ # are mixed in.
+
+ def _aggregate_simple(
+ self,
+ group_time_effects: Dict,
+ influence_func_info: Dict,
+ df: Optional[pd.DataFrame],
+ unit: Optional[str],
+ precomputed: Optional["PrecomputedData"] = None,
+ ) -> Tuple[float, float, Optional[int]]: ...
+
+ def _aggregate_event_study(
+ self,
+ group_time_effects: Dict,
+ influence_func_info: Dict,
+ groups: List[Any],
+ time_periods: List[Any],
+ balance_e: Optional[int] = None,
+ df: Optional[pd.DataFrame] = None,
+ unit: Optional[str] = None,
+ precomputed: Optional["PrecomputedData"] = None,
+ ) -> Any: ...
+
+ def _aggregate_by_group(
+ self,
+ group_time_effects: Dict,
+ influence_func_info: Dict,
+ groups: List[Any],
+ precomputed: Optional["PrecomputedData"] = None,
+ df: Optional[pd.DataFrame] = None,
+ unit: Optional[str] = None,
+ ) -> Dict[Any, Dict[str, Any]]: ...
+
+ def _run_multiplier_bootstrap(
+ self,
+ group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]],
+ influence_func_info: Dict[Tuple[Any, Any], Dict[str, Any]],
+ aggregate: Optional[str],
+ balance_e: Optional[int],
+ treatment_groups: List[Any],
+ time_periods: List[Any],
+ df: Any = None,
+ unit: Optional[str] = None,
+ precomputed: Any = None,
+ cband: bool = True,
+ ) -> Any: ...
+
+ def _fit_staggered_core(
+ self,
+ data: pd.DataFrame,
+ outcome: str,
+ unit: str,
+ time: str,
+ first_treat: str,
+ partition: str,
+ covariates: Optional[List[str]] = None,
+ aggregate: Optional[str] = None,
+ balance_e: Optional[int] = None,
+ survey_design: Optional["SurveyDesign"] = None,
+ *,
+ estimator_name: str,
+ partition_label: str,
+ _frame_offset: int = 0,
+ ) -> StaggeredTripleDiffResults:
+ """Fit the staggered triple-difference estimator.
+
+ The estimation body is the 3.x `StaggeredTripleDifference.fit` verbatim;
+ see the module docstring for the three knobs.
+ """
+ # Body-local alias so the relocated body stays byte-identical (the
+ # `time = post` precedent in triple_diff.py). The PARAMETER is
+ # `partition`; the third dimension is `eligibility` only internally.
+ eligibility = partition
+ self._warn_frame_offset = _frame_offset
+ try:
+ # Identification-assumption guard, enforced in the ENGINE so BOTH
+ # public surfaces fail closed on the same input. The merged
+ # constructor validates this eagerly; this call additionally covers
+ # the deprecated StaggeredTripleDifference (whose 3.x API SHAPE is
+ # frozen through removal - that freeze was never a licence to emit
+ # silently-biased numbers) and direct attribute mutation on either
+ # class, which bypasses __init__ and set_params alike.
+ validate_anticipation(self.anticipation)
+ from diff_diff.survey import (
+ _resolve_survey_for_fit,
+ _validate_unit_constant_survey,
+ compute_survey_metadata,
+ )
+
+ resolved_survey, survey_weights, survey_weight_type, survey_metadata = (
+ _resolve_survey_for_fit(survey_design, data, "analytical")
+ )
+
+ if resolved_survey is not None:
+ _validate_unit_constant_survey(data, unit, survey_design)
+ if resolved_survey.weight_type != "pweight":
+ raise ValueError(
+ f"{estimator_name} survey support requires "
+ f"weight_type='pweight', got '{resolved_survey.weight_type}'. "
+ f"The survey variance math assumes probability weights."
+ )
+ if aggregate is not None and aggregate not in [
+ "event_study",
+ "group",
+ "simple",
+ "all",
+ ]:
+ raise ValueError(
+ f"aggregate must be 'event_study', 'group', 'simple', or 'all', "
+ f"got '{aggregate}'"
+ )
+
+ df = data.copy()
+ self._validate_inputs(
+ df,
+ outcome,
+ unit,
+ time,
+ first_treat,
+ eligibility,
+ covariates,
+ partition_label=partition_label,
+ )
+
+ if self.cluster is not None:
+ warnings.warn(
+ "cluster parameter is accepted but cluster-robust analytical SEs "
+ "are not yet implemented for staggered DDD. Use n_bootstrap > 0 "
+ "for unit-level clustered inference via multiplier bootstrap.",
+ UserWarning,
+ stacklevel=2 + _frame_offset,
+ )
+
+ if first_treat != "first_treat":
+ df["first_treat"] = df[first_treat]
+ # Surface the inf → 0 recategorization the same way StaggeredDiD does
+ # (see `staggered.py:1508-1519`). Silently recoding inf would shift
+ # units between treated and never-treated pools with no signal
+ # (axis-E silent coercion under the Phase 2 audit).
+ _inf_mask = np.isposinf(df["first_treat"].values)
+ if _inf_mask.any():
+ n_inf_rows = int(_inf_mask.sum())
+ warnings.warn(
+ f"{n_inf_rows} row(s) have first_treat=inf; recoding to 0 "
+ f"(never-treated). Use first_treat=0 to suppress this warning.",
+ UserWarning,
+ stacklevel=2 + _frame_offset,
+ )
+ df["first_treat"] = df["first_treat"].replace([np.inf, float("inf")], 0)
+
+ # Negative cohorts belong to NEITHER population and must fail closed.
+ # _precompute_structures builds the treated set from `g > 0` and the
+ # never-enabled set from `g == 0`, so a unit encoded with the common
+ # `-1` never-treated convention (or `-inf`) is silently dropped from
+ # both - it still counts toward n_obs while contributing to no ATT
+ # comparison, and the fit returns a plausible finite estimate for a
+ # DIFFERENT population (measured: overall_att 3.29438582 -> 2.99470938
+ # with n_never_enabled 24 -> 0, no error and no warning).
+ # This is the same input axis the +inf branch above already defends,
+ # so silence here was a hole in an established contract, not an
+ # undocumented input. Raising rather than recoding: unlike +inf, a
+ # negative value has no unambiguous intent (-1 as "never" is a
+ # convention, not a limit), and guessing would be exactly the silent
+ # sample change being fixed.
+ _ft_vals = np.asarray(pd.to_numeric(df["first_treat"], errors="coerce"), dtype=float)
+ _neg_mask = _ft_vals < 0
+ if _neg_mask.any():
+ _bad = np.unique(_ft_vals[_neg_mask])
+ raise ValueError(
+ f"first_treat contains {int(_neg_mask.sum())} row(s) with negative "
+ f"cohort value(s) {sorted(_bad.tolist())}. Never-treated units must "
+ f"be encoded as 0 (or +inf, which is recoded to 0 with a warning); "
+ f"treated cohorts must be positive period labels. Negative values "
+ f"would be excluded from BOTH the treated and the comparison "
+ f"cohorts, silently estimating on a different population."
+ )
+
+ precomputed = self._precompute_structures(
+ df,
+ outcome,
+ unit,
+ time,
+ eligibility,
+ covariates,
+ resolved_survey=resolved_survey,
+ )
+
+ # Recompute survey metadata from unit-level resolved survey
+ if resolved_survey is not None and survey_metadata is not None:
+ resolved_survey_unit = precomputed.get("resolved_survey_unit")
+ if resolved_survey_unit is not None:
+ unit_w = resolved_survey_unit.weights
+ survey_metadata = compute_survey_metadata(resolved_survey_unit, unit_w)
+
+ # Survey df for t-distribution critical values
+ df_survey = precomputed.get("df_survey")
+ if (
+ df_survey is None
+ and resolved_survey is not None
+ and hasattr(resolved_survey, "uses_replicate_variance")
+ and resolved_survey.uses_replicate_variance
+ ):
+ df_survey = 0 # Forces NaN inference for undefined replicate df
+
+ treatment_groups = precomputed["treatment_groups"]
+ time_periods = precomputed["time_periods"]
+ all_units = precomputed["all_units"]
+ time_to_col = precomputed["time_to_col"]
+ unit_cohorts = precomputed["unit_cohorts"]
+ eligibility_per_unit = precomputed["eligibility_per_unit"]
+ n_units = len(all_units)
+
+ pscore_cache: Dict = {}
+
+ group_time_effects: Dict[Tuple, Dict[str, Any]] = {}
+ influence_func_info: Dict[Tuple, Dict[str, Any]] = {}
+ comparison_group_counts: Dict[Tuple, int] = {}
+ gmm_weights_store: Dict[Tuple, Dict] = {}
+ epv_diagnostics: Optional[Dict[Tuple, Dict[str, Any]]] = (
+ {} if (covariates and self.estimation_method in ("ipw", "dr")) else None
+ )
+
+ # Trackers for rank-deficient linalg solves across all (g, g_c, t)
+ # cells. `_compute_did_panel` appends to the OR-side tracker;
+ # `_compute_pscore` appends to the PS-side tracker. Both surface as
+ # ONE aggregate warning below rather than fanning out per cell.
+ self._lstsq_fallback_tracker: List[float] = []
+ self._ps_lstsq_fallback_tracker: List[float] = []
+
+ # A positive cohort with NO eligible treated units (every unit
+ # partition==0) cannot identify ATT(g,t): the DDD contrast needs the
+ # eligible-treated cell. Such a cohort silently contributed nothing
+ # while still being advertised in `groups`, so the aggregate narrowed
+ # with no statement of which cohort dropped out (the only warnings
+ # naming it describe its role as a COMPARISON cohort for other g).
+ # Warn per deficient cohort; the estimate itself is unchanged and
+ # remains valid for the cohorts that do identify.
+ for _g in treatment_groups:
+ _n_elig = int(np.sum((unit_cohorts == _g) & (eligibility_per_unit == 1)))
+ if _n_elig == 0:
+ _n_units_g = int(np.sum(unit_cohorts == _g))
+ warnings.warn(
+ f"Enabling cohort g={_g} has no eligible treated units "
+ f"(partition==1) among its {_n_units_g} unit(s), so ATT(g,t) is "
+ f"unidentified there and the cohort contributes to no aggregate. "
+ f"It is excluded from `groups`; the reported estimate covers the "
+ f"remaining cohort(s) only.",
+ UserWarning,
+ stacklevel=2 + _frame_offset,
+ )
+
+ for g in treatment_groups:
+ # In universal mode, skip the reference period (t == g-1-anticipation)
+ # so it's omitted from GT estimation. The event-study mixin injects
+ # a synthetic reference row with effect=0, matching CS behavior.
+ if self.base_period == "universal":
+ universal_base = g - 1 - self.anticipation
+ valid_periods = [t for t in time_periods if t != universal_base]
+ else:
+ valid_periods = time_periods
+
+ for t in valid_periods:
+ base_period_val = self._get_base_period(g, t)
+ if base_period_val is None:
+ continue
+ if base_period_val not in time_to_col:
+ warnings.warn(
+ f"Base period {base_period_val} for (g={g}, t={t}) is "
+ "outside the observed panel. Skipping this cell.",
+ UserWarning,
+ stacklevel=2 + _frame_offset,
+ )
+ continue
+ if t not in time_to_col:
+ continue
+
+ has_never_enabled = bool(np.any(unit_cohorts == 0))
+
+ if _is_never_treated(self.control_group):
+ # Only use never-enabled cohort as comparison
+ valid_gc = [0] if has_never_enabled else []
+ else:
+ # Use all valid comparison cohorts (not-yet-treated + never)
+ # Threshold accounts for anticipation: cohorts that start
+ # treatment within the anticipation window are contaminated.
+ nyt_threshold = max(t, base_period_val) + self.anticipation
+ valid_gc = [gc for gc in treatment_groups if gc > nyt_threshold and gc != g]
+ if has_never_enabled:
+ valid_gc = [0] + valid_gc
+
+ if not valid_gc:
+ warnings.warn(
+ f"No valid comparison groups for (g={g}, t={t}), skipping.",
+ UserWarning,
+ stacklevel=2 + _frame_offset,
+ )
+ continue
+
+ treated_mask = (unit_cohorts == g) & (eligibility_per_unit == 1)
+ n_treated = int(np.sum(treated_mask))
+ if n_treated == 0:
+ continue
+
+ att_vec = []
+ inf_raw = [] # unrescaled IFs
+ gc_labels = []
+ gc_cell_sizes = [] # size_gt_ctrl per surviving gc
+
+ for gc in valid_gc:
+ result = self._compute_ddd_gt_gc(
+ precomputed,
+ g,
+ gc,
+ t,
+ base_period_val,
+ covariates,
+ pscore_cache,
+ epv_diagnostics=epv_diagnostics,
+ )
+ if result is None:
+ continue
+ att_gc, inf_gc, size_gt_ctrl = result
+ if not np.isfinite(att_gc):
+ continue
+
+ att_vec.append(att_gc)
+ inf_raw.append(inf_gc)
+ gc_labels.append(gc)
+ gc_cell_sizes.append(size_gt_ctrl)
+
+ if not att_vec:
+ continue
+
+ # Compute size_gt from SURVIVING comparison cohorts only
+ # (not from all initially valid gc's)
+ surviving_units = treated_mask.copy()
+ for gc in gc_labels:
+ surviving_units |= (unit_cohorts == gc) | (unit_cohorts == g)
+ survey_w = precomputed.get("survey_weights")
+ if survey_w is not None:
+ size_gt = float(np.sum(survey_w[surviving_units]))
+ else:
+ size_gt = float(np.sum(surviving_units))
+
+ # Apply IF rescaling now that size_gt is known
+ inf_matrix = []
+ for inf_gc, size_gt_ctrl in zip(inf_raw, gc_cell_sizes):
+ if size_gt_ctrl > 0:
+ inf_gc = inf_gc * (size_gt / size_gt_ctrl)
+ inf_matrix.append(inf_gc)
+
+ att_gmm, inf_gmm, gmm_w, se_gt = self._combine_gmm(
+ np.array(att_vec),
+ np.array(inf_matrix),
+ n_units,
+ )
+
+ if not np.isfinite(att_gmm):
+ continue
+
+ # R's single-gc SE uses size_gt in denominator, not n_total.
+ # For multi-gc (GMM), the size_gt factor is already in Omega
+ # via the per-gc rescaling, so n_total is correct.
+ if len(gc_labels) == 1:
+ se_gt = float(np.sqrt(np.sum(inf_gmm**2) / size_gt**2))
+
+ if not np.isfinite(se_gt) or se_gt <= 0:
+ se_gt = np.nan
+
+ t_stat, p_value, conf_int = safe_inference(
+ att_gmm, se_gt, alpha=self.alpha, df=df_survey
+ )
+
+ # Rescale IF for mixin compatibility.
+ # R stores IF * (n/size_gt) in inf_func_mat, then uses
+ # SE = sqrt(sum(IF^2)/n^2) = sqrt(sum(psi^2)) with psi = IF/n.
+ # We need psi = IF_rescaled / n so mixin's sqrt(sum(psi^2)) works.
+ # IF is already at size_gt/size_gt_ctrl scale from above.
+ # Apply the final n/size_gt factor, then divide by n for mixin.
+ inf_gmm_rescaled = inf_gmm * (n_units / size_gt)
+ inf_gmm_scaled = inf_gmm_rescaled / n_units
+
+ # INVARIANT: np.where over boolean masks -> duplicate-free
+ # index arrays (fancy-+= scatter contract, see
+ # staggered_aggregation._combined_if_fast).
+ treated_idx = np.where(treated_mask)[0]
+ treated_inf = inf_gmm_scaled[treated_idx]
+ nonzero_mask = (inf_gmm_scaled != 0) & ~treated_mask
+ control_idx = np.where(nonzero_mask)[0]
+ control_inf = inf_gmm_scaled[control_idx]
+ n_control = int(np.sum(nonzero_mask))
+
+ group_time_effects[(g, t)] = {
+ "effect": att_gmm,
+ "se": se_gt,
+ "t_stat": t_stat,
+ "p_value": p_value,
+ "conf_int": conf_int,
+ "n_treated": n_treated,
+ "n_control": n_control,
+ }
+ influence_func_info[(g, t)] = {
+ "treated_idx": treated_idx,
+ "control_idx": control_idx,
+ "treated_inf": treated_inf,
+ "control_inf": control_inf,
+ }
+ comparison_group_counts[(g, t)] = len(gc_labels)
+ gmm_weights_store[(g, t)] = dict(zip(gc_labels, gmm_w.tolist()))
+
+ # Consolidated OR influence-function rank-deficiency warning.
+ # Finding #17 in the Phase 2 silent-failures audit: the per-pair OR
+ # solve at _compute_did_panel() previously fell back to lstsq with no
+ # signal, so near/fully singular X'WX in the covariate expansion went
+ # to the user as a normal result.
+ if self._lstsq_fallback_tracker and self.rank_deficient_action != "silent":
+ n_cells = len(self._lstsq_fallback_tracker)
+ finite_conds = [c for c in self._lstsq_fallback_tracker if np.isfinite(c)]
+ max_cond = max(finite_conds) if finite_conds else float("inf")
+ warnings.warn(
+ f"Rank-deficient X'WX detected in the outcome-regression "
+ f"influence-function step for {n_cells} (g, g_c, t) pair(s); "
+ f"dropped redundant direction(s) via a rank-guarded inverse. "
+ f"Max condition number of affected X'WX: {max_cond:.2e}. "
+ f"Standard errors use the identified covariate subset; consider "
+ f"dropping collinear covariates or using "
+ f"estimation_method='ipw' to avoid the OR projection.",
+ UserWarning,
+ stacklevel=2 + _frame_offset,
+ )
+
+ # Consolidated PS-Hessian rank-deficiency warning (sibling of the
+ # OR path above). `_compute_pscore` previously fell back from
+ # `np.linalg.inv(X'WX)` to `np.linalg.lstsq` with no signal, so
+ # a rank-deficient propensity-score design silently degraded
+ # IPW/DR influence-function corrections.
+ if self._ps_lstsq_fallback_tracker and self.rank_deficient_action != "silent":
+ n_cells = len(self._ps_lstsq_fallback_tracker)
+ finite_conds = [c for c in self._ps_lstsq_fallback_tracker if np.isfinite(c)]
+ max_cond = max(finite_conds) if finite_conds else float("inf")
+ warnings.warn(
+ f"Rank-deficient X'WX detected in the propensity-score "
+ f"Hessian for {n_cells} (g, g_c, t) pair(s); dropped redundant "
+ f"direction(s) via a rank-guarded inverse. Max condition number "
+ f"of affected X'WX: {max_cond:.2e}. IPW/DR influence-function "
+ f"corrections use the identified covariate subset; consider "
+ f"dropping collinear propensity-score covariates or using "
+ f"estimation_method='reg' to avoid the PS path.",
+ UserWarning,
+ stacklevel=2 + _frame_offset,
+ )
+
+ # Consolidated EPV summary warning
+ if epv_diagnostics:
+ low_epv = {k: v for k, v in epv_diagnostics.items() if v.get("is_low")}
+ if low_epv:
+ n_affected = len(low_epv)
+ n_total = len(epv_diagnostics)
+ min_entry = min(low_epv.values(), key=lambda v: v["epv"])
+ min_g = min(low_epv.keys(), key=lambda k: low_epv[k]["epv"])
+ warnings.warn(
+ f"Low Events Per Variable (EPV) detected in "
+ f"{n_affected} of {n_total} cohort-time cell(s). "
+ f"Minimum EPV: {min_entry['epv']:.1f} (cohort g={min_g[0]}). "
+ f"Consider estimation_method='reg' or fewer covariates. "
+ f"Call results.epv_summary() for per-cohort details.",
+ UserWarning,
+ stacklevel=2 + _frame_offset,
+ )
+
+ if not group_time_effects:
+ raise ValueError(
+ "No valid group-time effects could be computed. "
+ f"Check that the data has sufficient variation in treatment "
+ f"timing and {partition_label}."
+ )
+
+ # For aggregation: use eligible-treated-only cohort assignments so
+ # WIF weights match the point estimate weights (n_treated per cohort,
+ # i.e. P(S=g, Q=1)). This matches the paper's Eq 4.13 which defines
+ # aggregation weights over the treated population (G_i defined only
+ # for Q=1 units). Ineligible units get cohort=0 so they don't
+ # contribute to pg for any treatment group.
+ # Both precomputed["unit_cohorts"] AND df["first_treat"] must be
+ # zeroed for ineligible units because the WIF code reads both.
+ precomputed_agg = dict(precomputed)
+ cohorts_for_agg = precomputed["unit_cohorts"].copy()
+ cohorts_for_agg[eligibility_per_unit == 0] = 0
+ precomputed_agg["unit_cohorts"] = cohorts_for_agg
+
+ df_agg = df.copy()
+ df_agg.loc[df_agg[eligibility] == 0, "first_treat"] = 0
+
+ # Overall ATT via aggregation mixin
+ overall_att, overall_se, overall_effective_df = self._aggregate_simple(
+ group_time_effects, influence_func_info, df_agg, unit, precomputed_agg
+ )
+ # Preserve the ORIGINAL survey df before the simple-overall statistic mutates
+ # it below. The Eq. 4.14 overall (overall_att_es) must fall back to this
+ # original df, never to the simple overall's per-statistic replicate df — the
+ # two statistics can drop different replicate subsets, so reusing the simple
+ # overall's df would silently give overall_att_es the wrong p-value/CI.
+ df_survey_original = df_survey
+ # Use per-statistic effective df from replicate aggregation if available;
+ # otherwise fall back to the original df from the survey design.
+ if overall_effective_df is not None:
+ df_survey = overall_effective_df
+ if survey_metadata is not None:
+ survey_metadata.df_survey = df_survey
+ overall_t_stat, overall_p_value, overall_conf_int = safe_inference(
+ overall_att, overall_se, alpha=self.alpha, df=df_survey
+ )
+
+ # Aggregations
+ event_study_effects = None
+ group_effects = None
+ es_aggregation = None
+ if aggregate in ("event_study", "all"):
+ es_aggregation = self._aggregate_event_study(
+ group_time_effects,
+ influence_func_info,
+ treatment_groups,
+ time_periods,
+ balance_e,
+ df_agg,
+ unit,
+ precomputed_agg,
+ )
+ event_study_effects = es_aggregation.effects
+ if aggregate in ("group", "all"):
+ group_effects = self._aggregate_by_group(
+ group_time_effects,
+ influence_func_info,
+ treatment_groups,
+ precomputed_agg,
+ df_agg,
+ unit,
+ )
+
+ # Paper Eq. (4.14) overall ATT (event-study average): an opt-in summary
+ # alongside the default CS-simple ``overall_att``. ``_aggregate_event_study``
+ # RETURNS it on its aggregation object (it used to stash it on
+ # ``self._event_study_overall``); populated only when the event-study
+ # aggregation ran. Analytical inference here; the bootstrap block below
+ # overrides the SE when ``n_bootstrap > 0`` (mirroring ``overall_se``).
+ overall_att_es = None
+ overall_se_es = None
+ overall_t_stat_es = None
+ overall_p_value_es = None
+ overall_conf_int_es = None
+ # Whether the ANALYTICAL Eq. 4.14 SE was non-finite while its point estimate
+ # was finite (i.e. a contributing horizon's influence function was non-finite).
+ # Captured before any bootstrap override so the terminal warning below does not
+ # misdiagnose a bootstrap-side NaN SE (e.g. cluster-unidentified) as an
+ # analytical-IF failure (the bootstrap path emits its own warning).
+ analytical_overall_es_se_nonfinite = False
+ if aggregate in ("event_study", "all"):
+ es_overall = es_aggregation.overall if es_aggregation is not None else None
+ if es_overall is not None:
+ overall_att_es = es_overall["att"]
+ overall_se_es = es_overall["se"]
+ analytical_overall_es_se_nonfinite = bool(
+ np.isfinite(overall_att_es) and not np.isfinite(overall_se_es)
+ )
+ es_eff_df = es_overall.get("effective_df")
+ # Fall back to the ORIGINAL survey df, not the simple-overall's mutated
+ # per-statistic df (P1 fix): overall_att_es has its own replicate df.
+ df_for_es = es_eff_df if es_eff_df is not None else df_survey_original
+ overall_t_stat_es, overall_p_value_es, overall_conf_int_es = safe_inference(
+ overall_att_es, overall_se_es, alpha=self.alpha, df=df_for_es
+ )
+ else:
+ # Event-study aggregation was requested but yielded no post-treatment
+ # horizon: this is "requested but undefined" -> NaN + warning (the
+ # library's overall-aggregation contract, matching _aggregate_simple),
+ # distinct from "not requested" which leaves the fields None.
+ warnings.warn(
+ "Event-study aggregation was requested but no post-treatment "
+ "horizons are available for the Eq. 4.14 overall (overall_att_es); "
+ "returning NaN.",
+ UserWarning,
+ stacklevel=2 + _frame_offset,
+ )
+ overall_att_es = np.nan
+ overall_se_es = np.nan
+ overall_t_stat_es, overall_p_value_es, overall_conf_int_es = safe_inference(
+ np.nan, np.nan, alpha=self.alpha, df=df_survey
+ )
+
+ # Reject replicate-weight designs for bootstrap — replicate variance
+ # is an analytical alternative, not compatible with bootstrap
+ if (
+ self.n_bootstrap > 0
+ and resolved_survey is not None
+ and hasattr(resolved_survey, "uses_replicate_variance")
+ and resolved_survey.uses_replicate_variance
+ ):
+ raise NotImplementedError(
+ f"{estimator_name} bootstrap (n_bootstrap > 0) is not "
+ "supported with replicate-weight survey designs. Replicate "
+ "weights provide analytical variance; use n_bootstrap=0 instead."
+ )
+
+ # Bootstrap
+ bootstrap_results = None
+ cband_crit_value = None
+ if self.n_bootstrap > 0:
+ bootstrap_results = self._run_multiplier_bootstrap(
+ group_time_effects,
+ influence_func_info,
+ aggregate,
+ balance_e,
+ treatment_groups,
+ time_periods,
+ df_agg,
+ unit,
+ precomputed_agg,
+ self.cband,
+ )
+ if bootstrap_results is not None:
+ overall_se = bootstrap_results.overall_att_se
+ overall_t_stat, overall_p_value, overall_conf_int = safe_inference(
+ overall_att, overall_se, alpha=self.alpha, df=df_survey
+ )
+ overall_conf_int = bootstrap_results.overall_att_ci
+ overall_p_value = bootstrap_results.overall_att_p_value
+
+ # Mirror the override for the Eq. (4.14) event-study-average overall
+ # (only when event-study aggregation produced it). A NaN bootstrap SE
+ # (e.g. cluster-unidentified) correctly NaNs the inference.
+ if (
+ overall_att_es is not None
+ and bootstrap_results.overall_att_es_se is not None
+ ):
+ overall_se_es = bootstrap_results.overall_att_es_se
+ overall_t_stat_es, overall_p_value_es, overall_conf_int_es = safe_inference(
+ overall_att_es, overall_se_es, alpha=self.alpha, df=df_survey
+ )
+ overall_conf_int_es = bootstrap_results.overall_att_es_ci
+ overall_p_value_es = bootstrap_results.overall_att_es_p_value
+ if bootstrap_results.cband_crit_value is not None:
+ cband_crit_value = bootstrap_results.cband_crit_value
+
+ # Update group-time effects with bootstrap SEs
+ if bootstrap_results.group_time_ses:
+ for gt_key in group_time_effects:
+ if gt_key in bootstrap_results.group_time_ses:
+ group_time_effects[gt_key]["se"] = bootstrap_results.group_time_ses[
+ gt_key
+ ]
+ group_time_effects[gt_key]["conf_int"] = (
+ bootstrap_results.group_time_cis[gt_key]
+ )
+ group_time_effects[gt_key]["p_value"] = (
+ bootstrap_results.group_time_p_values[gt_key]
+ )
+ t_val, _, _ = safe_inference(
+ group_time_effects[gt_key]["effect"],
+ bootstrap_results.group_time_ses[gt_key],
+ alpha=self.alpha,
+ df=df_survey,
+ )
+ group_time_effects[gt_key]["t_stat"] = t_val
+
+ if event_study_effects and bootstrap_results.event_study_ses:
+ for e_key in event_study_effects:
+ if e_key in bootstrap_results.event_study_ses:
+ # ses/cis/p_values are populated together.
+ assert (
+ bootstrap_results.event_study_cis is not None
+ and bootstrap_results.event_study_p_values is not None
+ )
+ event_study_effects[e_key]["se"] = (
+ bootstrap_results.event_study_ses[e_key]
+ )
+ event_study_effects[e_key]["conf_int"] = (
+ bootstrap_results.event_study_cis[e_key]
+ )
+ event_study_effects[e_key]["p_value"] = (
+ bootstrap_results.event_study_p_values[e_key]
+ )
+ t_val, _, _ = safe_inference(
+ event_study_effects[e_key]["effect"],
+ bootstrap_results.event_study_ses[e_key],
+ alpha=self.alpha,
+ df=df_survey,
+ )
+ event_study_effects[e_key]["t_stat"] = t_val
+ if cband_crit_value is not None:
+ bs_se = bootstrap_results.event_study_ses[e_key]
+ eff = event_study_effects[e_key]["effect"]
+ event_study_effects[e_key]["cband_conf_int"] = (
+ eff - cband_crit_value * bs_se,
+ eff + cband_crit_value * bs_se,
+ )
+
+ # Update group effects with bootstrap SEs
+ if (
+ group_effects
+ and bootstrap_results.group_effect_ses is not None
+ and bootstrap_results.group_effect_cis is not None
+ and bootstrap_results.group_effect_p_values is not None
+ ):
+ grp_keys = [
+ g for g in group_effects if g in bootstrap_results.group_effect_ses
+ ]
+ for g_key in grp_keys:
+ group_effects[g_key]["se"] = bootstrap_results.group_effect_ses[g_key]
+ group_effects[g_key]["conf_int"] = bootstrap_results.group_effect_cis[
+ g_key
+ ]
+ group_effects[g_key]["p_value"] = (
+ bootstrap_results.group_effect_p_values[g_key]
+ )
+ t_val, _, _ = safe_inference(
+ group_effects[g_key]["effect"],
+ bootstrap_results.group_effect_ses[g_key],
+ alpha=self.alpha,
+ df=df_survey,
+ )
+ group_effects[g_key]["t_stat"] = t_val
+ # Bootstrap se/p/CI replaced the analytical ones, which
+ # is what the retained df described - keeping it would
+ # claim a t-reference that governed nothing.
+ group_effects[g_key]["df_used"] = None
+
+ # Eq. 4.14 overall: an ANALYTICAL non-finite SE under a finite point estimate
+ # (a contributing horizon's influence function is non-finite, or the variance is
+ # unidentified — e.g. a single-PSU/cluster design). Surface it — never NaN the SE
+ # silently. Gated on the analytical-origin flag (captured before the bootstrap
+ # override) and the final state still being non-finite: a bootstrap that supplies a
+ # finite SE rescues it (no warning), and a bootstrap that NaNs the SE for unrelated
+ # reasons (e.g. cluster-unidentified) is reported by the bootstrap's own warning,
+ # not misdiagnosed here as an analytical-IF failure.
+ if (
+ analytical_overall_es_se_nonfinite
+ and overall_se_es is not None
+ and not np.isfinite(overall_se_es)
+ ):
+ warnings.warn(
+ "Eq. 4.14 overall (overall_att_es) point estimate is defined but its "
+ "standard error is undefined (NaN): either a contributing post-treatment "
+ "event-study horizon has a non-finite influence function, or the variance "
+ "is unidentified (e.g. a single-PSU/cluster survey design). overall_se_es "
+ "and its inference fields are NaN.",
+ UserWarning,
+ stacklevel=2 + _frame_offset,
+ )
+
+ n_treated_units = int(np.sum((unit_cohorts > 0) & (eligibility_per_unit == 1)))
+ n_control_units = n_units - n_treated_units
+ n_never_enabled = int(np.sum(unit_cohorts == 0))
+ n_eligible = int(np.sum(eligibility_per_unit == 1))
+ n_ineligible = int(np.sum(eligibility_per_unit == 0))
+
+ self.results_ = StaggeredTripleDiffResults(
+ group_time_effects=group_time_effects,
+ overall_att=overall_att,
+ overall_se=overall_se,
+ overall_t_stat=overall_t_stat,
+ overall_p_value=overall_p_value,
+ overall_conf_int=overall_conf_int,
+ # Cohorts that actually produced at least one (g, t) cell. A
+ # cohort can drop out for several reasons (no eligible treated
+ # units, no valid comparison cohort, every base period outside
+ # the panel); reporting the INPUT roster made `groups` and
+ # `n_groups` claim coverage the estimate does not have. Derived
+ # from the realised cells so it is honest for every reason,
+ # order-preserving against the sorted input roster.
+ groups=[g for g in treatment_groups if g in {gt[0] for gt in group_time_effects}],
+ time_periods=time_periods,
+ n_obs=len(df),
+ n_treated_units=n_treated_units,
+ n_control_units=n_control_units,
+ n_never_enabled=n_never_enabled,
+ n_eligible=n_eligible,
+ n_ineligible=n_ineligible,
+ alpha=self.alpha,
+ control_group=self.control_group,
+ base_period=self.base_period,
+ anticipation=self.anticipation,
+ estimation_method=self.estimation_method,
+ event_study_effects=event_study_effects,
+ group_effects=group_effects,
+ bootstrap_results=bootstrap_results,
+ cband_crit_value=cband_crit_value,
+ pscore_trim=self.pscore_trim,
+ survey_metadata=survey_metadata,
+ comparison_group_counts=comparison_group_counts,
+ gmm_weights=gmm_weights_store,
+ epv_diagnostics=epv_diagnostics if epv_diagnostics else None,
+ epv_threshold=self.epv_threshold,
+ pscore_fallback=self.pscore_fallback,
+ overall_att_es=overall_att_es,
+ overall_se_es=overall_se_es,
+ overall_t_stat_es=overall_t_stat_es,
+ overall_p_value_es=overall_p_value_es,
+ overall_conf_int_es=overall_conf_int_es,
+ )
+ self.is_fitted_ = True
+ return self.results_
+ finally:
+ self._warn_frame_offset = 0
+
+ # ------------------------------------------------------------------
+ # Validation
+ # ------------------------------------------------------------------
+
+ def _validate_inputs(
+ self,
+ df: pd.DataFrame,
+ outcome: str,
+ unit: str,
+ time: str,
+ first_treat: str,
+ eligibility: str,
+ covariates: Optional[List[str]],
+ *,
+ partition_label: str = "eligibility",
+ ) -> None:
+ """Validate input data."""
+ required_cols = [outcome, unit, time, first_treat, eligibility]
+ if covariates:
+ required_cols.extend(covariates)
+ missing = [c for c in required_cols if c not in df.columns]
+ if missing:
+ raise ValueError(f"Missing columns: {missing}")
+
+ elig_vals = df[eligibility].dropna().unique()
+ if not set(elig_vals).issubset({0, 1, 0.0, 1.0}):
+ raise ValueError(
+ f"{partition_label.capitalize()} column '{eligibility}' must be binary (0/1). "
+ f"Found values: {sorted(elig_vals)}"
+ )
+ elig_by_unit = df.groupby(unit)[eligibility].nunique()
+ varying = elig_by_unit[elig_by_unit > 1]
+ if len(varying) > 0:
+ raise ValueError(
+ f"{partition_label.capitalize()} must be time-invariant within units. "
+ f"Found {len(varying)} units with varying {partition_label}."
+ )
+ for col in [outcome, first_treat, eligibility]:
+ if df[col].isna().any():
+ raise ValueError(f"Column '{col}' contains missing values.")
+
+ # Reject non-finite outcomes (Inf/-Inf)
+ if not np.all(np.isfinite(df[outcome])):
+ raise ValueError(
+ f"Column '{outcome}' contains non-finite values (Inf/-Inf). "
+ "All outcome values must be finite."
+ )
+
+ # Reject non-finite covariates
+ if covariates:
+ for cov in covariates:
+ if df[cov].isna().any():
+ raise ValueError(f"Covariate '{cov}' contains missing values.")
+ if not np.all(np.isfinite(df[cov])):
+ raise ValueError(f"Covariate '{cov}' contains non-finite values.")
+ if df[eligibility].nunique() < 2:
+ raise ValueError(
+ "Need both eligible (Q=1) and ineligible (Q=0) units. "
+ f"Only found Q={df[eligibility].unique()[0]}."
+ )
+
+ # Check unique (unit, time) pairs — no duplicate rows
+ dup = df.duplicated(subset=[unit, time], keep=False)
+ if dup.any():
+ raise ValueError(
+ f"Duplicate (unit, time) rows found. "
+ f"{int(dup.sum())} duplicates detected. Panel must have unique rows."
+ )
+
+ # Check balanced panel — every unit observed in exactly the global period set
+ global_periods = set(df[time].unique())
+ n_global_periods = len(global_periods)
+ unit_period_sets = df.groupby(unit)[time].apply(set)
+ mismatched = unit_period_sets[unit_period_sets != global_periods]
+ if len(mismatched) > 0:
+ raise ValueError(
+ "Unbalanced panel detected. All units must be observed in "
+ f"all {n_global_periods} periods. "
+ f"Found {len(mismatched)} units with different period sets."
+ )
+
+ # Check time-invariant first_treat
+ ft_by_unit = df.groupby(unit)[first_treat].nunique()
+ varying_ft = ft_by_unit[ft_by_unit > 1]
+ if len(varying_ft) > 0:
+ raise ValueError(
+ f"first_treat must be time-invariant within units. "
+ f"Found {len(varying_ft)} units with varying first_treat."
+ )
+
+ # Check time-invariant covariates
+ if covariates:
+ for cov in covariates:
+ cov_nunique = df.groupby(unit)[cov].nunique()
+ varying_cov = cov_nunique[cov_nunique > 1]
+ if len(varying_cov) > 0:
+ raise ValueError(
+ f"Covariate '{cov}' must be time-invariant within units. "
+ f"Found {len(varying_cov)} units with varying values."
+ )
+
+ # ------------------------------------------------------------------
+ # Precomputation
+ # ------------------------------------------------------------------
+
+ def _precompute_structures(
+ self,
+ df: pd.DataFrame,
+ outcome: str,
+ unit: str,
+ time: str,
+ eligibility: str,
+ covariates: Optional[List[str]],
+ resolved_survey=None,
+ ) -> PrecomputedData:
+ """Build precomputed structures for efficient computation."""
+ all_units = np.array(sorted(df[unit].unique()))
+ time_periods = sorted(df[time].unique())
+ n_units = len(all_units)
+ n_periods = len(time_periods)
+
+ unit_to_idx = {u: i for i, u in enumerate(all_units)}
+ time_to_col = {t: j for j, t in enumerate(time_periods)}
+
+ outcome_matrix = np.full((n_units, n_periods), np.nan)
+ for _, row in df.iterrows():
+ u_idx = unit_to_idx[row[unit]]
+ t_idx = time_to_col[row[time]]
+ outcome_matrix[u_idx, t_idx] = row[outcome]
+
+ unit_df = df.groupby(unit).first().reindex(all_units)
+ unit_cohorts = unit_df["first_treat"].values.astype(float)
+ eligibility_per_unit = unit_df[eligibility].values.astype(int)
+
+ treatment_groups = sorted([g for g in np.unique(unit_cohorts) if g > 0])
+
+ covariate_matrix = None
+ if covariates:
+ cov_wide = {}
+ for cov in covariates:
+ cov_vals = np.full(n_units, np.nan)
+ for u_id, idx in unit_to_idx.items():
+ u_data = df.loc[df[unit] == u_id, cov]
+ if len(u_data) > 0:
+ cov_vals[idx] = u_data.iloc[0]
+ cov_wide[cov] = cov_vals
+ covariate_matrix = np.column_stack(list(cov_wide.values()))
+
+ # Extract per-unit survey weights and collapse design to unit level
+ survey_weights_arr = None
+ resolved_survey_unit = None
+ if resolved_survey is not None:
+ from diff_diff.survey import collapse_survey_to_unit_level
+
+ survey_weights_arr = (
+ pd.Series(resolved_survey.weights, index=df.index)
+ .groupby(df[unit])
+ .first()
+ .reindex(all_units)
+ .values.astype(np.float64)
+ )
+ # Normalize to sum=n for aggregation/rescaling (matches pweight
+ # convention). Raw weights preserved in resolved_survey_unit for
+ # replicate w_r/w_full ratios — those are inherently scale-invariant.
+ sw_sum = np.sum(survey_weights_arr)
+ if sw_sum > 0:
+ survey_weights_arr = survey_weights_arr * (len(survey_weights_arr) / sw_sum)
+ resolved_survey_unit = collapse_survey_to_unit_level(
+ resolved_survey, df, unit, all_units
+ )
+
+ return {
+ "all_units": all_units,
+ "unit_to_idx": unit_to_idx,
+ "time_periods": time_periods,
+ "time_to_col": time_to_col,
+ "outcome_matrix": outcome_matrix,
+ "unit_cohorts": unit_cohorts,
+ "eligibility_per_unit": eligibility_per_unit,
+ "treatment_groups": treatment_groups,
+ "covariate_matrix": covariate_matrix,
+ "n_units": n_units,
+ "n_periods": n_periods,
+ "survey_weights": survey_weights_arr,
+ "resolved_survey_unit": resolved_survey_unit,
+ "df_survey": (
+ resolved_survey_unit.df_survey if resolved_survey_unit is not None else None
+ ),
+ }
+
+ # ------------------------------------------------------------------
+ # Base period
+ # ------------------------------------------------------------------
+
+ def _get_base_period(self, g: Any, t: Any) -> Optional[Any]:
+ """Determine base period for a (g, t) pair."""
+ if self.base_period == "universal":
+ return g - 1 - self.anticipation
+ else:
+ if t < g - self.anticipation:
+ return t - 1
+ else:
+ return g - 1 - self.anticipation
+
+ # ------------------------------------------------------------------
+ # Three-DiD DDD for one (g, g_c, t) triple
+ # ------------------------------------------------------------------
+
+ def _compute_ddd_gt_gc(
+ self,
+ precomputed: PrecomputedData,
+ g: Any,
+ g_c: Any,
+ t: Any,
+ base_period_val: Any,
+ covariates: Optional[List[str]],
+ pscore_cache: Dict,
+ epv_diagnostics: Optional[Dict] = None,
+ ) -> Optional[Tuple[float, np.ndarray, float]]:
+ """
+ Compute DDD ATT for one (g, g_c, t) triple.
+
+ Returns (att_ddd, inf_full_n_units, size_gt_ctrl) or None.
+ """
+ outcome_matrix = precomputed["outcome_matrix"]
+ time_to_col = precomputed["time_to_col"]
+ unit_cohorts = precomputed["unit_cohorts"]
+ eligibility_per_unit = precomputed["eligibility_per_unit"]
+ covariate_matrix = precomputed["covariate_matrix"]
+ n_units = precomputed["n_units"]
+ survey_weights = precomputed.get("survey_weights")
+
+ t_col = time_to_col[t]
+ b_col = time_to_col[base_period_val]
+
+ # Four sub-groups within this (g, g_c) cell
+ treated_mask = (unit_cohorts == g) & (eligibility_per_unit == 1) # subgroup 4
+ sub_a_mask = (unit_cohorts == g) & (eligibility_per_unit == 0) # subgroup 3
+ sub_b_mask = (unit_cohorts == g_c) & (eligibility_per_unit == 1) # subgroup 2
+ sub_c_mask = (unit_cohorts == g_c) & (eligibility_per_unit == 0) # subgroup 1
+
+ n_treated = int(np.sum(treated_mask))
+ n_a = int(np.sum(sub_a_mask))
+ n_b = int(np.sum(sub_b_mask))
+ n_c = int(np.sum(sub_c_mask))
+
+ # Check for empty subgroups (by count or by survey weight mass)
+ empty = []
+ if n_treated == 0:
+ empty.append(f"(S={g},Q=1)")
+ if n_a == 0:
+ empty.append(f"(S={g},Q=0)")
+ if n_b == 0:
+ empty.append(f"(S={g_c},Q=1)")
+ if n_c == 0:
+ empty.append(f"(S={g_c},Q=0)")
+ # Zero survey-weight mass after subpopulation filtering = effectively empty
+ if not empty and survey_weights is not None:
+ if np.sum(survey_weights[treated_mask]) <= 0:
+ empty.append(f"(S={g},Q=1,mass=0)")
+ if np.sum(survey_weights[sub_a_mask]) <= 0:
+ empty.append(f"(S={g},Q=0,mass=0)")
+ if np.sum(survey_weights[sub_b_mask]) <= 0:
+ empty.append(f"(S={g_c},Q=1,mass=0)")
+ if np.sum(survey_weights[sub_c_mask]) <= 0:
+ empty.append(f"(S={g_c},Q=0,mass=0)")
+ if empty:
+ warnings.warn(
+ f"Empty subgroup(s) {', '.join(empty)} for "
+ f"(g={g}, g_c={g_c}, t={t}). "
+ "Comparison unidentified, skipping.",
+ UserWarning,
+ stacklevel=3 + self._warn_frame_offset,
+ )
+ return None
+
+ if min(n_treated, n_a, n_b, n_c) < 5:
+ warnings.warn(
+ f"Small cell size for (g={g}, g_c={g_c}, t={t}). " "Estimates may be unreliable.",
+ UserWarning,
+ stacklevel=3 + self._warn_frame_offset,
+ )
+
+ # Outcome changes
+ delta_y_all = outcome_matrix[:, t_col] - outcome_matrix[:, b_col]
+ valid = np.isfinite(delta_y_all)
+ for m in [treated_mask, sub_a_mask, sub_b_mask, sub_c_mask]:
+ if not np.all(valid[m]):
+ return None
+
+ # Three pairwise DiDs, each on a 2-cell subset
+ # Collect per-DiD EPV diagnostics; merge worst into (g,t) key later
+ epv_diag_a = {} if epv_diagnostics is not None else None
+ epv_diag_b = {} if epv_diagnostics is not None else None
+ epv_diag_c = {} if epv_diagnostics is not None else None
+
+ # DiD_A: subgroup 4 vs 3 (treated-eligible vs treated-ineligible)
+ pair_a_mask = treated_mask | sub_a_mask
+ did_a = self._run_pairwise_did(
+ delta_y_all,
+ pair_a_mask,
+ treated_mask,
+ sub_a_mask,
+ covariate_matrix,
+ pscore_cache,
+ (g, g, 0, base_period_val),
+ survey_weights=survey_weights,
+ context_label=f"cohort g={g}, DiD_A (g_c={g_c})",
+ epv_diagnostics_out=epv_diag_a,
+ )
+
+ # DiD_B: subgroup 4 vs 2 (treated-eligible vs control-eligible)
+ pair_b_mask = treated_mask | sub_b_mask
+ did_b = self._run_pairwise_did(
+ delta_y_all,
+ pair_b_mask,
+ treated_mask,
+ sub_b_mask,
+ covariate_matrix,
+ pscore_cache,
+ (g, g_c, 1, base_period_val),
+ survey_weights=survey_weights,
+ context_label=f"cohort g={g}, DiD_B (g_c={g_c})",
+ epv_diagnostics_out=epv_diag_b,
+ )
+
+ # DiD_C: subgroup 4 vs 1 (treated-eligible vs control-ineligible)
+ pair_c_mask = treated_mask | sub_c_mask
+ did_c = self._run_pairwise_did(
+ delta_y_all,
+ pair_c_mask,
+ treated_mask,
+ sub_c_mask,
+ covariate_matrix,
+ pscore_cache,
+ (g, g_c, 0, base_period_val),
+ survey_weights=survey_weights,
+ context_label=f"cohort g={g}, DiD_C (g_c={g_c})",
+ epv_diagnostics_out=epv_diag_c,
+ )
+
+ # Merge per-DiD EPV diagnostics: keep the worst (lowest EPV) entry
+ # across all three DiDs for this g_c. If multiple g_c contribute to the
+ # same (g, t) cell, retain the overall minimum EPV across all g_c calls.
+ if epv_diagnostics is not None:
+ candidates = [d for d in [epv_diag_a, epv_diag_b, epv_diag_c] if d]
+ if candidates:
+ worst = min(candidates, key=lambda d: d.get("epv", float("inf")))
+ existing = epv_diagnostics.get((g, t))
+ if existing is None or worst.get("epv", float("inf")) < existing.get(
+ "epv", float("inf")
+ ):
+ epv_diagnostics[(g, t)] = worst
+
+ if did_a is None or did_b is None or did_c is None:
+ return None
+
+ att_a, inf_a = did_a
+ att_b, inf_b = did_b
+ att_c, inf_c = did_c
+
+ att_ddd = att_a + att_b - att_c
+
+ # Three-DiD IF combination: w_j = n_cell / n_pair_j (R's att_dr convention)
+ # With survey weights, use survey-weighted cell sizes
+ if survey_weights is not None:
+ sw_4 = float(np.sum(survey_weights[treated_mask]))
+ sw_3 = float(np.sum(survey_weights[sub_a_mask]))
+ sw_2 = float(np.sum(survey_weights[sub_b_mask]))
+ sw_1 = float(np.sum(survey_weights[sub_c_mask]))
+ n_cell_w = sw_4 + sw_3 + sw_2 + sw_1
+ n_pair_a_w = sw_4 + sw_3
+ n_pair_b_w = sw_4 + sw_2
+ n_pair_c_w = sw_4 + sw_1
+ w_3 = n_cell_w / n_pair_a_w if n_pair_a_w > 0 else 1.0
+ w_2 = n_cell_w / n_pair_b_w if n_pair_b_w > 0 else 1.0
+ w_1 = n_cell_w / n_pair_c_w if n_pair_c_w > 0 else 1.0
+ size_gt_ctrl = n_cell_w
+ else:
+ n_cell = n_treated + n_a + n_b + n_c
+ n_pair_a = n_treated + n_a
+ n_pair_b = n_treated + n_b
+ n_pair_c = n_treated + n_c
+ w_3 = n_cell / n_pair_a if n_pair_a > 0 else 1.0
+ w_2 = n_cell / n_pair_b if n_pair_b > 0 else 1.0
+ w_1 = n_cell / n_pair_c if n_pair_c > 0 else 1.0
+ size_gt_ctrl = float(n_cell)
+
+ # Scatter pair-level IFs into n_units-length vector
+ inf_full = np.zeros(n_units)
+ pair_a_idx = np.where(pair_a_mask)[0]
+ pair_b_idx = np.where(pair_b_mask)[0]
+ pair_c_idx = np.where(pair_c_mask)[0]
+
+ inf_full[pair_a_idx] += w_3 * inf_a
+ inf_full[pair_b_idx] += w_2 * inf_b
+ inf_full[pair_c_idx] -= w_1 * inf_c
+
+ return att_ddd, inf_full, size_gt_ctrl
+
+ # ------------------------------------------------------------------
+ # Pairwise DiD (matches R's compute_did)
+ # ------------------------------------------------------------------
+
+ def _run_pairwise_did(
+ self,
+ delta_y_all: np.ndarray,
+ pair_mask: np.ndarray,
+ treated_mask: np.ndarray,
+ control_mask: np.ndarray,
+ covariate_matrix: Optional[np.ndarray],
+ pscore_cache: Dict,
+ pscore_key: Any,
+ survey_weights: Optional[np.ndarray] = None,
+ context_label: str = "",
+ epv_diagnostics_out: Optional[dict] = None,
+ ) -> Optional[Tuple[float, np.ndarray]]:
+ """
+ Compute a single pairwise DiD ATT and IF on a 2-cell subset.
+
+ Matches R's triplediff::compute_did() formulation exactly:
+ Riesz/Hajek normalization, PS + OR IF corrections.
+
+ Returns (att, inf_func) where inf_func has length n_pair,
+ ordered by pair_mask indices. Returns None if insufficient data.
+ """
+ pair_idx = np.where(pair_mask)[0]
+ n_pair = len(pair_idx)
+ if n_pair == 0:
+ return None
+
+ delta_y = delta_y_all[pair_idx]
+ PA4 = treated_mask[pair_idx].astype(float)
+ PAa = control_mask[pair_idx].astype(float)
+ sw_pair = survey_weights[pair_idx] if survey_weights is not None else None
+
+ n_t = int(np.sum(PA4))
+ n_c = int(np.sum(PAa))
+ if n_t == 0 or n_c == 0:
+ return None
+
+ has_covariates = covariate_matrix is not None and self.estimation_method != "none"
+
+ # Build covariate matrix with intercept for the pair
+ covX = None
+ if has_covariates:
+ # The flag definition above guarantees this (mypy can't track it).
+ assert covariate_matrix is not None
+ X_pair = covariate_matrix[pair_idx]
+ covX = np.column_stack([np.ones(n_pair), X_pair])
+
+ # Compute nuisance parameters based on estimation method
+ pscore = None
+ hessian = None
+ or_delta = np.zeros(n_pair)
+
+ if self.estimation_method in ("ipw", "dr") and covX is not None:
+ pscore, hessian = self._compute_pscore(
+ PA4,
+ covX,
+ pscore_cache,
+ pscore_key,
+ survey_weights=sw_pair,
+ context_label=context_label,
+ epv_diagnostics_out=epv_diagnostics_out,
+ )
+
+ if self.estimation_method in ("reg", "dr") and covX is not None:
+ or_delta = self._compute_or(
+ delta_y,
+ PAa,
+ covX,
+ survey_weights=sw_pair,
+ )
+
+ # Compute ATT and IF (R's compute_did formulation)
+ return self._compute_did_panel(
+ delta_y,
+ PA4,
+ PAa,
+ covX,
+ pscore,
+ hessian,
+ or_delta,
+ survey_weights=sw_pair,
+ )
+
+ # ------------------------------------------------------------------
+ # Core DR/IPW/RA computation (matches R's compute_did exactly)
+ # ------------------------------------------------------------------
+
+ def _compute_did_panel(
+ self,
+ delta_y: np.ndarray,
+ PA4: np.ndarray,
+ PAa: np.ndarray,
+ covX: Optional[np.ndarray],
+ pscore: Optional[np.ndarray],
+ hessian: Optional[np.ndarray],
+ or_delta: np.ndarray,
+ survey_weights: Optional[np.ndarray] = None,
+ ) -> Tuple[float, np.ndarray]:
+ """
+ Pairwise DiD ATT and influence function.
+ Matches R's triplediff::compute_did() line-by-line.
+
+ Parameters
+ ----------
+ delta_y : outcome changes for 2-cell subset (n_pair,)
+ PA4 : treated indicator (n_pair,)
+ PAa : control indicator (n_pair,)
+ covX : covariate matrix with intercept (n_pair, p) or None
+ pscore : propensity scores (n_pair,) or None
+ hessian : (X'WX)^{-1} * n_pair or None
+ or_delta : OR predictions (n_pair,), zeros if no covariates
+ survey_weights : per-observation survey weights (n_pair,) or None
+
+ Returns
+ -------
+ (att, inf_func) where inf_func has length n_pair.
+ """
+ n_pair = len(delta_y)
+ est = self.estimation_method
+
+ # Riesz representers (R lines 243-250)
+ if est == "reg" or pscore is None:
+ w_treat = PA4.copy()
+ w_control = PAa.copy()
+ else:
+ w_treat = PA4.copy()
+ pscore_safe = np.clip(pscore, self.pscore_trim, 1 - self.pscore_trim)
+ w_control = pscore_safe * PAa / (1 - pscore_safe)
+
+ # Incorporate survey weights into Riesz representers
+ if survey_weights is not None:
+ w_treat = w_treat * survey_weights
+ w_control = w_control * survey_weights
+
+ # DR ATT via Hajek normalization (R lines 251-256)
+ resid = delta_y - or_delta
+ riesz_treat = w_treat * resid
+ riesz_control = w_control * resid
+
+ mean_w_treat = np.mean(w_treat)
+ mean_w_control = np.mean(w_control)
+
+ if mean_w_treat <= 0 or mean_w_control <= 0:
+ return float("nan"), np.zeros(n_pair)
+
+ att_treat = np.mean(riesz_treat) / mean_w_treat
+ att_control = np.mean(riesz_control) / mean_w_control
+ dr_att = att_treat - att_control
+
+ # Base IF (R lines 302-304)
+ inf_treat_did = riesz_treat - w_treat * att_treat
+ inf_control_did = riesz_control - w_control * att_control
+
+ # PS correction (R lines 262-273) — IPW and DR only
+ inf_control_pscore = 0.0
+ if est != "reg" and hessian is not None and covX is not None:
+ M2 = np.mean((w_control * (resid - att_control))[:, None] * covX, axis=0)
+ if survey_weights is not None:
+ score_ps = survey_weights[:, None] * (PA4 - pscore_safe)[:, None] * covX
+ else:
+ score_ps = (PA4 - pscore_safe)[:, None] * covX
+ asy_lin_rep_ps = score_ps @ hessian
+ inf_control_pscore = asy_lin_rep_ps @ M2
+
+ # OR correction (R lines 278-300) — reg and DR only
+ inf_treat_or = 0.0
+ inf_cont_or = 0.0
+ if est != "ipw" and covX is not None:
+ M1 = np.mean(w_treat[:, None] * covX, axis=0)
+ M3 = np.mean(w_control[:, None] * covX, axis=0)
+
+ if survey_weights is not None:
+ or_x = (PAa * survey_weights)[:, None] * covX
+ or_ex = (PAa * survey_weights * resid)[:, None] * covX
+ else:
+ or_x = PAa[:, None] * covX
+ or_ex = (PAa * resid)[:, None] * covX
+ XpX = or_x.T @ covX / n_pair
+
+ # Rank-guarded inverse: a near-singular X'WX (constant/collinear
+ # covariate) does not raise LinAlgError, so np.linalg.solve would
+ # return a garbage inverse that inflates the OR influence function.
+ # _rank_guarded_inv truncates redundant directions (finite SE on the
+ # identified subset) and is the sole owner of the tracker append.
+ XpX_inv, _, _ = _rank_guarded_inv(
+ XpX, tracker=getattr(self, "_lstsq_fallback_tracker", None)
+ )
+ asy_linear_or = (XpX_inv @ or_ex.T).T
+
+ inf_treat_or = -(asy_linear_or @ M1)
+ inf_cont_or = -(asy_linear_or @ M3)
+
+ # Final IF assembly (R lines 307-310)
+ inf_control = (inf_control_did + inf_control_pscore + inf_cont_or) / mean_w_control
+ inf_treat = (inf_treat_did + inf_treat_or) / mean_w_treat
+ inf_func = inf_treat - inf_control
+
+ return float(dr_att), inf_func
+
+ # ------------------------------------------------------------------
+ # Nuisance parameter computation
+ # ------------------------------------------------------------------
+
+ def _compute_pscore(
+ self,
+ PA4: np.ndarray,
+ covX: np.ndarray,
+ pscore_cache: Dict,
+ pscore_key: Any,
+ survey_weights: Optional[np.ndarray] = None,
+ context_label: str = "",
+ epv_diagnostics_out: Optional[dict] = None,
+ ) -> Tuple[np.ndarray, np.ndarray]:
+ """Fit logistic P(PA4=1|X). Returns (pscore, hessian).
+
+ hessian = (X'WX)^{-1} * n_pair, matching R's convention.
+ When survey_weights is provided, IRLS uses survey-weighted
+ working weights and the hessian accounts for survey weights.
+ """
+ cached = pscore_cache.get(pscore_key)
+ n_pair = len(PA4)
+
+ if cached is not None:
+ beta_logistic, cached_diag = cached
+ z = np.dot(covX, beta_logistic)
+ z = np.clip(z, -500, 500)
+ pscore = 1 / (1 + np.exp(-z))
+ if epv_diagnostics_out is not None and cached_diag:
+ epv_diagnostics_out.update(cached_diag)
+ else:
+ X_no_intercept = covX[:, 1:] # solve_logit adds its own intercept
+ diag = {}
+ try:
+ beta_logistic, pscore = solve_logit(
+ X_no_intercept,
+ PA4,
+ rank_deficient_action=self.rank_deficient_action,
+ weights=survey_weights,
+ epv_threshold=self.epv_threshold,
+ context_label=context_label,
+ diagnostics_out=diag,
+ )
+ _check_propensity_diagnostics(pscore, self.pscore_trim)
+ # Zero-fill NaN coefficients (from rank-deficient columns)
+ # before caching, so cache reuse doesn't propagate NaN.
+ # Cache alongside EPV diagnostics for replay on cache hits.
+ beta_clean = np.where(np.isfinite(beta_logistic), beta_logistic, 0.0)
+ pscore_cache[pscore_key] = (beta_clean, diag)
+ except (np.linalg.LinAlgError, ValueError):
+ if self.pscore_fallback == "error" or self.rank_deficient_action == "error":
+ raise
+ ctx = f" for {context_label}" if context_label else ""
+ warnings.warn(
+ f"Propensity score estimation failed{ctx}. "
+ f"Falling back to unconditional propensity "
+ f"(propensity model ignores covariates; outcome "
+ f"regression still uses them for DR). "
+ f"Consider estimation_method='reg' to avoid "
+ f"propensity scores entirely.",
+ UserWarning,
+ stacklevel=5 + self._warn_frame_offset,
+ )
+ # Use survey-weighted treated share when weights available
+ if survey_weights is not None:
+ pos = survey_weights > 0
+ if np.any(pos):
+ p_uc = np.average(PA4[pos], weights=survey_weights[pos])
+ else:
+ p_uc = np.mean(PA4)
+ else:
+ p_uc = np.mean(PA4)
+ pscore = np.full(n_pair, p_uc)
+ pscore = np.clip(pscore, self.pscore_trim, 1 - self.pscore_trim)
+ # No hessian for unconditional fallback
+ return pscore, None
+ if epv_diagnostics_out is not None and diag:
+ epv_diagnostics_out.update(diag)
+
+ pscore = np.clip(pscore, 1e-6, 1 - 1e-6)
+
+ # Hessian: (X'WX)^{-1} * n (matching R's compute_pscore)
+ W = pscore * (1 - pscore)
+ if survey_weights is not None:
+ W = W * survey_weights
+ XWX = covX.T @ (W[:, None] * covX)
+ # Rank-guarded inverse (sibling of the OR-side guard in
+ # _compute_did_panel). A near-singular X'WX (constant/collinear
+ # propensity covariate) does not raise LinAlgError, so the old
+ # np.linalg.inv returned a garbage inverse that inflated IPW/DR
+ # influence-function corrections. The helper truncates redundant
+ # directions and is the sole owner of the PS-Hessian tracker append.
+ XWX_inv, _, _ = _rank_guarded_inv(
+ XWX, tracker=getattr(self, "_ps_lstsq_fallback_tracker", None)
+ )
+ hessian = XWX_inv * n_pair
+
+ return pscore, hessian
+
+ def _compute_or(
+ self,
+ delta_y: np.ndarray,
+ PAa: np.ndarray,
+ covX: np.ndarray,
+ survey_weights: Optional[np.ndarray] = None,
+ ) -> np.ndarray:
+ """Fit OLS on control outcome changes. Returns or_delta for all pair units.
+
+ Honors self.rank_deficient_action for collinear covariates. The outcome
+ regression is fit through the shared scale-robust solver
+ (``solve_ols`` -> column-equilibrated SVD; matches TripleDifference and
+ R's lm()/QR), with optional WLS via ``solve_ols(weights=...)``.
+ """
+ from diff_diff.linalg import solve_ols as _solve_ols
+
+ control_mask = PAa > 0
+ n_c = int(np.sum(control_mask))
+ if n_c == 0:
+ return np.zeros(len(delta_y))
+
+ X_control = covX[control_mask]
+ y_control = delta_y[control_mask]
+ sw_control = survey_weights[control_mask] if survey_weights is not None else None
+
+ # Outcome regression via the shared scale-robust solver (`solve_ols` ->
+ # column-equilibrated SVD/gelsd; matches TripleDifference and R's
+ # lm()/QR). Replaces the prior `cho_solve(X'X)` cache fast path, which
+ # was NOT scale-equilibrated (a large-scale covariate would corrupt the
+ # OR fit via the normal-equations Cholesky). We only need beta for the
+ # OR prediction, so zero the NaN (dropped-column) coefficients.
+ beta, _, _ = _solve_ols(
+ X_control,
+ y_control,
+ rank_deficient_action=self.rank_deficient_action,
+ weights=sw_control,
+ return_vcov=False,
+ )
+ beta = np.where(np.isnan(beta), 0.0, beta)
+
+ return covX @ beta
+
+ # ------------------------------------------------------------------
+ # GMM-optimal combination (matches R's att_gt GMM procedure)
+ # ------------------------------------------------------------------
+
+ def _combine_gmm(
+ self,
+ att_vec: np.ndarray,
+ inf_func_matrix: np.ndarray,
+ n_units: int,
+ ) -> Tuple[float, np.ndarray, np.ndarray, float]:
+ """
+ Combine comparison-group-specific estimates via GMM-optimal weights.
+
+ Returns (att_gmm, inf_gmm, weights, se_gmm).
+ """
+ k = len(att_vec)
+
+ if k == 1:
+ att_gmm = float(att_vec[0])
+ inf_gmm = inf_func_matrix[0].copy()
+ # R's SE: sqrt(sum(IF^2) / n^2)
+ se_gmm = float(np.sqrt(np.sum(inf_gmm**2) / n_units**2))
+ return att_gmm, inf_gmm, np.array([1.0]), se_gmm
+
+ # R: OMEGA <- cov(inf_mat_local) — sample covariance, ddof=1
+ Omega = np.cov(inf_func_matrix)
+
+ ones = np.ones(k)
+ try:
+ Omega_inv = np.linalg.inv(Omega)
+ except np.linalg.LinAlgError:
+ warnings.warn(
+ "Singular covariance matrix in GMM combination. " "Using pseudoinverse.",
+ UserWarning,
+ stacklevel=3 + self._warn_frame_offset,
+ )
+ Omega_inv = np.linalg.pinv(Omega)
+
+ denom = float(ones @ Omega_inv @ ones)
+ if denom <= 0 or not np.isfinite(denom):
+ weights = np.full(k, 1.0 / k)
+ att_gmm = float(weights @ att_vec)
+ inf_gmm = weights @ inf_func_matrix
+ se_gmm = float(np.sqrt(np.sum(inf_gmm**2) / n_units**2))
+ else:
+ weights = (Omega_inv @ ones) / denom
+ att_gmm = float(weights @ att_vec)
+ inf_gmm = weights @ inf_func_matrix
+ # R: gmm_se <- sqrt(1 / (n * sum(inv_OMEGA)))
+ se_gmm = float(np.sqrt(1.0 / (n_units * denom)))
+
+ return att_gmm, inf_gmm, weights, se_gmm
diff --git a/diff_diff/guides/llms-autonomous.txt b/diff_diff/guides/llms-autonomous.txt
index 1f0a366b..bf7b251c 100644
--- a/diff_diff/guides/llms-autonomous.txt
+++ b/diff_diff/guides/llms-autonomous.txt
@@ -577,6 +577,9 @@ some states and some demographic subgroups within states):
- `TripleDifference` - classic two-period DDD.
- `StaggeredTripleDifference` - staggered DDD, robust to cohort-mixing.
+ DEPRECATED 3.9 (removed 4.0): use
+ `TripleDifference().fit(..., unit=, time=, first_treat=, partition=)`,
+ the same engine. `control_group` takes `not_yet_treated`/`never_treated`.
Triple-difference is not automatically detected by `profile_panel`;
it requires the caller to identify the third comparison axis. If a
diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt
index 57d5ff0f..66321805 100644
--- a/diff_diff/guides/llms-full.txt
+++ b/diff_diff/guides/llms-full.txt
@@ -1336,6 +1336,14 @@ plot_bacon(results)
### StaggeredTripleDifference
+DEPRECATED in 3.9, removed in 4.0 (ledger row M-013). Use
+`TripleDifference().fit(..., unit=, time=, first_treat=, partition=)` - the same
+engine, so the numbers are identical. Vocabulary on the merged surface:
+`eligibility=` is `partition=`, and `control_group` takes the underscored values
+`not_yet_treated`/`never_treated` (this class keeps R's compact `notyettreated`/
+`nevertreated` until removal). `cluster=` raises on the merged surface instead of
+being accepted-and-ignored. Alias `SDDD` is deprecated with the class.
+
Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD estimator for designs with two eligibility criteria and staggered treatment timing.
```python
diff --git a/diff_diff/guides/llms-practitioner.txt b/diff_diff/guides/llms-practitioner.txt
index 9e100f39..4015f887 100644
--- a/diff_diff/guides/llms-practitioner.txt
+++ b/diff_diff/guides/llms-practitioner.txt
@@ -173,7 +173,8 @@ Use this decision tree to select the appropriate estimator:
```
Is this a triple-difference (DDD) design? (Two criteria: e.g., policy + eligibility)
|-- YES, simultaneous (2x2x2): TripleDifference (DDD)
-|-- YES, staggered timing: StaggeredTripleDifference (SDDD)
+|-- YES, staggered timing: TripleDifference(...).fit(..., first_treat=)
+| (StaggeredTripleDifference/SDDD deprecated 3.9)
|
Is treatment continuous (doses/intensities)?
Route by ESTIMAND first, not by whether never-treated units happen
diff --git a/diff_diff/guides/llms.txt b/diff_diff/guides/llms.txt
index c41b3d7b..e2d8be19 100644
--- a/diff_diff/guides/llms.txt
+++ b/diff_diff/guides/llms.txt
@@ -68,14 +68,14 @@ The site is organized into 5 sections, each with a landing page:
- [SpilloverDiD](https://diff-diff.readthedocs.io/en/stable/api/spillover.html): Butts (2021) ring-indicator spillover-aware DiD identifying direct effect on treated + per-ring spillover-on-control; reuses `conley_coords` for ring construction; handles non-staggered and staggered timing; supports `SurveyDesign(weights, strata, psu, fpc)` under `vcov_type="hc1"` with optional `cluster=` for CR1 via Gerber (2026) Binder TSL (Wave E.1) and under `vcov_type="conley"` via a panel-aware stratified-Conley sandwich on per-period PSU totals (Wave E.2 cross-sectional `conley_lag_cutoff=0`) extended in Wave E.2 follow-up to `conley_lag_cutoff > 0` via panel-block composition with within-PSU serial Bartlett HAC (Newey-West 1987 separable form; `lag>0` requires an effective PSU via explicit `survey_design.psu` or injected `cluster=`), both composed with the Wave D Gardner GMM correction; `SurveyDesign.subpopulation()` preserves full-design `n_psu` / `df_survey` via zero-padded scores at the meat-helper boundary (Wave E.3, R `svyrecvar(subset())` form) (replicate weights queued as follow-up)
- [SyntheticDiD](https://diff-diff.readthedocs.io/en/stable/api/estimators.html): Synthetic DiD combining standard DiD and synthetic control methods for few treated units
- [SyntheticControl](https://diff-diff.readthedocs.io/en/stable/api/synthetic_control.html): Abadie, Diamond & Hainmueller (2010) classic synthetic control for ONE treated unit — donor-weight counterfactual, predictor-importance V via nested / cv (out-of-sample, ADH 2015; needs predictors spanning both train/val windows, so default single-period lags are rejected) / inverse-variance (1/Var on raw predictors, bypasses standardize) / custom, gap path + pre-RMSPE; no analytical SE (inference fields NaN), significance via in-space placebo permutation inference (`in_space_placebo()`, post/pre RMSPE-ratio, p = rank/(n_placebos+1)); ADH-2015 §4 robustness: `leave_one_out()` donor-robustness + `in_time_placebo()` backdating placebo; confidence sets by test inversion (Firpo-Possebom 2018 §4): `test_sharp_null()` + `confidence_set(family="constant"|"linear")` re-rank the placebo gaps into a confidence set for the effect path (the analytical `conf_int` stays NaN); conformal inference (Chernozhukov-Wüthrich-Zhu 2021): `conformal_test()` (joint sharp-null p-value), `conformal_confidence_intervals()` (pointwise per-period CIs), `conformal_average_effect()` (average-effect CI) — fit their OWN constrained-LS proxy under the null on all periods and permute residuals over time (moving-block / iid schemes; `conf_int` still NaN). Alias `SCM`.
-- [TripleDifference](https://diff-diff.readthedocs.io/en/stable/api/triple_diff.html): Triple difference (DDD) estimator for designs requiring two criteria for treatment eligibility
+- [TripleDifference](https://diff-diff.readthedocs.io/en/stable/api/triple_diff.html): Triple difference (DDD) estimator for designs requiring two criteria for treatment eligibility. Since 3.9 it serves BOTH DDD designs from one signature: the 2x2x2 design as before, and staggered adoption when `first_treat=` is supplied (with the unit id, the calendar period column and `partition=`; the staggered params are keyword-only, and mixing the two designs' params raises). `control_group` takes 'not_yet_treated'/'never_treated'
- [ContinuousDiD](https://diff-diff.readthedocs.io/en/stable/api/continuous_did.html): Callaway, Goodman-Bacon & Sant'Anna (2024) continuous treatment DiD with dose-response curves
- [HeterogeneousAdoptionDiD](https://diff-diff.readthedocs.io/en/stable/api/had.html): de Chaisemartin, Ciccia, D'Haultfœuille & Knau (2026) for designs where **no unit remains untreated**; local-linear estimator at the dose support boundary returning Weighted Average Slope (WAS) on Design 1' (`d̲=0` / QUG) or `WAS_{d̲}` on Design 1 (`d̲>0`, continuous-near-d̲ or mass-point), with multi-period event-study extension (last-treatment cohort, pointwise CIs; the mode is panel-inferred since 3.9 - two periods -> overall WAS, more -> event-study - and post-fit `results.aggregate('simple')`/`.aggregate('event_study')` are pure views, rows M-027/M-139). **Panel-only** in this release (repeated cross-sections rejected by the validator). Alias `HAD`.
- [RegressionDiscontinuity](https://diff-diff.readthedocs.io/en/stable/api/regression_discontinuity.html): Calonico, Cattaneo & Titiunik (2014) sharp AND fuzzy regression discontinuity with robust bias-corrected inference, parity-targeting R rdrobust 4.0.0 (all 10 data-driven bandwidth selectors, mass-point handling, three-row conventional/bias-corrected/robust output; canonical `att` = the bias-corrected estimate with a coherent robust CI - rdrobust's printed headline is `att_conventional`). Fuzzy via `fit(..., takeup=...)`: local Wald ratio (complier LATE for binary take-up under monotonicity; ratio-of-jumps otherwise - the `estimand` field says which), first-stage `first_stage*` block, weak-first-stage warning. Covariate adjustment via `fit(..., covariates=[...])` (CCFT 2019 additive common-coefficient, R's `covs=`): SAME estimand, precision only; requires covariate balance at the cutoff (testable: fit each covariate as the outcome); covariate-aware bandwidths; collinear columns dropped with a warning (`covs_drop`). Cluster-robust variance is a documented follow-up. Alias `RDD`.
- [StackedDiD](https://diff-diff.readthedocs.io/en/stable/api/stacked_did.html): Wing, Freedman & Hollingsworth (2024) stacked DiD with Q-weights and sub-experiments; optional covariate balancing (`balance="entropy"`, Ustyuzhanin 2026)
- [EfficientDiD](https://diff-diff.readthedocs.io/en/stable/api/efficient_did.html): Chen, Sant'Anna & Xie (2025) efficient DiD with optimal weighting for tighter SEs
- [TROP](https://diff-diff.readthedocs.io/en/stable/api/trop.html): Triply Robust Panel estimator (Athey et al. 2025) with nuclear norm factor adjustment (absorbing by default; `non_absorbing=True` for on/off treatment, method='local')
-- [StaggeredTripleDifference](https://diff-diff.readthedocs.io/en/stable/api/staggered.html#staggeredtripledifference): Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD with group-time ATT
+- [StaggeredTripleDifference](https://diff-diff.readthedocs.io/en/stable/api/staggered.html#staggeredtripledifference): Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD with group-time ATT. DEPRECATED in 3.9, removed in 4.0 - use `TripleDifference` with `first_treat=` (supplying the unit id, the calendar period column and `partition=`) - the same engine (`eligibility=` is `partition=` there; `control_group` takes the underscored values). Alias `SDDD` deprecated with it
- [WooldridgeDiD](https://diff-diff.readthedocs.io/en/stable/api/wooldridge_etwfe.html): Wooldridge (2023, 2025) ETWFE — saturated OLS, logit/Poisson QMLE (ASF-based ATT). Alias: ETWFE
- [LPDiD](https://diff-diff.readthedocs.io/en/stable/api/lpdid.html): Dube, Girardi, Jorda & Taylor (2025) Local Projections DiD: per-horizon long-difference event study on clean controls (no negative weighting); variance- or equally-weighted ATT, premean differencing, pooled pre/post, fast. Absorbing by default; non-absorbing (reversible) treatment via `non_absorbing="first_entry"` (Eq. 12) or `"effect_stabilization"` (Eq. 13, window `L`). Complex-survey designs (pweight + stratified-PSU TSL SEs) on the default path via `fit(survey_design=...)`.
- [ChangesInChanges](https://diff-diff.readthedocs.io/en/stable/api/changes_in_changes.html): Athey & Imbens (2006) nonlinear/distributional DiD for the 2x2 design: recovers the treated group's full counterfactual outcome distribution and quantile treatment effects (ATT + QTE grid) via the CDF transformation `F_10(F_00^{-1}(F_01(y)))`; invariant to monotone outcome transformations (unconditional fits; the covariate QR branch is not); bootstrap inference (panel or repeated cross-section resampling); point parity with R `qte::CiC()`, including its covariate branch (`covariates=` -> per-cell linear quantile regression, Melly-Santangelo-style conditional CiC). Continuous outcomes, numeric covariates. Alias `CiC`.
diff --git a/diff_diff/power.py b/diff_diff/power.py
index 69e1fda2..bba947fe 100644
--- a/diff_diff/power.py
+++ b/diff_diff/power.py
@@ -29,6 +29,10 @@
from scipy import stats
from diff_diff.results_base import Diagnostic
+from diff_diff.utils import (
+ STAGGERED_DDD_CTOR_PARAMS,
+ staggered_ddd_ctor_offenders,
+)
# Maximum sample size returned when effect is too small to detect
# (e.g., zero effect or extremely small relative to noise)
@@ -601,6 +605,69 @@ def _first(r: Any, *attrs: str, default: Any = _nan) -> Any:
# staggered cohort data, not factor-model or 2x2x2 data).
_SURVEY_UNSUPPORTED = frozenset({"TROP", "SyntheticDiD", "TripleDifference"})
+# TripleDifference serves two designs from one class since 3.9 (row M-013), but
+# power only ever fits the 2x2x2 one: both registered DDD profiles build
+# (group, partition, post) fit kwargs. A staggered-configured estimator would
+# therefore be simulated under the wrong design, so reject it at the front door.
+# The ROSTER of staggered-only constructor params is real information (it is not
+# derivable from the signature alone - the 2x2x2 params live there too), but the
+# DEFAULT VALUES are not: hard-coding them would let a future constructor-default
+# change silently reclassify an otherwise-default estimator as staggered-configured
+# and reject a legitimate 2x2x2 power run. So the roster is explicit and the values
+# are read off the estimator's OWN signature at call time. Deriving per-instance
+# rather than importing TripleDifference also preserves this module's deliberate
+# no-estimator-import design (dispatch is by ``type(estimator).__name__``).
+# Row M-013's staggered-only constructor roster and its default derivation live
+# in utils so this module and TripleDifference.fit() cannot drift apart: fit()
+# rejects these in 2x2x2 mode, and power rejects a staggered-configured
+# estimator at the front door. Both must mean the same thing by "staggered
+# configuration". See utils.STAGGERED_DDD_CTOR_PARAMS for why bootstrap_weights,
+# seed and cband are excluded (inert without n_bootstrap > 0, accepted by both).
+_DDD_STAGGERED_CTOR_PARAMS = STAGGERED_DDD_CTOR_PARAMS
+# The mode TRIGGER is a fit param, and estimator_kwargs IS the fit-kwargs
+# channel, so these can genuinely flip the merged class into staggered mode
+# against 2x2x2 data - the one path that reaches the trigger at all.
+#
+# The two tuples encode DIFFERENT tests, mirroring `TripleDifference.fit()`
+# exactly. `first_treat` and `unit` are sentinel-defaulted there
+# (`x is not NOT_SUPPLIED`), so SUPPLYING them at all is the signal - an
+# explicit `first_treat=None` still selects staggered mode and then fails on a
+# missing column. `aggregate`/`balance_e` default to None and fit only rejects a
+# NON-None value, so keying on presence here would reject
+# `estimator_kwargs={"aggregate": None}` - a config `fit()` accepts - and break
+# the documented "legal to fit implies legal to simulate" boundary.
+_DDD_STAGGERED_FIT_KEYS_BY_PRESENCE = ("first_treat", "unit")
+_DDD_STAGGERED_FIT_KEYS_BY_VALUE = ("aggregate", "balance_e")
+# Retained as the union for callers/tests that want the whole roster.
+_DDD_STAGGERED_FIT_KEYS = _DDD_STAGGERED_FIT_KEYS_BY_PRESENCE + _DDD_STAGGERED_FIT_KEYS_BY_VALUE
+
+
+def _reject_staggered_ddd_config(estimator: Any, est_kwargs: Dict[str, Any]) -> None:
+ """Reject a staggered-configured TripleDifference at the power front door.
+
+ Two arms, because the two config channels are different things: constructor
+ values live on the passed INSTANCE, while ``estimator_kwargs`` is forwarded
+ to ``fit()``. Note there is no working staggered route through the power
+ surface today - the custom ``data_generator`` path still builds fit kwargs
+ from the registered profile whenever one exists - so the message must not
+ advertise one.
+ """
+ if type(estimator).__name__ != "TripleDifference":
+ return
+ offenders = staggered_ddd_ctor_offenders(estimator)
+ offenders += [k for k in _DDD_STAGGERED_FIT_KEYS_BY_PRESENCE if k in est_kwargs]
+ offenders += [k for k in _DDD_STAGGERED_FIT_KEYS_BY_VALUE if est_kwargs.get(k) is not None]
+ if not offenders:
+ return
+ raise ValueError(
+ f"Power analysis for TripleDifference covers the 2x2x2 design only, but "
+ f"{', '.join(sorted(set(offenders)))} configure(s) the staggered DDD mode. Both "
+ f"registered DDD data generators produce 2x2x2 data and fit with "
+ f"(group=, partition=, post=), so a staggered configuration would be simulated "
+ f"under the wrong design. Staggered-DDD power is not supported yet (tracked in "
+ f"TODO.md); drop the staggered configuration to run the 2x2x2 analysis."
+ )
+
def _check_staggered_dgp_compat(
estimator: Any,
@@ -2209,6 +2276,10 @@ def simulate_power(
data_gen_kwargs = data_generator_kwargs or {}
est_kwargs = estimator_kwargs or {}
+ # Row M-013: TripleDifference is two designs behind one class; power fits
+ # the 2x2x2 one only.
+ _reject_staggered_ddd_config(estimator, est_kwargs)
+
# Block survey_design in estimator_kwargs when survey_config is active.
# Custom survey design overrides go through SurveyPowerConfig.survey_design.
if use_survey_dgp and "survey_design" in est_kwargs:
@@ -2886,6 +2957,12 @@ def simulate_mde(
estimator_name = type(estimator).__name__
search_path: List[Dict[str, float]] = []
+ # Row M-013: reject a staggered-configured TripleDifference here rather than
+ # letting it surface from inside a simulation replicate. This runs BEFORE the
+ # per-entry-point setup below, which is why it cannot simply ride on
+ # simulate_power's guard.
+ _reject_staggered_ddd_config(estimator, estimator_kwargs or {})
+
# Compute effective N for DDD (N is fixed throughout MDE search). Only the
# cross-sectional 2x2x2 DGP (n_periods <= 2) rounds n_units to a multiple of
# 8; the panel DGP (n_periods > 2) maps n_units directly, so report None.
@@ -3106,6 +3183,12 @@ def simulate_sample_size(
estimator_name = type(estimator).__name__
search_path: List[Dict[str, float]] = []
+ # Row M-013: reject a staggered-configured TripleDifference here rather than
+ # letting it surface from inside a simulation replicate. This runs BEFORE the
+ # per-entry-point setup below, which is why it cannot simply ride on
+ # simulate_power's guard.
+ _reject_staggered_ddd_config(estimator, estimator_kwargs or {})
+
# Determine min_n from registry. DDD splits cross-sectional (n_periods <= 2,
# 2x2x2 factorial) from panel (n_periods > 2, generate_ddd_panel_data).
registry = _get_registry()
diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py
index 82d37747..bc9ffbb6 100644
--- a/diff_diff/staggered.py
+++ b/diff_diff/staggered.py
@@ -7,7 +7,7 @@
import bisect
import warnings
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Tuple
import numpy as np
import pandas as pd
@@ -531,6 +531,11 @@ class CallawaySantAnna(
multiple time periods. Journal of Econometrics, 225(2), 200-230.
"""
+ # Names this estimator in the shared bootstrap mixin's user-facing
+ # warnings. The mixin is shared with the other hosts, so a hard-coded
+ # literal there would misname whichever surface was actually fit.
+ _BOOTSTRAP_LABEL: ClassVar[str] = "CallawaySantAnna"
+
def __init__(
self,
control_group: str = "never_treated",
diff --git a/diff_diff/staggered_bootstrap.py b/diff_diff/staggered_bootstrap.py
index 180f65a7..235476b8 100644
--- a/diff_diff/staggered_bootstrap.py
+++ b/diff_diff/staggered_bootstrap.py
@@ -8,7 +8,7 @@
import warnings
from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Iterator, List, Optional, Tuple
import numpy as np
@@ -128,6 +128,15 @@ class CallawaySantAnnaBootstrapMixin:
"""
# Type hints for attributes accessed from the main class
+ # Host-supplied estimator name for user-facing bootstrap warnings. Declared
+ # here because mypy type-checks the mixin independently of its hosts, so
+ # `self._BOOTSTRAP_LABEL` would otherwise be [attr-defined]. ClassVar, not a
+ # bare annotation: the hosts set it as class-level constant data, and an
+ # instance-variable declaration here would make each of those a
+ # "cannot override instance variable with class variable" [misc] error.
+ # (Contrast `_warn_frame_offset`, which IS assigned via instance and so must
+ # NOT be a ClassVar.)
+ _BOOTSTRAP_LABEL: ClassVar[str]
n_bootstrap: int
bootstrap_weights: str
alpha: float
@@ -194,14 +203,18 @@ def _run_multiplier_bootstrap(
CSBootstrapResults
Bootstrap inference results.
"""
- # Warn about low bootstrap iterations
+ # Warn about low bootstrap iterations. This site is USER-attributed, and
+ # the DDD engine reaches it through one or two extra frames depending on
+ # which surface was called, so it consults the offset the engine mirrors
+ # onto the instance for the duration of a fit. CallawaySantAnna never
+ # sets the attribute, so its attribution is bit-identical to 3.x.
if self.n_bootstrap < 50:
warnings.warn(
f"n_bootstrap={self.n_bootstrap} is low. Consider n_bootstrap >= 199 "
"for reliable inference. Percentile confidence intervals and p-values "
"may be unreliable with few iterations.",
UserWarning,
- stacklevel=3,
+ stacklevel=3 + getattr(self, "_warn_frame_offset", 0),
)
rng = np.random.default_rng(self.seed)
@@ -389,7 +402,7 @@ def _agg_mass(gt):
import warnings as _warnings
_warnings.warn(
- f"CallawaySantAnna bootstrap with survey/cluster design "
+ f"{self._BOOTSTRAP_LABEL} bootstrap with survey/cluster design "
f"has only {len(psu_ids)} PSU(s); bootstrap variance is "
"unidentified. All bootstrap inference fields "
"(overall_se, group_time_ses, event_study_ses, "
diff --git a/diff_diff/staggered_triple_diff.py b/diff_diff/staggered_triple_diff.py
index d60641bc..9f2aafaf 100644
--- a/diff_diff/staggered_triple_diff.py
+++ b/diff_diff/staggered_triple_diff.py
@@ -7,20 +7,19 @@
Core pairwise DiD computation matches R's triplediff::compute_did() exactly
(Riesz/Hajek normalization, separate M1/M3 OR corrections, hessian = (X'WX)^{-1}*n).
+
+The estimation engine itself lives in `_staggered_triple_diff_engine.py`, shared
+verbatim with `TripleDifference`'s staggered mode (ledger row M-013). This module
+keeps the deprecated class's 3.x surface: its own constructor (including R's
+compact `control_group` spellings) and its own `fit` signature (including the
+`eligibility` parameter name), both frozen until the 4.0 removal.
"""
import warnings
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
-
-import numpy as np
-import pandas as pd
+from typing import TYPE_CHECKING, ClassVar, List, Optional
from diff_diff._base import BaseEstimator
-from diff_diff.linalg import (
- _check_propensity_diagnostics,
- _rank_guarded_inv,
- solve_logit,
-)
+from diff_diff._staggered_triple_diff_engine import _StaggeredTripleDiffEngineMixin
from diff_diff.staggered_aggregation import (
CallawaySantAnnaAggregationMixin,
)
@@ -28,9 +27,11 @@
CallawaySantAnnaBootstrapMixin,
)
from diff_diff.staggered_triple_diff_results import StaggeredTripleDiffResults
-from diff_diff.utils import safe_inference, validate_n_bootstrap
+from diff_diff.utils import validate_n_bootstrap
if TYPE_CHECKING:
+ import pandas as pd
+
from diff_diff.survey import SurveyDesign
__all__ = [
@@ -38,11 +39,22 @@
"StaggeredTripleDiffResults",
]
-# Type alias for pre-computed structures
-PrecomputedData = Dict[str, Any]
+# StaggeredTripleDifference 3.9 deprecation message (row M-013; the SDDD alias
+# is the same class object, so constructing via the alias emits this too -
+# row M-064). Pinned verbatim by tests/test_v4_merge_ddd.py and the targeted
+# pytest filter in pyproject.toml.
+_SDDD_DEPRECATION_MSG = (
+ "StaggeredTripleDifference is deprecated and will be removed in 4.0; use "
+ "TripleDifference().fit(..., unit=, time=, first_treat=, partition=) "
+ "instead - the same engine, so the numbers are unchanged. Two vocabulary "
+ "changes on the merged surface: the eligibility= parameter is named "
+ "partition=, and control_group takes the underscored values "
+ "'not_yet_treated'/'never_treated'. The SDDD alias is deprecated with it."
+)
class StaggeredTripleDifference(
+ _StaggeredTripleDiffEngineMixin,
CallawaySantAnnaBootstrapMixin,
CallawaySantAnnaAggregationMixin,
BaseEstimator,
@@ -50,6 +62,14 @@ class StaggeredTripleDifference(
"""
Staggered Triple Difference (DDD) estimator.
+ .. deprecated:: 3.9
+ Deprecated in 3.9 and removed in 4.0 (ledger row M-013). Use
+ ``TripleDifference().fit(..., unit=, time=, first_treat=, partition=)``,
+ which runs this exact engine. The ``eligibility=`` parameter is named
+ ``partition=`` there, and ``control_group`` takes the underscored
+ values ``"not_yet_treated"``/``"never_treated"``. The ``SDDD`` alias is
+ deprecated with the class (row M-064).
+
Computes group-time average treatment effects ATT(g,t) for settings
with staggered adoption and a binary eligibility dimension, using the
three-DiD decomposition of Ortiz-Villavicencio & Sant'Anna (2025).
@@ -96,6 +116,11 @@ class StaggeredTripleDifference(
Triple Differences Estimators." arXiv:2505.09942v3.
"""
+ # Names this estimator in the shared bootstrap mixin's user-facing
+ # warnings. The mixin is shared with the other hosts, so a hard-coded
+ # literal there would misname whichever surface was actually fit.
+ _BOOTSTRAP_LABEL: ClassVar[str] = "StaggeredTripleDifference"
+
def __init__(
self,
estimation_method: str = "dr",
@@ -113,6 +138,11 @@ def __init__(
epv_threshold: float = 10,
pscore_fallback: str = "error",
):
+ # Row M-013. Emitted per construction (not once): a fitted instance
+ # never re-warns, but each new one does, and set_params re-emits via
+ # BaseEstimator's transactional probe re-init - the documented side
+ # effect MultiPeriodDiD's shim also has.
+ warnings.warn(_SDDD_DEPRECATION_MSG, FutureWarning, stacklevel=2)
if estimation_method not in ["dr", "ipw", "reg"]:
raise ValueError(
f"estimation_method must be 'dr', 'ipw', or 'reg', " f"got '{estimation_method}'"
@@ -164,15 +194,13 @@ def __init__(
self.is_fitted_ = False
self.results_: Optional[StaggeredTripleDiffResults] = None
- # get_params/set_params come from BaseEstimator.
-
# ------------------------------------------------------------------
# fit()
# ------------------------------------------------------------------
def fit(
self,
- data: pd.DataFrame,
+ data: "pd.DataFrame",
outcome: str,
unit: str,
time: str,
@@ -219,1444 +247,21 @@ def fit(
-------
StaggeredTripleDiffResults
"""
- from diff_diff.survey import (
- _resolve_survey_for_fit,
- _validate_unit_constant_survey,
- compute_survey_metadata,
- )
-
- resolved_survey, survey_weights, survey_weight_type, survey_metadata = (
- _resolve_survey_for_fit(survey_design, data, "analytical")
- )
-
- if resolved_survey is not None:
- _validate_unit_constant_survey(data, unit, survey_design)
- if resolved_survey.weight_type != "pweight":
- raise ValueError(
- f"StaggeredTripleDifference survey support requires "
- f"weight_type='pweight', got '{resolved_survey.weight_type}'. "
- f"The survey variance math assumes probability weights."
- )
- if aggregate is not None and aggregate not in [
- "event_study",
- "group",
- "simple",
- "all",
- ]:
- raise ValueError(
- f"aggregate must be 'event_study', 'group', 'simple', or 'all', "
- f"got '{aggregate}'"
- )
-
- df = data.copy()
- self._validate_inputs(df, outcome, unit, time, first_treat, eligibility, covariates)
-
- if self.cluster is not None:
- warnings.warn(
- "cluster parameter is accepted but cluster-robust analytical SEs "
- "are not yet implemented for staggered DDD. Use n_bootstrap > 0 "
- "for unit-level clustered inference via multiplier bootstrap.",
- UserWarning,
- stacklevel=2,
- )
-
- if first_treat != "first_treat":
- df["first_treat"] = df[first_treat]
- # Surface the inf → 0 recategorization the same way StaggeredDiD does
- # (see `staggered.py:1508-1519`). Silently recoding inf would shift
- # units between treated and never-treated pools with no signal
- # (axis-E silent coercion under the Phase 2 audit).
- _inf_mask = np.isposinf(df["first_treat"].values)
- if _inf_mask.any():
- n_inf_rows = int(_inf_mask.sum())
- warnings.warn(
- f"{n_inf_rows} row(s) have first_treat=inf; recoding to 0 "
- f"(never-treated). Use first_treat=0 to suppress this warning.",
- UserWarning,
- stacklevel=2,
- )
- df["first_treat"] = df["first_treat"].replace([np.inf, float("inf")], 0)
-
- precomputed = self._precompute_structures(
- df,
+ # The engine is shared with TripleDifference's staggered mode; this
+ # class supplies its own vocabulary and its own frame depth
+ # (user -> fit -> core), so warnings attribute exactly as in 3.x.
+ return self._fit_staggered_core(
+ data,
outcome,
unit,
time,
+ first_treat,
eligibility,
- covariates,
- resolved_survey=resolved_survey,
- )
-
- # Recompute survey metadata from unit-level resolved survey
- if resolved_survey is not None and survey_metadata is not None:
- resolved_survey_unit = precomputed.get("resolved_survey_unit")
- if resolved_survey_unit is not None:
- unit_w = resolved_survey_unit.weights
- survey_metadata = compute_survey_metadata(resolved_survey_unit, unit_w)
-
- # Survey df for t-distribution critical values
- df_survey = precomputed.get("df_survey")
- if (
- df_survey is None
- and resolved_survey is not None
- and hasattr(resolved_survey, "uses_replicate_variance")
- and resolved_survey.uses_replicate_variance
- ):
- df_survey = 0 # Forces NaN inference for undefined replicate df
-
- treatment_groups = precomputed["treatment_groups"]
- time_periods = precomputed["time_periods"]
- all_units = precomputed["all_units"]
- time_to_col = precomputed["time_to_col"]
- unit_cohorts = precomputed["unit_cohorts"]
- eligibility_per_unit = precomputed["eligibility_per_unit"]
- n_units = len(all_units)
-
- pscore_cache: Dict = {}
-
- group_time_effects: Dict[Tuple, Dict[str, Any]] = {}
- influence_func_info: Dict[Tuple, Dict[str, Any]] = {}
- comparison_group_counts: Dict[Tuple, int] = {}
- gmm_weights_store: Dict[Tuple, Dict] = {}
- epv_diagnostics: Optional[Dict[Tuple, Dict[str, Any]]] = (
- {} if (covariates and self.estimation_method in ("ipw", "dr")) else None
- )
-
- # Trackers for rank-deficient linalg solves across all (g, g_c, t)
- # cells. `_compute_did_panel` appends to the OR-side tracker;
- # `_compute_pscore` appends to the PS-side tracker. Both surface as
- # ONE aggregate warning below rather than fanning out per cell.
- self._lstsq_fallback_tracker: List[float] = []
- self._ps_lstsq_fallback_tracker: List[float] = []
-
- for g in treatment_groups:
- # In universal mode, skip the reference period (t == g-1-anticipation)
- # so it's omitted from GT estimation. The event-study mixin injects
- # a synthetic reference row with effect=0, matching CS behavior.
- if self.base_period == "universal":
- universal_base = g - 1 - self.anticipation
- valid_periods = [t for t in time_periods if t != universal_base]
- else:
- valid_periods = time_periods
-
- for t in valid_periods:
- base_period_val = self._get_base_period(g, t)
- if base_period_val is None:
- continue
- if base_period_val not in time_to_col:
- warnings.warn(
- f"Base period {base_period_val} for (g={g}, t={t}) is "
- "outside the observed panel. Skipping this cell.",
- UserWarning,
- stacklevel=2,
- )
- continue
- if t not in time_to_col:
- continue
-
- has_never_enabled = bool(np.any(unit_cohorts == 0))
-
- if self.control_group == "nevertreated":
- # Only use never-enabled cohort as comparison
- valid_gc = [0] if has_never_enabled else []
- else:
- # Use all valid comparison cohorts (not-yet-treated + never)
- # Threshold accounts for anticipation: cohorts that start
- # treatment within the anticipation window are contaminated.
- nyt_threshold = max(t, base_period_val) + self.anticipation
- valid_gc = [gc for gc in treatment_groups if gc > nyt_threshold and gc != g]
- if has_never_enabled:
- valid_gc = [0] + valid_gc
-
- if not valid_gc:
- warnings.warn(
- f"No valid comparison groups for (g={g}, t={t}), skipping.",
- UserWarning,
- stacklevel=2,
- )
- continue
-
- treated_mask = (unit_cohorts == g) & (eligibility_per_unit == 1)
- n_treated = int(np.sum(treated_mask))
- if n_treated == 0:
- continue
-
- att_vec = []
- inf_raw = [] # unrescaled IFs
- gc_labels = []
- gc_cell_sizes = [] # size_gt_ctrl per surviving gc
-
- for gc in valid_gc:
- result = self._compute_ddd_gt_gc(
- precomputed,
- g,
- gc,
- t,
- base_period_val,
- covariates,
- pscore_cache,
- epv_diagnostics=epv_diagnostics,
- )
- if result is None:
- continue
- att_gc, inf_gc, size_gt_ctrl = result
- if not np.isfinite(att_gc):
- continue
-
- att_vec.append(att_gc)
- inf_raw.append(inf_gc)
- gc_labels.append(gc)
- gc_cell_sizes.append(size_gt_ctrl)
-
- if not att_vec:
- continue
-
- # Compute size_gt from SURVIVING comparison cohorts only
- # (not from all initially valid gc's)
- surviving_units = treated_mask.copy()
- for gc in gc_labels:
- surviving_units |= (unit_cohorts == gc) | (unit_cohorts == g)
- survey_w = precomputed.get("survey_weights")
- if survey_w is not None:
- size_gt = float(np.sum(survey_w[surviving_units]))
- else:
- size_gt = float(np.sum(surviving_units))
-
- # Apply IF rescaling now that size_gt is known
- inf_matrix = []
- for inf_gc, size_gt_ctrl in zip(inf_raw, gc_cell_sizes):
- if size_gt_ctrl > 0:
- inf_gc = inf_gc * (size_gt / size_gt_ctrl)
- inf_matrix.append(inf_gc)
-
- att_gmm, inf_gmm, gmm_w, se_gt = self._combine_gmm(
- np.array(att_vec),
- np.array(inf_matrix),
- n_units,
- )
-
- if not np.isfinite(att_gmm):
- continue
-
- # R's single-gc SE uses size_gt in denominator, not n_total.
- # For multi-gc (GMM), the size_gt factor is already in Omega
- # via the per-gc rescaling, so n_total is correct.
- if len(gc_labels) == 1:
- se_gt = float(np.sqrt(np.sum(inf_gmm**2) / size_gt**2))
-
- if not np.isfinite(se_gt) or se_gt <= 0:
- se_gt = np.nan
-
- t_stat, p_value, conf_int = safe_inference(
- att_gmm, se_gt, alpha=self.alpha, df=df_survey
- )
-
- # Rescale IF for mixin compatibility.
- # R stores IF * (n/size_gt) in inf_func_mat, then uses
- # SE = sqrt(sum(IF^2)/n^2) = sqrt(sum(psi^2)) with psi = IF/n.
- # We need psi = IF_rescaled / n so mixin's sqrt(sum(psi^2)) works.
- # IF is already at size_gt/size_gt_ctrl scale from above.
- # Apply the final n/size_gt factor, then divide by n for mixin.
- inf_gmm_rescaled = inf_gmm * (n_units / size_gt)
- inf_gmm_scaled = inf_gmm_rescaled / n_units
-
- # INVARIANT: np.where over boolean masks -> duplicate-free
- # index arrays (fancy-+= scatter contract, see
- # staggered_aggregation._combined_if_fast).
- treated_idx = np.where(treated_mask)[0]
- treated_inf = inf_gmm_scaled[treated_idx]
- nonzero_mask = (inf_gmm_scaled != 0) & ~treated_mask
- control_idx = np.where(nonzero_mask)[0]
- control_inf = inf_gmm_scaled[control_idx]
- n_control = int(np.sum(nonzero_mask))
-
- group_time_effects[(g, t)] = {
- "effect": att_gmm,
- "se": se_gt,
- "t_stat": t_stat,
- "p_value": p_value,
- "conf_int": conf_int,
- "n_treated": n_treated,
- "n_control": n_control,
- }
- influence_func_info[(g, t)] = {
- "treated_idx": treated_idx,
- "control_idx": control_idx,
- "treated_inf": treated_inf,
- "control_inf": control_inf,
- }
- comparison_group_counts[(g, t)] = len(gc_labels)
- gmm_weights_store[(g, t)] = dict(zip(gc_labels, gmm_w.tolist()))
-
- # Consolidated OR influence-function rank-deficiency warning.
- # Finding #17 in the Phase 2 silent-failures audit: the per-pair OR
- # solve at _compute_did_panel() previously fell back to lstsq with no
- # signal, so near/fully singular X'WX in the covariate expansion went
- # to the user as a normal result.
- if self._lstsq_fallback_tracker and self.rank_deficient_action != "silent":
- n_cells = len(self._lstsq_fallback_tracker)
- finite_conds = [c for c in self._lstsq_fallback_tracker if np.isfinite(c)]
- max_cond = max(finite_conds) if finite_conds else float("inf")
- warnings.warn(
- f"Rank-deficient X'WX detected in the outcome-regression "
- f"influence-function step for {n_cells} (g, g_c, t) pair(s); "
- f"dropped redundant direction(s) via a rank-guarded inverse. "
- f"Max condition number of affected X'WX: {max_cond:.2e}. "
- f"Standard errors use the identified covariate subset; consider "
- f"dropping collinear covariates or using "
- f"estimation_method='ipw' to avoid the OR projection.",
- UserWarning,
- stacklevel=2,
- )
-
- # Consolidated PS-Hessian rank-deficiency warning (sibling of the
- # OR path above). `_compute_pscore` previously fell back from
- # `np.linalg.inv(X'WX)` to `np.linalg.lstsq` with no signal, so
- # a rank-deficient propensity-score design silently degraded
- # IPW/DR influence-function corrections.
- if self._ps_lstsq_fallback_tracker and self.rank_deficient_action != "silent":
- n_cells = len(self._ps_lstsq_fallback_tracker)
- finite_conds = [c for c in self._ps_lstsq_fallback_tracker if np.isfinite(c)]
- max_cond = max(finite_conds) if finite_conds else float("inf")
- warnings.warn(
- f"Rank-deficient X'WX detected in the propensity-score "
- f"Hessian for {n_cells} (g, g_c, t) pair(s); dropped redundant "
- f"direction(s) via a rank-guarded inverse. Max condition number "
- f"of affected X'WX: {max_cond:.2e}. IPW/DR influence-function "
- f"corrections use the identified covariate subset; consider "
- f"dropping collinear propensity-score covariates or using "
- f"estimation_method='reg' to avoid the PS path.",
- UserWarning,
- stacklevel=2,
- )
-
- # Consolidated EPV summary warning
- if epv_diagnostics:
- low_epv = {k: v for k, v in epv_diagnostics.items() if v.get("is_low")}
- if low_epv:
- n_affected = len(low_epv)
- n_total = len(epv_diagnostics)
- min_entry = min(low_epv.values(), key=lambda v: v["epv"])
- min_g = min(low_epv.keys(), key=lambda k: low_epv[k]["epv"])
- warnings.warn(
- f"Low Events Per Variable (EPV) detected in "
- f"{n_affected} of {n_total} cohort-time cell(s). "
- f"Minimum EPV: {min_entry['epv']:.1f} (cohort g={min_g[0]}). "
- f"Consider estimation_method='reg' or fewer covariates. "
- f"Call results.epv_summary() for per-cohort details.",
- UserWarning,
- stacklevel=2,
- )
-
- if not group_time_effects:
- raise ValueError(
- "No valid group-time effects could be computed. "
- "Check that the data has sufficient variation in treatment "
- "timing and eligibility."
- )
-
- # For aggregation: use eligible-treated-only cohort assignments so
- # WIF weights match the point estimate weights (n_treated per cohort,
- # i.e. P(S=g, Q=1)). This matches the paper's Eq 4.13 which defines
- # aggregation weights over the treated population (G_i defined only
- # for Q=1 units). Ineligible units get cohort=0 so they don't
- # contribute to pg for any treatment group.
- # Both precomputed["unit_cohorts"] AND df["first_treat"] must be
- # zeroed for ineligible units because the WIF code reads both.
- precomputed_agg = dict(precomputed)
- cohorts_for_agg = precomputed["unit_cohorts"].copy()
- cohorts_for_agg[eligibility_per_unit == 0] = 0
- precomputed_agg["unit_cohorts"] = cohorts_for_agg
-
- df_agg = df.copy()
- df_agg.loc[df_agg[eligibility] == 0, "first_treat"] = 0
-
- # Overall ATT via aggregation mixin
- overall_att, overall_se, overall_effective_df = self._aggregate_simple(
- group_time_effects, influence_func_info, df_agg, unit, precomputed_agg
+ covariates=covariates,
+ aggregate=aggregate,
+ balance_e=balance_e,
+ survey_design=survey_design,
+ estimator_name="StaggeredTripleDifference",
+ partition_label="eligibility",
+ _frame_offset=1,
)
- # Preserve the ORIGINAL survey df before the simple-overall statistic mutates
- # it below. The Eq. 4.14 overall (overall_att_es) must fall back to this
- # original df, never to the simple overall's per-statistic replicate df — the
- # two statistics can drop different replicate subsets, so reusing the simple
- # overall's df would silently give overall_att_es the wrong p-value/CI.
- df_survey_original = df_survey
- # Use per-statistic effective df from replicate aggregation if available;
- # otherwise fall back to the original df from the survey design.
- if overall_effective_df is not None:
- df_survey = overall_effective_df
- if survey_metadata is not None:
- survey_metadata.df_survey = df_survey
- overall_t_stat, overall_p_value, overall_conf_int = safe_inference(
- overall_att, overall_se, alpha=self.alpha, df=df_survey
- )
-
- # Aggregations
- event_study_effects = None
- group_effects = None
- es_aggregation = None
- if aggregate in ("event_study", "all"):
- es_aggregation = self._aggregate_event_study(
- group_time_effects,
- influence_func_info,
- treatment_groups,
- time_periods,
- balance_e,
- df_agg,
- unit,
- precomputed_agg,
- )
- event_study_effects = es_aggregation.effects
- if aggregate in ("group", "all"):
- group_effects = self._aggregate_by_group(
- group_time_effects,
- influence_func_info,
- treatment_groups,
- precomputed_agg,
- df_agg,
- unit,
- )
-
- # Paper Eq. (4.14) overall ATT (event-study average): an opt-in summary
- # alongside the default CS-simple ``overall_att``. ``_aggregate_event_study``
- # RETURNS it on its aggregation object (it used to stash it on
- # ``self._event_study_overall``); populated only when the event-study
- # aggregation ran. Analytical inference here; the bootstrap block below
- # overrides the SE when ``n_bootstrap > 0`` (mirroring ``overall_se``).
- overall_att_es = None
- overall_se_es = None
- overall_t_stat_es = None
- overall_p_value_es = None
- overall_conf_int_es = None
- # Whether the ANALYTICAL Eq. 4.14 SE was non-finite while its point estimate
- # was finite (i.e. a contributing horizon's influence function was non-finite).
- # Captured before any bootstrap override so the terminal warning below does not
- # misdiagnose a bootstrap-side NaN SE (e.g. cluster-unidentified) as an
- # analytical-IF failure (the bootstrap path emits its own warning).
- analytical_overall_es_se_nonfinite = False
- if aggregate in ("event_study", "all"):
- es_overall = es_aggregation.overall if es_aggregation is not None else None
- if es_overall is not None:
- overall_att_es = es_overall["att"]
- overall_se_es = es_overall["se"]
- analytical_overall_es_se_nonfinite = bool(
- np.isfinite(overall_att_es) and not np.isfinite(overall_se_es)
- )
- es_eff_df = es_overall.get("effective_df")
- # Fall back to the ORIGINAL survey df, not the simple-overall's mutated
- # per-statistic df (P1 fix): overall_att_es has its own replicate df.
- df_for_es = es_eff_df if es_eff_df is not None else df_survey_original
- overall_t_stat_es, overall_p_value_es, overall_conf_int_es = safe_inference(
- overall_att_es, overall_se_es, alpha=self.alpha, df=df_for_es
- )
- else:
- # Event-study aggregation was requested but yielded no post-treatment
- # horizon: this is "requested but undefined" -> NaN + warning (the
- # library's overall-aggregation contract, matching _aggregate_simple),
- # distinct from "not requested" which leaves the fields None.
- warnings.warn(
- "Event-study aggregation was requested but no post-treatment "
- "horizons are available for the Eq. 4.14 overall (overall_att_es); "
- "returning NaN.",
- UserWarning,
- stacklevel=2,
- )
- overall_att_es = np.nan
- overall_se_es = np.nan
- overall_t_stat_es, overall_p_value_es, overall_conf_int_es = safe_inference(
- np.nan, np.nan, alpha=self.alpha, df=df_survey
- )
-
- # Reject replicate-weight designs for bootstrap — replicate variance
- # is an analytical alternative, not compatible with bootstrap
- if (
- self.n_bootstrap > 0
- and resolved_survey is not None
- and hasattr(resolved_survey, "uses_replicate_variance")
- and resolved_survey.uses_replicate_variance
- ):
- raise NotImplementedError(
- "StaggeredTripleDifference bootstrap (n_bootstrap > 0) is not "
- "supported with replicate-weight survey designs. Replicate "
- "weights provide analytical variance; use n_bootstrap=0 instead."
- )
-
- # Bootstrap
- bootstrap_results = None
- cband_crit_value = None
- if self.n_bootstrap > 0:
- bootstrap_results = self._run_multiplier_bootstrap(
- group_time_effects,
- influence_func_info,
- aggregate,
- balance_e,
- treatment_groups,
- time_periods,
- df_agg,
- unit,
- precomputed_agg,
- self.cband,
- )
- if bootstrap_results is not None:
- overall_se = bootstrap_results.overall_att_se
- overall_t_stat, overall_p_value, overall_conf_int = safe_inference(
- overall_att, overall_se, alpha=self.alpha, df=df_survey
- )
- overall_conf_int = bootstrap_results.overall_att_ci
- overall_p_value = bootstrap_results.overall_att_p_value
-
- # Mirror the override for the Eq. (4.14) event-study-average overall
- # (only when event-study aggregation produced it). A NaN bootstrap SE
- # (e.g. cluster-unidentified) correctly NaNs the inference.
- if overall_att_es is not None and bootstrap_results.overall_att_es_se is not None:
- overall_se_es = bootstrap_results.overall_att_es_se
- overall_t_stat_es, overall_p_value_es, overall_conf_int_es = safe_inference(
- overall_att_es, overall_se_es, alpha=self.alpha, df=df_survey
- )
- overall_conf_int_es = bootstrap_results.overall_att_es_ci
- overall_p_value_es = bootstrap_results.overall_att_es_p_value
- if bootstrap_results.cband_crit_value is not None:
- cband_crit_value = bootstrap_results.cband_crit_value
-
- # Update group-time effects with bootstrap SEs
- if bootstrap_results.group_time_ses:
- for gt_key in group_time_effects:
- if gt_key in bootstrap_results.group_time_ses:
- group_time_effects[gt_key]["se"] = bootstrap_results.group_time_ses[
- gt_key
- ]
- group_time_effects[gt_key]["conf_int"] = (
- bootstrap_results.group_time_cis[gt_key]
- )
- group_time_effects[gt_key]["p_value"] = (
- bootstrap_results.group_time_p_values[gt_key]
- )
- t_val, _, _ = safe_inference(
- group_time_effects[gt_key]["effect"],
- bootstrap_results.group_time_ses[gt_key],
- alpha=self.alpha,
- df=df_survey,
- )
- group_time_effects[gt_key]["t_stat"] = t_val
-
- if event_study_effects and bootstrap_results.event_study_ses:
- for e_key in event_study_effects:
- if e_key in bootstrap_results.event_study_ses:
- # ses/cis/p_values are populated together.
- assert (
- bootstrap_results.event_study_cis is not None
- and bootstrap_results.event_study_p_values is not None
- )
- event_study_effects[e_key]["se"] = bootstrap_results.event_study_ses[
- e_key
- ]
- event_study_effects[e_key]["conf_int"] = (
- bootstrap_results.event_study_cis[e_key]
- )
- event_study_effects[e_key]["p_value"] = (
- bootstrap_results.event_study_p_values[e_key]
- )
- t_val, _, _ = safe_inference(
- event_study_effects[e_key]["effect"],
- bootstrap_results.event_study_ses[e_key],
- alpha=self.alpha,
- df=df_survey,
- )
- event_study_effects[e_key]["t_stat"] = t_val
- if cband_crit_value is not None:
- bs_se = bootstrap_results.event_study_ses[e_key]
- eff = event_study_effects[e_key]["effect"]
- event_study_effects[e_key]["cband_conf_int"] = (
- eff - cband_crit_value * bs_se,
- eff + cband_crit_value * bs_se,
- )
-
- # Update group effects with bootstrap SEs
- if (
- group_effects
- and bootstrap_results.group_effect_ses is not None
- and bootstrap_results.group_effect_cis is not None
- and bootstrap_results.group_effect_p_values is not None
- ):
- grp_keys = [g for g in group_effects if g in bootstrap_results.group_effect_ses]
- for g_key in grp_keys:
- group_effects[g_key]["se"] = bootstrap_results.group_effect_ses[g_key]
- group_effects[g_key]["conf_int"] = bootstrap_results.group_effect_cis[g_key]
- group_effects[g_key]["p_value"] = bootstrap_results.group_effect_p_values[
- g_key
- ]
- t_val, _, _ = safe_inference(
- group_effects[g_key]["effect"],
- bootstrap_results.group_effect_ses[g_key],
- alpha=self.alpha,
- df=df_survey,
- )
- group_effects[g_key]["t_stat"] = t_val
- # Bootstrap se/p/CI replaced the analytical ones, which
- # is what the retained df described - keeping it would
- # claim a t-reference that governed nothing.
- group_effects[g_key]["df_used"] = None
-
- # Eq. 4.14 overall: an ANALYTICAL non-finite SE under a finite point estimate
- # (a contributing horizon's influence function is non-finite, or the variance is
- # unidentified — e.g. a single-PSU/cluster design). Surface it — never NaN the SE
- # silently. Gated on the analytical-origin flag (captured before the bootstrap
- # override) and the final state still being non-finite: a bootstrap that supplies a
- # finite SE rescues it (no warning), and a bootstrap that NaNs the SE for unrelated
- # reasons (e.g. cluster-unidentified) is reported by the bootstrap's own warning,
- # not misdiagnosed here as an analytical-IF failure.
- if (
- analytical_overall_es_se_nonfinite
- and overall_se_es is not None
- and not np.isfinite(overall_se_es)
- ):
- warnings.warn(
- "Eq. 4.14 overall (overall_att_es) point estimate is defined but its "
- "standard error is undefined (NaN): either a contributing post-treatment "
- "event-study horizon has a non-finite influence function, or the variance "
- "is unidentified (e.g. a single-PSU/cluster survey design). overall_se_es "
- "and its inference fields are NaN.",
- UserWarning,
- stacklevel=2,
- )
-
- n_treated_units = int(np.sum((unit_cohorts > 0) & (eligibility_per_unit == 1)))
- n_control_units = n_units - n_treated_units
- n_never_enabled = int(np.sum(unit_cohorts == 0))
- n_eligible = int(np.sum(eligibility_per_unit == 1))
- n_ineligible = int(np.sum(eligibility_per_unit == 0))
-
- self.results_ = StaggeredTripleDiffResults(
- group_time_effects=group_time_effects,
- overall_att=overall_att,
- overall_se=overall_se,
- overall_t_stat=overall_t_stat,
- overall_p_value=overall_p_value,
- overall_conf_int=overall_conf_int,
- groups=treatment_groups,
- time_periods=time_periods,
- n_obs=len(df),
- n_treated_units=n_treated_units,
- n_control_units=n_control_units,
- n_never_enabled=n_never_enabled,
- n_eligible=n_eligible,
- n_ineligible=n_ineligible,
- alpha=self.alpha,
- control_group=self.control_group,
- base_period=self.base_period,
- anticipation=self.anticipation,
- estimation_method=self.estimation_method,
- event_study_effects=event_study_effects,
- group_effects=group_effects,
- bootstrap_results=bootstrap_results,
- cband_crit_value=cband_crit_value,
- pscore_trim=self.pscore_trim,
- survey_metadata=survey_metadata,
- comparison_group_counts=comparison_group_counts,
- gmm_weights=gmm_weights_store,
- epv_diagnostics=epv_diagnostics if epv_diagnostics else None,
- epv_threshold=self.epv_threshold,
- pscore_fallback=self.pscore_fallback,
- overall_att_es=overall_att_es,
- overall_se_es=overall_se_es,
- overall_t_stat_es=overall_t_stat_es,
- overall_p_value_es=overall_p_value_es,
- overall_conf_int_es=overall_conf_int_es,
- )
- self.is_fitted_ = True
- return self.results_
-
- # ------------------------------------------------------------------
- # Validation
- # ------------------------------------------------------------------
-
- def _validate_inputs(
- self,
- df: pd.DataFrame,
- outcome: str,
- unit: str,
- time: str,
- first_treat: str,
- eligibility: str,
- covariates: Optional[List[str]],
- ) -> None:
- """Validate input data."""
- required_cols = [outcome, unit, time, first_treat, eligibility]
- if covariates:
- required_cols.extend(covariates)
- missing = [c for c in required_cols if c not in df.columns]
- if missing:
- raise ValueError(f"Missing columns: {missing}")
-
- elig_vals = df[eligibility].dropna().unique()
- if not set(elig_vals).issubset({0, 1, 0.0, 1.0}):
- raise ValueError(
- f"Eligibility column '{eligibility}' must be binary (0/1). "
- f"Found values: {sorted(elig_vals)}"
- )
- elig_by_unit = df.groupby(unit)[eligibility].nunique()
- varying = elig_by_unit[elig_by_unit > 1]
- if len(varying) > 0:
- raise ValueError(
- f"Eligibility must be time-invariant within units. "
- f"Found {len(varying)} units with varying eligibility."
- )
- for col in [outcome, first_treat, eligibility]:
- if df[col].isna().any():
- raise ValueError(f"Column '{col}' contains missing values.")
-
- # Reject non-finite outcomes (Inf/-Inf)
- if not np.all(np.isfinite(df[outcome])):
- raise ValueError(
- f"Column '{outcome}' contains non-finite values (Inf/-Inf). "
- "All outcome values must be finite."
- )
-
- # Reject non-finite covariates
- if covariates:
- for cov in covariates:
- if df[cov].isna().any():
- raise ValueError(f"Covariate '{cov}' contains missing values.")
- if not np.all(np.isfinite(df[cov])):
- raise ValueError(f"Covariate '{cov}' contains non-finite values.")
- if df[eligibility].nunique() < 2:
- raise ValueError(
- "Need both eligible (Q=1) and ineligible (Q=0) units. "
- f"Only found Q={df[eligibility].unique()[0]}."
- )
-
- # Check unique (unit, time) pairs — no duplicate rows
- dup = df.duplicated(subset=[unit, time], keep=False)
- if dup.any():
- raise ValueError(
- f"Duplicate (unit, time) rows found. "
- f"{int(dup.sum())} duplicates detected. Panel must have unique rows."
- )
-
- # Check balanced panel — every unit observed in exactly the global period set
- global_periods = set(df[time].unique())
- n_global_periods = len(global_periods)
- unit_period_sets = df.groupby(unit)[time].apply(set)
- mismatched = unit_period_sets[unit_period_sets != global_periods]
- if len(mismatched) > 0:
- raise ValueError(
- "Unbalanced panel detected. All units must be observed in "
- f"all {n_global_periods} periods. "
- f"Found {len(mismatched)} units with different period sets."
- )
-
- # Check time-invariant first_treat
- ft_by_unit = df.groupby(unit)[first_treat].nunique()
- varying_ft = ft_by_unit[ft_by_unit > 1]
- if len(varying_ft) > 0:
- raise ValueError(
- f"first_treat must be time-invariant within units. "
- f"Found {len(varying_ft)} units with varying first_treat."
- )
-
- # Check time-invariant covariates
- if covariates:
- for cov in covariates:
- cov_nunique = df.groupby(unit)[cov].nunique()
- varying_cov = cov_nunique[cov_nunique > 1]
- if len(varying_cov) > 0:
- raise ValueError(
- f"Covariate '{cov}' must be time-invariant within units. "
- f"Found {len(varying_cov)} units with varying values."
- )
-
- # ------------------------------------------------------------------
- # Precomputation
- # ------------------------------------------------------------------
-
- def _precompute_structures(
- self,
- df: pd.DataFrame,
- outcome: str,
- unit: str,
- time: str,
- eligibility: str,
- covariates: Optional[List[str]],
- resolved_survey=None,
- ) -> PrecomputedData:
- """Build precomputed structures for efficient computation."""
- all_units = np.array(sorted(df[unit].unique()))
- time_periods = sorted(df[time].unique())
- n_units = len(all_units)
- n_periods = len(time_periods)
-
- unit_to_idx = {u: i for i, u in enumerate(all_units)}
- time_to_col = {t: j for j, t in enumerate(time_periods)}
-
- outcome_matrix = np.full((n_units, n_periods), np.nan)
- for _, row in df.iterrows():
- u_idx = unit_to_idx[row[unit]]
- t_idx = time_to_col[row[time]]
- outcome_matrix[u_idx, t_idx] = row[outcome]
-
- unit_df = df.groupby(unit).first().reindex(all_units)
- unit_cohorts = unit_df["first_treat"].values.astype(float)
- eligibility_per_unit = unit_df[eligibility].values.astype(int)
-
- treatment_groups = sorted([g for g in np.unique(unit_cohorts) if g > 0])
-
- covariate_matrix = None
- if covariates:
- cov_wide = {}
- for cov in covariates:
- cov_vals = np.full(n_units, np.nan)
- for u_id, idx in unit_to_idx.items():
- u_data = df.loc[df[unit] == u_id, cov]
- if len(u_data) > 0:
- cov_vals[idx] = u_data.iloc[0]
- cov_wide[cov] = cov_vals
- covariate_matrix = np.column_stack(list(cov_wide.values()))
-
- # Extract per-unit survey weights and collapse design to unit level
- survey_weights_arr = None
- resolved_survey_unit = None
- if resolved_survey is not None:
- from diff_diff.survey import collapse_survey_to_unit_level
-
- survey_weights_arr = (
- pd.Series(resolved_survey.weights, index=df.index)
- .groupby(df[unit])
- .first()
- .reindex(all_units)
- .values.astype(np.float64)
- )
- # Normalize to sum=n for aggregation/rescaling (matches pweight
- # convention). Raw weights preserved in resolved_survey_unit for
- # replicate w_r/w_full ratios — those are inherently scale-invariant.
- sw_sum = np.sum(survey_weights_arr)
- if sw_sum > 0:
- survey_weights_arr = survey_weights_arr * (len(survey_weights_arr) / sw_sum)
- resolved_survey_unit = collapse_survey_to_unit_level(
- resolved_survey, df, unit, all_units
- )
-
- return {
- "all_units": all_units,
- "unit_to_idx": unit_to_idx,
- "time_periods": time_periods,
- "time_to_col": time_to_col,
- "outcome_matrix": outcome_matrix,
- "unit_cohorts": unit_cohorts,
- "eligibility_per_unit": eligibility_per_unit,
- "treatment_groups": treatment_groups,
- "covariate_matrix": covariate_matrix,
- "n_units": n_units,
- "n_periods": n_periods,
- "survey_weights": survey_weights_arr,
- "resolved_survey_unit": resolved_survey_unit,
- "df_survey": (
- resolved_survey_unit.df_survey if resolved_survey_unit is not None else None
- ),
- }
-
- # ------------------------------------------------------------------
- # Base period
- # ------------------------------------------------------------------
-
- def _get_base_period(self, g: Any, t: Any) -> Optional[Any]:
- """Determine base period for a (g, t) pair."""
- if self.base_period == "universal":
- return g - 1 - self.anticipation
- else:
- if t < g - self.anticipation:
- return t - 1
- else:
- return g - 1 - self.anticipation
-
- # ------------------------------------------------------------------
- # Three-DiD DDD for one (g, g_c, t) triple
- # ------------------------------------------------------------------
-
- def _compute_ddd_gt_gc(
- self,
- precomputed: PrecomputedData,
- g: Any,
- g_c: Any,
- t: Any,
- base_period_val: Any,
- covariates: Optional[List[str]],
- pscore_cache: Dict,
- epv_diagnostics: Optional[Dict] = None,
- ) -> Optional[Tuple[float, np.ndarray, float]]:
- """
- Compute DDD ATT for one (g, g_c, t) triple.
-
- Returns (att_ddd, inf_full_n_units, size_gt_ctrl) or None.
- """
- outcome_matrix = precomputed["outcome_matrix"]
- time_to_col = precomputed["time_to_col"]
- unit_cohorts = precomputed["unit_cohorts"]
- eligibility_per_unit = precomputed["eligibility_per_unit"]
- covariate_matrix = precomputed["covariate_matrix"]
- n_units = precomputed["n_units"]
- survey_weights = precomputed.get("survey_weights")
-
- t_col = time_to_col[t]
- b_col = time_to_col[base_period_val]
-
- # Four sub-groups within this (g, g_c) cell
- treated_mask = (unit_cohorts == g) & (eligibility_per_unit == 1) # subgroup 4
- sub_a_mask = (unit_cohorts == g) & (eligibility_per_unit == 0) # subgroup 3
- sub_b_mask = (unit_cohorts == g_c) & (eligibility_per_unit == 1) # subgroup 2
- sub_c_mask = (unit_cohorts == g_c) & (eligibility_per_unit == 0) # subgroup 1
-
- n_treated = int(np.sum(treated_mask))
- n_a = int(np.sum(sub_a_mask))
- n_b = int(np.sum(sub_b_mask))
- n_c = int(np.sum(sub_c_mask))
-
- # Check for empty subgroups (by count or by survey weight mass)
- empty = []
- if n_treated == 0:
- empty.append(f"(S={g},Q=1)")
- if n_a == 0:
- empty.append(f"(S={g},Q=0)")
- if n_b == 0:
- empty.append(f"(S={g_c},Q=1)")
- if n_c == 0:
- empty.append(f"(S={g_c},Q=0)")
- # Zero survey-weight mass after subpopulation filtering = effectively empty
- if not empty and survey_weights is not None:
- if np.sum(survey_weights[treated_mask]) <= 0:
- empty.append(f"(S={g},Q=1,mass=0)")
- if np.sum(survey_weights[sub_a_mask]) <= 0:
- empty.append(f"(S={g},Q=0,mass=0)")
- if np.sum(survey_weights[sub_b_mask]) <= 0:
- empty.append(f"(S={g_c},Q=1,mass=0)")
- if np.sum(survey_weights[sub_c_mask]) <= 0:
- empty.append(f"(S={g_c},Q=0,mass=0)")
- if empty:
- warnings.warn(
- f"Empty subgroup(s) {', '.join(empty)} for "
- f"(g={g}, g_c={g_c}, t={t}). "
- "Comparison unidentified, skipping.",
- UserWarning,
- stacklevel=3,
- )
- return None
-
- if min(n_treated, n_a, n_b, n_c) < 5:
- warnings.warn(
- f"Small cell size for (g={g}, g_c={g_c}, t={t}). " "Estimates may be unreliable.",
- UserWarning,
- stacklevel=3,
- )
-
- # Outcome changes
- delta_y_all = outcome_matrix[:, t_col] - outcome_matrix[:, b_col]
- valid = np.isfinite(delta_y_all)
- for m in [treated_mask, sub_a_mask, sub_b_mask, sub_c_mask]:
- if not np.all(valid[m]):
- return None
-
- # Three pairwise DiDs, each on a 2-cell subset
- # Collect per-DiD EPV diagnostics; merge worst into (g,t) key later
- epv_diag_a = {} if epv_diagnostics is not None else None
- epv_diag_b = {} if epv_diagnostics is not None else None
- epv_diag_c = {} if epv_diagnostics is not None else None
-
- # DiD_A: subgroup 4 vs 3 (treated-eligible vs treated-ineligible)
- pair_a_mask = treated_mask | sub_a_mask
- did_a = self._run_pairwise_did(
- delta_y_all,
- pair_a_mask,
- treated_mask,
- sub_a_mask,
- covariate_matrix,
- pscore_cache,
- (g, g, 0, base_period_val),
- survey_weights=survey_weights,
- context_label=f"cohort g={g}, DiD_A (g_c={g_c})",
- epv_diagnostics_out=epv_diag_a,
- )
-
- # DiD_B: subgroup 4 vs 2 (treated-eligible vs control-eligible)
- pair_b_mask = treated_mask | sub_b_mask
- did_b = self._run_pairwise_did(
- delta_y_all,
- pair_b_mask,
- treated_mask,
- sub_b_mask,
- covariate_matrix,
- pscore_cache,
- (g, g_c, 1, base_period_val),
- survey_weights=survey_weights,
- context_label=f"cohort g={g}, DiD_B (g_c={g_c})",
- epv_diagnostics_out=epv_diag_b,
- )
-
- # DiD_C: subgroup 4 vs 1 (treated-eligible vs control-ineligible)
- pair_c_mask = treated_mask | sub_c_mask
- did_c = self._run_pairwise_did(
- delta_y_all,
- pair_c_mask,
- treated_mask,
- sub_c_mask,
- covariate_matrix,
- pscore_cache,
- (g, g_c, 0, base_period_val),
- survey_weights=survey_weights,
- context_label=f"cohort g={g}, DiD_C (g_c={g_c})",
- epv_diagnostics_out=epv_diag_c,
- )
-
- # Merge per-DiD EPV diagnostics: keep the worst (lowest EPV) entry
- # across all three DiDs for this g_c. If multiple g_c contribute to the
- # same (g, t) cell, retain the overall minimum EPV across all g_c calls.
- if epv_diagnostics is not None:
- candidates = [d for d in [epv_diag_a, epv_diag_b, epv_diag_c] if d]
- if candidates:
- worst = min(candidates, key=lambda d: d.get("epv", float("inf")))
- existing = epv_diagnostics.get((g, t))
- if existing is None or worst.get("epv", float("inf")) < existing.get(
- "epv", float("inf")
- ):
- epv_diagnostics[(g, t)] = worst
-
- if did_a is None or did_b is None or did_c is None:
- return None
-
- att_a, inf_a = did_a
- att_b, inf_b = did_b
- att_c, inf_c = did_c
-
- att_ddd = att_a + att_b - att_c
-
- # Three-DiD IF combination: w_j = n_cell / n_pair_j (R's att_dr convention)
- # With survey weights, use survey-weighted cell sizes
- if survey_weights is not None:
- sw_4 = float(np.sum(survey_weights[treated_mask]))
- sw_3 = float(np.sum(survey_weights[sub_a_mask]))
- sw_2 = float(np.sum(survey_weights[sub_b_mask]))
- sw_1 = float(np.sum(survey_weights[sub_c_mask]))
- n_cell_w = sw_4 + sw_3 + sw_2 + sw_1
- n_pair_a_w = sw_4 + sw_3
- n_pair_b_w = sw_4 + sw_2
- n_pair_c_w = sw_4 + sw_1
- w_3 = n_cell_w / n_pair_a_w if n_pair_a_w > 0 else 1.0
- w_2 = n_cell_w / n_pair_b_w if n_pair_b_w > 0 else 1.0
- w_1 = n_cell_w / n_pair_c_w if n_pair_c_w > 0 else 1.0
- size_gt_ctrl = n_cell_w
- else:
- n_cell = n_treated + n_a + n_b + n_c
- n_pair_a = n_treated + n_a
- n_pair_b = n_treated + n_b
- n_pair_c = n_treated + n_c
- w_3 = n_cell / n_pair_a if n_pair_a > 0 else 1.0
- w_2 = n_cell / n_pair_b if n_pair_b > 0 else 1.0
- w_1 = n_cell / n_pair_c if n_pair_c > 0 else 1.0
- size_gt_ctrl = float(n_cell)
-
- # Scatter pair-level IFs into n_units-length vector
- inf_full = np.zeros(n_units)
- pair_a_idx = np.where(pair_a_mask)[0]
- pair_b_idx = np.where(pair_b_mask)[0]
- pair_c_idx = np.where(pair_c_mask)[0]
-
- inf_full[pair_a_idx] += w_3 * inf_a
- inf_full[pair_b_idx] += w_2 * inf_b
- inf_full[pair_c_idx] -= w_1 * inf_c
-
- return att_ddd, inf_full, size_gt_ctrl
-
- # ------------------------------------------------------------------
- # Pairwise DiD (matches R's compute_did)
- # ------------------------------------------------------------------
-
- def _run_pairwise_did(
- self,
- delta_y_all: np.ndarray,
- pair_mask: np.ndarray,
- treated_mask: np.ndarray,
- control_mask: np.ndarray,
- covariate_matrix: Optional[np.ndarray],
- pscore_cache: Dict,
- pscore_key: Any,
- survey_weights: Optional[np.ndarray] = None,
- context_label: str = "",
- epv_diagnostics_out: Optional[dict] = None,
- ) -> Optional[Tuple[float, np.ndarray]]:
- """
- Compute a single pairwise DiD ATT and IF on a 2-cell subset.
-
- Matches R's triplediff::compute_did() formulation exactly:
- Riesz/Hajek normalization, PS + OR IF corrections.
-
- Returns (att, inf_func) where inf_func has length n_pair,
- ordered by pair_mask indices. Returns None if insufficient data.
- """
- pair_idx = np.where(pair_mask)[0]
- n_pair = len(pair_idx)
- if n_pair == 0:
- return None
-
- delta_y = delta_y_all[pair_idx]
- PA4 = treated_mask[pair_idx].astype(float)
- PAa = control_mask[pair_idx].astype(float)
- sw_pair = survey_weights[pair_idx] if survey_weights is not None else None
-
- n_t = int(np.sum(PA4))
- n_c = int(np.sum(PAa))
- if n_t == 0 or n_c == 0:
- return None
-
- has_covariates = covariate_matrix is not None and self.estimation_method != "none"
-
- # Build covariate matrix with intercept for the pair
- covX = None
- if has_covariates:
- # The flag definition above guarantees this (mypy can't track it).
- assert covariate_matrix is not None
- X_pair = covariate_matrix[pair_idx]
- covX = np.column_stack([np.ones(n_pair), X_pair])
-
- # Compute nuisance parameters based on estimation method
- pscore = None
- hessian = None
- or_delta = np.zeros(n_pair)
-
- if self.estimation_method in ("ipw", "dr") and covX is not None:
- pscore, hessian = self._compute_pscore(
- PA4,
- covX,
- pscore_cache,
- pscore_key,
- survey_weights=sw_pair,
- context_label=context_label,
- epv_diagnostics_out=epv_diagnostics_out,
- )
-
- if self.estimation_method in ("reg", "dr") and covX is not None:
- or_delta = self._compute_or(
- delta_y,
- PAa,
- covX,
- survey_weights=sw_pair,
- )
-
- # Compute ATT and IF (R's compute_did formulation)
- return self._compute_did_panel(
- delta_y,
- PA4,
- PAa,
- covX,
- pscore,
- hessian,
- or_delta,
- survey_weights=sw_pair,
- )
-
- # ------------------------------------------------------------------
- # Core DR/IPW/RA computation (matches R's compute_did exactly)
- # ------------------------------------------------------------------
-
- def _compute_did_panel(
- self,
- delta_y: np.ndarray,
- PA4: np.ndarray,
- PAa: np.ndarray,
- covX: Optional[np.ndarray],
- pscore: Optional[np.ndarray],
- hessian: Optional[np.ndarray],
- or_delta: np.ndarray,
- survey_weights: Optional[np.ndarray] = None,
- ) -> Tuple[float, np.ndarray]:
- """
- Pairwise DiD ATT and influence function.
- Matches R's triplediff::compute_did() line-by-line.
-
- Parameters
- ----------
- delta_y : outcome changes for 2-cell subset (n_pair,)
- PA4 : treated indicator (n_pair,)
- PAa : control indicator (n_pair,)
- covX : covariate matrix with intercept (n_pair, p) or None
- pscore : propensity scores (n_pair,) or None
- hessian : (X'WX)^{-1} * n_pair or None
- or_delta : OR predictions (n_pair,), zeros if no covariates
- survey_weights : per-observation survey weights (n_pair,) or None
-
- Returns
- -------
- (att, inf_func) where inf_func has length n_pair.
- """
- n_pair = len(delta_y)
- est = self.estimation_method
-
- # Riesz representers (R lines 243-250)
- if est == "reg" or pscore is None:
- w_treat = PA4.copy()
- w_control = PAa.copy()
- else:
- w_treat = PA4.copy()
- pscore_safe = np.clip(pscore, self.pscore_trim, 1 - self.pscore_trim)
- w_control = pscore_safe * PAa / (1 - pscore_safe)
-
- # Incorporate survey weights into Riesz representers
- if survey_weights is not None:
- w_treat = w_treat * survey_weights
- w_control = w_control * survey_weights
-
- # DR ATT via Hajek normalization (R lines 251-256)
- resid = delta_y - or_delta
- riesz_treat = w_treat * resid
- riesz_control = w_control * resid
-
- mean_w_treat = np.mean(w_treat)
- mean_w_control = np.mean(w_control)
-
- if mean_w_treat <= 0 or mean_w_control <= 0:
- return float("nan"), np.zeros(n_pair)
-
- att_treat = np.mean(riesz_treat) / mean_w_treat
- att_control = np.mean(riesz_control) / mean_w_control
- dr_att = att_treat - att_control
-
- # Base IF (R lines 302-304)
- inf_treat_did = riesz_treat - w_treat * att_treat
- inf_control_did = riesz_control - w_control * att_control
-
- # PS correction (R lines 262-273) — IPW and DR only
- inf_control_pscore = 0.0
- if est != "reg" and hessian is not None and covX is not None:
- M2 = np.mean((w_control * (resid - att_control))[:, None] * covX, axis=0)
- if survey_weights is not None:
- score_ps = survey_weights[:, None] * (PA4 - pscore_safe)[:, None] * covX
- else:
- score_ps = (PA4 - pscore_safe)[:, None] * covX
- asy_lin_rep_ps = score_ps @ hessian
- inf_control_pscore = asy_lin_rep_ps @ M2
-
- # OR correction (R lines 278-300) — reg and DR only
- inf_treat_or = 0.0
- inf_cont_or = 0.0
- if est != "ipw" and covX is not None:
- M1 = np.mean(w_treat[:, None] * covX, axis=0)
- M3 = np.mean(w_control[:, None] * covX, axis=0)
-
- if survey_weights is not None:
- or_x = (PAa * survey_weights)[:, None] * covX
- or_ex = (PAa * survey_weights * resid)[:, None] * covX
- else:
- or_x = PAa[:, None] * covX
- or_ex = (PAa * resid)[:, None] * covX
- XpX = or_x.T @ covX / n_pair
-
- # Rank-guarded inverse: a near-singular X'WX (constant/collinear
- # covariate) does not raise LinAlgError, so np.linalg.solve would
- # return a garbage inverse that inflates the OR influence function.
- # _rank_guarded_inv truncates redundant directions (finite SE on the
- # identified subset) and is the sole owner of the tracker append.
- XpX_inv, _, _ = _rank_guarded_inv(
- XpX, tracker=getattr(self, "_lstsq_fallback_tracker", None)
- )
- asy_linear_or = (XpX_inv @ or_ex.T).T
-
- inf_treat_or = -(asy_linear_or @ M1)
- inf_cont_or = -(asy_linear_or @ M3)
-
- # Final IF assembly (R lines 307-310)
- inf_control = (inf_control_did + inf_control_pscore + inf_cont_or) / mean_w_control
- inf_treat = (inf_treat_did + inf_treat_or) / mean_w_treat
- inf_func = inf_treat - inf_control
-
- return float(dr_att), inf_func
-
- # ------------------------------------------------------------------
- # Nuisance parameter computation
- # ------------------------------------------------------------------
-
- def _compute_pscore(
- self,
- PA4: np.ndarray,
- covX: np.ndarray,
- pscore_cache: Dict,
- pscore_key: Any,
- survey_weights: Optional[np.ndarray] = None,
- context_label: str = "",
- epv_diagnostics_out: Optional[dict] = None,
- ) -> Tuple[np.ndarray, np.ndarray]:
- """Fit logistic P(PA4=1|X). Returns (pscore, hessian).
-
- hessian = (X'WX)^{-1} * n_pair, matching R's convention.
- When survey_weights is provided, IRLS uses survey-weighted
- working weights and the hessian accounts for survey weights.
- """
- cached = pscore_cache.get(pscore_key)
- n_pair = len(PA4)
-
- if cached is not None:
- beta_logistic, cached_diag = cached
- z = np.dot(covX, beta_logistic)
- z = np.clip(z, -500, 500)
- pscore = 1 / (1 + np.exp(-z))
- if epv_diagnostics_out is not None and cached_diag:
- epv_diagnostics_out.update(cached_diag)
- else:
- X_no_intercept = covX[:, 1:] # solve_logit adds its own intercept
- diag = {}
- try:
- beta_logistic, pscore = solve_logit(
- X_no_intercept,
- PA4,
- rank_deficient_action=self.rank_deficient_action,
- weights=survey_weights,
- epv_threshold=self.epv_threshold,
- context_label=context_label,
- diagnostics_out=diag,
- )
- _check_propensity_diagnostics(pscore, self.pscore_trim)
- # Zero-fill NaN coefficients (from rank-deficient columns)
- # before caching, so cache reuse doesn't propagate NaN.
- # Cache alongside EPV diagnostics for replay on cache hits.
- beta_clean = np.where(np.isfinite(beta_logistic), beta_logistic, 0.0)
- pscore_cache[pscore_key] = (beta_clean, diag)
- except (np.linalg.LinAlgError, ValueError):
- if self.pscore_fallback == "error" or self.rank_deficient_action == "error":
- raise
- ctx = f" for {context_label}" if context_label else ""
- warnings.warn(
- f"Propensity score estimation failed{ctx}. "
- f"Falling back to unconditional propensity "
- f"(propensity model ignores covariates; outcome "
- f"regression still uses them for DR). "
- f"Consider estimation_method='reg' to avoid "
- f"propensity scores entirely.",
- UserWarning,
- stacklevel=5,
- )
- # Use survey-weighted treated share when weights available
- if survey_weights is not None:
- pos = survey_weights > 0
- if np.any(pos):
- p_uc = np.average(PA4[pos], weights=survey_weights[pos])
- else:
- p_uc = np.mean(PA4)
- else:
- p_uc = np.mean(PA4)
- pscore = np.full(n_pair, p_uc)
- pscore = np.clip(pscore, self.pscore_trim, 1 - self.pscore_trim)
- # No hessian for unconditional fallback
- return pscore, None
- if epv_diagnostics_out is not None and diag:
- epv_diagnostics_out.update(diag)
-
- pscore = np.clip(pscore, 1e-6, 1 - 1e-6)
-
- # Hessian: (X'WX)^{-1} * n (matching R's compute_pscore)
- W = pscore * (1 - pscore)
- if survey_weights is not None:
- W = W * survey_weights
- XWX = covX.T @ (W[:, None] * covX)
- # Rank-guarded inverse (sibling of the OR-side guard in
- # _compute_did_panel). A near-singular X'WX (constant/collinear
- # propensity covariate) does not raise LinAlgError, so the old
- # np.linalg.inv returned a garbage inverse that inflated IPW/DR
- # influence-function corrections. The helper truncates redundant
- # directions and is the sole owner of the PS-Hessian tracker append.
- XWX_inv, _, _ = _rank_guarded_inv(
- XWX, tracker=getattr(self, "_ps_lstsq_fallback_tracker", None)
- )
- hessian = XWX_inv * n_pair
-
- return pscore, hessian
-
- def _compute_or(
- self,
- delta_y: np.ndarray,
- PAa: np.ndarray,
- covX: np.ndarray,
- survey_weights: Optional[np.ndarray] = None,
- ) -> np.ndarray:
- """Fit OLS on control outcome changes. Returns or_delta for all pair units.
-
- Honors self.rank_deficient_action for collinear covariates. The outcome
- regression is fit through the shared scale-robust solver
- (``solve_ols`` -> column-equilibrated SVD; matches TripleDifference and
- R's lm()/QR), with optional WLS via ``solve_ols(weights=...)``.
- """
- from diff_diff.linalg import solve_ols as _solve_ols
-
- control_mask = PAa > 0
- n_c = int(np.sum(control_mask))
- if n_c == 0:
- return np.zeros(len(delta_y))
-
- X_control = covX[control_mask]
- y_control = delta_y[control_mask]
- sw_control = survey_weights[control_mask] if survey_weights is not None else None
-
- # Outcome regression via the shared scale-robust solver (`solve_ols` ->
- # column-equilibrated SVD/gelsd; matches TripleDifference and R's
- # lm()/QR). Replaces the prior `cho_solve(X'X)` cache fast path, which
- # was NOT scale-equilibrated (a large-scale covariate would corrupt the
- # OR fit via the normal-equations Cholesky). We only need beta for the
- # OR prediction, so zero the NaN (dropped-column) coefficients.
- beta, _, _ = _solve_ols(
- X_control,
- y_control,
- rank_deficient_action=self.rank_deficient_action,
- weights=sw_control,
- return_vcov=False,
- )
- beta = np.where(np.isnan(beta), 0.0, beta)
-
- return covX @ beta
-
- # ------------------------------------------------------------------
- # GMM-optimal combination (matches R's att_gt GMM procedure)
- # ------------------------------------------------------------------
-
- def _combine_gmm(
- self,
- att_vec: np.ndarray,
- inf_func_matrix: np.ndarray,
- n_units: int,
- ) -> Tuple[float, np.ndarray, np.ndarray, float]:
- """
- Combine comparison-group-specific estimates via GMM-optimal weights.
-
- Returns (att_gmm, inf_gmm, weights, se_gmm).
- """
- k = len(att_vec)
-
- if k == 1:
- att_gmm = float(att_vec[0])
- inf_gmm = inf_func_matrix[0].copy()
- # R's SE: sqrt(sum(IF^2) / n^2)
- se_gmm = float(np.sqrt(np.sum(inf_gmm**2) / n_units**2))
- return att_gmm, inf_gmm, np.array([1.0]), se_gmm
-
- # R: OMEGA <- cov(inf_mat_local) — sample covariance, ddof=1
- Omega = np.cov(inf_func_matrix)
-
- ones = np.ones(k)
- try:
- Omega_inv = np.linalg.inv(Omega)
- except np.linalg.LinAlgError:
- warnings.warn(
- "Singular covariance matrix in GMM combination. " "Using pseudoinverse.",
- UserWarning,
- stacklevel=3,
- )
- Omega_inv = np.linalg.pinv(Omega)
-
- denom = float(ones @ Omega_inv @ ones)
- if denom <= 0 or not np.isfinite(denom):
- weights = np.full(k, 1.0 / k)
- att_gmm = float(weights @ att_vec)
- inf_gmm = weights @ inf_func_matrix
- se_gmm = float(np.sqrt(np.sum(inf_gmm**2) / n_units**2))
- else:
- weights = (Omega_inv @ ones) / denom
- att_gmm = float(weights @ att_vec)
- inf_gmm = weights @ inf_func_matrix
- # R: gmm_se <- sqrt(1 / (n * sum(inv_OMEGA)))
- se_gmm = float(np.sqrt(1.0 / (n_units * denom)))
-
- return att_gmm, inf_gmm, weights, se_gmm
diff --git a/diff_diff/triple_diff.py b/diff_diff/triple_diff.py
index 6a28b802..41cd5524 100644
--- a/diff_diff/triple_diff.py
+++ b/diff_diff/triple_diff.py
@@ -29,7 +29,7 @@
import warnings
from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Tuple, Union
import numpy as np
import pandas as pd
@@ -41,10 +41,19 @@
resolve_renamed_kwarg,
warn_deprecated_kwarg,
)
+from diff_diff._staggered_triple_diff_engine import _StaggeredTripleDiffEngineMixin
from diff_diff.linalg import _rank_guarded_inv, solve_logit, solve_ols
from diff_diff.results import _format_survey_block, _get_significance_stars
from diff_diff.results_base import BaseResults
-from diff_diff.utils import safe_inference
+from diff_diff.staggered_aggregation import CallawaySantAnnaAggregationMixin
+from diff_diff.staggered_bootstrap import CallawaySantAnnaBootstrapMixin
+from diff_diff.staggered_triple_diff_results import StaggeredTripleDiffResults
+from diff_diff.utils import (
+ safe_inference,
+ staggered_ddd_ctor_offenders,
+ validate_anticipation,
+ validate_n_bootstrap,
+)
if TYPE_CHECKING:
from diff_diff.survey import SurveyDesign
@@ -384,7 +393,12 @@ def epv_summary(self, show_all: bool = False) -> pd.DataFrame:
# =============================================================================
-class TripleDifference(BaseEstimator):
+class TripleDifference(
+ _StaggeredTripleDiffEngineMixin,
+ CallawaySantAnnaBootstrapMixin,
+ CallawaySantAnnaAggregationMixin,
+ BaseEstimator,
+):
"""
Triple Difference (DDD) estimator.
@@ -455,10 +469,44 @@ class TripleDifference(BaseEstimator):
When ``rank_deficient_action="error"``, errors are always
re-raised regardless of this setting.
+ control_group : str, default="not_yet_treated"
+ Comparison cohort for staggered mode: ``"not_yet_treated"`` (units
+ whose enabling period is still in the future) or ``"never_treated"``
+ (the never-enabled cohort only). The compact R spellings
+ ``"notyettreated"``/``"nevertreated"`` are accepted only by the
+ deprecated ``StaggeredTripleDifference`` and die with it at 4.0.
+ anticipation : int, default=0
+ Number of periods before the enabling period in which units may
+ already respond. Must be a non-negative integer. Shifts each cohort's
+ base period earlier and excludes cohorts entering treatment within the
+ window from the comparison group.
+ base_period : str, default="varying"
+ ``"varying"`` uses consecutive comparisons - each pre-treatment period
+ ``t`` is compared to ``t - 1``, while post-treatment periods use the
+ fixed cohort reference ``g - 1 - anticipation``. ``"universal"`` uses
+ ``g - 1 - anticipation`` for every period, pre and post alike.
+ The two therefore differ only on the PRE-treatment effects, which is
+ what makes ``"varying"`` the pre-trend-readable default.
+ n_bootstrap : int, default=0
+ Multiplier-bootstrap replications for staggered mode. ``0`` means the
+ analytical influence-function SEs (the library-wide ``0 = off``
+ convention, row M-081). Clustered inference in staggered mode is only
+ available through this path - ``cluster=`` is rejected there.
+ bootstrap_weights : str, default="rademacher"
+ Multiplier distribution: ``"rademacher"``, ``"mammen"`` or ``"webb"``.
+ seed : int, optional
+ Seed for the multiplier bootstrap.
+ cband : bool, default=True
+ Whether to compute simultaneous (sup-t) confidence bands alongside the
+ pointwise ones.
+
Attributes
----------
- results_ : TripleDifferenceResults
- Estimation results after calling fit().
+ results_ : TripleDifferenceResults or StaggeredTripleDiffResults
+ Estimation results after calling fit(). The 2x2x2 mode returns
+ ``TripleDifferenceResults``; the staggered mode returns
+ ``StaggeredTripleDiffResults`` (the two containers unify at 4.0, row
+ M-014).
is_fitted_ : bool
Whether the model has been fitted.
@@ -515,6 +563,22 @@ class TripleDifference(BaseEstimator):
This is weaker than requiring separate parallel trends for two DiDs,
as biases can cancel out in the differencing.
+ **Which parameters belong to which mode.** ``control_group``,
+ ``anticipation``, ``base_period``, ``n_bootstrap``, ``bootstrap_weights``,
+ ``seed`` and ``cband`` apply to the STAGGERED mode only
+ (``fit(..., first_treat=...)``).
+
+ The first four raise a ``ValueError`` if given a non-default value and then
+ used to fit the 2x2x2 design, rather than being silently ignored: that
+ engine has one pre/post contrast, so there is no comparison-cohort choice,
+ no anticipation window, no base-period rule and no multiplier bootstrap.
+
+ The other three - ``bootstrap_weights``, ``seed`` and ``cband`` - are
+ ACCEPTED and inert in 2x2x2 mode. They act only through ``n_bootstrap > 0``,
+ which that mode already rejects, so they are unreachable by construction
+ rather than silently ignored. The power helpers apply the same boundary, so
+ an estimator that is legal to ``fit()`` is legal to simulate.
+
References
----------
.. [1] Ortiz-Villavicencio, M., & Sant'Anna, P. H. C. (2025).
@@ -525,6 +589,11 @@ class TripleDifference(BaseEstimator):
American Economic Review, 84(3), 622-641.
"""
+ # Names this estimator in the shared bootstrap mixin's user-facing
+ # warnings. The mixin is shared with the other hosts, so a hard-coded
+ # literal there would misname whichever surface was actually fit.
+ _BOOTSTRAP_LABEL: ClassVar[str] = "TripleDifference"
+
_PARAM_ATTR_ALIASES = {"robust": "_robust_arg"}
_DERIVED_CONFIG_ATTRS = ("robust",)
@@ -539,11 +608,67 @@ def __init__(
rank_deficient_action: str = "warn",
epv_threshold: float = 10,
pscore_fallback: str = "error",
+ control_group: str = "not_yet_treated",
+ anticipation: int = 0,
+ base_period: str = "varying",
+ n_bootstrap: int = 0,
+ bootstrap_weights: str = "rademacher",
+ seed: Optional[int] = None,
+ cband: bool = True,
):
if estimation_method not in ("dr", "reg", "ipw"):
raise ValueError(
f"estimation_method must be 'dr', 'reg', or 'ipw', " f"got '{estimation_method}'"
)
+ # Staggered-mode value contracts (M-013). The merged class uses the
+ # UNDERSCORED library vocabulary from birth; R's compact spellings
+ # ("notyettreated"/"nevertreated") die with StaggeredTripleDifference
+ # at 4.0 and are deliberately NOT accepted here.
+ if control_group not in ("never_treated", "not_yet_treated"):
+ raise ValueError(
+ f"control_group must be 'never_treated' or 'not_yet_treated', "
+ f"got '{control_group}'"
+ )
+ # `anticipation` feeds BOTH the base-period rule and the comparison-cohort
+ # threshold, so an out-of-domain value silently changes the estimand
+ # rather than failing: anticipation=-1 makes the universal base period
+ # `g` (an ALREADY-TREATED period) and relaxes the not-yet-treated
+ # threshold to max(t, base) - 1, admitting cohorts treated at the
+ # evaluation period as "clean" controls. Neither condition is
+ # observable in the output. `bool` is rejected too - `True` would
+ # otherwise coerce to a silent one-period window. (The sibling
+ # estimators taking `anticipation` are not yet uniformly validated;
+ # aligning the family is tracked in TODO.md.)
+ validate_anticipation(anticipation)
+ if base_period not in ("varying", "universal"):
+ raise ValueError(f"base_period must be 'varying' or 'universal', got '{base_period}'")
+ if bootstrap_weights not in ("rademacher", "mammen", "webb"):
+ raise ValueError(
+ f"bootstrap_weights must be 'rademacher', 'mammen', or 'webb', "
+ f"got '{bootstrap_weights}'"
+ )
+ validate_n_bootstrap(n_bootstrap)
+ # pscore_trim gains the staggered engine's range check (row M-142). It
+ # was previously unvalidated here, but the value feeds
+ # np.clip(pscore, trim, 1 - trim) in both engines: trim=0 disables the
+ # overlap guard that keeps the 1/(1-p) IPW/DR weights finite, and
+ # trim >= 0.5 inverts the clip bounds.
+ # The TYPE guard precedes the range check so the documented ValueError is
+ # what users actually see: a bare `0 < x < 0.5` raises an incidental
+ # TypeError on None/str/complex/list, an ambiguous-truth ValueError on a
+ # multi-element array, and - worst - ACCEPTS a 1-element array, storing an
+ # ndarray as the parameter. Same shape as validate_n_bootstrap: reject
+ # bool (True would read as a 1.0 trim), then non-real-scalar, then
+ # non-finite, then the range.
+ if isinstance(pscore_trim, bool) or not isinstance(
+ pscore_trim, (int, float, np.integer, np.floating)
+ ):
+ raise ValueError(
+ f"pscore_trim must be a real number in (0, 0.5), got {pscore_trim!r} "
+ f"(type {type(pscore_trim).__name__})"
+ )
+ if not np.isfinite(pscore_trim) or not 0 < pscore_trim < 0.5:
+ raise ValueError(f"pscore_trim must be in (0, 0.5), got {pscore_trim}")
if rank_deficient_action not in ["warn", "error", "silent"]:
raise ValueError(
f"rank_deficient_action must be 'warn', 'error', or 'silent', "
@@ -580,23 +705,46 @@ def __init__(
self.rank_deficient_action = rank_deficient_action
self.epv_threshold = epv_threshold
self.pscore_fallback = pscore_fallback
+ # Staggered-mode config (M-013). Inert in 2x2x2 mode, and fit() rejects
+ # any non-default value there rather than letting it pass unused.
+ self.control_group = control_group
+ self.anticipation = anticipation
+ self.base_period = base_period
+ self.n_bootstrap = n_bootstrap
+ self.bootstrap_weights = bootstrap_weights
+ self.seed = seed
+ self.cband = cband
self.is_fitted_ = False
- self.results_: Optional[TripleDifferenceResults] = None
+ self.results_: Optional[Union[TripleDifferenceResults, StaggeredTripleDiffResults]] = None
def fit(
self,
data: pd.DataFrame,
outcome: str,
- group: str,
- partition: str,
+ group: Any = NOT_SUPPLIED,
+ partition: Any = NOT_SUPPLIED,
post: Any = NOT_SUPPLIED,
covariates: Optional[List[str]] = None,
survey_design=None,
time: Any = NOT_SUPPLIED,
- ) -> TripleDifferenceResults:
+ *,
+ unit: Any = NOT_SUPPLIED,
+ first_treat: Any = NOT_SUPPLIED,
+ aggregate: Optional[str] = None,
+ balance_e: Optional[int] = None,
+ ) -> Union[TripleDifferenceResults, StaggeredTripleDiffResults]:
"""
- Fit the Triple Difference model.
+ Fit the Triple Difference model, in either of its two designs.
+
+ The 2x2x2 design takes ``(group, partition, post)``; the staggered
+ design takes ``(unit, time, first_treat, partition)``. ``first_treat=``
+ selects the staggered engine - mixing the two parameter sets is an
+ error, never a guess. This mirrors the reference implementation, whose
+ ``triplediff::ddd()`` serves both designs from one signature.
+
+ The staggered-only parameters are keyword-only, so a staggered call
+ written positionally cannot silently bind to the 2x2x2 slots.
Parameters
----------
@@ -605,13 +753,14 @@ def fit(
outcome : str
Name of the outcome variable column.
group : str
- Name of the group indicator column (0/1).
+ 2x2x2 mode. Name of the group indicator column (0/1).
1 = treated group (e.g., states that enacted policy).
0 = control group.
partition : str
- Name of the partition/eligibility indicator column (0/1).
- 1 = eligible partition (e.g., women, targeted demographic).
- 0 = ineligible partition.
+ BOTH modes. Name of the partition/eligibility indicator column
+ (0/1). 1 = eligible partition (e.g., women, targeted demographic);
+ 0 = ineligible partition. In staggered mode it must be
+ time-invariant within units.
post : str
Name of the post-period indicator column (0/1).
1 = post-treatment period.
@@ -625,24 +774,119 @@ def fit(
provided, uses survey weights for estimation and Taylor Series
Linearization (TSL) for variance estimation. Supported with
all estimation methods ("reg", "ipw", "dr").
+ time : str
+ In STAGGERED mode, the calendar period column (keyword). In 2x2x2
+ mode it is the deprecated alias for ``post`` (row M-031) and warns
+ with ``FutureWarning``. From 4.0 it means the calendar column only
+ (row M-085).
+ unit : str, keyword-only
+ Staggered mode. Unit identifier column.
+ first_treat : str, keyword-only
+ Staggered mode, and the MODE TRIGGER. Column giving each unit's
+ enabling period; 0 or ``np.inf`` marks never-enabled units.
+ aggregate : str, optional, keyword-only
+ Staggered mode. ``"event_study"``, ``"group"``, ``"simple"`` or
+ ``"all"``.
+ balance_e : int, optional, keyword-only
+ Staggered mode. Event time to balance the event study on.
Returns
-------
- TripleDifferenceResults
- Object containing estimation results.
+ TripleDifferenceResults or StaggeredTripleDiffResults
+ 2x2x2 mode returns ``TripleDifferenceResults``; staggered mode
+ returns ``StaggeredTripleDiffResults``. The two containers unify at
+ 4.0 (row M-014).
Raises
------
ValueError
- If required columns are missing or data validation fails.
+ If required columns are missing, data validation fails, or the two
+ designs' parameters are mixed.
NotImplementedError
If survey_design is used with wild_bootstrap inference.
-
- The keyword-only ``time`` parameter is a deprecated alias for
- ``post`` (row M-031); it warns with ``FutureWarning``. From 4.0,
- ``time=`` on the merged staggered interface means the CALENDAR
- column only (row M-085) - the 2x2x2 post dummy is ``post=``.
"""
+ # ---- Step 0: mode-INDEPENDENT value validation -------------------
+ # A bad `aggregate` value must not be masked by a mode error (3(a)'s
+ # `spec` precedent). Consequence, deliberate: fit(aggregate="bogus")
+ # WITHOUT first_treat reports the invalid value, not the mode error.
+ # `_validate_vcov_type` is deliberately NOT hoisted here - moving it
+ # ahead of the M-031 shim would make a direct-attribute-mutated
+ # vcov_type raise BEFORE the rename FutureWarning that fires today.
+ if aggregate is not None and aggregate not in ("event_study", "group", "simple", "all"):
+ raise ValueError(
+ f"aggregate must be 'event_study', 'group', 'simple', or 'all', "
+ f"got '{aggregate}'"
+ )
+
+ # ---- Step 1: staggered mode ---------------------------------------
+ # Raw sentinels forwarded: the 2x2x2 rename shim must not run, because
+ # `time=` is the calendar column on this path, not a deprecated alias.
+ if first_treat is not NOT_SUPPLIED:
+ return self._fit_staggered(
+ data,
+ outcome,
+ unit=unit,
+ time=time,
+ first_treat=first_treat,
+ partition=partition,
+ group=group,
+ post=post,
+ covariates=covariates,
+ aggregate=aggregate,
+ balance_e=balance_e,
+ survey_design=survey_design,
+ )
+
+ # ---- Step 2: 2x2x2 mode, staggered-only FIT params ----------------
+ _staggered_only = [
+ name
+ for name, supplied in (
+ ("unit", unit is not NOT_SUPPLIED),
+ ("aggregate", aggregate is not None),
+ ("balance_e", balance_e is not None),
+ )
+ if supplied
+ ]
+ if _staggered_only:
+ raise ValueError(
+ f"{', '.join(_staggered_only)} require(s) first_treat= (the staggered DDD "
+ f"mode); the 2x2x2 TripleDifference design estimates a single ATT from the "
+ f"(group=, partition=, post=) triple and has no adoption cohorts to "
+ f"aggregate over. For staggered adoption pass fit(..., unit=, "
+ f"time=, first_treat=, "
+ f"partition=)."
+ )
+
+ # ---- Step 3: 2x2x2 mode, staggered-only CONSTRUCTOR params --------
+ # The mode is only known at fit time, so this coherence check cannot
+ # live in __init__. bootstrap_weights/seed/cband are deliberately NOT
+ # listed: they only take effect through n_bootstrap > 0, which is
+ # rejected here, so they are unreachable-by-construction rather than
+ # silently ignored (documented in REGISTRY and the M-013 notes).
+ # Roster AND defaults come from the shared helper (derived from this
+ # class's own __init__ signature), not a local copy: the power entry
+ # points apply the same boundary, and two hand-maintained lists would
+ # eventually disagree about what "staggered-configured" means.
+ _staggered_ctor = staggered_ddd_ctor_offenders(self)
+ if _staggered_ctor:
+ raise ValueError(
+ f"TripleDifference({', '.join(n + '=' for n in _staggered_ctor)}) applies to "
+ f"the staggered DDD mode only (fit(..., first_treat=)): the 2x2x2 engine has "
+ f"one pre/post contrast, so there is no comparison-cohort choice, no "
+ f"anticipation window, no base-period rule and no multiplier bootstrap - its "
+ f"SEs come from the efficient influence function (cluster= for CR1). Drop the "
+ f"parameter(s), or pass first_treat= to fit the staggered design."
+ )
+
+ # ---- Step 4: 2x2x2 mode, the M-031 rename shim --------------------
+ # group/partition are sentinel-defaulted now (they must be omissible in
+ # staggered mode), so their missing-argument TypeErrors are restored
+ # here. They run BEFORE the rename shim: after it, fit(df, "y",
+ # time="t") would newly emit the M-031 FutureWarning before raising the
+ # missing-group TypeError (today it raises immediately, with no
+ # warning, and under -W error the surfaced exception type would flip).
+ require_arg("TripleDifference.fit", "group", group)
+ require_arg("TripleDifference.fit", "partition", partition)
post = resolve_renamed_kwarg(
"TripleDifference.fit",
"time",
@@ -863,6 +1107,102 @@ def fit(
self.is_fitted_ = True
return self.results_
+ def _fit_staggered(
+ self,
+ data: pd.DataFrame,
+ outcome: str,
+ *,
+ unit: Any,
+ time: Any,
+ first_treat: Any,
+ partition: Any,
+ group: Any,
+ post: Any,
+ covariates: Optional[List[str]],
+ aggregate: Optional[str],
+ balance_e: Optional[int],
+ survey_design,
+ ) -> StaggeredTripleDiffResults:
+ """Staggered branch of the merged fit (row M-013).
+
+ Reached only when ``first_treat=`` was supplied. Rejects the other
+ design's parameters rather than ignoring them, then hands off to the
+ engine shared with the deprecated ``StaggeredTripleDifference``.
+ """
+ # Step 1: the 2x2x2 params are rejected, not ignored. A positional 3rd
+ # or 5th argument lands in group/post and is caught here.
+ _2x2x2_only = [
+ name for name, value in (("group", group), ("post", post)) if value is not NOT_SUPPLIED
+ ]
+ if _2x2x2_only:
+ raise ValueError(
+ f"{', '.join(_2x2x2_only)} belong(s) to the 2x2x2 mode and cannot be combined "
+ f"with first_treat=: staggered DDD reads the treated cohort from first_treat= "
+ f"(0 or np.inf = never enabled) and the pre/post contrast from time= (the "
+ f"calendar column), so there is no 0/1 group dummy and no 0/1 post dummy. "
+ f"Positional arguments 3-5 are the 2x2x2 (group, partition, post) triple - in "
+ f"staggered mode pass unit=, time=, first_treat= and partition= as keywords."
+ )
+
+ # Step 2: the reverse constructor guard, the mirror of the 2x2x2 one.
+ # The two arms have DIFFERENT reachability. `robust=` is constructible
+ # (it only warns), so that arm is reached normally. `vcov_type` is
+ # validated eagerly in __init__ and by set_params' probe re-init, so a
+ # bad value never survives construction - its only live route is direct
+ # attribute mutation, the same bypass __init__'s comment documents. That
+ # makes this arm the staggered mode's direct-mutation guard, not dead
+ # code, and it is why no separate _validate_vcov_type call is needed.
+ _2x2x2_ctor = [
+ name
+ for name, non_default in (
+ ("robust", self._robust_arg is not None),
+ ("vcov_type", self.vcov_type != "hc1"),
+ )
+ if non_default
+ ]
+ if _2x2x2_ctor:
+ raise ValueError(
+ f"TripleDifference({', '.join(n + '=' for n in _2x2x2_ctor)}) applies to the "
+ f"2x2x2 mode only: the staggered engine's SEs come from the GMM-combined "
+ f"influence function (analytical) or the multiplier bootstrap "
+ f"(n_bootstrap > 0), so there is no sandwich family to select."
+ )
+
+ # Step 3: required staggered columns (first_treat is present already).
+ require_arg("TripleDifference.fit", "unit", unit)
+ require_arg("TripleDifference.fit", "time", time)
+ require_arg("TripleDifference.fit", "partition", partition)
+
+ # Step 4: cluster= is a CONSTRUCTOR param and the staggered engine has
+ # no clustered analytical path. The deprecated class accepts-then-warns
+ # -then-ignores it; this surface is new, so it raises from birth rather
+ # than shipping an accepted-but-unhonored parameter (3(a)'s precedent).
+ if self.cluster is not None:
+ raise ValueError(
+ "cluster= is not supported in staggered DDD mode: cluster-robust ANALYTICAL "
+ "SEs are not implemented for the staggered engine. Use n_bootstrap > 0 for "
+ "unit-level clustered inference via the multiplier bootstrap, or fit the "
+ "2x2x2 design, where cluster= gives Liang-Zeger CR1 standard errors."
+ )
+
+ results = self._fit_staggered_core(
+ data,
+ outcome,
+ unit,
+ time,
+ first_treat,
+ partition,
+ covariates=covariates,
+ aggregate=aggregate,
+ balance_e=balance_e,
+ survey_design=survey_design,
+ estimator_name="TripleDifference",
+ partition_label="partition",
+ _frame_offset=2,
+ )
+ self.results_ = results
+ return results
+
def _validate_data(
self,
data: pd.DataFrame,
diff --git a/diff_diff/utils.py b/diff_diff/utils.py
index f1d5ce9d..ea7ae4d4 100644
--- a/diff_diff/utils.py
+++ b/diff_diff/utils.py
@@ -2,6 +2,7 @@
Utility functions for difference-in-differences estimation.
"""
+import inspect
import os
import warnings
from dataclasses import dataclass, field
@@ -490,7 +491,8 @@ def validate_n_bootstrap(n_bootstrap: Any) -> None:
accepts numpy integers, rejects bool/None/float/negative). ``0`` means
bootstrap off wherever a ``> 0`` gate exists — the zero-default
estimators (CallawaySantAnna, SunAbraham, EfficientDiD, ImputationDiD,
- TwoStageDiD, WooldridgeDiD, ContinuousDiD, StaggeredTripleDifference)
+ TwoStageDiD, WooldridgeDiD, ContinuousDiD, StaggeredTripleDifference,
+ TripleDifference)
and every analytical lane. On the DiD/TWFE wild-bootstrap lane 0 never
meant off (the routing consults only ``inference=``); their fit-level
floor rejects ``n_bootstrap < 2`` under ``inference="wild_bootstrap"``.
@@ -501,6 +503,81 @@ def validate_n_bootstrap(n_bootstrap: Any) -> None:
raise ValueError(f"n_bootstrap must be a non-negative integer, got '{n_bootstrap}'")
+# The staggered-only TripleDifference constructor params that genuinely select
+# staggered behavior (row M-013). Lives here because TWO independent consumers
+# need the same boundary and must not drift apart: TripleDifference.fit(), which
+# rejects them in 2x2x2 mode, and the power entry points, which reject a
+# staggered-configured estimator at the front door. `bootstrap_weights`, `seed`
+# and `cband` are deliberately EXCLUDED - they act only through n_bootstrap > 0,
+# which 2x2x2 mode already rejects, so both surfaces accept them as inert.
+STAGGERED_DDD_CTOR_PARAMS = (
+ "control_group",
+ "anticipation",
+ "base_period",
+ "n_bootstrap",
+)
+
+
+def staggered_ddd_ctor_defaults(estimator: Any) -> Dict[str, Any]:
+ """Read :data:`STAGGERED_DDD_CTOR_PARAMS` defaults off ``estimator``'s class.
+
+ Derived from the live ``__init__`` signature rather than hard-coded, so a
+ future default change cannot make the two mode-detection boundaries
+ disagree — or, worse, reclassify an otherwise-default estimator as
+ staggered-configured and reject a legitimate 2x2x2 run.
+
+ Params absent from the signature, or present with no default, are skipped:
+ a param with no default has no "non-default value" to detect, and treating
+ ``Parameter.empty`` as the default would flag every instance.
+ """
+ try:
+ params = inspect.signature(type(estimator).__init__).parameters
+ except (TypeError, ValueError): # pragma: no cover - exotic callables
+ return {}
+ return {
+ name: params[name].default
+ for name in STAGGERED_DDD_CTOR_PARAMS
+ if name in params and params[name].default is not inspect.Parameter.empty
+ }
+
+
+def staggered_ddd_ctor_offenders(estimator: Any) -> List[str]:
+ """Names in :data:`STAGGERED_DDD_CTOR_PARAMS` set to a non-default value."""
+ return [
+ name
+ for name, default in staggered_ddd_ctor_defaults(estimator).items()
+ if getattr(estimator, name, default) != default
+ ]
+
+
+def validate_anticipation(anticipation: Any) -> None:
+ """Raise ValueError unless ``anticipation`` is a non-negative integer.
+
+ An out-of-domain anticipation window does not fail loudly on its own — it
+ silently changes the ESTIMAND. On the staggered DDD engine the value feeds
+ both the base-period rule and the comparison-cohort threshold, so
+ ``anticipation=-1`` makes the universal base period ``g`` (an
+ ALREADY-TREATED period) and relaxes the not-yet-treated threshold to
+ ``max(t, base) - 1``, admitting cohorts treated at the evaluation period as
+ clean controls. Neither condition is observable in the output.
+
+ ``bool`` is rejected: ``True`` would otherwise coerce to a silent
+ one-period window. Numpy integers are accepted, matching
+ :func:`validate_n_bootstrap`, whose shape this follows.
+
+ Adopted by ``TripleDifference`` and the shared staggered engine (so the
+ deprecated ``StaggeredTripleDifference`` fails closed too). The remaining
+ estimators taking ``anticipation`` are tracked for alignment in TODO.md.
+ """
+ if isinstance(anticipation, bool) or not isinstance(anticipation, (int, np.integer)):
+ raise ValueError(
+ f"anticipation must be a non-negative integer; got "
+ f"{anticipation!r} (type {type(anticipation).__name__})."
+ )
+ if anticipation < 0:
+ raise ValueError(f"anticipation must be a non-negative integer; got {anticipation!r}.")
+
+
def resolve_tail_df(
df_convention: str,
*,
diff --git a/docs/api/staggered.rst b/docs/api/staggered.rst
index 78542cb7..ebfe1731 100644
--- a/docs/api/staggered.rst
+++ b/docs/api/staggered.rst
@@ -137,6 +137,14 @@ StaggeredTripleDifference
Ortiz-Villavicencio & Sant'Anna (2025) staggered triple-difference (DDD) estimator
with group-time ATT identification under heterogeneous treatment timing.
+.. deprecated:: 3.9
+ Removed in 4.0 (ledger row M-013). Use
+ :class:`~diff_diff.TripleDifference` with
+ ``fit(..., unit=, time=, first_treat=, partition=)``, which runs the same
+ engine. ``eligibility=`` is named ``partition=`` there, and ``control_group``
+ takes the underscored values ``"not_yet_treated"``/``"never_treated"``. The
+ ``SDDD`` alias is deprecated with the class.
+
.. autoclass:: diff_diff.StaggeredTripleDifference
:no-index:
:members:
diff --git a/docs/api/triple_diff.rst b/docs/api/triple_diff.rst
index 46e19731..8d4a5ee8 100644
--- a/docs/api/triple_diff.rst
+++ b/docs/api/triple_diff.rst
@@ -29,6 +29,39 @@ TripleDifference
Main estimator class for Triple Difference designs.
+Since 3.9 this class serves BOTH triple-difference designs (ledger row M-013):
+the 2x2x2 design, and the staggered-adoption design that
+:class:`~diff_diff.StaggeredTripleDifference` used to own. ``first_treat=``
+selects the staggered engine; mixing the two designs' parameters raises rather
+than guessing. The estimation cores are unchanged - both surfaces share one
+engine, so the staggered numbers are identical to the deprecated class's.
+
+.. code-block:: python
+
+ from diff_diff import TripleDifference
+
+ # 2x2x2 design (unchanged)
+ ddd = TripleDifference(estimation_method="dr")
+ res = ddd.fit(df, outcome="y", group="state", partition="eligible", post="post")
+
+ # staggered adoption - the staggered params are keyword-only
+ sddd = TripleDifference(estimation_method="dr", control_group="not_yet_treated")
+ res = sddd.fit(
+ df,
+ outcome="y",
+ partition="eligible",
+ unit="id",
+ time="period",
+ first_treat="enacted",
+ aggregate="event_study",
+ )
+
+2x2x2 mode returns :class:`~diff_diff.TripleDifferenceResults`; staggered mode
+returns :class:`~diff_diff.StaggeredTripleDiffResults` (the containers unify at
+4.0, row M-014). ``cluster=`` gives Liang-Zeger CR1 in 2x2x2 mode and raises in
+staggered mode, where unit-level clustering is available through
+``n_bootstrap > 0``.
+
.. autoclass:: diff_diff.TripleDifference
:no-index:
:members:
diff --git a/docs/choosing_estimator.rst b/docs/choosing_estimator.rst
index 10d9349f..56519d96 100644
--- a/docs/choosing_estimator.rst
+++ b/docs/choosing_estimator.rst
@@ -16,7 +16,7 @@ Start here and follow the questions:
- **No** → Go to question 1
- **Yes, simultaneous treatment (2×2×2)** → Use :class:`~diff_diff.TripleDifference`
- - **Yes, with staggered timing** → Use :class:`~diff_diff.StaggeredTripleDifference`
+ - **Yes, with staggered timing** → Use :class:`~diff_diff.TripleDifference` with ``first_treat=`` (:class:`~diff_diff.StaggeredTripleDifference` is deprecated in 3.9 and runs the same engine until its 4.0 removal)
1. **Is treatment continuous?** (Units receive different doses or intensities)
@@ -844,12 +844,12 @@ estimation. The depth of support varies by estimator and variance method:
- Full (TSL)
- Full (analytical)
- Group-level (warning)
- * - ``TripleDifference``
+ * - ``TripleDifference`` (2x2x2 mode)
- pweight only
- Full
- Full (analytical)
- --
- * - ``StaggeredTripleDifference``
+ * - ``TripleDifference`` (staggered mode, ``first_treat=``)
- pweight only
- Full
- Full
diff --git a/docs/dev-status.md b/docs/dev-status.md
index b94f7ce7..8f31340f 100644
--- a/docs/dev-status.md
+++ b/docs/dev-status.md
@@ -35,7 +35,7 @@ Target: ideally < 1000 lines per module; modules ≥3000 lines are candidates fo
| `estimators.py` | 2441 | Monitor |
| `continuous_did.py` | 2459 | Monitor |
| `sun_abraham.py` | 2314 | Monitor |
-| `triple_diff.py` | 2231 | Monitor |
+| `triple_diff.py` | 2533 | Monitor — grew with the phase-3(b) DDD facade (merged constructor + dispatch + staggered branch) |
| `wooldridge.py` | 2192 | Monitor |
| `practitioner.py` | 2113 | Monitor — grew with per-estimator handlers (was 1511 on 2026-07-13) |
| `efficient_did.py` | 1729 | Acceptable — dropped below 2000 when the M-023 aggregate() migration extracted the aggregation mixin into `efficient_did_aggregation.py` (~520 lines, below this table's floor) |
@@ -45,7 +45,8 @@ Target: ideally < 1000 lines per module; modules ≥3000 lines are candidates fo
| `pretrends.py` | 1879 | Acceptable |
| `prep.py` | 1878 | Acceptable |
| `efficient_did_covariates.py` | 1818 | Acceptable |
-| `staggered_triple_diff.py` | 1680 | Acceptable |
+| `_staggered_triple_diff_engine.py` | 1635 | Acceptable — the shared staggered DDD engine, relocated here in phase 3(b) (row M-013) |
+| `staggered_triple_diff.py` | 262 | Acceptable — shrank from 1680 when the engine moved out; now the deprecated class's frozen 3.x surface |
| `trop_local.py` | 1662 | Acceptable |
| `lpdid.py` | 1607 | Acceptable |
| `stacked_did.py` | 1589 | Acceptable |
diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml
index 2165fda6..e1a9268b 100644
--- a/docs/doc-deps.yaml
+++ b/docs/doc-deps.yaml
@@ -29,6 +29,14 @@ groups:
staggered_triple_diff:
- diff_diff/staggered_triple_diff.py
- diff_diff/staggered_triple_diff_results.py
+ # The estimation engine relocated here in phase 3(b) (row M-013) and is
+ # shared with TripleDifference's staggered mode. Private (leading
+ # underscore), so tests/test_doc_deps_integrity.py exempts it from needing
+ # its own `sources:` key - but it is mapped as a GROUP MEMBER anyway,
+ # because /docs-check and /docs-impact have no private carve-out and,
+ # more importantly, the whole staggered methodology would otherwise live
+ # in a module with no path to REGISTRY.md.
+ - diff_diff/_staggered_triple_diff_engine.py
trop:
- diff_diff/trop.py
- diff_diff/trop_global.py
@@ -226,6 +234,12 @@ sources:
- path: diff_diff/guides/llms.txt
section: "Estimators"
type: user_guide
+ # Hand-edit-only surfaces before phase 3(b): neither was mapped, so
+ # /docs-impact could not surface them when this estimator changed.
+ - path: docs/survey-roadmap.md
+ type: user_guide
+ - path: docs/methodology/papers/ortiz-villavicencio-santanna-2025-review.md
+ type: methodology
# ── SunAbraham ──────────────────────────────────────────────────────
diff --git a/docs/index.rst b/docs/index.rst
index b7c81ffa..9294a49b 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -162,7 +162,7 @@ Supported Estimators
* - :class:`~diff_diff.TripleDifference`
- Triple difference (DDD) estimator
* - :class:`~diff_diff.StaggeredTripleDifference`
- - Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD with group-time ATT
+ - Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD with group-time ATT (deprecated 3.9 - use :class:`~diff_diff.TripleDifference` with ``first_treat=``)
* - :class:`~diff_diff.ContinuousDiD`
- Callaway, Goodman-Bacon & Sant'Anna (2024) continuous-treatment dose-response DiD
* - :class:`~diff_diff.HeterogeneousAdoptionDiD`
diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md
index d7d914be..2dabcdbd 100644
--- a/docs/methodology/REGISTRY.md
+++ b/docs/methodology/REGISTRY.md
@@ -3102,8 +3102,170 @@ contract changes.
---
+
+### Staggered mode (3.9, ledger row M-013)
+
+Since 3.9 `TripleDifference` serves BOTH DDD designs from one signature, mirroring
+the reference implementation (`triplediff::ddd()` does the same). The estimand,
+estimator and variance of each design are unchanged - the merge is API-only, and
+the staggered estimation core is the relocated `StaggeredTripleDifference` engine,
+shared verbatim.
+
+- **Dispatch.** `first_treat=` selects the staggered engine; the 2x2x2 engine
+ serves `(group, partition, post)`. Mixing the two parameter sets is an ERROR in
+ both directions, constructor params included, never a silent guess. The
+ staggered-only fit params (`unit`, `first_treat`, `aggregate`, `balance_e`) are
+ keyword-only, so a positionally-written staggered call cannot bind to the 2x2x2
+ slots; positional slots 1-8 are unchanged.
+- **Note (`time=` semantics):** the dual role is resolved by dispatch ORDER, not
+ value inspection. In staggered mode `time=` is the calendar column and emits NO
+ rename warning; in 2x2x2 mode it stays the deprecated alias for `post=` (row
+ M-031) and warns. At 4.0 it means the calendar column only (row M-085).
+- **Note (vocabulary):** the third dimension is `partition` in both modes, and
+ `control_group` takes the underscored values `"not_yet_treated"` /
+ `"never_treated"`. R's compact spellings are accepted only by the deprecated
+ `StaggeredTripleDifference` and die with it.
+- **Note (`cluster=` raises in staggered mode):** cluster-robust ANALYTICAL SEs
+ are not implemented for the staggered engine, so on this surface `cluster=`
+ raises a `ValueError` steering to `n_bootstrap > 0` (unit-level clustering via
+ the multiplier bootstrap). The deprecated class keeps its 3.x
+ accepted-then-warned-then-ignored `UserWarning`. **Deviation from the dying
+ class, deliberate:** the raise is live from the mode's 3.9 birth, because a new
+ surface needs no deprecation window and should not accept a parameter it does
+ not honor. In 2x2x2 mode `cluster=` continues to give Liang-Zeger CR1.
+- **Note (unreachable-by-construction params):** `bootstrap_weights`, `seed` and
+ `cband` pass silently in 2x2x2 mode. They take effect only through
+ `n_bootstrap > 0`, which 2x2x2 mode rejects outright, so they are unreachable
+ rather than silently ignored. `control_group`, `anticipation`, `base_period` and
+ `n_bootstrap` themselves all raise in 2x2x2 mode. **The power surface applies
+ the SAME boundary:** `simulate_power` / `simulate_mde` / `simulate_sample_size`
+ reject only those four and accept the three inert ones, so an estimator that is
+ legal to `fit()` is legal to simulate. That equivalence extends to the FIT
+ kwargs, and the test differs per param because `fit()`'s does: `first_treat`
+ and `unit` are sentinel-defaulted, so SUPPLYING them at all selects staggered
+ mode (an explicit `first_treat=None` still does, then fails on a missing
+ column), while `aggregate`/`balance_e` default to `None` and only a NON-`None`
+ value is staggered-only. Power therefore keys on presence for the first pair
+ and on value for the second — keying on presence throughout would reject
+ `estimator_kwargs={"aggregate": None}`, a configuration `fit()` accepts. An earlier 3(b) revision rejected all
+ seven at the power front door, which made `TripleDifference(seed=7)` fittable
+ but not simulatable - `seed` is set habitually for reproducibility and the
+ `simulate_*` helpers take their own separate `seed=`.
+- **Note (`pscore_trim` validation tightened, ledger row M-142):** the merged
+ constructor adopts the staggered engine's `0 < pscore_trim < 0.5` rule;
+ `TripleDifference`'s `pscore_trim` was previously unvalidated. The value feeds
+ `np.clip(pscore, trim, 1 - trim)` in both engines, so `trim=0` disabled the
+ overlap guard that keeps the `1/(1-p)` IPW/DR weights finite and `trim >= 0.5`
+ inverted the clip bounds. `TripleDifference(pscore_trim=0)` therefore changes
+ from accepted to a loud `ValueError`. **Sibling divergence, recorded rather than
+ silently tolerated:** `ContinuousDiD` still validates `0.0 <= pscore_trim < 0.5`
+ and so still admits `0`; aligning it is out of scope for a DDD merge and is
+ tracked as a `TODO.md` row.
+- **Note (the `triple_difference()` wrapper stays 2x2x2-only):** the deprecated
+ functional wrapper is deliberately NOT extended to staggered mode. It reaches
+ only the 2x2x2 design and forwards its own `time=` as `post=`. Three reasons,
+ all program-level rather than local: the wrapper is itself deprecated in 3.9
+ and removed at 4.0 (row M-075), so widening it would mint new removal
+ obligations for params introduced solely to be deleted; phase 3(a) set the
+ precedent by extending no wrapper with its new event-study mode; and the
+ wrapper's own deprecation message already steers to the class, which is where
+ staggered mode lives. A caller who needs staggered DDD migrates to
+ `TripleDifference(...).fit(..., first_treat=...)` - the migration the row
+ exists to produce. Pinned by
+ `tests/test_v4_merge_ddd.py::TestGateHConsumers::test_deprecated_wrapper_stays_2x2x2_only`.
+- **Note (bootstrap `seed` is estimator state, not results provenance):** the
+ staggered container does NOT carry the `seed` that generated its bootstrap
+ SEs, p-values and sup-t bands; it is readable from the estimator
+ (`est.seed` / `get_params()`), not from the returned result. `seed` IS fully
+ propagated in the operative sense - it reaches the multiplier draws, so the
+ same seed reproduces an SE bit-exactly and a different seed moves it - and it
+ appears in `get_params()`/`set_params()`. What is absent is a results-object
+ FIELD. This is deliberate and pre-dates the 3.9 merge: `seed` was already a
+ `StaggeredTripleDifference` constructor param that its container never
+ recorded, and `CallawaySantAnnaResults` has exactly the same gap through the
+ shared `CallawaySantAnnaBootstrapMixin`. Adding the field to one container
+ alone would leave the two siblings inconsistent, and
+ `StaggeredTripleDiffResults` is already slated for the 4.0 container port
+ (row M-014), so the fix is sequenced there and tracked as a `TODO.md` row
+ covering BOTH containers. Estimators that do expose it
+ (`ContinuousDiDResults`, `EfficientDiDResults`, `SyntheticDiDResults`) are
+ the precedent for the eventual shape.
+- **Note (`anticipation` domain validated from birth):** `anticipation` is one of
+ the seven staggered-only constructor params introduced on `TripleDifference` in
+ 3.9, and it is validated as a non-negative integer (`bool` rejected) from that
+ birth - so this is a new param's input contract, not a tightening of an existing
+ one, and it needs no lifecycle row. The guard is load-bearing rather than
+ cosmetic: the value feeds BOTH the base-period rule and the comparison-cohort
+ threshold, so `anticipation=-1` would make the universal base period `g` - an
+ ALREADY-TREATED period - and relax the not-yet-treated threshold to
+ `max(t, base) - 1`, admitting cohorts treated at the evaluation period as clean
+ controls. Neither condition is observable in the output, so an unvalidated
+ window is a silent estimand change (see the no-silent-failures policy).
+ `bool` is rejected because `True` would otherwise coerce to a silent
+ one-period window. The check is the shared `utils.validate_anticipation`, and
+ it is ALSO called from the staggered engine, so the deprecated
+ `StaggeredTripleDifference` fails closed too — its frozen 3.x API shape was
+ never a licence to emit silently-biased numbers from the same engine. The
+ split is deliberate: on the dying class CONSTRUCTION still succeeds (its
+ signature contract is untouched) and FIT raises, because this is an
+ identification guard rather than an API change. The engine-level call also
+ covers direct attribute mutation, which bypasses `__init__` and `set_params`
+ alike. **Sibling divergence, recorded rather than silently tolerated:** among
+ the remaining estimators taking `anticipation`, only `spillover.py` and
+ `wooldridge.py` validate it (and neither rejects `bool`); aligning the family
+ behind the shared helper is tracked as a `TODO.md` row.
+- **Note (degenerate enabling cohorts are reported, not hidden):** a positive
+ `first_treat` cohort whose units are ALL `partition == 0` cannot identify
+ `ATT(g,t)` — the DDD contrast needs the eligible-treated cell — so it
+ contributes to no aggregate. That was previously invisible: the cohort was
+ still listed in `results.groups`/`n_groups`, and the only warnings naming it
+ described its role as a COMPARISON cohort for other `g`, never as a treated
+ cohort that dropped out. The estimate quietly covered fewer cohorts than the
+ metadata advertised. Now a `UserWarning` names the cohort, its unit count and
+ the consequence, and `groups` is derived from the cohorts that actually
+ produced a `(g, t)` cell — so it is honest for every drop-out reason (no
+ eligible treated units, no valid comparison cohort, all base periods outside
+ the panel), not just this one. **The estimate is deliberately unchanged:** it
+ remains valid for the cohorts that do identify, and this engine already
+ warns-and-skips rather than raising for unidentified CELLS, so raising for a
+ cohort would have been the odd one out. Pre-existing behavior on
+ `StaggeredTripleDifference`; the fix is in the shared engine, so both
+ surfaces report identically.
+- **Note (`first_treat` cohort encoding is fail-closed on negatives):** the
+ contract is `0` for never-treated (or `+inf`, recoded to `0` with a
+ `UserWarning`) and positive period labels for treated cohorts. Negative values
+ — including the common `-1` never-treated convention and `-inf` — raise a
+ `ValueError` naming the offending values. They are NOT recoded: unlike `+inf`,
+ a negative value carries no unambiguous intent, and guessing would be the
+ silent sample change the guard exists to stop. Why it is a guard and not a
+ documentation matter: `_precompute_structures` builds the treated set from
+ `g > 0` and the never-enabled set from `g == 0`, so a negatively-encoded unit
+ belonged to NEITHER — still counted toward `n_obs` while contributing to no
+ ATT comparison — and the fit returned a plausible finite estimate for a
+ different population (measured on the 3(b) fixture: `overall_att`
+ 3.29438582 → 2.99470938 with `n_never_enabled` 24 → 0, no error and no
+ warning). The guard lives in the shared engine, so BOTH surfaces fail closed;
+ the behavior it replaces was pre-existing on `StaggeredTripleDifference`.
+- **Note (results containers):** 2x2x2 mode returns `TripleDifferenceResults` and
+ staggered mode returns `StaggeredTripleDiffResults` through 3.9; `fit`'s return
+ type is a `Union`. The one-shape unification lands at 4.0 with row M-014, so no
+ downstream consumer changes in 3.9.
+
## StaggeredTripleDifference
+- **Note:** Deprecated in 3.9, removed in 4.0 (ledger row M-013). The successor is
+ `TripleDifference().fit(..., unit=, time=, first_treat=, partition=)`, which runs
+ this EXACT engine - the estimation core moved verbatim to
+ `diff_diff/_staggered_triple_diff_engine.py` and both classes mix it in, so the
+ numbers are identical by construction (pinned bit-exactly in
+ `tests/test_v4_merge_ddd.py`). Two vocabulary differences on the merged surface:
+ `eligibility=` is named `partition=`, and `control_group` takes the underscored
+ values `"not_yet_treated"`/`"never_treated"` (R's compact spellings stay on this
+ class and die with it). One behavior difference: `cluster=` RAISES on the merged
+ surface instead of being accepted-then-warned-then-ignored, because a new surface
+ should not ship a parameter it does not honor. Everything below remains the 3.x
+ contract for the deprecated class through its removal.
+
**Primary source:** [Ortiz-Villavicencio, M., & Sant'Anna, P.H.C. (2025). Better Understanding Triple Differences Estimators. arXiv:2505.09942v3.](https://arxiv.org/abs/2505.09942v3). Paper review on file: `docs/methodology/papers/ortiz-villavicencio-santanna-2025-review.md`.
**Key implementation requirements:**
diff --git a/docs/methodology/papers/ortiz-villavicencio-santanna-2025-review.md b/docs/methodology/papers/ortiz-villavicencio-santanna-2025-review.md
index 3b5387b7..155860fd 100644
--- a/docs/methodology/papers/ortiz-villavicencio-santanna-2025-review.md
+++ b/docs/methodology/papers/ortiz-villavicencio-santanna-2025-review.md
@@ -220,7 +220,7 @@ The paper recommends practitioners **favor the optimal GMM DR DDD estimator** th
### Relation to Existing diff-diff Estimators
- **`TripleDifference`** (library, **Complete**) implements the **two-period DDD estimand and three-DR-DiD-component decomposition** of Eq. (3.5)/(4.1) with CS-2021-style influence-function SEs — this estimand and decomposition are what the paper is the primary source for. The library's path is **repeated cross-section** (`triplediff::ddd(panel=FALSE)`); the paper assumes a balanced panel and leaves repeated-cross-section / unbalanced data handling to future work (§7), so that data-path detail follows the companion `triplediff` package rather than the paper directly.
-- **`StaggeredTripleDifference`** (library, **In Progress**) implements the staggered / multi-period case — group-time `ATT(g,t)` via Eq. (4.1), the optimal-GMM combination across comparison cohorts (Eqs. 4.11–4.12), event-study via the CS aggregation mixin (Eq. 4.13), IF-based SEs, and a multiplier bootstrap for simultaneous bands. R reference `triplediff::ddd(panel = TRUE)` + `agg_ddd()`.
+- **`StaggeredTripleDifference`** (library, **Complete**; DEPRECATED in 3.9, removed in 4.0 - the staggered case is served by `TripleDifference().fit(..., first_treat=)` since the phase-3(b) merge, ledger row M-013, running this same engine) implements the staggered / multi-period case — group-time `ATT(g,t)` via Eq. (4.1), the optimal-GMM combination across comparison cohorts (Eqs. 4.11–4.12), event-study via the CS aggregation mixin (Eq. 4.13), IF-based SEs, and a multiplier bootstrap for simultaneous bands. R reference `triplediff::ddd(panel = TRUE)` + `agg_ddd()`.
- Both reuse Callaway–Sant'Anna machinery (the DR-DiD building block, the cohort-share event-study aggregation, and the multiplier-bootstrap inference of Remark 4.7).
*The following are points where the **library's implementation choices** differ from the **paper and/or the R `triplediff` package** — for example, the comparison-cohort admissibility rule (a) differs from the paper and (b) matches R. They are recorded here for the implementation record and are formalized as REGISTRY deviations separately; the methodology summary above sources only from the paper.*
diff --git a/docs/survey-roadmap.md b/docs/survey-roadmap.md
index 222ce7cc..056c4833 100644
--- a/docs/survey-roadmap.md
+++ b/docs/survey-roadmap.md
@@ -79,8 +79,10 @@ Weighted `solve_logit()` in `linalg.py` — survey weights enter IRLS as
workflow (strata, PSU, FPC, replicates, subpopulation, DEFF)
- **7d.** HonestDiD + survey: survey df and event-study VCV propagated
to sensitivity analysis with t-distribution critical values
-- **7e.** StaggeredTripleDifference survey support (only implementation
- in R or Python with design-based DDD variance)
+- **7e.** Staggered DDD survey support (only implementation in R or Python
+ with design-based DDD variance). Reached via `TripleDifference` with
+ `first_treat=` since 3.9; `StaggeredTripleDifference` is deprecated
+ (row M-013) but runs the same engine until its 4.0 removal.
### Phase 8: Survey Maturity (v2.8.3-v2.8.4)
diff --git a/docs/tutorials/06_power_analysis.ipynb b/docs/tutorials/06_power_analysis.ipynb
index 917bd3e2..0b685e13 100644
--- a/docs/tutorials/06_power_analysis.ipynb
+++ b/docs/tutorials/06_power_analysis.ipynb
@@ -425,7 +425,18 @@
{
"cell_type": "markdown",
"id": "6qpu05hi18s",
- "source": "### Triple Difference\n\n`TripleDifference` power routes by `n_periods`:\n\n- **`n_periods ≤ 2`** → the cross-sectional 2×2×2 factorial DGP (`generate_ddd_data`, group × partition × time). Sample sizes are **rounded via `n_per_cell = max(2, n_units // 8)`**, so the minimum effective N is 16 (2 units per cell × 8 cells), and the `effective_n_units` field tracks any rounding. `simulate_sample_size()` snaps `required_n` to multiples of 8.\n- **`n_periods > 2`** → the **panel** DGP (`generate_ddd_panel_data`), which honors `n_periods`/`treatment_period`. `n_units` maps directly to panel units (no rounding; `effective_n_units` is `None`), and `simulate_sample_size()` searches a continuous grid. The panel DGP has within-unit serial correlation, so construct the estimator as **`TripleDifference(cluster=\"unit\")`** — otherwise the unclustered SEs overstate power and a `UserWarning` is emitted. `treatment_fraction` is inert (balanced 2×2×2); pass `group_frac`/`partition_frac` via `data_generator_kwargs` to vary the split.\n\nThe example below uses the panel path (`n_periods=6`):",
+ "source": [
+ "### Triple Difference\n",
+ "\n",
+ "`TripleDifference` power routes by `n_periods`:\n",
+ "\n",
+ "- **`n_periods ≤ 2`** → the cross-sectional 2×2×2 factorial DGP (`generate_ddd_data`, group × partition × time). Sample sizes are **rounded via `n_per_cell = max(2, n_units // 8)`**, so the minimum effective N is 16 (2 units per cell × 8 cells), and the `effective_n_units` field tracks any rounding. `simulate_sample_size()` snaps `required_n` to multiples of 8.\n",
+ "- **`n_periods > 2`** → the **panel** DGP (`generate_ddd_panel_data`), which honors `n_periods`/`treatment_period`. `n_units` maps directly to panel units (no rounding; `effective_n_units` is `None`), and `simulate_sample_size()` searches a continuous grid. The panel DGP has within-unit serial correlation, so construct the estimator as **`TripleDifference(cluster=\"unit\")`** — otherwise the unclustered SEs overstate power and a `UserWarning` is emitted. `treatment_fraction` is inert (balanced 2×2×2); pass `group_frac`/`partition_frac` via `data_generator_kwargs` to vary the split.\n",
+ "\n",
+ "The example below uses the panel path (`n_periods=6`):\n",
+ "\n",
+ "Since 3.9 `TripleDifference` also serves the staggered-adoption DDD design (`fit(..., first_treat=)`), but **power analysis covers the 2x2x2 design only**: both registered DDD generators emit 2x2x2 data, so a staggered-configured estimator is rejected at the front door rather than simulated under the wrong design."
+ ],
"metadata": {}
},
{
@@ -575,7 +586,43 @@
{
"cell_type": "code",
"id": "v06p7ubbj9p",
- "source": "def my_dgp(n_units, n_periods, treatment_effect, treatment_fraction,\n treatment_period, noise_sd, seed=None):\n \"\"\"Custom DGP with heterogeneous unit effects.\"\"\"\n rng = np.random.default_rng(seed)\n n_treat = int(n_units * treatment_fraction)\n\n rows = []\n for i in range(n_units):\n unit_fe = rng.normal(0, 3) # heterogeneous unit effect\n treated_unit = i < n_treat\n for t in range(n_periods):\n post = int(t >= treatment_period)\n effect = treatment_effect * post if treated_unit else 0.0\n y = unit_fe + 2.0 * t + effect + rng.normal(0, noise_sd)\n rows.append({\n \"unit\": i, \"period\": t, \"outcome\": y,\n \"ever_treated\": int(treated_unit), \"post\": post,\n })\n return pd.DataFrame(rows)\n\n# Use the custom DGP with simulate_power\ncustom_results = simulate_power(\n estimator=DifferenceInDifferences(),\n n_units=80,\n n_periods=4,\n treatment_effect=4.0,\n sigma=3.0,\n n_simulations=100,\n seed=42,\n progress=False,\n data_generator=my_dgp,\n estimator_kwargs={\"outcome\": \"outcome\", \"treatment\": \"ever_treated\", \"time\": \"post\"},\n)\n\nprint(custom_results.summary())",
+ "source": [
+ "def my_dgp(n_units, n_periods, treatment_effect, treatment_fraction,\n",
+ " treatment_period, noise_sd, seed=None):\n",
+ " \"\"\"Custom DGP with heterogeneous unit effects.\"\"\"\n",
+ " rng = np.random.default_rng(seed)\n",
+ " n_treat = int(n_units * treatment_fraction)\n",
+ "\n",
+ " rows = []\n",
+ " for i in range(n_units):\n",
+ " unit_fe = rng.normal(0, 3) # heterogeneous unit effect\n",
+ " treated_unit = i < n_treat\n",
+ " for t in range(n_periods):\n",
+ " post = int(t >= treatment_period)\n",
+ " effect = treatment_effect * post if treated_unit else 0.0\n",
+ " y = unit_fe + 2.0 * t + effect + rng.normal(0, noise_sd)\n",
+ " rows.append({\n",
+ " \"unit\": i, \"period\": t, \"outcome\": y,\n",
+ " \"ever_treated\": int(treated_unit), \"post\": post,\n",
+ " })\n",
+ " return pd.DataFrame(rows)\n",
+ "\n",
+ "# Use the custom DGP with simulate_power\n",
+ "custom_results = simulate_power(\n",
+ " estimator=DifferenceInDifferences(),\n",
+ " n_units=80,\n",
+ " n_periods=4,\n",
+ " treatment_effect=4.0,\n",
+ " sigma=3.0,\n",
+ " n_simulations=100,\n",
+ " seed=42,\n",
+ " progress=False,\n",
+ " data_generator=my_dgp,\n",
+ " estimator_kwargs={\"outcome\": \"outcome\", \"treatment\": \"ever_treated\", \"post\": \"post\"},\n",
+ ")\n",
+ "\n",
+ "print(custom_results.summary())"
+ ],
"metadata": {},
"execution_count": null,
"outputs": []
@@ -679,4 +726,4 @@
},
"nbformat": 4,
"nbformat_minor": 5
-}
\ No newline at end of file
+}
diff --git a/docs/tutorials/08_triple_diff.ipynb b/docs/tutorials/08_triple_diff.ipynb
index 73dec636..084403e5 100644
--- a/docs/tutorials/08_triple_diff.ipynb
+++ b/docs/tutorials/08_triple_diff.ipynb
@@ -367,6 +367,84 @@
" print(f\" {cell}: {mean:.4f}\")"
]
},
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Staggered adoption\n",
+ "\n",
+ "Since 3.9 the same class also serves the **staggered-adoption** DDD design, where\n",
+ "groups become treatment-enabled in different periods, while each unit's\n",
+ "`partition` (eligibility) status stays time-invariant - the staggering is in\n",
+ "the enabling cohort, not in the third dimension. Supply `first_treat=`\n",
+ "(each unit's enabling period, with `0` or `np.inf` for never-enabled units;\n",
+ "negative values are rejected) and\n",
+ "the estimator switches engines: instead of one ATT it returns group-time\n",
+ "`ATT(g,t)` effects, combined across comparison cohorts by optimal GMM weighting\n",
+ "(Ortiz-Villavicencio & Sant'Anna 2025, Eqs. 4.11-4.12).\n",
+ "\n",
+ "The staggered parameters are keyword-only, and mixing the two designs' parameters\n",
+ "raises rather than guessing which one you meant.\n",
+ "\n",
+ "> This replaces `StaggeredTripleDifference`, which is deprecated in 3.9 and\n",
+ "> removed in 4.0. It runs the same engine, so the numbers are unchanged; on the\n",
+ "> merged surface the third dimension is named `partition=` (not `eligibility=`)\n",
+ "> and `control_group` takes `\"not_yet_treated\"` / `\"never_treated\"`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from diff_diff.prep import generate_staggered_ddd_data\n",
+ "\n",
+ "panel = generate_staggered_ddd_data(\n",
+ " n_units=120, n_periods=6, cohort_periods=[3, 4], seed=42\n",
+ ")\n",
+ "print(panel.head())\n",
+ "\n",
+ "staggered = TripleDifference(\n",
+ " estimation_method=\"dr\", control_group=\"not_yet_treated\"\n",
+ ")\n",
+ "staggered_results = staggered.fit(\n",
+ " panel,\n",
+ " outcome=\"outcome\",\n",
+ " partition=\"eligibility\",\n",
+ " unit=\"unit\",\n",
+ " time=\"period\",\n",
+ " first_treat=\"first_treat\",\n",
+ " aggregate=\"event_study\",\n",
+ ")\n",
+ "print(staggered_results.summary())"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Group-time effects, and the event-study aggregation\n",
+ "es = staggered_results.to_dataframe(level=\"event_study\")\n",
+ "print(es)\n",
+ "\n",
+ "# The 2x2x2 parameters are rejected in staggered mode, not silently ignored\n",
+ "try:\n",
+ " staggered.fit(\n",
+ " panel,\n",
+ " outcome=\"outcome\",\n",
+ " partition=\"eligibility\",\n",
+ " unit=\"unit\",\n",
+ " time=\"period\",\n",
+ " first_treat=\"first_treat\",\n",
+ " post=\"period\",\n",
+ " )\n",
+ "except ValueError as exc:\n",
+ " print(f\"\\nValueError: {exc}\")"
+ ]
+ },
{
"cell_type": "markdown",
"metadata": {},
diff --git a/docs/tutorials/16_survey_did.ipynb b/docs/tutorials/16_survey_did.ipynb
index b61276e9..552b16d9 100644
--- a/docs/tutorials/16_survey_did.ipynb
+++ b/docs/tutorials/16_survey_did.ipynb
@@ -1087,7 +1087,67 @@
"cell_type": "markdown",
"id": "cell-35-f1ef376c",
"metadata": {},
- "source": "## Which Estimators Support Survey Design?\n\n`diff-diff` supports survey design across all estimators, though the level of support varies:\n\n| Estimator | Weights | Strata/PSU/FPC (TSL) | Replicate Weights | Survey-Aware Bootstrap |\n|-----------|---------|---------------------|-------------------|------------------------|\n| **DifferenceInDifferences** | Full | Full | -- | -- |\n| **TwoWayFixedEffects** | Full | Full | -- | -- |\n| **MultiPeriodDiD** | Full | Full | -- | -- |\n| **CallawaySantAnna** | pweight only | Full | Full | Multiplier at PSU |\n| **TripleDifference** | pweight only | Full | Full (analytical) | -- |\n| **StaggeredTripleDifference** | pweight only | Full | Full | Multiplier at PSU |\n| **SunAbraham** | Full | Full | -- | Rao-Wu rescaled |\n| **StackedDiD** | pweight only | Full (pweight only) | -- | -- |\n| **ImputationDiD** | pweight only | Partial (no FPC) | -- | Multiplier at PSU |\n| **TwoStageDiD** | pweight only | Partial (no FPC) | -- | Multiplier at PSU |\n| **ContinuousDiD** | Full | Full | Full (analytical) | Multiplier at PSU |\n| **EfficientDiD** | Full | Full | Full (analytical) | Multiplier at PSU |\n| **SyntheticDiD** | pweight only | Full (all three variance methods) | -- | Hybrid pairs-bootstrap + Rao-Wu (bootstrap); stratified permutation (placebo); PSU-LOO (jackknife) |\n| **TROP** | pweight only | -- | -- | Rao-Wu rescaled |\n| **BaconDecomposition** | Diagnostic | Diagnostic | -- | -- |\n\n**Legend:**\n- **Full**: All weight types (pweight/fweight/aweight) + strata/PSU/FPC + Taylor Series Linearization variance\n- **Full (pweight only)**: Full TSL support with strata/PSU/FPC, but only accepts `pweight` weight type (`fweight`/`aweight` rejected because Q-weight composition changes their semantics)\n- **Partial (no FPC)**: Weights + strata (for df) + PSU (for clustering); FPC raises `NotImplementedError`\n- **pweight only** (Weights column): Only `pweight` accepted; `fweight`/`aweight` raise an error\n- **pweight only** (TSL column): Sampling weights for point estimates; no strata/PSU/FPC design elements\n- **Full (all three variance methods)** (SyntheticDiD TSL column): Strata/PSU/FPC supported on all three `variance_method` choices — `bootstrap` via weighted Frank-Wolfe + Rao-Wu, `placebo` via stratified permutation + weighted FW, `jackknife` via PSU-level LOO with stratum aggregation. Replicate-weight designs remain rejected (pre-existing limitation).\n- **Diagnostic**: Weighted descriptive statistics only (no inference)\n- **--**: Not supported\n\n**Note on SyntheticDiD:** all three variance methods now support full strata/PSU/FPC designs.\n\n- **Bootstrap** (PR #355) composes per-draw Rao-Wu rescaled weights with a weighted Frank-Wolfe variant of `_sc_weight_fw`. Each draw solves `min ||A·diag(rw)·ω - b||² + ζ²·Σ rw_i ω_i²` and composes `ω_eff = rw·ω/Σ(rw·ω)` for the SDID estimator. Pweight-only fits use the constant per-control survey weight as `rw`; full designs use Rao-Wu rescaling per draw.\n- **Placebo** uses a stratified permutation allocator: pseudo-treated indices are drawn from controls *within each stratum* containing actual treated units; weighted FW re-estimates ω and λ per draw with per-control survey weights flowing into both loss and regularization. SE follows Arkhangelsky Algorithm 4. The allocator requires at least `n_treated_h` controls per treated-containing stratum; fit-time guards raise targeted `ValueError` on infeasible configurations.\n- **Jackknife** uses PSU-level leave-one-out with stratum aggregation: `SE² = Σ_h (1-f_h)·(n_h-1)/n_h·Σ_{j∈h}(τ̂_{(h,j)} - τ̄_h)²` (Rust & Rao 1996). FPC folded via `(1-f_h)`; strata with fewer than 2 PSUs are silently skipped. Known anti-conservatism with few PSUs per stratum — for tight SE calibration in that regime, prefer `variance_method=\"bootstrap\"`.\n\nSee `docs/methodology/REGISTRY.md` §SyntheticDiD `Note (survey + bootstrap / placebo / jackknife composition)` for the full objectives, allocator asymmetry rationale (placebo ignores PSU axis, jackknife respects it), and validation details.\n\n**Note:** `EfficientDiD` supports `covariates` and `survey_design` simultaneously. The doubly-robust (DR) path threads survey weights through WLS outcome regression, weighted sieve propensity ratios, and survey-weighted kernel smoothing.\n\nFor full details, see `docs/survey-roadmap.md`."
+ "source": [
+ "## Which Estimators Support Survey Design?\n",
+ "\n",
+ "`diff-diff` supports survey design across all estimators, though the level of support varies:\n",
+ "\n",
+ "| Estimator | Weights | Strata/PSU/FPC (TSL) | Replicate Weights | Survey-Aware Bootstrap |\n",
+ "|-----------|---------|---------------------|-------------------|------------------------|\n",
+ "| **DifferenceInDifferences** | Full | Full | -- | -- |\n",
+ "| **TwoWayFixedEffects** | Full | Full | -- | -- |\n",
+ "| **MultiPeriodDiD** | Full | Full | -- | -- |\n",
+ "| **CallawaySantAnna** | pweight only | Full | Full | Multiplier at PSU |\n",
+ "| **TripleDifference** — 2x2x2 mode (`group=`, `partition=`, `post=`) | pweight only | Full | Full (analytical) | -- |\n",
+ "| **TripleDifference** — staggered mode (`first_treat=`) | pweight only | Full | Full | Multiplier at PSU |\n",
+ "| **SunAbraham** | Full | Full | -- | Rao-Wu rescaled |\n",
+ "| **StackedDiD** | pweight only | Full (pweight only) | -- | -- |\n",
+ "| **ImputationDiD** | pweight only | Partial (no FPC) | -- | Multiplier at PSU |\n",
+ "| **TwoStageDiD** | pweight only | Partial (no FPC) | -- | Multiplier at PSU |\n",
+ "| **ContinuousDiD** | Full | Full | Full (analytical) | Multiplier at PSU |\n",
+ "| **EfficientDiD** | Full | Full | Full (analytical) | Multiplier at PSU |\n",
+ "| **SyntheticDiD** | pweight only | Full (all three variance methods) | -- | Hybrid pairs-bootstrap + Rao-Wu (bootstrap); stratified permutation (placebo); PSU-LOO (jackknife) |\n",
+ "| **TROP** | pweight only | -- | -- | Rao-Wu rescaled |\n",
+ "| **BaconDecomposition** | Diagnostic | Diagnostic | -- | -- |\n",
+ "\n",
+ "**Legend:**\n",
+ "- **Full**: All weight types (pweight/fweight/aweight) + strata/PSU/FPC + Taylor Series Linearization variance\n",
+ "- **Full (pweight only)**: Full TSL support with strata/PSU/FPC, but only accepts `pweight` weight type (`fweight`/`aweight` rejected because Q-weight composition changes their semantics)\n",
+ "- **Partial (no FPC)**: Weights + strata (for df) + PSU (for clustering); FPC raises `NotImplementedError`\n",
+ "- **pweight only** (Weights column): Only `pweight` accepted; `fweight`/`aweight` raise an error\n",
+ "- **pweight only** (TSL column): Sampling weights for point estimates; no strata/PSU/FPC design elements\n",
+ "- **Full (all three variance methods)** (SyntheticDiD TSL column): Strata/PSU/FPC supported on all three `variance_method` choices — `bootstrap` via weighted Frank-Wolfe + Rao-Wu, `placebo` via stratified permutation + weighted FW, `jackknife` via PSU-level LOO with stratum aggregation. Replicate-weight designs remain rejected (pre-existing limitation).\n",
+ "- **Diagnostic**: Weighted descriptive statistics only (no inference)\n",
+ "- **--**: Not supported\n",
+ "\n",
+ "**Migration note:** `StaggeredTripleDifference` is deprecated in 3.9 and removed in\n",
+ "4.0. It is not a separate capability - it runs the SAME staggered engine as the\n",
+ "row above, reached by passing `first_treat=` to **`fit()`** (it is a fit\n",
+ "argument, not a constructor one):\n",
+ "\n",
+ "```python\n",
+ "TripleDifference().fit(\n",
+ " data, outcome='y', unit='unit', time='period',\n",
+ " first_treat='first_treat', partition='eligibility',\n",
+ " survey_design=design,\n",
+ ")\n",
+ "```\n",
+ "\n",
+ "Pass `partition=` instead of `eligibility=`, and use the underscored\n",
+ "`control_group` values (`\"not_yet_treated\"` / `\"never_treated\"`).\n",
+ "\n",
+ "**Note on SyntheticDiD:** all three variance methods now support full strata/PSU/FPC designs.\n",
+ "\n",
+ "- **Bootstrap** (PR #355) composes per-draw Rao-Wu rescaled weights with a weighted Frank-Wolfe variant of `_sc_weight_fw`. Each draw solves `min ||A·diag(rw)·ω - b||² + ζ²·Σ rw_i ω_i²` and composes `ω_eff = rw·ω/Σ(rw·ω)` for the SDID estimator. Pweight-only fits use the constant per-control survey weight as `rw`; full designs use Rao-Wu rescaling per draw.\n",
+ "- **Placebo** uses a stratified permutation allocator: pseudo-treated indices are drawn from controls *within each stratum* containing actual treated units; weighted FW re-estimates ω and λ per draw with per-control survey weights flowing into both loss and regularization. SE follows Arkhangelsky Algorithm 4. The allocator requires at least `n_treated_h` controls per treated-containing stratum; fit-time guards raise targeted `ValueError` on infeasible configurations.\n",
+ "- **Jackknife** uses PSU-level leave-one-out with stratum aggregation: `SE² = Σ_h (1-f_h)·(n_h-1)/n_h·Σ_{j∈h}(τ̂_{(h,j)} - τ̄_h)²` (Rust & Rao 1996). FPC folded via `(1-f_h)`; strata with fewer than 2 PSUs are silently skipped. Known anti-conservatism with few PSUs per stratum — for tight SE calibration in that regime, prefer `variance_method=\"bootstrap\"`.\n",
+ "\n",
+ "See `docs/methodology/REGISTRY.md` §SyntheticDiD `Note (survey + bootstrap / placebo / jackknife composition)` for the full objectives, allocator asymmetry rationale (placebo ignores PSU axis, jackknife respects it), and validation details.\n",
+ "\n",
+ "**Note:** `EfficientDiD` supports `covariates` and `survey_design` simultaneously. The doubly-robust (DR) path threads survey weights through WLS outcome regression, weighted sieve propensity ratios, and survey-weighted kernel smoothing.\n",
+ "\n",
+ "For full details, see `docs/survey-roadmap.md`."
+ ]
},
{
"cell_type": "markdown",
@@ -1436,4 +1496,4 @@
},
"nbformat": 4,
"nbformat_minor": 5
-}
\ No newline at end of file
+}
diff --git a/docs/v4-deprecations.yaml b/docs/v4-deprecations.yaml
index bfc81f8f..63ffc393 100644
--- a/docs/v4-deprecations.yaml
+++ b/docs/v4-deprecations.yaml
@@ -171,11 +171,12 @@ rows:
introduced_in: "3.9"
deprecated_in: "3.9"
removed_in: "4.0"
- status: planned
- phase: 3
+ status: shimmed
+ phase: 5
warning: FutureWarning
- code_refs: [diff_diff/triple_diff.py, diff_diff/staggered_triple_diff.py, diff_diff/__init__.py]
- notes: "Facade dispatch: 2x2x2 engine vs staggered engine, both internally unchanged (R-parity preserved). Partition param unified as 'partition'. The compact 'notyettreated' control_group spelling rides this class removal - no value shim on the dying class; the unified class uses underscored values from birth."
+ code_refs: [diff_diff/triple_diff.py, diff_diff/staggered_triple_diff.py, diff_diff/_staggered_triple_diff_engine.py, diff_diff/staggered_bootstrap.py, diff_diff/power.py, diff_diff/__init__.py]
+ test_ref: tests/test_v4_merge_ddd.py
+ notes: "Facade dispatch: 2x2x2 engine vs staggered engine, both internally unchanged (R-parity preserved). Partition param unified as 'partition'. The compact 'notyettreated' control_group spelling rides this class removal - no value shim on the dying class; the unified class uses underscored values from birth. SHIPPED 3.9 (Phase 3(b)). Mechanics: (1) the estimation core moved VERBATIM to the private diff_diff/_staggered_triple_diff_engine.py, mixed into both classes, so neither surface forks the math - the pre-merge oracles reproduce bit-identically on both engines (tests/_capture_v4_merge_ddd_oracles.py). (2) The core takes three knobs: estimator_name (messages never name the deprecated class on the merged surface), partition_label (applied to BOTH sentences of the time-invariance error - the tail 'varying eligibility' would otherwise leak onto the partition= interface), and _frame_offset, mirrored onto self._warn_frame_offset for the four helper depths plus the shared bootstrap mixin, so warning attribution is unchanged on both paths. (3) fit dispatches on first_treat=; the staggered-only fit params are KEYWORD-ONLY (a positionally-written staggered call cannot bind to the 2x2x2 slots), while positional slots 1-8 including the deprecated time= are preserved byte-for-byte, so nothing narrows. (4) Mixing the two designs' params raises in BOTH directions, constructor params included (control_group/anticipation/base_period/n_bootstrap in 2x2x2; robust/vcov_type in staggered). bootstrap_weights/seed/cband pass silently in 2x2x2 because they are unreachable without n_bootstrap > 0, which is itself rejected - recorded here and in REGISTRY rather than left to discovery. The power front door (simulate_power/simulate_mde/simulate_sample_size) applies the SAME boundary, rejecting only control_group/anticipation/base_period/n_bootstrap: an earlier revision had it reject all seven, which made TripleDifference(seed=7) legal to fit and illegal to simulate. The phase also added TWO fail-closed identification guards in the SHARED engine, so both surfaces (the deprecated class included) reject them: negative/-inf first_treat cohort values, which previously fell into neither the treated nor the comparison population and silently estimated on a different sample (overall_att 3.29438582 -> 2.99470938, n_never_enabled 24 -> 0 on the 3(b) fixture); and a non-negative-integer anticipation window via the new shared utils.validate_anticipation. The dying class's API SHAPE stays frozen - construction still succeeds, fit raises - because these are identification guards, not signature changes. (5) cluster= RAISES in staggered mode (the deprecated class keeps its accepted-then-warned-then-ignored UserWarning): a new surface should not ship an accepted-but-unhonored param, the 3(a) precedent. (6) Returns StaggeredTripleDiffResults through 3.9; the container unification is M-014's job at 4.0. (7) pscore_trim validation tightened, rowed separately as M-142. Fit-time aggregate=/balance_e= are carried on the surviving class and rowed as M-140/M-141."
- id: M-014
kind: class
group: merge-ddd
@@ -186,7 +187,7 @@ rows:
status: planned
phase: 5
code_refs: [diff_diff/staggered_triple_diff_results.py, diff_diff/__init__.py]
- notes: "Deprecation rides parent [M-013]. Successor = unified TripleDifference results shape (degenerate single-ATT view for 2x2x2)."
+ notes: "Deprecation rides parent [M-013]. Successor = unified TripleDifference results shape (degenerate single-ATT view for 2x2x2). 4.0 REMOVAL NOTE: since Phase 3(b) this container has a SECOND construction site - TripleDifference's staggered mode returns it too (the facade is API-only in 3.9; user-approved 2026-08-08). The 4.0 PR that unifies the shape must therefore decouple _fit_staggered_core's construction site as well as the deprecated class's, and it needs a disposition for TripleDifferenceResults, which has no row of its own. M-140/M-141's removed_in: 4.0 machine-checks that the unified container gains the post-fit aggregate() this release, so the two obligations land together."
- id: M-015
kind: class
group: merge-qdid
@@ -768,9 +769,9 @@ rows:
deprecated_in: "3.9"
removed_in: "4.0"
status: planned
- phase: 3
+ phase: 5
code_refs: [diff_diff/__init__.py]
- notes: "Dies with its class [M-013]; deprecation warning rides the parent class shim (same object). Migrate the old_target locator when [M-013] flips to removed (cross-row rule)."
+ notes: "Dies with its class [M-013]; deprecation warning rides the parent class shim (same object). Migrate the old_target locator when [M-013] flips to removed (cross-row rule). 3.9 FutureWarning live via the parent [M-013] shim since Phase 3(b) - SDDD stays a plain module global (NOT an entry in __init__'s _DEPRECATED_ALIASES: that table is read only by the module __getattr__, which never fires for a real global, so an entry would be dead code). Status stays 'planned' because this row's own transition is the 4.0 removal; the parent's message names the alias explicitly. Pinned by tests/test_v4_merge_ddd.py."
- id: M-062
kind: alias
group: alias-table
@@ -878,8 +879,8 @@ rows:
phase: 5
test_ref: tests/test_v4_wrapper_shims.py
warning: FutureWarning
- code_refs: [diff_diff/triple_diff.py, diff_diff/__init__.py, diff_diff/guides/llms-full.txt, docs/api/triple_diff.rst, docs/tutorials/08_triple_diff.ipynb]
- notes: "Use TripleDifference. 3.9 shim shipped (2(d) PR-A): FutureWarning at call + docstring deprecation note; the autofunction API page stays through 3.9."
+ code_refs: [diff_diff/triple_diff.py, diff_diff/__init__.py, diff_diff/guides/llms-full.txt, docs/api/triple_diff.rst, docs/tutorials/08_triple_diff.ipynb, docs/methodology/REGISTRY.md]
+ notes: "Use TripleDifference. 3.9 shim shipped (2(d) PR-A): FutureWarning at call + docstring deprecation note; the autofunction API page stays through 3.9. Phase 3(b) DELIBERATELY did not widen it to the merged staggered mode ([M-013]): the wrapper reaches only the 2x2x2 design, because growing a surface that is removed at 4.0 would mint fresh removal obligations for params introduced solely to be deleted, and phase 3(a) set the precedent by extending no wrapper with its new event-study mode. A caller needing staggered DDD migrates to the class - the migration this row exists to produce. Pinned by tests/test_v4_merge_ddd.py::TestGateHConsumers::test_deprecated_wrapper_stays_2x2x2_only and REGISTRY-noted."
- id: M-076
kind: function
group: function-wrappers
@@ -1679,3 +1680,41 @@ rows:
test_ref: tests/test_aggregate_contract.py
code_refs: [diff_diff/had_pretests.py, diff_diff/had.py]
notes: "The workflow twin of [M-027], added as a pre-cut amendment (next free id - the reserved 2(b) pool is spent/earmarked: M-118/M-119/M-120 claimed, M-116/M-121 intentionally unused). did_had_pretest_workflow's aggregate= routed WHICH pretest battery runs on the same 'overall'/'event_study' vocabulary the fit param used; it never calls fit(), so the successor is the SAME panel-shape inference (_infer_aggregate_mode, shared with fit so the two surfaces cannot drift), applied after the workflow's own alias reconciliation. A plain workflow call never warns; supplying ANY value warns once, then the legacy routing runs unchanged (invalid values still raise ValueError after the warning). `new` is null: there is no successor symbol - the HADPretestReport.aggregate FIELD survives as honest output metadata recording which battery ran (SURFACE_ALLOWLIST entry in tests/test_naming_guard.py; the MODES survive 4.0, only the param dies). The mode vocabulary dies from both public surfaces at 4.0."
+
+ - id: M-140
+ kind: param
+ group: merge-ddd
+ old: "diff_diff:TripleDifference.fit[aggregate]"
+ new: "diff_diff:TripleDifferenceResults.aggregate"
+ introduced_in: null
+ deprecated_in: null
+ removed_in: "4.0"
+ status: planned
+ phase: 5
+ code_refs: [diff_diff/triple_diff.py]
+ notes: "The fit-time aggregate= carried onto the SURVIVING class by the M-013 merge, and the one documented exception to section 6's aggregate-postfit program (M-020..M-027): the staggered DDD container has no post-fit aggregate() to steer users to (no AggregationMixin on StaggeredTripleDiffResults), so fit-time routing ships here as the only route rather than as a deprecated one. Hence deprecated_in and warning are null - nothing is being steered away from yet. introduced_in is DELIBERATELY null even though the param ships in 3.9: collect_due_problems fails any row past its introduced_in still at 'planned', 'shimmed' would require the successor locator to resolve (it cannot until the 4.0 container port), and 'removed' trips the early-removal guard - the ledger's status vocabulary cannot express 'introduced now, successor not built yet', so the row is modelled purely as a scheduled REMOVAL. removed_in: 4.0 is load-bearing twice: it satisfies the phase-table direction-2 citation predicate at phase 5, and it machine-checks the container port at the 4.0 bump, which is M-014's job in the same release. Section 6 carries the matching dated amendment."
+ - id: M-141
+ kind: param
+ group: merge-ddd
+ old: "diff_diff:TripleDifference.fit[balance_e]"
+ new: "diff_diff:TripleDifferenceResults.aggregate[balance_e]"
+ introduced_in: null
+ deprecated_in: null
+ removed_in: "4.0"
+ status: planned
+ phase: 5
+ code_refs: [diff_diff/triple_diff.py]
+ notes: "Twin of [M-140] for the balance_e knob - see its notes for the null-version rationale. Rowed separately rather than riding M-140's notes, per the M-117..M-120 precedent and M-020's own note that prose-only balance_e tracking was exactly the un-rowed-obligation class the gating-completeness amendment closed."
+ - id: M-142
+ kind: behavior
+ group: merge-ddd
+ old: "diff_diff:TripleDifference[pscore_trim]"
+ new: null
+ introduced_in: "3.9"
+ deprecated_in: null
+ removed_in: null
+ status: done
+ phase: 3
+ test_ref: tests/test_v4_merge_ddd.py
+ code_refs: [diff_diff/triple_diff.py]
+ notes: "Input-validation tightening shipped with the M-013 merge: TripleDifference's pscore_trim was previously UNVALIDATED, and the merged constructor adopts the staggered engine's 0 < x < 0.5 rule. Not cosmetic - the value feeds np.clip(pscore, trim, 1 - trim) in both engines, so trim=0 disables the overlap guard that keeps the 1/(1-p) IPW/DR weights finite and trim >= 0.5 inverts the clip bounds. TripleDifference(pscore_trim=0) therefore changes from accepted to a loud ValueError; no in-repo caller passed it. Same shape as [M-096] (a 3.9 validation tightening on a previously-unvalidated/silently-degrading param, rowed in the PR that shipped it), and status 'done' is terminal so the row is exempt from both phase-table directions. SIBLING DIVERGENCE, deliberate and recorded in REGISTRY: ContinuousDiD still validates 0.0 <= pscore_trim < 0.5, i.e. it admits 0 - aligning it is out of scope for a DDD merge and carries a TODO.md row instead of silent drift."
diff --git a/docs/v4-design.md b/docs/v4-design.md
index 7010cee4..a414a732 100644
--- a/docs/v4-design.md
+++ b/docs/v4-design.md
@@ -368,10 +368,23 @@ mechanics and the test triple in tests/test_v4_merge_mpd.py.)
```python
ddd = TripleDifference()
ddd.fit(df, outcome, group, partition, post) # 2x2x2 (RC engine)
-ddd.fit(df, outcome, unit, time, first_treat, partition, # staggered (panel engine)
- covariates=None, ...)
+ddd.fit(df, outcome, partition=..., unit=..., time=..., # staggered (panel engine)
+ first_treat=..., covariates=None, ...)
```
+**(Amended 2026-08-08, Phase 3(b) ship.)** The staggered snippet is KEYWORD
+form, not positional. Positional slots 1-8 belong to the 2x2x2 design and are
+preserved byte-for-byte (`data, outcome, group, partition, post, covariates,
+survey_design, time`) - including `time` at slot 8, so nothing narrows and no
+compat row is needed. The four staggered-only params (`unit`, `first_treat`,
+`aggregate`, `balance_e`) are KEYWORD-ONLY: a staggered call written
+positionally would otherwise bind `group="unit"`, `partition="period"`,
+`post="first_treat"` and land in 2x2x2 mode with a confusing downstream error
+instead of a mode error. `time=` keeps its dual role, resolved by dispatch
+ORDER rather than by value inspection: the staggered branch is taken before the
+M-031 rename shim runs, so `time=` is the calendar column there and emits no
+rename warning.
+
**Routing semantics.** One class, two engines, both internally UNCHANGED in
this program (R-parity preserved; engine unification is a possible 4.x
internal refactor, explicitly out of scope). Dispatch is by signature shape:
@@ -386,13 +399,30 @@ both modes (the paper's and R package's vocabulary); the staggered engine's
**Inference note.** The two engines keep their existing inference stacks
(analytical influence-function SEs on the 2x2x2 engine; multiplier bootstrap +
GMM weighting on the staggered engine). `cluster=` analytical SEs remain
-staggered-mode-unsupported (bootstrap required), as today - documented in the
-class docstring, tracked as post-4.0 backlog.
+staggered-mode-unsupported (bootstrap required) - documented in the class
+docstring, tracked as post-4.0 backlog. **(Amended 2026-08-08, Phase 3(b) ship,
+user-approved.)** On the MERGED surface `cluster=` now RAISES a `ValueError`
+steering to `n_bootstrap > 0`, rather than the deprecated class's
+accepted-then-warned-then-ignored `UserWarning`, which is preserved unchanged on
+that class. Same reasoning 3(a) ratified for MPD's wild-bootstrap fallback: the
+raise is live from the mode's 3.9 BIRTH, because a new surface needs no
+deprecation window and should not ship a parameter it does not honor.
**Results.** One results shape: the staggered container's structure with the
2x2x2 case as the degenerate single-ATT view (canonical quintet always
populated; group-time table empty in 2x2x2 mode). StaggeredTripleDiffResults
-dies [M-014].
+dies [M-014]. **(Amended 2026-08-08, Phase 3(b) ship, user-approved.)** That end
+state lands at 4.0, not in 3.9. The 3.9 facade is API-only: 2x2x2 mode returns
+`TripleDifferenceResults`, staggered mode returns `StaggeredTripleDiffResults`,
+and `fit`'s return type widens to a `Union` (the `TwoWayFixedEffects` precedent
+from 3(a)). Rationale: 3(a) could consume a finished container because
+`EventStudyResults` shipped in Phase 2, whereas unifying here would have to
+INVENT one - the two containers disagree on storage (native quintet vs
+`overall_*` + property aliases), nine 2x2x2 fields have no staggered
+counterpart, `to_dataframe`/`epv_summary`/`summary` diverge, and changing the
+2x2x2 return type is an unrowed breaking change to a surface with no
+deprecation row. [M-014] keeps the obligation and now records the second
+construction site.
**Deprecation choreography.** 3.9: staggered params ship on TripleDifference;
`time` -> `post` in 2x2x2 mode [M-031] (the `time` NAME persists as the
@@ -470,6 +500,18 @@ strongest norm (`did::aggte`, `etwfe::emfx`, Stata `estat aggregation`).
`fit(aggregate=)` is deprecated in 3.9 and removed in 4.0
([M-020]..[M-027]); `balance_e` moves to `aggregate()` alongside it.
+**(Amended 2026-08-08, Phase 3(b) ship - ONE documented exception.)** The DDD
+staggered mode is the single surface where fit-time `aggregate=`/`balance_e=`
+ship as CANONICAL rather than deprecated, because it has no post-fit successor
+to steer to: `StaggeredTripleDiffResults` carries no `AggregationMixin`, and
+`results_base.py`'s absent-surface hint for it literally reads "refit with
+aggregate='event_study'". The M-013 merge therefore carries both params onto
+the surviving `TripleDifference`, rowed as [M-140] [M-141] with
+`deprecated_in: null` (nothing to steer away from yet) and `removed_in: "4.0"`,
+which machine-checks that the 4.0 container unification ([M-014]) delivers the
+post-fit `aggregate()` in the same release the params die. The rule above is
+otherwise unchanged: no OTHER estimator may add a fit-time `aggregate=`.
+
**Vocabulary.** Closed set: `"simple"`, `"event_study"`, `"group"`,
`"calendar"`, plus per-estimator documented extras where the estimand demands
them (ContinuousDiD adds `"dose"` [M-025]; HAD's `"overall"` maps to
@@ -781,9 +823,9 @@ above; anything only one PR cares about stays in that PR's plan.**
|---|---|---|
| 1 (this PR) | - | Spec + matrix + enforcement test + support edits |
| 2: contract foundations | 3.9 | (a) results base + unified event-study representation [M-092] + to_dict completion + the Diagnostic marker base on the diagnostic result roster [M-091] (section 3.5); (b) `aggregate()` + fit(aggregate=) shims [M-020..M-027] [M-139] (M-020's shim already shipped; M-139 is the HAD workflow twin, a pre-cut amendment); (c) param renames [M-030..M-047] [M-084] [M-086..M-089] + their results-field mirrors [M-094] [M-095] (section 8 rule 9) + the public-function completeness sweep [M-097..M-113] (section 8 rule 10) + the dCDH results mirror [M-114] + the fourth `robust` site [M-115] + the 2(c)-ii missed-rename amendments [M-136..M-138] (LPDiD `level` value; the two post-dummy diagnostics params) + BaseEstimator mixin + ContinuousDiD covariates move; (d) alias introduction [M-062] (the Spillover introduction is cancelled [M-063]) + the alias-diet `__getattr__` warning shim [M-135] + wrapper deprecations [M-070..M-077] + the two inference-surface policies: `n_bootstrap` semantic unification [M-081] and the wild-cluster-bootstrap roster guard [M-096]; shipped insertions (all done): the aggregate contract [M-122], the ETWFE reference-period family [M-123] [M-124] [M-125], and the variance-consolidation program [M-126] [M-127] |
-| 3: merges | 3.9 | (a) TWFE event-study mode [M-010] + EventStudy warn [M-060] + the fit `time`->`post` rename [M-082] (gates: section 4.1's equivalence/divergence/pooled-parity test triple) (shipped: tests/test_v4_merge_mpd.py; consumer ports incl. HonestDiD/PreTrendsPower calendar routes); (b) TripleDifference facade [M-013] + the SDDD alias [M-064]; (c) CiC method= [M-015] |
+| 3: merges | 3.9 | (a) TWFE event-study mode [M-010] + EventStudy warn [M-060] + the fit `time`->`post` rename [M-082] (gates: section 4.1's equivalence/divergence/pooled-parity test triple) (shipped: tests/test_v4_merge_mpd.py; consumer ports incl. HonestDiD/PreTrendsPower calendar routes); (b) TripleDifference facade [M-013] + the SDDD alias [M-064] (shipped: tests/test_v4_merge_ddd.py; the engine relocation into the private shared mixin, the keyword-only staggered fit params, and the pscore_trim tightening [M-142]. The fit-time aggregate=/balance_e= carve-out rows are scheduled REMOVALS and so are cited in the phase-5 cell, not here - their only lifecycle version is removed_in 4.0); (c) CiC method= [M-015] |
| 4: release + soak | 3.9 cut | Migration guide written (skeleton: section 10); maintainer cuts 3.9; maint/3.8 rule active |
-| 5: enforcement | 4.0 | Removals [M-010..M-015, M-020..M-027, M-139, M-030, M-032..M-047 old names, M-060, M-061, M-064, M-070..M-077, M-084, M-086..M-089, M-001..M-003, M-117, M-118, M-119, M-120] + the alias diet [M-132]..[M-134] + the amendment's old names [M-094] [M-095] [M-097..M-115] [M-136..M-138] (incl. their consumer migrations and the `clean_control` serialized reporting key); M-031's old `time` name persists as the merged class's calendar column, so it is deliberately absent from the removal roster (its 4.0 enforcement is the M-085 behavior entry below); property window: [M-016] property-flips at 4.0 (removal at 5.0); storage flips [M-050..M-058]; default policies [M-004..M-006, M-128..M-131, M-080]; merged-class behavior enforcements [M-083] [M-085]; warning retirement [M-007]; fastpath go/no-go [M-008]; diagnostic-family docs/roster reorganization [M-090]; sentinel retirement [M-093]; docs/llms.txt/README refresh |
+| 5: enforcement | 4.0 | Removals [M-010..M-015, M-020..M-027, M-139, M-030, M-032..M-047 old names, M-060, M-061, M-064, M-070..M-077, M-084, M-086..M-089, M-001..M-003, M-117, M-118, M-119, M-120, M-140, M-141] + the alias diet [M-132]..[M-134] + the amendment's old names [M-094] [M-095] [M-097..M-115] [M-136..M-138] (incl. their consumer migrations and the `clean_control` serialized reporting key); M-031's old `time` name persists as the merged class's calendar column, so it is deliberately absent from the removal roster (its 4.0 enforcement is the M-085 behavior entry below); property window: [M-016] property-flips at 4.0 (removal at 5.0); storage flips [M-050..M-058]; default policies [M-004..M-006, M-128..M-131, M-080]; merged-class behavior enforcements [M-083] [M-085]; warning retirement [M-007]; fastpath go/no-go [M-008]; diagnostic-family docs/roster reorganization [M-090]; sentinel retirement [M-093]; docs/llms.txt/README refresh |
| 6: front door | 4.1 | `event_study(data, outcome, unit, time, first_treat, estimator=...)` comparison entry point over the staggered family (sketch only; specified in its own plan) |
Citation semantic for the table: a cell may cite a row whose current `phase`
diff --git a/pyproject.toml b/pyproject.toml
index 355dc174..852650e7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -123,14 +123,20 @@ markers = [
"slow: marks tests as slow (run `pytest -m 'not slow'` to exclude, or `pytest -m slow` to run only slow tests)",
"realdata: marks tests that validate against real survey datasets (API, NHANES, RECS)",
]
-# The ~140 MPD behavior-suite constructions stay noise-free through the 3.9
-# deprecation window (row M-010; the suites pin MPD behavior until its 4.0
-# removal - REMOVE this filter with the class). pytest.warns overrides it,
-# so the dedicated deprecation pins in tests/test_v4_merge_mpd.py still
-# fire; note a test-local `warnings.simplefilter(...)` RESETS the filter
+# The ~140 MPD and ~103 StaggeredTripleDifference behavior-suite constructions
+# stay noise-free through the 3.9 deprecation window (rows M-010 / M-013; the
+# suites pin each class's 3.x behavior until its 4.0 removal - REMOVE each
+# filter with its class). pytest.warns overrides them, so the dedicated
+# deprecation pins in tests/test_v4_merge_mpd.py and tests/test_v4_merge_ddd.py
+# still fire; note a test-local `warnings.simplefilter(...)` RESETS the filter
# list, so record-based tests must select warnings by message, never index.
+# The SDDD suites already do: several DO count and index (e.g.
+# test_staggered_triple_diff.py's rank-guard assertions), but every one of them
+# substring-filters the record list first, so the new FutureWarning never
+# enters a counted list.
filterwarnings = [
"ignore:MultiPeriodDiD is deprecated:FutureWarning",
+ "ignore:StaggeredTripleDifference is deprecated:FutureWarning",
]
[tool.black]
diff --git a/tests/_capture_v4_merge_ddd_oracles.py b/tests/_capture_v4_merge_ddd_oracles.py
new file mode 100644
index 00000000..6ee490a7
--- /dev/null
+++ b/tests/_capture_v4_merge_ddd_oracles.py
@@ -0,0 +1,244 @@
+"""Step 0: capture the pre-merge oracles for tests/test_v4_merge_ddd.py.
+
+MUST run on the UNMODIFIED tree, under DIFF_DIFF_BACKEND=python.
+
+Two oracles:
+ A) 2x2x2 engine (TripleDifference.fit today) - guards the rewritten fit
+ prologue/signature/dispatch.
+ B) staggered engine (StaggeredTripleDifference.fit today) - guards the
+ ~1700-line relocation into the private engine mixin. Nothing else in CI
+ pins this in absolute terms (the SDDD suites carry no committed numeric
+ pins; the R-golden lane skips when the gitignored CSVs are absent).
+
+The +-cluster axis of oracle A uses generate_ddd_panel_data, NOT
+generate_ddd_data: the cross-sectional generator emits `unit_id` incremented
+once per ROW, so every cluster would be a singleton (and there is no `unit`
+column at all). The panel generator has real repeated units - the same shape
+tests/test_prep.py:1438 already fits with cluster="unit".
+
+Oracle B covers every branch the relocation MOVED, not just the happy path:
+the DR base config and its bootstrap, the three nuisance models under
+covariates (_compute_pscore / _compute_or), the never-treated comparison fork
+(_is_never_treated) and the survey-pweight path. Re-running this file on a
+tree where the engine has already moved reproduces the same literals - that is
+the point of the gate, but it is NOT a substitute for capturing pre-move: once
+relocated there is no independent path left to disagree with.
+"""
+
+import json
+
+import numpy as np
+
+from diff_diff import StaggeredTripleDifference, TripleDifference
+from diff_diff.prep_dgp import (
+ generate_ddd_data,
+ generate_ddd_panel_data,
+ generate_staggered_ddd_data,
+)
+from diff_diff.survey import SurveyDesign
+
+CROSS_KW = dict(
+ n_per_cell=200,
+ treatment_effect=2.0,
+ group_effect=1.0,
+ partition_effect=0.5,
+ time_effect=0.7,
+ noise_sd=1.0,
+ add_covariates=True,
+ seed=42,
+)
+PANEL_KW = dict(n_units=80, n_periods=4, treatment_period=2, noise_sd=1.0, seed=42)
+STAG_KW = dict(n_units=96, n_periods=6, cohort_periods=[3, 4], seed=42)
+# add_covariates=True feeds x1/x2 into the OUTCOME, so the covariate lanes are a
+# different DGP draw than STAG_KW - they get their own fixture and their own
+# oracle keys rather than being compared against the plain-data literals.
+STAG_COV_KW = dict(STAG_KW, add_covariates=True)
+
+COVS = ["age", "education"]
+STAG_COVS = ["x1", "x2"]
+
+
+def _f(x):
+ """JSON-safe float (NaN -> None sentinel string handled by the caller)."""
+ if x is None:
+ return None
+ x = float(x)
+ return "NaN" if np.isnan(x) else x
+
+
+def quintet(r):
+ ci = getattr(r, "conf_int", None)
+ return {
+ "att": _f(r.att),
+ "se": _f(r.se),
+ "t_stat": _f(r.t_stat),
+ "p_value": _f(r.p_value),
+ "conf_int_lower": _f(ci[0]) if ci is not None else None,
+ "conf_int_upper": _f(ci[1]) if ci is not None else None,
+ }
+
+
+def capture_2x2x2():
+ out = {}
+ df = generate_ddd_data(**CROSS_KW)
+ for method in ("dr", "reg", "ipw"):
+ r = TripleDifference(estimation_method=method).fit(
+ df, outcome="outcome", group="group", partition="partition", post="time"
+ )
+ out[f"cross_{method}"] = dict(
+ quintet(r),
+ n_obs=int(r.n_obs),
+ n_treated_eligible=int(r.n_treated_eligible),
+ n_treated_ineligible=int(r.n_treated_ineligible),
+ n_control_eligible=int(r.n_control_eligible),
+ n_control_ineligible=int(r.n_control_ineligible),
+ vcov_type=r.vcov_type,
+ cluster_name=r.cluster_name,
+ n_clusters=r.n_clusters,
+ )
+ r = TripleDifference(estimation_method="dr").fit(
+ df,
+ outcome="outcome",
+ group="group",
+ partition="partition",
+ post="time",
+ covariates=COVS,
+ )
+ out["cross_dr_cov"] = dict(quintet(r), n_obs=int(r.n_obs))
+
+ # survey pweight lane (cross-sectional)
+ dfw = df.copy()
+ rng = np.random.default_rng(7)
+ dfw["w"] = rng.uniform(0.5, 2.0, size=len(dfw))
+ r = TripleDifference(estimation_method="reg").fit(
+ dfw,
+ outcome="outcome",
+ group="group",
+ partition="partition",
+ post="time",
+ survey_design=SurveyDesign(weights="w"),
+ )
+ out["cross_survey_pweight_reg"] = dict(quintet(r), n_obs=int(r.n_obs))
+
+ # cluster lane MUST use the panel generator (see module docstring)
+ pdf = generate_ddd_panel_data(**PANEL_KW)
+ r = TripleDifference(estimation_method="dr", cluster="unit").fit(
+ pdf, outcome="outcome", group="group", partition="partition", post="post"
+ )
+ out["panel_dr_cluster_unit"] = dict(
+ quintet(r),
+ n_obs=int(r.n_obs),
+ cluster_name=r.cluster_name,
+ n_clusters=int(r.n_clusters) if r.n_clusters is not None else None,
+ )
+ return out
+
+
+def _gt_table(r):
+ """Sorted (g,t) -> effect/se, key-stringified for JSON."""
+ return {
+ f"{g}|{t}": {"effect": _f(v["effect"]), "se": _f(v["se"])}
+ for (g, t), v in sorted(r.group_time_effects.items())
+ }
+
+
+def capture_staggered():
+ out = {}
+ df = generate_staggered_ddd_data(**STAG_KW)
+ fit_cols = dict(
+ outcome="outcome",
+ unit="unit",
+ time="period",
+ first_treat="first_treat",
+ eligibility="eligibility",
+ )
+ base = StaggeredTripleDifference(estimation_method="dr").fit(df, aggregate="all", **fit_cols)
+ out["stag_dr_all"] = {
+ "overall": quintet(base),
+ "overall_att_es": _f(base.overall_att_es),
+ "overall_se_es": _f(base.overall_se_es),
+ "n_obs": int(base.n_obs),
+ "n_treated_units": int(base.n_treated_units),
+ "n_never_enabled": int(base.n_never_enabled),
+ "group_time": _gt_table(base),
+ }
+ boot = StaggeredTripleDifference(
+ estimation_method="dr", n_bootstrap=49, seed=7, cband=True
+ ).fit(df, aggregate="all", **fit_cols)
+ out["stag_dr_all_boot49"] = {
+ "overall": quintet(boot),
+ "cband_crit_value": _f(boot.cband_crit_value),
+ "group_time": _gt_table(boot),
+ }
+
+ # --- Branch coverage for the RELOCATED nuisance/comparison/survey code ----
+ # The two lanes above ride DR-without-covariates only. The mixin also moved
+ # _compute_pscore, _compute_or and the never-treated comparison branch, and
+ # those are otherwise pinned only by comparing two callers of the SAME moved
+ # code - a shared transcription slip would keep that parity green. Each lane
+ # below is an ABSOLUTE pin on one of the moved branches.
+
+ # NOTE: dr/ipw/reg are numerically IDENTICAL on covariate-free data (the
+ # propensity score is constant and the outcome regression is a bare mean, so
+ # all three collapse to the same simple DiD - verified at capture time).
+ # Committing three identical literal blocks would look like nuisance-model
+ # coverage while providing none, so the no-covariate lanes are NOT captured
+ # per method; the test asserts that convergence live instead, and the
+ # discriminating pins are the *_cov lanes below.
+
+ # never-treated comparison branch (the _is_never_treated fork). The compact
+ # spelling is the DYING class's vocabulary; the merged surface passes
+ # "never_treated" and must land on the same numbers.
+ r = StaggeredTripleDifference(estimation_method="dr", control_group="nevertreated").fit(
+ df, aggregate="all", **fit_cols
+ )
+ out["stag_dr_nevertreated"] = {
+ "overall": quintet(r),
+ "overall_att_es": _f(r.overall_att_es),
+ "overall_se_es": _f(r.overall_se_es),
+ "group_time": _gt_table(r),
+ }
+
+ # covariate lanes - the DGP differs (x1/x2 enter the outcome), so these pin
+ # against their own fixture, not against stag_dr_all.
+ dfc = generate_staggered_ddd_data(**STAG_COV_KW)
+ for method in ("dr", "ipw", "reg"):
+ r = StaggeredTripleDifference(estimation_method=method).fit(
+ dfc, aggregate="all", covariates=STAG_COVS, **fit_cols
+ )
+ out[f"stag_{method}_cov"] = {
+ "overall": quintet(r),
+ "overall_att_es": _f(r.overall_att_es),
+ "overall_se_es": _f(r.overall_se_es),
+ "group_time": _gt_table(r),
+ }
+
+ # survey pweight lane through the staggered engine.
+ dfw = df.copy()
+ rng = np.random.default_rng(11)
+ per_unit = {
+ u: w
+ for u, w in zip(
+ sorted(dfw["unit"].unique()), rng.uniform(0.5, 2.0, size=dfw["unit"].nunique())
+ )
+ }
+ dfw["w"] = dfw["unit"].map(per_unit)
+ r = StaggeredTripleDifference(estimation_method="dr").fit(
+ dfw, aggregate="all", survey_design=SurveyDesign(weights="w"), **fit_cols
+ )
+ out["stag_dr_survey_pweight"] = {
+ "overall": quintet(r),
+ "overall_att_es": _f(r.overall_att_es),
+ "overall_se_es": _f(r.overall_se_es),
+ "group_time": _gt_table(r),
+ }
+ return out
+
+
+if __name__ == "__main__":
+ import warnings
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ payload = {"2x2x2": capture_2x2x2(), "staggered": capture_staggered()}
+ print(json.dumps(payload, indent=1, sort_keys=True))
diff --git a/tests/test_base_estimator.py b/tests/test_base_estimator.py
index 064e1087..4d04a2c2 100644
--- a/tests/test_base_estimator.py
+++ b/tests/test_base_estimator.py
@@ -176,9 +176,10 @@ def test_init_signature_matches_get_params(cls):
# (3.9 class merges, v4-design section 4.1). The round-trip test's
# warnings-as-errors filter exists to catch UNEXPECTED warnings from a
# re-init; these messages are expected by contract and are ignored inside
-# the error filter. Forward home for the phase-3 siblings (SDDD, QDiD).
+# the error filter. Forward home for the remaining phase-3 sibling (QDiD).
DEPRECATED_CLASS_WARNINGS = {
"MultiPeriodDiD": r"MultiPeriodDiD is deprecated",
+ "StaggeredTripleDifference": r"StaggeredTripleDifference is deprecated",
}
diff --git a/tests/test_naming_guard.py b/tests/test_naming_guard.py
index a3d21f7b..d5689597 100644
--- a/tests/test_naming_guard.py
+++ b/tests/test_naming_guard.py
@@ -784,7 +784,14 @@ def _live_family_code_refs(tok):
# 3.9 class merges, v4-design section 4.1). Their effective __init__ is no
# longer the base's function object, but calls under the shim's name still
# read every base constructor param - so they stay in the base's
-# init-sharing group. Forward home for the phase-3 siblings (SDDD, QDiD).
+# init-sharing group. Forward home for the remaining phase-3 sibling (QDiD).
+#
+# StaggeredTripleDifference was EVALUATED for this table in phase 3(b) and is
+# deliberately absent: it is not a forwarding shim. It keeps its own __init__
+# (R's compact control_group spellings, frozen until the 4.0 removal) and only
+# SHARES an engine mixin with TripleDifference, so identity grouping already
+# sees it correctly. Adding it would wrongly inject it into TripleDifference's
+# init-sharing group and widen that class's param-row consumer set.
_FORWARDING_INIT_SHIMS = {
"MultiPeriodDiD": "DifferenceInDifferences",
}
diff --git a/tests/test_v4_inference_policy.py b/tests/test_v4_inference_policy.py
index 16600289..e963d534 100644
--- a/tests/test_v4_inference_policy.py
+++ b/tests/test_v4_inference_policy.py
@@ -44,6 +44,7 @@
QDiD,
StaggeredTripleDifference,
SunAbraham,
+ TripleDifference,
TwoStageDiD,
TwoWayFixedEffects,
WooldridgeDiD,
@@ -96,6 +97,7 @@ def _floor_msg(n: int) -> str:
WooldridgeDiD,
ContinuousDiD,
StaggeredTripleDifference,
+ TripleDifference,
ChangesInChanges,
QDiD,
]
diff --git a/tests/test_v4_matrix.py b/tests/test_v4_matrix.py
index c8912a0f..d52bb691 100644
--- a/tests/test_v4_matrix.py
+++ b/tests/test_v4_matrix.py
@@ -126,11 +126,13 @@
# reserved pool) = 118, plus 2b PR-3b's Imputation/TwoStage balance_e rows
# (M-118, M-119, claimed from the reserved pool) = 120, plus 2b PR-4's
# HAD workflow-aggregate row (M-139, next free id - the reserved pool is
-# spent/earmarked) = 121.
+# spent/earmarked) = 121; + the phase-3(b) DDD merge rows (M-140/M-141 carry
+# fit-time aggregate=/balance_e= onto the surviving TripleDifference, M-142 the
+# pscore_trim tightening) = 124.
# Ids are never reused and terminal rows are never deleted, so the ledger
# only grows - raise the floor when rows are added; a lower parse count
# means scanner/format drift or an illegal row deletion.
-ROW_COUNT_FLOOR = 121
+ROW_COUNT_FLOOR = 124
# Committed snapshot of the shipped id set ("ids are never deleted or reused"
# contract - a delete-one-add-one edit keeps the count above the floor but trips
@@ -182,6 +184,7 @@
(136, 138),
(118, 119),
(139, 139),
+ (140, 142),
]
EXPECTED_INITIAL_IDS = frozenset(
f"M-{n:03d}" for lo, hi in _INITIAL_ID_RANGES for n in range(lo, hi + 1)
@@ -584,10 +587,10 @@ def test_initial_ids_never_deleted():
Phase 1 + diagnostic-family + M-092/M-093 + M-094..M-096 + the
M-097..M-115 public-function completeness sweep + M-117..M-120/M-122 +
M-123/M-124 + M-125 + M-126 + M-127..M-131 + M-132..M-135 +
- M-136..M-138 + M-139)."""
+ M-136..M-138 + M-139 + M-140..M-142)."""
missing = sorted(EXPECTED_INITIAL_IDS - set(_ROW_IDS))
assert not missing, f"ledger rows deleted (ids are permanent): {missing}"
- assert len(EXPECTED_INITIAL_IDS) == 121
+ assert len(EXPECTED_INITIAL_IDS) == 124
def test_version_tuple_pads_to_three_components():
diff --git a/tests/test_v4_merge_ddd.py b/tests/test_v4_merge_ddd.py
new file mode 100644
index 00000000..67e4dc64
--- /dev/null
+++ b/tests/test_v4_merge_ddd.py
@@ -0,0 +1,1746 @@
+"""Phase 3(b) merge gates: TripleDifference absorbs StaggeredTripleDifference.
+
+Ledger rows under test: M-013 (the class merge), M-064 (the SDDD alias warning
+riding the class), M-140/M-141 (fit-time aggregate=/balance_e= carried on the
+surviving class), M-142 (the pscore_trim validation tightening).
+
+TOLERANCE DOCTRINE
+------------------
+Every merge-PARITY gate is BIT-EXACT (``assert_array_equal``): both sides run
+the same engine in the same process, so a needed tolerance IS the finding.
+
+Two deliberate exceptions, both comparing against literals captured in a
+DIFFERENT process:
+
+* the committed ORACLES below, and
+* the R-parity lane (that suite's own tolerances).
+
+The oracles are NOT portable across backends - ``TripleDifference`` routes OLS
+through ``solve_ols`` (``diff_diff/triple_diff.py``), which dispatches to Rust
+or NumPy, and the repo's own equivalence claim is only ``decimal=8``
+(``tests/test_rust_backend.py``), while CI runs the suite under BOTH
+``DIFF_DIFF_BACKEND=rust`` and ``=python`` (``.github/workflows/rust-test.yml``).
+So oracle assertions use ``assert_allclose(rtol=1e-9, atol=1e-12)`` AND are
+guarded to the python backend.
+
+ORACLE PROVENANCE
+-----------------
+Captured on the UNMODIFIED pre-merge tree, before any source edit::
+
+ DIFF_DIFF_BACKEND=python python3 capture_oracles.py
+
+commit be694f493ac4e1640cdcccb5bdf7997a5b1f24dd.
+
+Why two oracles, not one. ORACLE_2X2X2 guards the rewritten fit prologue,
+signature and dispatch. ORACLE_STAGGERED guards the ~1700-line relocation of
+the staggered engine into ``diff_diff/_staggered_triple_diff_engine.py``:
+nothing else in CI pins that engine in ABSOLUTE terms, because the parity gates
+compare two callers of the same relocated core, the four SDDD suites carry no
+committed numeric pins, and the R-golden lane skips whenever the gitignored
+CSVs are absent (``.gitignore``). Without it a transcription slip in the move
+would be invisible.
+
+Two capture-config notes worth keeping:
+
+* The +-cluster oracle lane uses ``generate_ddd_panel_data``, NOT
+ ``generate_ddd_data``: the cross-sectional generator emits ``unit_id``
+ incremented once per ROW (and no ``unit`` column at all), so every cluster
+ would be a singleton. The panel generator has real repeated units - the shape
+ ``tests/test_prep.py`` already fits with ``cluster="unit"``.
+* ``cross_dr``/``cross_reg``/``cross_ipw`` agree to ~1e-13: on a saturated DDD
+ with no covariates all three estimators collapse to the same cell-mean
+ contrast. Those three lanes pin that the prologue still ROUTES correctly;
+ they are not independent evidence about method-specific math (the
+ ``_cov``/survey/cluster lanes carry that).
+* The same collapse applies to the STAGGERED engine, which is why there is no
+ per-method covariate-free staggered oracle: three identical literal blocks
+ would look like nuisance-model coverage while providing none. The convergence
+ is asserted live (``test_methods_converge_without_covariates``) and the
+ discriminating pins are the ``stag_{dr,ipw,reg}_cov`` lanes, whose mutual
+ distinctness is itself asserted so the parametrization cannot quietly stop
+ discriminating.
+
+ORACLE_STAGGERED covers, in absolute terms, every branch the relocation moved:
+the DR base config and its bootstrap, the three nuisance models under
+covariates (``_compute_pscore`` / ``_compute_or``), the never-treated
+comparison fork (``_is_never_treated``), and the survey-pweight path. The
+covariate lanes use ``STAG_COV_KW`` (``add_covariates=True``), a DIFFERENT draw
+than ``STAG_KW`` because x1/x2 enter the outcome - they pin against their own
+fixture, never against ``stag_dr_all``.
+"""
+
+import os
+import re
+import warnings
+
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose, assert_array_equal
+
+from diff_diff import SDDD, StaggeredTripleDifference, TripleDifference
+from diff_diff.prep_dgp import (
+ generate_ddd_data,
+ generate_ddd_panel_data,
+ generate_staggered_ddd_data,
+)
+from diff_diff.staggered_triple_diff import _SDDD_DEPRECATION_MSG
+from diff_diff.staggered_triple_diff_results import StaggeredTripleDiffResults
+from diff_diff.survey import SurveyDesign
+from diff_diff.triple_diff import TripleDifferenceResults
+
+# Oracle assertions compare literals captured in a DIFFERENT process, so they
+# are tolerance-based AND python-backend-guarded (see the tolerance doctrine).
+ORACLE_RTOL = 1e-9
+ORACLE_ATOL = 1e-12
+_BACKEND_IS_PYTHON = os.environ.get("DIFF_DIFF_BACKEND", "").lower() == "python"
+requires_python_backend = pytest.mark.skipif(
+ not _BACKEND_IS_PYTHON,
+ reason="committed oracle literals were captured under DIFF_DIFF_BACKEND=python; "
+ "solve_ols dispatches to Rust or NumPy and the repo's cross-backend claim is "
+ "only decimal=8, looser than this gate's tolerance",
+)
+
+# captured at be694f493ac4e1640cdcccb5bdf7997a5b1f24dd
+ORACLE_2X2X2 = {
+ "cross_dr": {
+ "att": 2.0192439829287876,
+ "cluster_name": None,
+ "conf_int_lower": 1.3227242700327153,
+ "conf_int_upper": 2.71576369582486,
+ "n_clusters": None,
+ "n_control_eligible": 400,
+ "n_control_ineligible": 400,
+ "n_obs": 1600,
+ "n_treated_eligible": 400,
+ "n_treated_ineligible": 400,
+ "p_value": 1.5412648419886862e-08,
+ "se": 0.35510355075673344,
+ "t_stat": 5.686352554418942,
+ "vcov_type": "hc1",
+ },
+ "cross_dr_cov": {
+ "att": 2.0051590625940046,
+ "conf_int_lower": 1.6213078544981028,
+ "conf_int_upper": 2.3890102706899063,
+ "n_obs": 1600,
+ "p_value": 6.681912366997255e-24,
+ "se": 0.19569696877741513,
+ "t_stat": 10.246244870939538,
+ },
+ "cross_ipw": {
+ "att": 2.0192439829287636,
+ "cluster_name": None,
+ "conf_int_lower": 1.3227242700326913,
+ "conf_int_upper": 2.715763695824836,
+ "n_clusters": None,
+ "n_control_eligible": 400,
+ "n_control_ineligible": 400,
+ "n_obs": 1600,
+ "n_treated_eligible": 400,
+ "n_treated_ineligible": 400,
+ "p_value": 1.5412648419892834e-08,
+ "se": 0.35510355075673344,
+ "t_stat": 5.686352554418875,
+ "vcov_type": "hc1",
+ },
+ "cross_reg": {
+ "att": 2.0192439829288134,
+ "cluster_name": None,
+ "conf_int_lower": 1.322724270032741,
+ "conf_int_upper": 2.7157636958248856,
+ "n_clusters": None,
+ "n_control_eligible": 400,
+ "n_control_ineligible": 400,
+ "n_obs": 1600,
+ "n_treated_eligible": 400,
+ "n_treated_ineligible": 400,
+ "p_value": 1.5412648419880546e-08,
+ "se": 0.35510355075673344,
+ "t_stat": 5.686352554419014,
+ "vcov_type": "hc1",
+ },
+ "cross_survey_pweight_reg": {
+ "att": 2.1916660209559,
+ "conf_int_lower": 1.324875700947568,
+ "conf_int_upper": 3.0584563409642316,
+ "n_obs": 1600,
+ "p_value": 7.820014665174102e-07,
+ "se": 0.4419133298063041,
+ "t_stat": 4.9594928985657285,
+ },
+ "panel_dr_cluster_unit": {
+ "att": 1.9176442049480977,
+ "cluster_name": "unit",
+ "conf_int_lower": 1.0160086827192953,
+ "conf_int_upper": 2.8192797271769,
+ "n_clusters": 80,
+ "n_obs": 320,
+ "p_value": 3.7155532583036706e-05,
+ "se": 0.4582420854203778,
+ "t_stat": 4.184784126034402,
+ },
+}
+
+ORACLE_STAGGERED = {
+ "stag_dr_all": {
+ "group_time": {
+ "3.0|2": {"effect": -0.027719840746928036, "se": 0.2811010312694456},
+ "3.0|3": {"effect": 2.762999666832223, "se": 0.2838203199927689},
+ "3.0|4": {"effect": 2.946059062614397, "se": 0.395161527831989},
+ "3.0|5": {"effect": 3.0227256449202673, "se": 0.3155992489776067},
+ "3.0|6": {"effect": 3.426518594109598, "se": 0.3635995672681256},
+ "4.0|2": {"effect": -0.0027171020705352477, "se": 0.32263986406339495},
+ "4.0|3": {"effect": -0.544147723760115, "se": 0.3356048332466622},
+ "4.0|4": {"effect": 3.5587114889457214, "se": 0.3634683141995581},
+ "4.0|5": {"effect": 3.597702813259514, "se": 0.34627545455921666},
+ "4.0|6": {"effect": 3.7459834980563302, "se": 0.39030264674197895},
+ },
+ "n_never_enabled": 24,
+ "n_obs": 576,
+ "n_treated_units": 36,
+ "overall": {
+ "att": 3.2943858241054356,
+ "conf_int_lower": 2.850906376698114,
+ "conf_int_upper": 3.737865271512757,
+ "p_value": 5.076360843907034e-48,
+ "se": 0.22626918193672466,
+ "t_stat": 14.559586930520208,
+ },
+ "overall_att_es": 3.3109024203559563,
+ "overall_se_es": 0.23391014483341283,
+ },
+ "stag_dr_all_boot49": {
+ "cband_crit_value": 2.7121511569356116,
+ "group_time": {
+ "3.0|2": {"effect": -0.027719840746928036, "se": 0.2722181669173392},
+ "3.0|3": {"effect": 2.762999666832223, "se": 0.28424292019694336},
+ "3.0|4": {"effect": 2.946059062614397, "se": 0.3731732550535591},
+ "3.0|5": {"effect": 3.0227256449202673, "se": 0.3110908492941959},
+ "3.0|6": {"effect": 3.426518594109598, "se": 0.33836191268742344},
+ "4.0|2": {"effect": -0.0027171020705352477, "se": 0.31692583310553385},
+ "4.0|3": {"effect": -0.544147723760115, "se": 0.3762241748319726},
+ "4.0|4": {"effect": 3.5587114889457214, "se": 0.3607135604501477},
+ "4.0|5": {"effect": 3.597702813259514, "se": 0.4049317945954622},
+ "4.0|6": {"effect": 3.7459834980563302, "se": 0.39946787358947394},
+ },
+ "overall": {
+ "att": 3.2943858241054356,
+ "conf_int_lower": 2.973695470025734,
+ "conf_int_upper": 3.82236144438262,
+ "p_value": 0.02,
+ "se": 0.23513162889644051,
+ "t_stat": 14.010815301910695,
+ },
+ },
+ "stag_dr_nevertreated": {
+ "group_time": {
+ "3.0|2": {"effect": -0.036712929377447884, "se": 0.3156231893405813},
+ "3.0|3": {"effect": 2.4505616566073596, "se": 0.3418289853541229},
+ "3.0|4": {"effect": 2.946059062614397, "se": 0.395161527831989},
+ "3.0|5": {"effect": 3.0227256449202673, "se": 0.3155992489776067},
+ "3.0|6": {"effect": 3.426518594109598, "se": 0.3635995672681256},
+ "4.0|2": {"effect": -0.02225820203848114, "se": 0.36226342694022373},
+ "4.0|3": {"effect": -0.544147723760115, "se": 0.3356048332466622},
+ "4.0|4": {"effect": 3.5587114889457214, "se": 0.3634683141995581},
+ "4.0|5": {"effect": 3.597702813259514, "se": 0.34627545455921666},
+ "4.0|6": {"effect": 3.7459834980563302, "se": 0.39030264674197895},
+ },
+ "overall": {
+ "att": 3.249751822644741,
+ "conf_int_lower": 2.8069280007214155,
+ "conf_int_upper": 3.6925756445680666,
+ "p_value": 6.560096242110517e-47,
+ "se": 0.22593467299208733,
+ "t_stat": 14.383590529101983,
+ },
+ "overall_att_es": 3.2718476690778484,
+ "overall_se_es": 0.23479955028064609,
+ },
+ "stag_dr_cov": {
+ "group_time": {
+ "3.0|2": {"effect": 0.30104528952593945, "se": 0.3288203170665867},
+ "3.0|3": {"effect": 3.0785075003936813, "se": 0.2880034245519164},
+ "3.0|4": {"effect": 2.7415492276631688, "se": 0.3403067398792659},
+ "3.0|5": {"effect": 2.701920247082487, "se": 0.32408769717572766},
+ "3.0|6": {"effect": 3.1354001973198953, "se": 0.3767773891775244},
+ "4.0|2": {"effect": 0.09285109959308976, "se": 0.3024820435924664},
+ "4.0|3": {"effect": 0.2869519248812908, "se": 0.323592049203466},
+ "4.0|4": {"effect": 2.7543305469828474, "se": 0.3709804006699708},
+ "4.0|5": {"effect": 2.7349820440287345, "se": 0.2930394677448108},
+ "4.0|6": {"effect": 2.6886356893894936, "se": 0.3633732146038028},
+ },
+ "overall": {
+ "att": 2.8336179218371864,
+ "conf_int_lower": 2.4550853638028727,
+ "conf_int_upper": 3.2121504798715,
+ "p_value": 9.758456975554044e-49,
+ "se": 0.19313240499321932,
+ "t_stat": 14.671892694220174,
+ },
+ "overall_att_es": 2.871340706272525,
+ "overall_se_es": 0.20149333462640176,
+ },
+ "stag_ipw_cov": {
+ "group_time": {
+ "3.0|2": {"effect": 0.30697667162647824, "se": 0.3247400010524773},
+ "3.0|3": {"effect": 3.1252746382126366, "se": 0.29781018754636845},
+ "3.0|4": {"effect": 2.730696345860042, "se": 0.3223912049564754},
+ "3.0|5": {"effect": 2.668735095865907, "se": 0.31800505390579015},
+ "3.0|6": {"effect": 3.0764584177360605, "se": 0.3616484472028155},
+ "4.0|2": {"effect": 0.10289546605785345, "se": 0.30357100920990315},
+ "4.0|3": {"effect": 0.29457053449196857, "se": 0.3264538981435691},
+ "4.0|4": {"effect": 2.725589827482862, "se": 0.4021132659327127},
+ "4.0|5": {"effect": 2.70592308497532, "se": 0.3077614474407017},
+ "4.0|6": {"effect": 2.6604495789311056, "se": 0.3767586528071246},
+ },
+ "overall": {
+ "att": 2.8133038555805623,
+ "conf_int_lower": 2.4218667691120874,
+ "conf_int_upper": 3.204740942049037,
+ "p_value": 4.59862408907509e-45,
+ "se": 0.19971646905559512,
+ "t_stat": 14.086489055629269,
+ },
+ "overall_att_es": 2.8461981758499997,
+ "overall_se_es": 0.20425902876397498,
+ },
+ "stag_reg_cov": {
+ "group_time": {
+ "3.0|2": {"effect": 0.2776783805103399, "se": 0.33365048698863964},
+ "3.0|3": {"effect": 3.076572394024985, "se": 0.28703632625261727},
+ "3.0|4": {"effect": 2.649951806165864, "se": 0.37479998504995804},
+ "3.0|5": {"effect": 2.7012789558536996, "se": 0.31653025014644537},
+ "3.0|6": {"effect": 3.0420008720472342, "se": 0.3881434808314953},
+ "4.0|2": {"effect": 0.08041142066229977, "se": 0.3049263704997503},
+ "4.0|3": {"effect": 0.2858927581065889, "se": 0.3264180220652882},
+ "4.0|4": {"effect": 2.769419984274427, "se": 0.36003067748880446},
+ "4.0|5": {"effect": 2.733860817994398, "se": 0.2947015467495497},
+ "4.0|6": {"effect": 2.691465318642148, "se": 0.36282711170283977},
+ },
+ "overall": {
+ "att": 2.8092214498575365,
+ "conf_int_lower": 2.4288798195291026,
+ "conf_int_upper": 3.1895630801859705,
+ "p_value": 1.7084282921997317e-47,
+ "se": 0.1940554180222292,
+ "t_stat": 14.476387613850276,
+ },
+ "overall_att_es": 2.8383188776312487,
+ "overall_se_es": 0.20378953973090694,
+ },
+ "stag_dr_survey_pweight": {
+ "group_time": {
+ "3.0|2": {"effect": -0.15168667565630167, "se": 0.2957758221514469},
+ "3.0|3": {"effect": 2.6028992505775195, "se": 0.32935262794723136},
+ "3.0|4": {"effect": 2.6888855004876224, "se": 0.3918630493389082},
+ "3.0|5": {"effect": 2.9347276471371493, "se": 0.3423394186475062},
+ "3.0|6": {"effect": 3.400335285909308, "se": 0.4190520610387603},
+ "4.0|2": {"effect": 0.15529987882584567, "se": 0.30936054767650695},
+ "4.0|3": {"effect": -0.7550528967872525, "se": 0.32574440796417875},
+ "4.0|4": {"effect": 3.494966537949268, "se": 0.3614205225792733},
+ "4.0|5": {"effect": 3.717983004091062, "se": 0.3845593915382934},
+ "4.0|6": {"effect": 3.8371645373489773, "se": 0.39151964978005016},
+ },
+ "overall": {
+ "att": 3.2649376305989826,
+ "conf_int_lower": 2.789054729251909,
+ "conf_int_upper": 3.7408205319460563,
+ "p_value": 4.606519827221944e-24,
+ "se": 0.23970918564291063,
+ "t_stat": 13.620411006955264,
+ },
+ "overall_att_es": 3.28298050259862,
+ "overall_se_es": 0.25007641854869433,
+ },
+}
+
+
+# ---------------------------------------------------------------------------
+# Fixtures and helpers
+# ---------------------------------------------------------------------------
+
+CROSS_KW = dict(
+ n_per_cell=200,
+ treatment_effect=2.0,
+ group_effect=1.0,
+ partition_effect=0.5,
+ time_effect=0.7,
+ noise_sd=1.0,
+ add_covariates=True,
+ seed=42,
+)
+PANEL_KW = dict(n_units=80, n_periods=4, treatment_period=2, noise_sd=1.0, seed=42)
+STAG_KW = dict(n_units=96, n_periods=6, cohort_periods=[3, 4], seed=42)
+# add_covariates=True feeds x1/x2 into the OUTCOME, so this is a different draw
+# than STAG_KW - the covariate oracles pin against it, not against stag_dr_all.
+STAG_COV_KW = dict(STAG_KW, add_covariates=True)
+COVS = ["age", "education"]
+STAG_COVS = ["x1", "x2"]
+
+
+# The dying class emits a FutureWarning per construction; the pyproject filter
+# keeps the legacy suites quiet, but a test-local simplefilter RESETS the filter
+# list, so build it through this helper inside any recording block.
+def _sddd(**kwargs):
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore", FutureWarning)
+ return StaggeredTripleDifference(**kwargs)
+
+
+@pytest.fixture(scope="module")
+def cross():
+ return generate_ddd_data(**CROSS_KW)
+
+
+@pytest.fixture(scope="module")
+def panel():
+ return generate_ddd_panel_data(**PANEL_KW)
+
+
+@pytest.fixture(scope="module")
+def stag():
+ return generate_staggered_ddd_data(**STAG_KW)
+
+
+@pytest.fixture(scope="module")
+def stag_cov():
+ return generate_staggered_ddd_data(**STAG_COV_KW)
+
+
+@pytest.fixture(scope="module")
+def stag_survey(stag):
+ """Staggered panel with a per-UNIT pweight (weights must not vary within a
+ unit - a within-unit-varying weight is a different design entirely)."""
+ df = stag.copy()
+ rng = np.random.default_rng(11)
+ units = sorted(df["unit"].unique())
+ per_unit = dict(zip(units, rng.uniform(0.5, 2.0, size=len(units))))
+ df["w"] = df["unit"].map(per_unit)
+ return df
+
+
+C_COLS = dict(outcome="outcome", group="group", partition="partition", post="time")
+S_NEW = dict(outcome="outcome", unit="unit", time="period", first_treat="first_treat")
+S_OLD = dict(outcome="outcome", unit="unit", time="period", first_treat="first_treat")
+
+
+def _fit_new(df, **kwargs):
+ """Staggered fit through the MERGED surface."""
+ return TripleDifference(**kwargs.pop("ctor", {})).fit(
+ df, partition="eligibility", **S_NEW, **kwargs
+ )
+
+
+def _fit_old(df, **kwargs):
+ """Staggered fit through the DEPRECATED surface."""
+ return _sddd(**kwargs.pop("ctor", {})).fit(df, eligibility="eligibility", **S_OLD, **kwargs)
+
+
+def _eq_with_nans(a, b):
+ """NaN-mask equality first, then the finite subset - so a NaN-vs-0.0
+ regression cannot be masked by nan_to_num-style comparison."""
+ a, b = np.asarray(a, dtype=float), np.asarray(b, dtype=float)
+ assert_array_equal(np.isnan(a), np.isnan(b))
+ m = ~np.isnan(a)
+ assert_array_equal(a[m], b[m])
+
+
+def _eq_dicts_with_nans(a, b, keys=("effect", "se")):
+ """Identical key sets in identical ITERATION order, then values.
+
+ Order matters: the bootstrap draw order depends on cell order, so a
+ reordering would silently change seeded results.
+ """
+ assert list(a.keys()) == list(b.keys())
+ for k in a:
+ for f in keys:
+ _eq_with_nans(a[k][f], b[k][f])
+
+
+def _quintet(r):
+ ci = r.conf_int
+ return np.array([r.att, r.se, r.t_stat, r.p_value, ci[0], ci[1]], dtype=float)
+
+
+# ---------------------------------------------------------------------------
+# Gate A - the 2x2x2 engine is unchanged (committed pre-merge oracle)
+# ---------------------------------------------------------------------------
+
+
+@requires_python_backend
+class TestGateA2x2x2Unchanged:
+ """The rewritten fit prologue/signature/dispatch moved no 2x2x2 numbers.
+
+ There is no surviving in-process pre-merge path to compare against: direct
+ construction and the triple_difference() wrapper both route through the
+ modified class, so a shared regression would be invisible to a
+ self-comparison. Hence the committed literals.
+ """
+
+ @pytest.mark.parametrize("method", ["dr", "reg", "ipw"])
+ def test_cross_sectional_lanes_match_oracle(self, cross, method):
+ exp = ORACLE_2X2X2[f"cross_{method}"]
+ r = TripleDifference(estimation_method=method).fit(cross, **C_COLS)
+ assert_allclose(
+ _quintet(r),
+ [
+ exp["att"],
+ exp["se"],
+ exp["t_stat"],
+ exp["p_value"],
+ exp["conf_int_lower"],
+ exp["conf_int_upper"],
+ ],
+ rtol=ORACLE_RTOL,
+ atol=ORACLE_ATOL,
+ )
+ assert r.n_obs == exp["n_obs"]
+ assert r.n_treated_eligible == exp["n_treated_eligible"]
+ assert r.n_control_ineligible == exp["n_control_ineligible"]
+ assert r.vcov_type == exp["vcov_type"]
+ assert r.cluster_name == exp["cluster_name"]
+ assert r.n_clusters == exp["n_clusters"]
+
+ def test_covariate_lane_matches_oracle(self, cross):
+ exp = ORACLE_2X2X2["cross_dr_cov"]
+ r = TripleDifference(estimation_method="dr").fit(cross, covariates=COVS, **C_COLS)
+ assert_allclose(r.att, exp["att"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(r.se, exp["se"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+
+ def test_survey_pweight_lane_matches_oracle(self, cross):
+ exp = ORACLE_2X2X2["cross_survey_pweight_reg"]
+ dfw = cross.copy()
+ dfw["w"] = np.random.default_rng(7).uniform(0.5, 2.0, size=len(dfw))
+ r = TripleDifference(estimation_method="reg").fit(
+ dfw, survey_design=SurveyDesign(weights="w"), **C_COLS
+ )
+ assert_allclose(r.att, exp["att"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(r.se, exp["se"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+
+ def test_cluster_lane_matches_oracle(self, panel):
+ """The cluster axis runs on the PANEL generator deliberately: the
+ cross-sectional one emits unit_id incremented once per ROW (and no
+ `unit` column), so every cluster would be a singleton."""
+ exp = ORACLE_2X2X2["panel_dr_cluster_unit"]
+ r = TripleDifference(estimation_method="dr", cluster="unit").fit(
+ panel, outcome="outcome", group="group", partition="partition", post="post"
+ )
+ assert_allclose(r.att, exp["att"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(r.se, exp["se"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert r.cluster_name == "unit"
+ assert r.n_clusters == exp["n_clusters"] > 1
+
+ def test_post_still_lands_in_slot_5(self, cross):
+ """Positional slot 5 is `post`; the M-031 pin passes it positionally."""
+ with warnings.catch_warnings():
+ warnings.simplefilter("error") # positional canonical form must not warn
+ positional = TripleDifference().fit(cross, "outcome", "group", "partition", "time")
+ keyword = TripleDifference().fit(cross, **C_COLS)
+ _eq_with_nans(_quintet(positional), _quintet(keyword))
+
+ def test_time_alias_still_warns_and_routes_identically(self, cross):
+ with pytest.warns(FutureWarning, match="calendar column only"):
+ aliased = TripleDifference().fit(
+ cross, outcome="outcome", group="group", partition="partition", time="time"
+ )
+ _eq_with_nans(_quintet(aliased), _quintet(TripleDifference().fit(cross, **C_COLS)))
+
+
+# ---------------------------------------------------------------------------
+# Gate B' - the RELOCATED staggered engine against its own committed oracle
+# ---------------------------------------------------------------------------
+
+
+@requires_python_backend
+class TestGateBPrimeStaggeredOracle:
+ """The engine that actually moved, pinned in ABSOLUTE terms.
+
+ Gate B below compares two callers of the same relocated core, so it cannot
+ see a transcription slip in the move itself; the SDDD suites carry no
+ committed numeric pins; and the R-golden lane skips whenever the gitignored
+ CSVs are absent. This gate is the one that would catch it.
+ """
+
+ def test_base_config_matches_oracle(self, stag):
+ exp = ORACLE_STAGGERED["stag_dr_all"]
+ r = _fit_new(stag, aggregate="all")
+ assert_allclose(r.overall_att, exp["overall"]["att"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(r.overall_se, exp["overall"]["se"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(r.overall_att_es, exp["overall_att_es"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert r.n_obs == exp["n_obs"]
+ assert r.n_never_enabled == exp["n_never_enabled"]
+ got = {f"{g}|{t}": v for (g, t), v in sorted(r.group_time_effects.items())}
+ assert set(got) == set(exp["group_time"])
+ for k, v in exp["group_time"].items():
+ assert_allclose(got[k]["effect"], v["effect"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(got[k]["se"], v["se"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+
+ def test_bootstrap_config_matches_oracle(self, stag):
+ exp = ORACLE_STAGGERED["stag_dr_all_boot49"]
+ r = _fit_new(stag, aggregate="all", ctor=dict(n_bootstrap=49, seed=7, cband=True))
+ assert_allclose(r.overall_att, exp["overall"]["att"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(r.overall_se, exp["overall"]["se"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(
+ r.cband_crit_value, exp["cband_crit_value"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL
+ )
+
+ # -- branch coverage for the rest of the relocated engine -----------------
+ # The two lanes above ride DR-without-covariates. _compute_pscore,
+ # _compute_or, the never-treated comparison fork and the survey path all
+ # moved too, and Gate B cannot see a slip in any of them (both of its sides
+ # run the moved code). These are absolute pins on each.
+
+ def test_never_treated_branch_matches_oracle(self, stag):
+ """The _is_never_treated fork. The merged surface passes the UNDERSCORED
+ spelling; the oracle was captured through the dying class's compact one,
+ so this also pins the vocabulary bridge onto real numbers."""
+ exp = ORACLE_STAGGERED["stag_dr_nevertreated"]
+ r = _fit_new(stag, aggregate="all", ctor=dict(control_group="never_treated"))
+ self._assert_against(r, exp)
+ # ...and it must actually be a DIFFERENT comparison group, else a bridge
+ # that collapsed every spelling to one branch would still pass above.
+ base = ORACLE_STAGGERED["stag_dr_all"]["overall"]["att"]
+ assert exp["overall"]["att"] != base
+
+ @pytest.mark.parametrize("method", ["dr", "ipw", "reg"])
+ def test_covariate_lanes_match_oracle(self, stag_cov, method):
+ """Covariates are what separate the three nuisance models, so this is
+ the lane that would actually catch a _compute_pscore/_compute_or slip."""
+ exp = ORACLE_STAGGERED[f"stag_{method}_cov"]
+ r = _fit_new(
+ stag_cov,
+ aggregate="all",
+ covariates=STAG_COVS,
+ ctor=dict(estimation_method=method),
+ )
+ self._assert_against(r, exp)
+
+ def test_covariate_lanes_are_mutually_distinct(self):
+ """Guard on the guard: if the three covariate oracles ever coincide, the
+ parametrized test above silently stops discriminating between the
+ nuisance models even while passing."""
+ atts = {ORACLE_STAGGERED[f"stag_{m}_cov"]["overall"]["att"] for m in ("dr", "ipw", "reg")}
+ assert len(atts) == 3
+
+ def test_methods_converge_without_covariates(self, stag):
+ """Why there is no per-method no-covariate oracle: with no covariates the
+ propensity score is constant and the outcome regression is a bare mean,
+ so dr/ipw/reg collapse to the same estimator. Asserted live rather than
+ committed as three identical literal blocks that would LOOK like
+ nuisance-model coverage while providing none."""
+ fits = [
+ _fit_new(stag, aggregate="all", ctor=dict(estimation_method=m))
+ for m in ("dr", "ipw", "reg")
+ ]
+ for r in fits[1:]:
+ assert_allclose(r.overall_att, fits[0].overall_att, rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(r.overall_se, fits[0].overall_se, rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+
+ def test_survey_pweight_matches_oracle(self, stag_survey):
+ exp = ORACLE_STAGGERED["stag_dr_survey_pweight"]
+ r = _fit_new(stag_survey, aggregate="all", survey_design=SurveyDesign(weights="w"))
+ self._assert_against(r, exp)
+ # the weights must actually bite
+ assert exp["overall"]["att"] != ORACLE_STAGGERED["stag_dr_all"]["overall"]["att"]
+
+ @staticmethod
+ def _assert_against(r, exp):
+ assert_allclose(r.overall_att, exp["overall"]["att"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(r.overall_se, exp["overall"]["se"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(r.overall_att_es, exp["overall_att_es"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ got = {f"{g}|{t}": v for (g, t), v in sorted(r.group_time_effects.items())}
+ assert set(got) == set(exp["group_time"])
+ for k, v in exp["group_time"].items():
+ assert_allclose(got[k]["effect"], v["effect"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+ assert_allclose(got[k]["se"], v["se"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+
+
+# ---------------------------------------------------------------------------
+# Gate B - merged staggered mode == deprecated class, bit-exact
+# ---------------------------------------------------------------------------
+
+# ONE AXIS AT A TIME against a fixed base, not a Cartesian product: the full
+# cross-product is ~5.8k configs x two fits x bootstrap draws, which is not
+# runnable. Named interaction cells cover the couplings that matter.
+_BASE = dict(estimation_method="dr")
+_AXES = [
+ ("method_ipw", dict(estimation_method="ipw"), {}),
+ ("method_reg", dict(estimation_method="reg"), {}),
+ ("base_period_universal", dict(base_period="universal"), {}),
+ ("anticipation_1", dict(anticipation=1), {}),
+ ("cband_off", dict(n_bootstrap=49, seed=7, cband=False), {}),
+ ("weights_mammen", dict(n_bootstrap=49, seed=7, bootstrap_weights="mammen"), {}),
+ ("weights_webb", dict(n_bootstrap=49, seed=7, bootstrap_weights="webb"), {}),
+ ("aggregate_none", {}, dict(aggregate=None)),
+ ("aggregate_event_study", {}, dict(aggregate="event_study")),
+ ("aggregate_group", {}, dict(aggregate="group")),
+ ("aggregate_simple", {}, dict(aggregate="simple")),
+ ("balance_e", {}, dict(aggregate="event_study", balance_e=1)),
+ ("covariates", {}, dict(covariates=["x1"])),
+ # named interaction cells
+ ("boot_x_cband", dict(n_bootstrap=49, seed=7, cband=True), dict(aggregate="all")),
+ ("universal_x_anticipation", dict(base_period="universal", anticipation=1), {}),
+]
+
+_STAG_FIELDS = [
+ "overall_att",
+ "overall_se",
+ "overall_t_stat",
+ "overall_p_value",
+ "overall_att_es",
+ "overall_se_es",
+ "overall_t_stat_es",
+ "overall_p_value_es",
+ "n_obs",
+ "n_treated_units",
+ "n_control_units",
+ "n_never_enabled",
+ "n_eligible",
+ "n_ineligible",
+ "cband_crit_value",
+]
+
+
+def _assert_staggered_parity(old, new):
+ assert type(new) is type(old) is StaggeredTripleDiffResults
+ for f in _STAG_FIELDS:
+ a, b = getattr(old, f), getattr(new, f)
+ if a is None or b is None:
+ assert a is b, f"{f}: {a!r} vs {b!r}"
+ else:
+ _eq_with_nans(a, b)
+ _eq_with_nans(
+ np.asarray(old.overall_conf_int, dtype=float), np.asarray(new.overall_conf_int, dtype=float)
+ )
+ _eq_dicts_with_nans(old.group_time_effects, new.group_time_effects)
+ for attr in ("event_study_effects", "group_effects"):
+ a, b = getattr(old, attr), getattr(new, attr)
+ assert (a is None) == (b is None)
+ if a is not None:
+ _eq_dicts_with_nans(a, b, keys=("effect", "se"))
+
+
+class TestGateBStaggeredParity:
+ """Every lane: merged staggered mode is BIT-EXACT vs the dying class."""
+
+ @pytest.mark.parametrize("label,ctor,fit_kw", _AXES, ids=[a[0] for a in _AXES])
+ def test_axis_parity(self, stag, label, ctor, fit_kw):
+ full_ctor = {**_BASE, **ctor}
+ base_fit = dict(aggregate="all")
+ base_fit.update(fit_kw)
+ if "covariates" in fit_kw:
+ df = generate_staggered_ddd_data(**STAG_KW, add_covariates=True)
+ base_fit["covariates"] = ["x1"]
+ else:
+ df = stag
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ old = _fit_old(df, ctor=dict(full_ctor), **base_fit)
+ new = _fit_new(df, ctor=dict(full_ctor), **base_fit)
+ _assert_staggered_parity(old, new)
+
+ @pytest.mark.parametrize(
+ "new_value,old_value",
+ [("not_yet_treated", "notyettreated"), ("never_treated", "nevertreated")],
+ )
+ def test_control_group_vocabulary_parity(self, stag, new_value, old_value):
+ """Each class takes only ITS OWN spelling, so a parity pair is the
+ underscored value on the new surface vs the compact one on the old."""
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ old = _fit_old(stag, ctor=dict(control_group=old_value), aggregate="all")
+ new = _fit_new(stag, ctor=dict(control_group=new_value), aggregate="all")
+ _assert_staggered_parity(old, new)
+
+ def test_control_group_values_are_SEMANTICALLY_different(self, stag):
+ """Guards a _is_never_treated that collapses both spellings.
+
+ Parity alone cannot catch it: both surfaces route through the same
+ helper, so a collapsing helper keeps every parity gate green while
+ silently switching the comparison-group semantics. Nothing else in the
+ repo differentiates the two values (the R-golden lane skips in CI).
+ """
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ nyt = _fit_new(stag, ctor=dict(control_group="not_yet_treated"), aggregate="all")
+ nev = _fit_new(stag, ctor=dict(control_group="never_treated"), aggregate="all")
+ assert nyt.overall_att != nev.overall_att, (
+ "not_yet_treated and never_treated produced identical ATTs - the "
+ "comparison-group branch is not reading control_group"
+ )
+
+ def test_first_treat_inf_recode_warning_parity(self, stag):
+ df = stag.copy()
+ df["first_treat"] = df["first_treat"].astype(float)
+ df.loc[df["first_treat"] == 0, "first_treat"] = np.inf
+ msgs = []
+ for fit in (_fit_old, _fit_new):
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ fit(df, aggregate="simple")
+ msgs.append([str(x.message) for x in w if "first_treat=inf" in str(x.message)])
+ assert msgs[0] == msgs[1] and msgs[0], msgs
+
+ def test_bootstrap_x_replicate_survey_raises_on_BOTH_surfaces(self, stag):
+ """Excluded from the parity matrix by construction: this combination
+ raises before any parity field exists, so it is exception-parity."""
+ df = stag.copy()
+ rng = np.random.default_rng(3)
+ df["w"] = 1.0
+ # Replicate weights must be UNIT-constant for a panel estimator, so
+ # draw per unit and map onto rows.
+ units = df["unit"].unique()
+ for i in range(1, 4):
+ per_unit = dict(zip(units, rng.uniform(0.5, 1.5, size=len(units))))
+ df[f"rw{i}"] = df["unit"].map(per_unit)
+ design = SurveyDesign(
+ weights="w",
+ replicate_weights=[f"rw{i}" for i in range(1, 4)],
+ replicate_method="JK1",
+ )
+ for fit, name in ((_fit_old, "StaggeredTripleDifference"), (_fit_new, "TripleDifference")):
+ with pytest.raises(NotImplementedError) as exc:
+ fit(df, ctor=dict(n_bootstrap=49, seed=7), survey_design=design)
+ assert name in str(exc.value), f"{name} not named: {exc.value}"
+
+
+# ---------------------------------------------------------------------------
+# Gate C - the rejection matrix, both directions
+# ---------------------------------------------------------------------------
+
+
+class TestGateCRejectionMatrix:
+ def test_staggered_rejects_2x2x2_fit_params(self, stag):
+ for kw in ({"post": "period"}, {"group": "eligibility"}):
+ with pytest.raises(ValueError, match=r"belong\(s\) to the 2x2x2 mode"):
+ _fit_new(stag, **kw)
+
+ @pytest.mark.parametrize("kw", [{"unit": "unit"}, {"aggregate": "all"}, {"balance_e": 1}])
+ def test_2x2x2_rejects_staggered_fit_params(self, cross, kw):
+ with pytest.raises(ValueError, match=r"require\(s\) first_treat="):
+ TripleDifference().fit(cross, **C_COLS, **kw)
+
+ @pytest.mark.parametrize(
+ "param,value",
+ [
+ ("control_group", "never_treated"),
+ ("anticipation", 1),
+ ("base_period", "universal"),
+ ("n_bootstrap", 9),
+ ],
+ )
+ def test_2x2x2_rejects_staggered_ctor_params(self, cross, param, value):
+ with pytest.raises(ValueError, match="applies to the staggered DDD mode only"):
+ TripleDifference(**{param: value}).fit(cross, **C_COLS)
+
+ def test_staggered_rejects_robust(self, stag):
+ with pytest.warns(FutureWarning, match="robust"):
+ est = TripleDifference(robust=True)
+ with pytest.raises(ValueError, match="applies to the 2x2x2 mode only"):
+ est.fit(stag, partition="eligibility", **S_NEW)
+
+ def test_staggered_rejects_mutated_vcov_type(self, stag):
+ """vcov_type is validated eagerly in __init__ and by set_params' probe
+ re-init, so its ONLY live route is direct attribute mutation - the
+ bypass __init__'s comment documents. This arm is that guard."""
+ est = TripleDifference()
+ est.vcov_type = "classical"
+ with pytest.raises(ValueError, match="applies to the 2x2x2 mode only"):
+ est.fit(stag, partition="eligibility", **S_NEW)
+
+ def test_staggered_rejects_cluster(self, stag):
+ with pytest.raises(ValueError, match="not supported in staggered DDD mode"):
+ TripleDifference(cluster="unit").fit(stag, partition="eligibility", **S_NEW)
+
+ @pytest.mark.parametrize(
+ "param,value",
+ [
+ ("bootstrap_weights", "webb"),
+ ("seed", 7),
+ ("cband", False),
+ ("alpha", 0.10),
+ ("pscore_trim", 0.02),
+ ("rank_deficient_action", "silent"),
+ ("epv_threshold", 5),
+ ("pscore_fallback", "unconditional"),
+ ],
+ )
+ def test_shared_and_unreachable_params_pass_silently_in_2x2x2(self, cross, param, value):
+ """The six shared params must stay silent, and the three
+ bootstrap satellites are unreachable-by-construction (n_bootstrap > 0
+ is itself rejected) rather than silently ignored."""
+ r = TripleDifference(**{param: value}).fit(cross, **C_COLS)
+ assert np.isfinite(r.att)
+
+ def test_bad_aggregate_VALUE_precedes_the_mode_error(self, cross):
+ """Deliberate ordering (3(a)'s `spec` precedent): a bad value is
+ reported as a bad value even in the wrong mode."""
+ with pytest.raises(ValueError, match="aggregate must be"):
+ TripleDifference().fit(cross, aggregate="bogus", **C_COLS)
+
+ @pytest.mark.parametrize("missing", ["group", "partition", "post"])
+ def test_2x2x2_required_args_named_individually(self, cross, missing):
+ kw = dict(C_COLS)
+ kw.pop(missing)
+ with pytest.raises(TypeError, match=f"missing required argument: '{missing}'"):
+ TripleDifference().fit(cross, **kw)
+
+ def test_missing_group_raises_without_the_rename_warning(self, cross):
+ """require_arg for group/partition runs BEFORE the M-031 shim: after
+ it, this call would newly emit the rename FutureWarning before raising,
+ and under -W error the surfaced exception type would flip."""
+ with warnings.catch_warnings():
+ warnings.simplefilter("error")
+ with pytest.raises(TypeError, match="missing required argument: 'group'"):
+ TripleDifference().fit(cross, outcome="outcome", time="time")
+
+ @pytest.mark.parametrize("missing", ["unit", "time", "partition"])
+ def test_staggered_required_args_named_individually(self, stag, missing):
+ kw = dict(partition="eligibility", **S_NEW)
+ kw.pop(missing)
+ with pytest.raises(TypeError, match=f"missing required argument: '{missing}'"):
+ TripleDifference().fit(stag, **kw)
+
+ def test_staggered_params_are_keyword_only(self, stag):
+ """A positionally-written staggered call must be a clean TypeError, not
+ a silent bind into the 2x2x2 slots."""
+ with pytest.raises(TypeError):
+ TripleDifference().fit(
+ stag, "outcome", "unit", "period", "first_treat", "eligibility", "extra", "x", "y"
+ )
+
+
+# ---------------------------------------------------------------------------
+# Gate D - time= semantics across the two modes
+# ---------------------------------------------------------------------------
+
+
+class TestGateDTimeSemantics:
+ def test_staggered_time_emits_no_rename_warning(self, stag):
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ _fit_new(stag, aggregate="simple")
+ assert not [x for x in w if "calendar column only" in str(x.message)]
+
+ def test_staggered_time_is_bit_exact_vs_the_dying_surface(self, stag):
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ _assert_staggered_parity(
+ _fit_old(stag, aggregate="all"), _fit_new(stag, aggregate="all")
+ )
+
+ def test_2x2x2_both_post_and_time_raises(self, cross):
+ with pytest.raises(ValueError, match="pass only post="):
+ TripleDifference().fit(
+ cross,
+ outcome="outcome",
+ group="group",
+ partition="partition",
+ post="time",
+ time="time",
+ )
+
+ def test_m085_is_not_live_in_3_9(self, cross):
+ """2x2x2 time= WARNS; it does not raise. The 4.0 flip is M-085's job."""
+ with pytest.warns(FutureWarning):
+ r = TripleDifference().fit(
+ cross, outcome="outcome", group="group", partition="partition", time="time"
+ )
+ assert np.isfinite(r.att)
+
+
+# ---------------------------------------------------------------------------
+# Gate E - deprecation choreography (M-013 / M-064)
+# ---------------------------------------------------------------------------
+
+
+class TestGateEDeprecation:
+ def test_construction_warns_with_the_pinned_message(self):
+ with pytest.warns(FutureWarning, match=re.escape(_SDDD_DEPRECATION_MSG)):
+ StaggeredTripleDifference()
+
+ def test_sddd_alias_is_the_same_class_and_warns_once(self):
+ assert SDDD is StaggeredTripleDifference
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ SDDD()
+ fired = [
+ x
+ for x in w
+ if x.category is FutureWarning
+ and "StaggeredTripleDifference is deprecated" in str(x.message)
+ ]
+ assert len(fired) == 1, [str(x.message) for x in w]
+ assert "SDDD alias is deprecated with it" in str(fired[0].message)
+
+ def test_warns_and_still_works(self, stag):
+ with pytest.warns(FutureWarning):
+ est = StaggeredTripleDifference()
+ r = est.fit(stag, eligibility="eligibility", **S_OLD, aggregate="simple")
+ assert np.isfinite(r.overall_att)
+
+ def test_successor_does_not_warn(self, stag):
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ TripleDifference()
+ assert not [x for x in w if x.category is FutureWarning]
+
+ def test_deprecation_attributes_to_the_caller(self):
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ StaggeredTripleDifference()
+ assert w[0].filename.endswith("test_v4_merge_ddd.py")
+
+ def test_set_params_re_emits(self):
+ """Documented side effect of BaseEstimator's transactional probe
+ re-init, shared with MultiPeriodDiD's shim."""
+ est = _sddd()
+ with pytest.warns(FutureWarning, match="StaggeredTripleDifference is deprecated"):
+ est.set_params(alpha=0.10)
+
+
+# ---------------------------------------------------------------------------
+# Gate F - introspection, the new validation contracts, cross-mode refit
+# ---------------------------------------------------------------------------
+
+_MERGED_PARAMS = {
+ "estimation_method",
+ "robust",
+ "cluster",
+ "vcov_type",
+ "alpha",
+ "pscore_trim",
+ "rank_deficient_action",
+ "epv_threshold",
+ "pscore_fallback",
+ "control_group",
+ "anticipation",
+ "base_period",
+ "n_bootstrap",
+ "bootstrap_weights",
+ "seed",
+ "cband",
+}
+
+
+class TestGateFIntrospectionAndValidation:
+ def test_merged_param_set_is_pinned(self):
+ assert set(TripleDifference().get_params()) == _MERGED_PARAMS
+
+ def test_dying_class_param_set_is_frozen(self):
+ assert set(_sddd().get_params()) == {
+ "estimation_method",
+ "control_group",
+ "alpha",
+ "anticipation",
+ "base_period",
+ "n_bootstrap",
+ "bootstrap_weights",
+ "seed",
+ "cband",
+ "pscore_trim",
+ "cluster",
+ "rank_deficient_action",
+ "epv_threshold",
+ "pscore_fallback",
+ }
+
+ def test_round_trip_both_classes(self):
+ est = TripleDifference(control_group="never_treated", n_bootstrap=9)
+ assert TripleDifference(**est.get_params()).get_params() == est.get_params()
+ old = _sddd(control_group="nevertreated")
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore", FutureWarning)
+ assert StaggeredTripleDifference(**old.get_params()).get_params() == old.get_params()
+
+ def test_set_params_round_trip(self):
+ est = TripleDifference()
+ assert est.set_params(alpha=0.10) is est
+ assert est.get_params()["alpha"] == 0.10
+
+ @pytest.mark.parametrize("bad", [0, 0.5, -0.1, 0.9, float("nan"), float("inf")])
+ def test_pscore_trim_boundaries_rejected(self, bad):
+ with pytest.raises(ValueError, match=r"pscore_trim must be in \(0, 0.5\)"):
+ TripleDifference(pscore_trim=bad)
+
+ @pytest.mark.parametrize(
+ "bad", [None, "0.01", 1j, [0.01], (0.01,), np.array([0.01]), True, False]
+ )
+ def test_pscore_trim_non_scalar_rejected(self, bad):
+ """The TYPE guard, not just the range one. A bare `0 < x < 0.5` raises an
+ incidental TypeError on None/str/complex/list, an ambiguous-truth error on
+ a multi-element array, and ACCEPTS a 1-element array - storing an ndarray
+ as the parameter, which then rides into np.clip(pscore, trim, 1 - trim)."""
+ with pytest.raises(ValueError, match="pscore_trim must be"):
+ TripleDifference(pscore_trim=bad)
+
+ def test_pscore_trim_multi_element_array_rejected(self):
+ with pytest.raises(ValueError, match="pscore_trim must be"):
+ TripleDifference(pscore_trim=np.array([0.01, 0.02]))
+
+ @pytest.mark.parametrize("good", [0.01, 0.49, 0.001, np.float64(0.02), np.float32(0.02)])
+ def test_pscore_trim_interior_accepted(self, good):
+ assert TripleDifference(pscore_trim=good).pscore_trim == good
+
+ def test_set_params_pscore_trim_is_transactional(self):
+ est = TripleDifference(pscore_trim=0.01)
+ with pytest.raises(ValueError):
+ est.set_params(pscore_trim=0)
+ assert est.pscore_trim == 0.01, "probe re-init must not mutate on failure"
+
+ # -- anticipation domain (silent-estimand guard) --------------------------
+
+ @pytest.mark.parametrize("bad", [-1, -5, 1.5, 0.5, "1", None, True, False, np.float64(2.0)])
+ def test_anticipation_domain_rejected(self, bad):
+ with pytest.raises(ValueError, match="anticipation must be a non-negative integer"):
+ TripleDifference(anticipation=bad)
+
+ @pytest.mark.parametrize("good", [0, 1, 3, np.int64(2)])
+ def test_anticipation_domain_accepted(self, good):
+ assert TripleDifference(anticipation=good).anticipation == good
+
+ def test_set_params_anticipation_is_transactional(self):
+ est = TripleDifference(anticipation=1)
+ with pytest.raises(ValueError, match="anticipation must be a non-negative integer"):
+ est.set_params(anticipation=-1)
+ assert est.anticipation == 1, "probe re-init must not mutate on failure"
+
+ def test_negative_anticipation_cannot_reach_the_engine(self, stag):
+ """The behavioral point of the guard, not just the raise: a negative
+ window would make the universal base period `g` - an ALREADY-TREATED
+ period - and relax the not-yet-treated threshold to max(t, base) - 1,
+ admitting cohorts treated at the evaluation period as clean controls.
+ Neither is visible in the output, so the constructor is the only place
+ it can be stopped."""
+ with pytest.raises(ValueError, match="anticipation must be a non-negative integer"):
+ TripleDifference(anticipation=-1, base_period="universal")
+ # the guard is EAGER: it fires at construction, so no fit can be reached
+ # with an out-of-domain window even via the staggered branch.
+ with pytest.raises(ValueError, match="anticipation must be a non-negative integer"):
+ TripleDifference(anticipation=-1).fit(
+ stag, partition="eligibility", **S_NEW, aggregate="simple"
+ )
+
+ @pytest.mark.parametrize("bad", [-1, 1.5, True])
+ def test_deprecated_sibling_also_fails_closed_on_anticipation(self, stag, bad):
+ """The dying class's frozen 3.x API SHAPE (param names, `eligibility=`,
+ the compact control_group vocabulary, accepted-then-ignored `cluster=`)
+ was never a licence to emit silently-biased numbers: it runs the SAME
+ engine, where a negative window selects an already-treated base period
+ and admits contaminated controls. So the guard lives in the engine and
+ BOTH surfaces fail closed.
+
+ The split is deliberate and pinned below: CONSTRUCTION still succeeds on
+ the deprecated class (its signature contract is untouched), while FIT
+ raises - the check is an identification guard, not an API change."""
+ est = _sddd(anticipation=bad)
+ assert est.anticipation == bad, "the dying class's constructor contract is unchanged"
+ with pytest.raises(ValueError, match="anticipation must be a non-negative integer"):
+ est.fit(stag, eligibility="eligibility", **S_OLD, aggregate="simple")
+
+ # -- first_treat cohort encoding (silent-population guard) ----------------
+
+ @pytest.mark.parametrize("sentinel", [-1, -np.inf, -0.5])
+ @pytest.mark.parametrize("surface", ["merged", "deprecated"])
+ def test_negative_cohorts_fail_closed_on_both_surfaces(self, stag, sentinel, surface):
+ """_precompute_structures builds the treated set from `g > 0` and the
+ never-enabled set from `g == 0`, so a unit encoded with the common `-1`
+ never-treated convention belonged to NEITHER: still counted in n_obs,
+ contributing to no ATT comparison, and the fit returned a plausible
+ finite estimate for a DIFFERENT population (measured before the guard:
+ overall_att 3.29438582 -> 2.99470938, n_never_enabled 24 -> 0, silently).
+
+ The +inf branch already defended this same input axis with an explicit
+ recode-and-warn, so silence on negatives was a hole in an established
+ contract. Raising rather than recoding: -1-as-never is a convention, not
+ an unambiguous limit like +inf, and guessing would BE the silent sample
+ change."""
+ df = stag.copy()
+ df["first_treat"] = df["first_treat"].astype(float)
+ df.loc[df["first_treat"] == 0, "first_treat"] = sentinel
+ fit = _fit_new if surface == "merged" else _fit_old
+ with pytest.raises(ValueError, match="negative cohort value"):
+ fit(df, aggregate="simple")
+
+ # -- degenerate enabling cohort (honest-reporting guard) ------------------
+
+ @staticmethod
+ def _strip_eligibility_from_cohort(df, cohort):
+ out = df.copy()
+ units = out.loc[out["first_treat"] == cohort, "unit"].unique()
+ out.loc[out["unit"].isin(units), "eligibility"] = 0
+ return out
+
+ @pytest.mark.parametrize("surface", ["merged", "deprecated"])
+ def test_degenerate_cohort_is_named_and_dropped_from_groups(self, stag, surface):
+ """A positive cohort whose units are ALL partition==0 cannot identify
+ ATT(g,t), so it contributes to no aggregate - but it was still advertised
+ in `groups`/`n_groups`, and no warning named it (the ones mentioning it
+ describe its role as a COMPARISON cohort for other g). The estimate is
+ deliberately unchanged: it is valid for the cohorts that do identify."""
+ df = self._strip_eligibility_from_cohort(stag, 4)
+ fit = _fit_new if surface == "merged" else _fit_old
+ with pytest.warns(UserWarning, match=r"cohort g=4\.0 has no eligible treated units"):
+ r = fit(df, aggregate="simple")
+ assert list(r.groups) == [3.0], "groups must not advertise a non-contributing cohort"
+ assert r.to_dict()["n_groups"] == 1
+ assert {g for g, _ in r.group_time_effects} == {3.0}
+
+ def test_degenerate_cohort_does_not_move_the_estimate(self, stag):
+ """The fix is reporting-only - pinned so a later 'improvement' cannot
+ quietly turn it into an estimand change."""
+ df = self._strip_eligibility_from_cohort(stag, 4)
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ new, old = _fit_new(df, aggregate="simple"), _fit_old(df, aggregate="simple")
+ _eq_with_nans([new.overall_att, new.overall_se], [old.overall_att, old.overall_se])
+ assert list(new.groups) == list(old.groups)
+
+ def test_healthy_cohorts_are_untouched(self, stag):
+ """Negative pin: no warning and no groups change on a well-posed panel."""
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ r = _fit_new(stag, aggregate="simple")
+ assert not [x for x in w if "no eligible treated units" in str(x.message)]
+ assert list(r.groups) == [3.0, 4.0]
+
+ def test_degenerate_cohort_under_survey_weights(self, stag_survey):
+ """Survey counterpart: rows exist for the cohort but carry no eligible
+ treated mass."""
+ df = self._strip_eligibility_from_cohort(stag_survey, 4)
+ with pytest.warns(UserWarning, match=r"cohort g=4\.0 has no eligible treated units"):
+ r = _fit_new(df, aggregate="simple", survey_design=SurveyDesign(weights="w"))
+ assert list(r.groups) == [3.0]
+
+ def test_positive_inf_still_recodes_with_a_warning(self, stag):
+ """The negative guard must not disturb the +inf lane it sits beside."""
+ df = stag.copy()
+ df["first_treat"] = df["first_treat"].astype(float)
+ df.loc[df["first_treat"] == 0, "first_treat"] = np.inf
+ with pytest.warns(UserWarning, match="recoding to 0"):
+ r = _fit_new(df, aggregate="simple")
+ plain = _fit_new(stag, aggregate="simple")
+ _eq_with_nans([r.overall_att, r.overall_se], [plain.overall_att, plain.overall_se])
+ assert r.n_never_enabled == plain.n_never_enabled
+
+ def test_valid_cohorts_are_untouched_by_the_guard(self, stag):
+ """Regression pin: the guard changes no accepted fit."""
+ r = _fit_new(stag, aggregate="all")
+ exp = ORACLE_STAGGERED["stag_dr_all"]
+ assert_allclose(r.overall_att, exp["overall"]["att"], rtol=ORACLE_RTOL, atol=ORACLE_ATOL)
+
+ def test_anticipation_guard_survives_attribute_mutation(self, stag):
+ """The engine-level guard also covers the __init__/set_params bypass."""
+ est = TripleDifference()
+ est.anticipation = -1 # neither validator sees this
+ with pytest.raises(ValueError, match="anticipation must be a non-negative integer"):
+ est.fit(stag, partition="eligibility", **S_NEW, aggregate="simple")
+
+ @pytest.mark.parametrize("compact", ["notyettreated", "nevertreated"])
+ def test_merged_class_rejects_the_compact_vocabulary(self, compact):
+ """The boundary Design 3's bridge and Gate B's pairing rest on."""
+ with pytest.raises(ValueError, match="control_group must be"):
+ TripleDifference(control_group=compact)
+
+ @pytest.mark.parametrize("underscored", ["not_yet_treated", "never_treated"])
+ def test_dying_class_rejects_the_underscored_vocabulary(self, underscored):
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore", FutureWarning)
+ with pytest.raises(ValueError, match="control_group must be"):
+ StaggeredTripleDifference(control_group=underscored)
+
+ def test_base_period_and_bootstrap_weights_contracts(self):
+ with pytest.raises(ValueError, match="base_period must be"):
+ TripleDifference(base_period="bogus")
+ with pytest.raises(ValueError, match="bootstrap_weights must be"):
+ TripleDifference(bootstrap_weights="bogus")
+
+ def test_n_bootstrap_type_contract(self):
+ for bad in (True, 2.5, -1, None):
+ with pytest.raises(ValueError, match="n_bootstrap"):
+ TripleDifference(n_bootstrap=bad)
+
+ def test_cross_mode_refit_on_one_instance(self, cross, stag):
+ """fit() writes mode-specific state (results_ changes TYPE, is_fitted_,
+ _replicate_n_valid is reset only in the 2x2x2 prologue), so a refit
+ across modes must match a fresh instance in both directions."""
+ est = TripleDifference()
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ stag_first = est.fit(stag, partition="eligibility", **S_NEW, aggregate="all")
+ assert isinstance(stag_first, StaggeredTripleDiffResults)
+ then_cross = est.fit(cross, **C_COLS)
+ assert isinstance(then_cross, TripleDifferenceResults)
+ fresh_cross = TripleDifference().fit(cross, **C_COLS)
+ _eq_with_nans(_quintet(then_cross), _quintet(fresh_cross))
+
+ back_to_stag = est.fit(stag, partition="eligibility", **S_NEW, aggregate="all")
+ _assert_staggered_parity(stag_first, back_to_stag)
+ assert est.is_fitted_
+
+ def test_frame_offset_is_restored_after_a_RAISING_staggered_fit(self, cross, stag):
+ """The try/finally must restore the offset, or every later warning on
+ the instance attributes to the wrong frame."""
+ est = TripleDifference()
+ bad = stag.copy()
+ bad.loc[bad.index[0], "eligibility"] = 1 - bad.loc[bad.index[0], "eligibility"]
+ with pytest.raises(ValueError):
+ est.fit(bad, partition="eligibility", **S_NEW)
+ assert est._warn_frame_offset == 0
+
+
+# ---------------------------------------------------------------------------
+# Gate G - warning attribution, message threading, the R lane
+# ---------------------------------------------------------------------------
+
+
+class TestGateGAttributionAndThreading:
+ def test_user_attributed_warnings_land_on_the_CALLER_on_both_surfaces(self, stag):
+ """The facade adds a frame; _frame_offset must absorb it. Uses the
+ base-period-outside-panel warning, which fires from the fit body."""
+ df = generate_staggered_ddd_data(**{**STAG_KW, "cohort_periods": [1, 3]})
+ for fit in (_fit_old, _fit_new):
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ fit(df, aggregate="simple")
+ fired = [x for x in w if "outside the observed panel" in str(x.message)]
+ assert fired, "expected the base-period warning"
+ assert all(x.filename.endswith("test_v4_merge_ddd.py") for x in fired), [
+ x.filename for x in fired
+ ]
+
+ def test_low_bootstrap_warning_attributes_to_the_caller_on_both_surfaces(self, stag):
+ """Shared with CallawaySantAnna via the bootstrap mixin, so it reads the
+ offset through getattr rather than taking it as an argument."""
+ for fit in (_fit_old, _fit_new):
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ fit(stag, ctor=dict(n_bootstrap=9, seed=1), aggregate="simple")
+ fired = [x for x in w if "is low" in str(x.message)]
+ assert fired and all(x.filename.endswith("test_v4_merge_ddd.py") for x in fired), [
+ (x.filename, str(x.message)[:40]) for x in fired
+ ]
+
+ def test_callawaysantanna_attribution_is_unchanged(self):
+ """CS never sets _warn_frame_offset, so the shared site's getattr
+ default keeps its attribution bit-identical to 3.x."""
+ from diff_diff import CallawaySantAnna
+ from diff_diff.prep_dgp import generate_staggered_data
+
+ df = generate_staggered_data(n_units=60, n_periods=5, seed=1)
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ CallawaySantAnna(n_bootstrap=9, seed=1).fit(
+ df, outcome="outcome", unit="unit", time="period", first_treat="first_treat"
+ )
+ fired = [x for x in w if "is low" in str(x.message)]
+ assert fired and all(x.filename.endswith("test_v4_merge_ddd.py") for x in fired)
+
+ def test_survey_pweight_error_names_the_surface_that_was_fit(self, stag):
+ df = stag.copy()
+ df["w"] = 1.0
+ design = SurveyDesign(weights="w", weight_type="fweight")
+ for fit, name in ((_fit_old, "StaggeredTripleDifference"), (_fit_new, "TripleDifference")):
+ with pytest.raises(ValueError) as exc:
+ fit(df, survey_design=design)
+ assert str(exc.value).startswith(f"{name} survey support requires")
+
+ def test_partition_vocabulary_never_leaks_on_the_merged_surface(self, stag):
+ """BOTH sentences of the time-invariance error are parameterized - the
+ tail 'varying eligibility' would otherwise say the wrong word."""
+ df = stag.copy()
+ df.loc[df.index[0], "eligibility"] = 1 - df.loc[df.index[0], "eligibility"]
+ with pytest.raises(ValueError) as exc:
+ _fit_new(df)
+ msg = str(exc.value)
+ assert "Partition must be time-invariant" in msg
+ assert "eligibility" not in msg, msg
+ with pytest.raises(ValueError) as exc_old:
+ _fit_old(df)
+ assert "Eligibility must be time-invariant" in str(exc_old.value)
+ assert "varying eligibility" in str(exc_old.value)
+
+ def test_no_valid_group_time_message_uses_the_right_vocabulary(self, stag):
+ """An anticipation window wider than the panel skips every cell, which
+ is the cheapest route to this message. (Filtering the partition instead
+ would trip the Q-notation validator first - that message is paper
+ vocabulary and is deliberately NOT parameterized.)"""
+ for fit, word, wrong in (
+ (_fit_new, "partition", "eligibility"),
+ (_fit_old, "eligibility", "partition"),
+ ):
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ with pytest.raises(ValueError, match="No valid group-time effects") as exc:
+ fit(stag, ctor=dict(anticipation=10))
+ assert word in str(exc.value)
+ assert wrong not in str(exc.value)
+
+
+# ---------------------------------------------------------------------------
+# Gate H - downstream consumers
+# ---------------------------------------------------------------------------
+
+
+class TestGateHConsumers:
+ def test_event_study_surface_identical_across_surfaces(self, stag):
+ from diff_diff.results_base import build_event_study_surface
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ old = build_event_study_surface(_fit_old(stag, aggregate="event_study"))
+ new = build_event_study_surface(_fit_new(stag, aggregate="event_study"))
+ _eq_with_nans(old.att, new.att)
+ _eq_with_nans(old.se, new.se)
+ assert old.source == new.source
+
+ def test_business_report_handles_both_control_group_vocabularies(self, stag):
+ from diff_diff.business_report import BusinessReport
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ new = BusinessReport(
+ _fit_new(stag, ctor=dict(control_group="not_yet_treated"), aggregate="simple")
+ ).to_dict()
+ old = BusinessReport(
+ _fit_old(stag, ctor=dict(control_group="notyettreated"), aggregate="simple")
+ ).to_dict()
+ assert new["sample"]["n_never_enabled"] == old["sample"]["n_never_enabled"]
+
+ def test_summary_renders_in_both_modes(self, cross, stag):
+ """summary()/print_summary() delegate to results_, so the merged class
+ renders the staggered container in staggered mode - new public
+ behavior that follows for free but should be pinned."""
+ est = TripleDifference()
+ est.fit(cross, **C_COLS)
+ assert "Triple Difference" in est.summary()
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ est.fit(stag, partition="eligibility", **S_NEW, aggregate="all")
+ text = est.summary()
+ assert "ATT" in text and len(text.splitlines()) > 10
+
+ def test_describe_target_parameter_unchanged(self, stag):
+ from diff_diff._reporting_helpers import describe_target_parameter
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ assert describe_target_parameter(
+ _fit_new(stag, aggregate="simple")
+ ) == describe_target_parameter(_fit_old(stag, aggregate="simple"))
+
+ def test_power_rejects_staggered_config_at_all_three_entry_points(self):
+ from diff_diff.power import simulate_mde, simulate_power, simulate_sample_size
+
+ for fn in (simulate_power, simulate_mde, simulate_sample_size):
+ with pytest.raises(ValueError, match="staggered DDD mode"):
+ fn(TripleDifference(control_group="never_treated"), n_simulations=1)
+ with pytest.raises(ValueError, match="staggered DDD mode"):
+ fn(TripleDifference(), n_simulations=1, estimator_kwargs={"first_treat": "g"})
+
+ def test_power_still_works_on_defaults(self):
+ from diff_diff.power import simulate_power
+
+ r = simulate_power(TripleDifference(), n_units=64, n_simulations=2, seed=1)
+ assert r is not None
+
+ @pytest.mark.parametrize(
+ "inert",
+ [dict(seed=7), dict(bootstrap_weights="mammen"), dict(cband=False)],
+ )
+ def test_power_accepts_the_inert_bootstrap_satellites(self, inert):
+ """The three params that fit() accepts in 2x2x2 mode as inert must not be
+ rejected by power. They take effect only through n_bootstrap > 0, which
+ 2x2x2 mode already rejects, so a non-default value cannot signal a
+ staggered configuration - and refusing them made the same estimator legal
+ to fit and illegal to simulate. `seed` is the one that bites in practice:
+ users set it habitually, and simulate_* take their own separate seed=."""
+ from diff_diff.power import simulate_mde, simulate_power, simulate_sample_size
+
+ # the fit really does accept it (the premise of this gate)
+ cross = generate_ddd_data(**CROSS_KW)
+ assert TripleDifference(**inert).fit(cross, **C_COLS) is not None
+
+ assert (
+ simulate_power(TripleDifference(**inert), n_units=64, n_simulations=2, seed=1)
+ is not None
+ )
+ assert simulate_mde(
+ TripleDifference(**inert),
+ n_units=64,
+ n_simulations=2,
+ seed=1,
+ max_steps=2,
+ progress=False,
+ )
+ assert simulate_sample_size(
+ TripleDifference(**inert),
+ n_simulations=2,
+ seed=1,
+ max_steps=2,
+ progress=False,
+ )
+
+ def test_deprecated_wrapper_stays_2x2x2_only(self, cross):
+ """`triple_difference()` is deliberately NOT extended to the staggered
+ mode. It is itself deprecated (row M-075, removed at 4.0), phase 3(a)
+ extended no wrapper with its new mode, and adding params to a dying
+ surface would mint fresh 4.0 removal obligations. The migration steer -
+ 'construct the estimator instead' - is the intended pressure. Pinned so
+ the omission reads as a decision."""
+ import inspect
+
+ from diff_diff import triple_difference
+
+ params = set(inspect.signature(triple_difference).parameters)
+ staggered_only = {
+ "unit",
+ "first_treat",
+ "aggregate",
+ "balance_e",
+ "control_group",
+ "anticipation",
+ "base_period",
+ "n_bootstrap",
+ "bootstrap_weights",
+ "seed",
+ "cband",
+ }
+ assert not (params & staggered_only), (
+ "the deprecated wrapper grew a staggered-mode param; it is slated for "
+ "4.0 removal and must not widen"
+ )
+
+ # ...and the 2x2x2 route it DOES own still works, warning exactly once
+ with pytest.warns(FutureWarning, match="triple_difference.. is deprecated"):
+ r = triple_difference(cross, "outcome", "group", "partition", "time")
+ direct = TripleDifference().fit(cross, **C_COLS)
+ _eq_with_nans(_quintet(r), _quintet(direct))
+
+ @pytest.mark.parametrize("fn_name", ["simulate_power", "simulate_mde", "simulate_sample_size"])
+ def test_power_accepts_explicit_none_aggregate_and_balance_e(self, cross, fn_name):
+ """`aggregate`/`balance_e` default to None and fit() rejects only a
+ NON-None value, so keying the power guard on KEY PRESENCE rejected
+ `estimator_kwargs={"aggregate": None}` - a config fit() accepts. That
+ broke the boundary the guard exists to uphold: legal to fit implies
+ legal to simulate."""
+ import diff_diff.power as power_mod
+
+ kw = {"aggregate": None, "balance_e": None}
+ assert TripleDifference().fit(cross, **C_COLS, **kw) is not None
+
+ fn = getattr(power_mod, fn_name)
+ extra = {"n_units": 48} if fn_name != "simulate_sample_size" else {}
+ if fn_name != "simulate_power":
+ extra.update(max_steps=2, progress=False)
+ assert fn(TripleDifference(), n_simulations=2, seed=1, estimator_kwargs=kw, **extra)
+
+ @pytest.mark.parametrize("key,value", [("aggregate", "simple"), ("balance_e", 1)])
+ def test_power_still_rejects_non_none_aggregate_and_balance_e(self, key, value):
+ from diff_diff.power import simulate_power
+
+ with pytest.raises(ValueError, match="staggered DDD mode"):
+ simulate_power(
+ TripleDifference(), n_units=48, n_simulations=1, estimator_kwargs={key: value}
+ )
+
+ @pytest.mark.parametrize("key", ["first_treat", "unit"])
+ @pytest.mark.parametrize("value", ["g", None])
+ def test_power_rejects_mode_selectors_by_presence(self, key, value):
+ """The other half of the same rule: `first_treat`/`unit` are
+ sentinel-defaulted in fit(), so SUPPLYING them at all is the signal -
+ an explicit None still selects staggered mode and then fails on a
+ missing column. Presence is the correct test for these two."""
+ from diff_diff.power import simulate_power
+
+ with pytest.raises(ValueError, match="staggered DDD mode"):
+ simulate_power(
+ TripleDifference(), n_units=48, n_simulations=1, estimator_kwargs={key: value}
+ )
+
+ @pytest.mark.parametrize("surface", ["merged", "deprecated"])
+ def test_bootstrap_warning_names_the_estimator_that_was_fit(self, stag, surface):
+ """The single-PSU/degenerate-design warning lives in the mixin shared by
+ CallawaySantAnna and BOTH DDD classes, so the hard-coded literal named
+ CallawaySantAnna no matter which surface was fit - a misleading diagnosis
+ on a fit that is failing closed.
+
+ Routed through a real degenerate fit rather than asserted against source
+ text: a source pin cannot see whether the interpolated value actually
+ reaches the user."""
+ df = stag.copy()
+ df["w"] = 1.0
+ df["psu"] = 0 # collapse to one PSU -> bootstrap variance unidentified
+ df["fpc_col"] = 10**6 # design-based variance so the PSU is retained
+ design = SurveyDesign(weights="w", psu="psu", fpc="fpc_col")
+ ctor = dict(estimation_method="reg", n_bootstrap=20, seed=7)
+ fit = _fit_new if surface == "merged" else _fit_old
+ expected = "TripleDifference" if surface == "merged" else "StaggeredTripleDifference"
+
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ fit(df, aggregate="event_study", survey_design=design, ctor=ctor)
+ hits = [
+ str(w.message)
+ for w in caught
+ if "bootstrap with survey/cluster design" in str(w.message)
+ ]
+ assert hits, "expected the single-PSU bootstrap warning"
+ assert all(m.startswith(f"{expected} bootstrap") for m in hits), hits
+ assert not any(m.startswith("CallawaySantAnna") for m in hits)
+
+ def test_bootstrap_label_is_declared_on_every_host(self):
+ """CallawaySantAnna shares the mixin and its message was already correct,
+ so its label must stay byte-identical."""
+ from diff_diff.staggered import CallawaySantAnna
+
+ assert CallawaySantAnna._BOOTSTRAP_LABEL == "CallawaySantAnna"
+ assert TripleDifference._BOOTSTRAP_LABEL == "TripleDifference"
+ assert StaggeredTripleDifference._BOOTSTRAP_LABEL == "StaggeredTripleDifference"
+
+ @pytest.mark.parametrize(
+ "staggered", ["control_group", "anticipation", "base_period", "n_bootstrap"]
+ )
+ def test_power_still_rejects_the_genuinely_staggered_four(self, staggered):
+ """The other side of the same boundary - narrowing the roster must not
+ have opened a hole for the params that DO select staggered behavior."""
+ from diff_diff.power import simulate_power
+
+ value = {
+ "control_group": "never_treated",
+ "anticipation": 1,
+ "base_period": "universal",
+ "n_bootstrap": 49,
+ }[staggered]
+ with pytest.raises(ValueError, match="staggered DDD mode"):
+ simulate_power(TripleDifference(**{staggered: value}), n_simulations=1)
+
+ def test_staggered_roster_reads_live_constructor_defaults(self):
+ """The guard compares each staggered-only param against its DEFAULT. If
+ those defaults were hard-coded, a future constructor-default change would
+ silently reclassify a plain TripleDifference() as staggered-configured
+ and reject a legitimate 2x2x2 power run."""
+ import inspect
+
+ from diff_diff.utils import STAGGERED_DDD_CTOR_PARAMS, staggered_ddd_ctor_defaults
+
+ sig = inspect.signature(TripleDifference.__init__).parameters
+ missing = [n for n in STAGGERED_DDD_CTOR_PARAMS if n not in sig]
+ assert not missing, f"roster names no longer on the constructor: {missing}"
+
+ derived = staggered_ddd_ctor_defaults(TripleDifference())
+ assert set(derived) == set(STAGGERED_DDD_CTOR_PARAMS)
+ for name, default in derived.items():
+ assert default == sig[name].default
+
+ # and a default-constructed estimator is never an offender
+ from diff_diff.power import _reject_staggered_ddd_config
+
+ _reject_staggered_ddd_config(TripleDifference(), {})
+
+ def test_fit_and_power_share_one_staggered_boundary(self):
+ """The defect this centralization fixes: fit() used to hardcode the four
+ defaults while power derived them, so a future default change could make
+ the two mode-detection boundaries disagree about what
+ 'staggered-configured' means. Both now read the same helper - asserted
+ by behavior on every rostered param, not just by shared imports."""
+ from diff_diff.power import simulate_power
+ from diff_diff.utils import STAGGERED_DDD_CTOR_PARAMS
+
+ cross = generate_ddd_data(**CROSS_KW)
+ non_default = {
+ "control_group": "never_treated",
+ "anticipation": 1,
+ "base_period": "universal",
+ "n_bootstrap": 49,
+ }
+ assert set(non_default) == set(STAGGERED_DDD_CTOR_PARAMS)
+ for name, value in non_default.items():
+ est = TripleDifference(**{name: value})
+ with pytest.raises(ValueError, match="staggered"):
+ est.fit(cross, **C_COLS)
+ with pytest.raises(ValueError, match="staggered DDD mode"):
+ simulate_power(TripleDifference(**{name: value}), n_simulations=1)
+
+
+# ---------------------------------------------------------------------------
+# Gate G (R lane) - the golden, plus its always-running substitute
+# ---------------------------------------------------------------------------
+
+
+class TestGateGRParity:
+ """The R-golden lane SKIPS in CI: the generated CSVs are gitignored
+ (.gitignore benchmarks/data/synthetic/*.csv), so `_load_r_data` calls
+ pytest.skip. It is therefore paired with a substitute that always runs and
+ proves the routing is identity-preserving on the same config.
+ """
+
+ def test_r_golden_through_the_merged_surface(self):
+ """Local-only. Reuses the methodology suite's loader and tolerances
+ rather than duplicating them; maps the golden's compact control_group
+ to the merged class's underscored vocabulary."""
+ import json
+
+ from tests.test_methodology_staggered_triple_diff import (
+ ATT_ATOL,
+ ATT_RTOL,
+ TestStaggeredDDDAggregation,
+ _load_r_data,
+ )
+
+ results_file = (
+ __import__("pathlib").Path(__file__).parent.parent
+ / "benchmarks"
+ / "data"
+ / "synthetic"
+ / "staggered_ddd_r_results.json"
+ )
+ if not results_file.exists():
+ pytest.skip("R golden results not present")
+ golden = json.loads(results_file.read_text())
+ key = "s42_dgp1_dr_nyt"
+ if key not in golden:
+ pytest.skip(f"scenario {key} absent from the golden file")
+ scenario = golden[key]
+ df = _load_r_data(42, 1) # skips when the gitignored CSV is absent
+
+ # No aggregate=: the default simple aggregation is what the golden's
+ # overall_att_simple records, and it is how the methodology suite fits.
+ r = TripleDifference(estimation_method="dr", control_group="not_yet_treated").fit(
+ df, partition="eligibility", **S_NEW
+ )
+ # Tolerances mirror the methodology suite EXACTLY rather than inventing
+ # new ones: per-(g,t) cells are pinned tightly, while the simple overall
+ # carries that suite's documented looser envelope (a weighting deviation
+ # on individual values - see TestStaggeredDDDAggregation's docstring).
+ assert_allclose(
+ r.overall_att,
+ scenario["overall_att_simple"],
+ rtol=TestStaggeredDDDAggregation.AGG_RTOL,
+ atol=TestStaggeredDDDAggregation.AGG_ATOL,
+ )
+ # the per-(g,t) table too, so the golden pins more than one scalar
+ py_gt = dict(sorted(r.group_time_effects.items()))
+ r_gt = list(zip(scenario["gt_groups"], scenario["gt_periods"]))
+ assert len(py_gt) == len(r_gt), f"{len(py_gt)} cells vs R's {len(r_gt)}"
+ for i, key_gt in enumerate(r_gt):
+ assert_allclose(
+ py_gt[key_gt]["effect"], scenario["gt_att"][i], rtol=ATT_RTOL, atol=ATT_ATOL
+ )
+
+ def test_always_running_substitute_on_the_same_config(self, stag):
+ """What the merge actually needs to establish: the routing is
+ identity-preserving on the golden's config, even when R data is absent."""
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ old = _fit_old(
+ stag,
+ ctor=dict(estimation_method="dr", control_group="notyettreated"),
+ aggregate="simple",
+ )
+ new = _fit_new(
+ stag,
+ ctor=dict(estimation_method="dr", control_group="not_yet_treated"),
+ aggregate="simple",
+ )
+ _assert_staggered_parity(old, new)