From 72cc39b13b596c9208e6dba0135a9e99f5493bb6 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 27 Aug 2026 20:04:47 -0400 Subject: [PATCH 1/4] feat(dml): DMLDiD survey-design and cluster support on both lanes (DML survey PR) survey_design= (pweight full-design TSL: weights/strata/PSU/FPC) on fit() and cluster= on the constructor, on the panel AND declared-RCS lanes. Declared designs weight the moment kernels (Hajek p-hat / lambda-hat / theta, weighted G_2lambda) and pass sample_weight into the nuisance learners (user learner objects must accept it by keyword - TypeError up front otherwise); cross-fitting switches to PSU-cohesive folds when the PSU is strictly coarser than the sampling unit, with the explicit Chang I_kz^c fold-composition guard on the RCS lane, an effective-fold-count reduction to n_psu (warned, provenance on DMLDiDResults.effective_n_folds / summary / to_dict / repr) when 2 <= n_psu < n_folds, and a stratified-fold fallback only for the single-PSU design (whose clustered variance is NaN either way). Per-cell SEs are PSU-gated exactly like CallawaySantAnna (_cluster_robust_se_from_per_gt_if; NaN propagates as the deliberate unidentified-variance signal on a RETAINED cell); analytical inference uses df = n_PSU - n_strata t-statistics; the survey multiplier bootstrap and <2-PSU NaN contracts activate through the inherited CS machinery, and aggregation kits now carry the producer's bootstrap label so post-fit replays brand warnings with the fitting estimator (the _BOOTSTRAP_LABEL ClassVar becomes a plain annotation across its five hosts). Bare cluster= keeps the kernels unweighted - it drives folds, variance, and df_inference only. zero_weight_mass joins the DMLDiD skip vocabulary; universal-base reference cells gate on the weighted cohort mass; aggregate('total') fails closed on declared-survey fits (is_survey_fit kit gate). A documented LIBRARY EXTENSION of Chang (2020)'s i.i.d. theory (REGISTRY DMLDiD survey Notes; Theorem 2's coverage claim does not carry over on the weighted-lambda RCS lane). Replicate-weight designs fail closed (TODO.md Actionable row). No v4 ledger row: additive kwargs, no behavior change on existing calls (precedent check recorded). Downstream consumer surfaces synced: reporting helpers' target-parameter description, MMM scale hint, practitioner learner-sensitivity template, BusinessReport narratives, DiagnosticReport design_effect roster comment, variance-conventions inventory (CI-pinned doc/test pair), survey-theory rosters, guides (llms/full/autonomous/practitioner incl. the signature-parity blocks), choosing_estimator + practitioner_decision_tree + survey-roadmap + migration-4.0 + v4-design prose, api/dml_did.rst, doc-deps.yaml mappings. Tests: new tests/test_survey_dml.py (62 tests - invariant battery, per-cell SE == compute_survey_if_variance cross-checks, cluster-wiring branches, fold cohesion/reduction/boundary pins, spy-learner weight-threading, zero_weight_mass skips, survey bootstrap + replay label + legacy-kit fallback, rejections, df threading/rendering, aggregation masses incl. the >2^53 cohort-key regression); weighted-slope tests in test_dr_scores.py (uniform == unweighted, finite-difference identity, validation rejections on the public AND internal entries); guide signature pins extended; existing DMLDiD/CS suites pass unchanged (no-survey fits bit-identical). --- CHANGELOG.md | 22 +- DEFERRED.md | 3 +- README.md | 2 +- TODO.md | 3 +- diff_diff/_dr_scores.py | 54 +- diff_diff/_reporting_helpers.py | 31 +- diff_diff/business_report.py | 13 +- diff_diff/diagnostic_report.py | 5 +- diff_diff/dml_did.py | 720 +++++++++++++++-- diff_diff/dml_did_results.py | 49 +- diff_diff/guides/llms-autonomous.txt | 14 +- diff_diff/guides/llms-full.txt | 5 +- diff_diff/guides/llms-practitioner.txt | 6 +- diff_diff/guides/llms.txt | 2 +- diff_diff/mmm.py | 5 +- diff_diff/practitioner.py | 17 +- diff_diff/staggered.py | 9 +- diff_diff/staggered_bootstrap.py | 17 +- diff_diff/staggered_results.py | 18 +- diff_diff/staggered_triple_diff.py | 4 +- diff_diff/triple_diff.py | 4 +- docs/api/dml_did.rst | 38 +- docs/choosing_estimator.rst | 17 +- docs/doc-deps.yaml | 6 + docs/index.rst | 2 +- docs/methodology/REGISTRY.md | 106 ++- docs/methodology/papers/chang-2020-review.md | 6 +- docs/methodology/survey-theory.md | 8 +- docs/methodology/variance-conventions.md | 4 +- docs/migration-4.0.md | 2 +- docs/practitioner_decision_tree.rst | 16 +- docs/survey-roadmap.md | 4 +- docs/v4-design.md | 8 +- tests/test_dml_did.py | 10 +- tests/test_dr_scores.py | 119 +++ tests/test_guides.py | 4 + tests/test_survey_dml.py | 781 +++++++++++++++++++ tests/test_variance_conventions.py | 11 +- 38 files changed, 1935 insertions(+), 210 deletions(-) create mode 100644 tests/test_survey_dml.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 018af339d..865ae8ab8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- **`DMLDiD` survey-design and cluster support (both lanes)**: `fit()` gains + `survey_design=` (pweight-only full-design TSL — weights/strata/PSU/FPC) + and the constructor gains `cluster=` (coarser-than-unit clustering). A + declared design weights the moment kernels (Hájek `p̂`/`λ̂`/`θ̂`, weighted + `Ĝ₂λ`), passes `sample_weight` into the nuisance learners (user learner + objects must accept it by keyword — rejected up front with `TypeError` + otherwise), switches cross-fitting to PSU-cohesive folds when the PSU is + strictly coarser than the sampling unit (with an explicit Chang + `I_kz^c` fold-composition guard on the RCS lane), and routes per-cell and + aggregate variances through the design-based kernels with + `df = n_PSU − n_strata` t-inference; the survey multiplier bootstrap and + the <2-PSU NaN contracts activate through the inherited CS machinery + (post-fit bootstrap replays now brand their warnings with the fitting + estimator's name). Bare `cluster=` keeps the kernels unweighted + (variance/folds/df only). A documented LIBRARY EXTENSION of Chang + (2020)'s i.i.d. theory (REGISTRY DMLDiD survey Notes; Theorem 2's + coverage claim does not carry over on the weighted-λ RCS lane). + Replicate-weight designs fail closed (TODO.md row); + `aggregate('total')` fails closed on declared-survey fits. - **`DMLDiD` — Chang (2020) double/debiased machine learning DiD** (staggered panel ATT(g,t)): per-(g,t)-cell DML2 cross-fitting of configurable nuisance learners (string names or any sklearn-style object; the adaptive polynomial @@ -28,8 +47,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 Theorem 2 variance (`chang_rcs_score_augmented` with the explicit `Ĝ₂λ(T − λ̂)` term). Per-observation influence functions; aggregation weights are fixed cohort row masses (WIF-consistent SEs); - `aggregate('total')` fails closed on RCS fits; no survey weights (weighted - RCS belongs to `CallawaySantAnna(panel=False, survey_design=...)`). + `aggregate('total')` fails closed on RCS fits. Validated by equation-level fixtures, oracle closed forms, derivative-identity checks, double robustness in both directions, and a committed `DoubleMLDIDCSBinary` characterization spike (no parity oracle diff --git a/DEFERRED.md b/DEFERRED.md index 25c394b72..5690c1c13 100644 --- a/DEFERRED.md +++ b/DEFERRED.md @@ -43,7 +43,7 @@ provenance and AI-review deviation-documentation: a row here (or in | `SpilloverDiD(vcov_type="conley", conley_lag_cutoff>0, survey_design=...)` no-effective-PSU serial Bartlett HAC: weights-only / strata-only designs without a cluster fallback raise `NotImplementedError` (each pseudo-PSU appears in one period, so the serial cross-period loop contributes zero). Needs a unit-level serial fallback derivation or routing through `conley_unit` with documented IF-allocator asymmetry. | `spillover.py`, `two_stage.py::_compute_stratified_serial_bartlett_meat` | Wave E.2 tail | Low | | `SpilloverDiD` data-driven `d_bar` selection (Butts 2021b / 2023 JUE Insight cross-validation). | `spillover.py` | follow-up | Low | | **`LPDiD` non-absorbing exit-event dynamics** (Dube et al. 2025 online Appendix C `eta_h^{g,n}`): the shipped `non_absorbing` modes estimate the **entry-effect** estimands (Eq. 12/13) only; separate dynamic event-studies for treatment switch-*offs* are not implemented. Needs the exit-event clean-sample derivation + estimand contract. | `lpdid.py`, REGISTRY | PR-C follow-up | Low | -| MMM interop, survey/RC totals for the four `aggregate('total')` adopters (CS/EfficientDiD/ImputationDiD/TwoStageDiD; the shipped v1 is panel non-survey only and fails closed with `NotImplementedError` on repeated-cross-section-routed fits, fits declaring a `survey_design=`, and the CS bare-`cluster=` divergent corner where the cohort-mass weighting disagrees with the complete-case count): needs (a) DESIGN-AWARE raw-mass recovery - analytic pweight/aweight resolved weights are normalized to sum = n and CS retains no raw record (fit-time RAW retention required there; CS accepts pweight only), while EDiD's analytic fweight fits and all replicate-design fits already retain raw scale in the kits and Imputation/TwoStage retain raw `SurveyMetadata.sum_weights` - and (b) the mass-uncertainty `att*dC` variance derivation (the shipped relay conditions on the realized mass and omits it). Never ship a survey/RC total from normalized masses. | `diff_diff/staggered_results.py`, `diff_diff/survey.py` | mmm-interop | Low | +| MMM interop, survey/RC totals for the five `aggregate('total')` adopters (CS/DMLDiD/EfficientDiD/ImputationDiD/TwoStageDiD; the shipped v1 is panel non-survey only and fails closed with `NotImplementedError` on repeated-cross-section-routed fits, fits declaring a `survey_design=`, and the CS/DMLDiD bare-`cluster=` divergent corner where the cohort-mass weighting disagrees with the complete-case count): needs (a) DESIGN-AWARE raw-mass recovery - analytic pweight/aweight resolved weights are normalized to sum = n and CS retains no raw record (fit-time RAW retention required there; CS accepts pweight only), while EDiD's analytic fweight fits and all replicate-design fits already retain raw scale in the kits and Imputation/TwoStage retain raw `SurveyMetadata.sum_weights` - and (b) the mass-uncertainty `att*dC` variance derivation (the shipped relay conditions on the realized mass and omits it). Never ship a survey/RC total from normalized masses. | `diff_diff/staggered_results.py`, `diff_diff/survey.py` | mmm-interop | Low | | **`LPDiD` non-absorbing interior-gap support**: non-absorbing modes require a gap-free panel within each unit's observed span and raise on interior time gaps (the `[t-L, t+h]` window conditions can't be verified across a gap). The absorbing path already reindexes interior gaps to the calendar grid; extending that fail-closed handling (per-window gap masking) to non-absorbing is deferred. | `lpdid.py::_prepare_panel` | PR-C follow-up | Low | ## Needs external reference (R / Stata / Julia) @@ -157,4 +157,3 @@ decisions (refactor waivers, perf trade-offs, test-infrastructure calls) are rec | **`ImputationDiD` SE vcov is already rank-guarded upstream.** Excluded from the structural rank-guard sweep: the lead/effect vcov comes from `solve_ols(..., return_vcov=True, rank_deficient_action=...)` at the OLS fit (`imputation.py:~2316`), which already drops rank-deficient columns. The only raw inverse (`solve(V_gamma, gamma)`, `imputation.py:~2530`) is the pretrends **Wald F-test statistic** with a safe `NaN` fallback — a test statistic, not a sandwich bread — so there is no garbage-SE exposure. No structural rank-guard needed. | `imputation.py` | structural-rank-guard / 2026-06-28 | | **TWFE HC2/HC2-BM full-dummy dedup: drift-prone duplication already resolved by the shared builder; full delegation waived.** The former Actionable row (origin: follow-up review, citing pre-#655 line numbers) asked to extract a shared dummy-construction helper or delegate TWFE's HC2/HC2-BM path to DiD's `fixed_effects=` branch. The shared-helper half SHIPPED in #655: both sites now delegate dummy construction, drop-first convention, FE column naming, and the duplicate-term backstop to the single `build_fe_dummy_blocks` (`utils.py`) + `validate_design_term_names` implementation (`twfe.py::fit` full-dummy branch; `estimators.py::DifferenceInDifferences.fit` `fixed_effects=` branch) — the FE-naming / survey-behavior drift risk the row targeted is gone. What remains per site is ~4 lines of genuinely estimator-specific design-matrix assembly (TWFE stacks `const`/`ATT`/covariates; DiD stacks its formula terms), which is not drift-prone duplication. The remaining full-delegation option — routing `TWFE.fit` through DiD machinery with TWFE-specific cluster-default threading — would touch TWFE's user-visible result surface (coefficient-dict keys, cluster-label conventions, warning text) for near-zero residual benefit; waived on cost/benefit. | `twfe.py::fit`, `estimators.py::DifferenceInDifferences.fit`, `utils.py::build_fe_dummy_blocks` | #655 / 2026-07-10 | | **Survey TSL SE intentionally counts genuine-subpopulation zero-weight PSUs (matches R, NOT a bug).** Recorded as the REGISTRY § "Subpopulation Analysis" TSL-meat Note (Lumley 2004 §3.4 full-design domain convention; R `survey::svyrecvar(subset())` parity); regression-locked by `tests/test_survey.py::TestZeroWeightPsuConventionWaiver`. | `survey.py` (`_compute_stratified_psu_meat`) | PR-B / 2026-06-30 | -| DMLDiD survey/cluster support: Chang (2020) assumes i.i.d. sampling, so `DMLDiD.fit()` accepts no `survey_design=`/`cluster=` (bare TypeError). The per-unit augmented-score influence function IS unit-level clustering (REGISTRY DMLDiD M-080 Note), so the 4.0 auto-cluster default flip is inert; what is missing is the COARSER-than-unit CR1 surface CS exposes and any design-based (survey) variance. The CS mixins were verified safe with survey keys absent (every read is `.get` with panel fallback), so the extension is additive: survey keys in `_precompute` AND `_precompute_rcs` (both design lanes) + kit label threading (`_BOOTSTRAP_LABEL` reaches the survey-bootstrap <2-PSU warning) + a clustered-score variance derivation consistent with the cross-fitting — per-UNIT influence functions on the panel lane, per-OBSERVATION on the declared-RCS lane (where clustering would group rows into design clusters). Needs a methodology decision on clustered cross-fitting (cluster-level folds vs unit folds with clustered scores). Note RCS data is typically survey data (BRFSS/ACS/CPS), so the RCS lane raises the priority of the survey half. | `diff_diff/dml_did.py` | DML PR-B1 | Low | diff --git a/README.md b/README.md index db347ed3b..c85a6da35 100644 --- a/README.md +++ b/README.md @@ -121,7 +121,7 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`. - [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 via `method="qdid"`; bootstrap inference; R qte parity. Alias `CiC` - [LWDiD](https://diff-diff.readthedocs.io/en/stable/api/lwdid.html) - Lee & Wooldridge (2025, 2026) rolling-transformation DiD: unit-specific demean/detrend converts panel to cross-section, staggered adoption, `estimation_method` in `reg`/`ipw`/`dr`/`psm` (the papers' RA/IPW/IPWRA plus propensity-score matching), exact small-N inference on the classical collapsed regression -- [DMLDiD](https://diff-diff.readthedocs.io/en/stable/api/dml_did.html) - Chang (2020) double/debiased machine learning DiD: staggered ATT(g,t) with cross-fitted ML nuisance learners (DML2) and Neyman-orthogonal scores, for flexible/high-dimensional covariate adjustment under conditional parallel trends; panel or declared repeated cross sections (`panel=False`) +- [DMLDiD](https://diff-diff.readthedocs.io/en/stable/api/dml_did.html) - Chang (2020) double/debiased machine learning DiD: staggered ATT(g,t) with cross-fitted ML nuisance learners (DML2) and Neyman-orthogonal scores, for flexible/high-dimensional covariate adjustment under conditional parallel trends; panel or declared repeated cross sections (`panel=False`); survey/cluster support on both lanes - [BaconDecomposition](https://diff-diff.readthedocs.io/en/stable/api/bacon.html) - Goodman-Bacon (2021) decomposition for diagnosing TWFE bias in staggered settings ## Diagnostics & Sensitivity diff --git a/TODO.md b/TODO.md index 69764c3f5..0a89c6d02 100644 --- a/TODO.md +++ b/TODO.md @@ -37,7 +37,8 @@ Related tracking surfaces: | `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 | | Opt-in singleton-group pruning for TwoWayFixedEffects (static + event-study mode; reghdfe parity): singleton units/periods are currently RETAINED class-wide - the within-demeaned row is zero so points are unchanged, but N/G/residual-df count it and CR1/finite-sample SEs shift (~0.41019 -> 0.40962 measured; REGISTRY "Deviation from R" Note, R5 review) - reghdfe iteratively drops singletons by default while fixest retains them (diff-diff matches fixest); an opt-in knob needs iterative unit+period pruning with consistent cluster/survey/replicate/Conley array subsetting and a default-flip decision protocol (moves published SEs) | `diff_diff/twfe.py`, `diff_diff/estimators.py`, `diff_diff/utils.py` | 3(a) R5 | Mid | Low | | Cohort-timing validation input for the simultaneous-adoption event-study family (TWFE `event_study=True` + MultiPeriodDiD through 3.9): an optional `first_treat=`/`cohort=` column so simultaneous adoption becomes checkable under the contract-valid time-invariant `D_i` indicator - today the staggered-adoption advisory derives timing from within-unit 0->1 transitions, so it can only fire on off-contract time-varying `D_it` input, and with valid `D_i` adoption timing is not observable in the inputs at all (REGISTRY "staggered-adoption detection limit" Notes, both sections); design questions: validate-only vs steering error, and interplay with the M-011 removal | `diff_diff/twfe.py`, `diff_diff/estimators.py` | 3(a) R2 | Mid | Medium | -| Bootstrapped ES REPLAY containers (CS and EfficientDiD) publish the analytical `survey_metadata.df_survey` provenance scalar beside percentile inference on survey fits — a cross-estimator provenance residual (shipped CS behaves identically; the per-row df/inference channels are correctly NaN, the scalar metadata field is the residual). Evaluate clearing/gating it on both estimators together (cross-surface twins) | `diff_diff/staggered_results.py`, `diff_diff/efficient_did_results.py`, `diff_diff/results_base.py` | EDiD-replay review | Quick | Low | +| Bootstrapped ES REPLAY containers (CS, DMLDiD and EfficientDiD) publish the analytical `survey_metadata.df_survey` provenance scalar beside percentile inference on survey fits — a cross-estimator provenance residual (shipped CS behaves identically; the per-row df/inference channels are correctly NaN, the scalar metadata field is the residual; DMLDiD joined with its 3.11 survey support). Evaluate clearing/gating it on all three estimators together (cross-surface twins) | `diff_diff/staggered_results.py`, `diff_diff/dml_did_results.py`, `diff_diff/efficient_did_results.py`, `diff_diff/results_base.py` | EDiD-replay review | Quick | Low | +| DMLDiD replicate-weight survey designs (currently fail closed with `NotImplementedError` pointing here): the IF-reweighting route (`compute_replicate_if_variance` on the augmented scores) is BETTER justified for DMLDiD than for CS — Neyman orthogonality makes the score first-order insensitive to nuisance perturbation — but needs per-cell replicate SE plumbing, the aggregate `_se_from_psi` replicate branch already activates on kit keys, and its own rejection/df test matrix (QR-rank df, `n_valid - 1` tightening) | `diff_diff/dml_did.py` | DML survey PR | Mid | Low | | ImputationDiD/TwoStageDiD `aggregate()` recompute levels on bootstrapped fits fail closed ('simple' relays since the M-027 per-level convergence; M-021/M-022); ImputationDiD's per-target psi machinery makes seeded replay tractable (the panel-backed kit retains everything the psi precompute reads), TwoStageDiD's per-level GMM scores are function-locals and would need retention | `diff_diff/imputation_results.py`, `diff_diff/two_stage_results.py`, `diff_diff/aggregation.py` | 2(b) PR-3b | Mid | Low | | ContinuousDiD `aggregate('event_study')` on bootstrapped fits fails closed (M-025); a seeded post-fit bootstrap-ES replay is tractable - the multiplier draws are seeded (`np.random.default_rng(self.seed)`) - but needs the FULL per-cell `_bootstrap_info` (bread/ee_treated/Psi_eval/dPsi_*/beta_pred) the pruned kit deliberately drops, so shipping it means a kit-payload change with its own memory contract | `diff_diff/continuous_did_aggregation.py`, `diff_diff/continuous_did_results.py` | 2(b) PR-3c | Mid | Low | | EfficientDiD, ImputationDiD, ContinuousDiD and HeterogeneousAdoptionDiD are the outstanding M-092 event-study df-provenance holes: the container's per-row df is all-NaN even on survey fits where a finite `_survey_df` governed the p-values (the container-level scalar `df_survey` IS exposed - the hole is the PER-ROW column only; no event_study_df/df_inference field; pre-existing, NOT a regression of the M-023 PR - today's builder output is identical). The kits now retain the scalar (ImputationDiD's since 2(b) PR-3b, ContinuousDiD's since 2(b) PR-3c - same shape: scalar `df_survey` exposed, per-row column all-NaN, identical to each fit-time surface); threading it into the per-row channel is a contained follow-up | `diff_diff/efficient_did_results.py`, `diff_diff/imputation_results.py`, `diff_diff/continuous_did_results.py`, `diff_diff/results_base.py` | 2(b) PR-3a | Quick | Low | diff --git a/diff_diff/_dr_scores.py b/diff_diff/_dr_scores.py index 806e9a226..a9324baa3 100644 --- a/diff_diff/_dr_scores.py +++ b/diff_diff/_dr_scores.py @@ -35,7 +35,7 @@ difference-in-differences models. The Econometrics Journal, 23(2), 177-191. """ -from typing import Tuple +from typing import Optional, Tuple import numpy as np @@ -343,6 +343,33 @@ def chang_rcs_score( return weight * ((T - lam_hat) * y - m2_hat) +def _validate_slope_weights( + weights: Optional[np.ndarray], n: int, context: str +) -> Optional[np.ndarray]: + """Validate optional survey weights for the lambda-slope family. + + Returns the coerced float64 weight vector, or ``None`` when no weights + were supplied. Shared by the public wrappers AND the internal + ``_with_slope`` entry (DMLDiD's RCS cell calls the latter directly, so + validating only the public surface would leave the estimator's own + call path unvalidated). + """ + if weights is None: + return None + w = np.asarray(weights, dtype=np.float64) + if w.ndim != 1: + raise ValueError(f"{context}: weights must be 1-dimensional") + if w.shape[0] != n: + raise ValueError(f"{context}: weights has length {w.shape[0]}, expected {n}") + if not np.all(np.isfinite(w)): + raise ValueError(f"{context}: weights contains non-finite values") + if np.any(w < 0): + raise ValueError(f"{context}: weights must be non-negative") + if not np.sum(w) > 0: + raise ValueError(f"{context}: weights must have a positive sum") + return w + + def chang_rcs_lambda_slope( y: np.ndarray, D: np.ndarray, @@ -351,6 +378,8 @@ def chang_rcs_lambda_slope( ps: np.ndarray, p_hat: float, lam_hat: float, + *, + weights: Optional[np.ndarray] = None, ) -> float: """Chang (2020) Case 2 lambda-slope estimator ``G_2lambda``. @@ -373,7 +402,8 @@ def chang_rcs_lambda_slope( y, D, T, m2_hat, ps = _validate_chang_rcs_inputs( y, D, T, m2_hat, ps, p_hat, lam_hat, "chang_rcs_lambda_slope" ) - return _chang_rcs_lambda_slope_validated(y, D, T, m2_hat, ps, p_hat, lam_hat) + weights = _validate_slope_weights(weights, y.shape[0], "chang_rcs_lambda_slope") + return _chang_rcs_lambda_slope_validated(y, D, T, m2_hat, ps, p_hat, lam_hat, weights=weights) def _chang_rcs_lambda_slope_validated( @@ -384,8 +414,11 @@ def _chang_rcs_lambda_slope_validated( ps: np.ndarray, p_hat: float, lam_hat: float, + *, + weights: Optional[np.ndarray] = None, ) -> float: - # Assumes inputs already coerced/validated by _validate_chang_rcs_inputs. + # Assumes inputs already coerced/validated by _validate_chang_rcs_inputs + # (and weights by _validate_slope_weights). odds = (D - ps) / (1.0 - ps) term1 = ( -((1.0 - 2.0 * lam_hat) / (lam_hat**2 * (1.0 - lam_hat) ** 2)) @@ -393,7 +426,9 @@ def _chang_rcs_lambda_slope_validated( * ((T - lam_hat) * y - m2_hat) ) term2 = -(y / (p_hat * lam_hat * (1.0 - lam_hat))) * odds - return float(np.mean(term1 + term2)) + if weights is None: + return float(np.mean(term1 + term2)) + return float(np.average(term1 + term2, weights=weights)) def chang_rcs_score_augmented( @@ -406,6 +441,8 @@ def chang_rcs_score_augmented( theta: float, p_hat: float, lam_hat: float, + *, + weights: Optional[np.ndarray] = None, ) -> np.ndarray: """Chang (2020) Case 2 augmented score with BOTH finite-dim corrections. @@ -425,7 +462,7 @@ def chang_rcs_score_augmented( omits the lambda term; see the committed characterization spike). """ return _chang_rcs_score_augmented_with_slope( - summand, D, T, y, m2_hat, ps, theta, p_hat, lam_hat + summand, D, T, y, m2_hat, ps, theta, p_hat, lam_hat, weights=weights )[0] @@ -439,6 +476,8 @@ def _chang_rcs_score_augmented_with_slope( theta: float, p_hat: float, lam_hat: float, + *, + weights: Optional[np.ndarray] = None, ) -> Tuple[np.ndarray, float]: """Internal variant returning ``(psi_bar, g2_lambda)``. @@ -460,5 +499,8 @@ def _chang_rcs_score_augmented_with_slope( y, D, T, m2_hat, ps = _validate_chang_rcs_inputs(y, D, T, m2_hat, ps, p_hat, lam_hat, context) if summand.shape[0] != y.shape[0]: raise ValueError(f"{context}: summand has length {summand.shape[0]}, expected {y.shape[0]}") - g2_lambda = _chang_rcs_lambda_slope_validated(y, D, T, m2_hat, ps, p_hat, lam_hat) + weights = _validate_slope_weights(weights, y.shape[0], context) + g2_lambda = _chang_rcs_lambda_slope_validated( + y, D, T, m2_hat, ps, p_hat, lam_hat, weights=weights + ) return summand - D * theta / p_hat + g2_lambda * (T - lam_hat), g2_lambda diff --git a/diff_diff/_reporting_helpers.py b/diff_diff/_reporting_helpers.py index 456d46c49..9c79d64a7 100644 --- a/diff_diff/_reporting_helpers.py +++ b/diff_diff/_reporting_helpers.py @@ -150,12 +150,37 @@ def describe_target_parameter(results: Any) -> Dict[str, Any]: } if name == "DMLDiDResults": + # Survey/cluster fits switch the aggregation weight source to cohort + # masses (survey masses on declared designs; all-ones cohort masses + # on bare cluster=). getattr-defaulted reads only: this function is + # exercised on attribute-less stubs by the exhaustiveness guard. + _has_design = ( + getattr(results, "survey_metadata", None) is not None + or getattr(results, "df_inference", None) is not None + ) if getattr(results, "panel", True) is False: + if _has_design: + weight_name = "survey-cohort-mass-weighted" + weight_clause = ( + "cell weights are SURVEY cohort masses (per-observation " + "design weights summed by cohort; the per-cell " + "complete-case ``n_treated`` is display-only; REGISTRY " + "DMLDiD survey Note)" + ) + else: + weight_name = "cohort-mass-weighted" + weight_clause = ( + "cell weights are FIXED cohort row masses (the CS-RCS " + "convention, WIF-consistent; the per-cell complete-case " + "``n_treated`` is display-only; REGISTRY DMLDiD RCS Note)" + ) + elif _has_design: weight_name = "cohort-mass-weighted" weight_clause = ( - "cell weights are FIXED cohort row masses (the CS-RCS " - "convention, WIF-consistent; the per-cell complete-case " - "``n_treated`` is display-only; REGISTRY DMLDiD RCS Note)" + "cell weights are cohort masses (survey masses on declared " + "designs; full-cohort counts on bare ``cluster=`` fits — " + "NOT the no-design per-cell complete-case ``n_treated``; " + "REGISTRY DMLDiD survey Note)" ) else: weight_name = "valid-treated-count-weighted" diff --git a/diff_diff/business_report.py b/diff_diff/business_report.py index 7d740da1b..a4d29027c 100644 --- a/diff_diff/business_report.py +++ b/diff_diff/business_report.py @@ -1576,6 +1576,15 @@ def _describe_assumption(estimator_name: str, results: Any = None) -> Dict[str, } if estimator_name == "DMLDiDResults": is_rcs = getattr(results, "panel", True) is False + _is_survey = getattr(results, "survey_metadata", None) is not None + _survey_caveat = ( + " Under the declared survey design, inference is the library's " + "design-based extension (Chang 2020 assumes i.i.d. sampling; " + "Theorem 2's coverage guarantee does not carry over — REGISTRY " + "DMLDiD survey Note)." + if _is_survey + else "" + ) if is_rcs: tail = ( "cross-fitting (DML2) removing own-observation overfitting. " @@ -1586,14 +1595,14 @@ def _describe_assumption(estimator_name: str, results: Any = None) -> Dict[str, "waves while outcomes are the period-specific potential " "outcomes — not data-checkable), and valid " "normal inference requires the nuisance learners to satisfy " - "Chang (2020)'s Case 2 rate conditions (Assumption 3.2(h))." + "Chang (2020)'s Case 2 rate conditions (Assumption 3.2(h))." + _survey_caveat ) else: tail = ( "cross-fitting (DML2) removing own-observation overfitting; " "valid normal inference additionally requires the nuisance " "learners to satisfy Chang (2020)'s rate conditions " - "(Assumption 3.1(f))." + "(Assumption 3.1(f))." + _survey_caveat ) return { "parallel_trends_variant": "conditional_on_covariates", diff --git a/diff_diff/diagnostic_report.py b/diff_diff/diagnostic_report.py index 367d3f89f..e79863122 100644 --- a/diff_diff/diagnostic_report.py +++ b/diff_diff/diagnostic_report.py @@ -125,8 +125,9 @@ # ``bacon`` refits BaconDecomposition from the panel independently of # the estimator (every staggered binary-treatment peer carries it); # ``design_effect`` follows the 17-of-18 convention (instance-gated on - # survey_metadata, which DMLDiD never sets -> informative skip). Only - # ``epv`` is excluded: no in-cell unpenalized-logit EPV diagnostics. + # survey_metadata: LIVE on declared survey_design= fits, informative + # skip otherwise). Only ``epv`` is excluded: no in-cell + # unpenalized-logit EPV diagnostics. "DMLDiDResults": frozenset( { "parallel_trends", diff --git a/diff_diff/dml_did.py b/diff_diff/dml_did.py index 0403b9917..7c3fe164d 100644 --- a/diff_diff/dml_did.py +++ b/diff_diff/dml_did.py @@ -18,10 +18,13 @@ ``DMLDiD`` writes the CallawaySantAnna per-(g,t) ``influence_func_info`` payload (per-sampling-unit entries ``psi_bar_i / n_cell`` — per unit on panel fits, per observation on RCS fits — so ``sqrt(sum(if**2))`` IS the -cell SE) and inherits the CS aggregation + multiplier-bootstrap mixins: -event study with sup-t bands, group/simple aggregation (plus total on -panel fits only; RCS fits fail ``total`` closed), and post-fit -``results.aggregate()`` with bootstrap replay. +cell SE on no-design fits; under ``survey_design=``/``cluster=`` the +payload is the weighted-IF analogue ``w_i * psi_bar_i / sum(w)`` and the +per-cell SE is the design-based CR1/weighted-IF variance instead) and +inherits the CS aggregation + multiplier-bootstrap mixins: event study +with sup-t bands, group/simple aggregation (plus total on panel +non-survey fits; RCS and declared-survey fits fail ``total`` closed), and +post-fit ``results.aggregate()`` with bootstrap replay. See docs/methodology/REGISTRY.md "DMLDiD" for equations, implementation Notes (global p-hat, D-stratified folds, pooled fold weighting, trimming, @@ -29,9 +32,10 @@ """ import decimal +import inspect import secrets import warnings -from typing import Any, ClassVar, Dict, Iterable, List, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union, cast import numpy as np import pandas as pd @@ -54,6 +58,7 @@ from diff_diff.linalg import _check_propensity_diagnostics from diff_diff.staggered import ( _build_aggregation_kit, + _cluster_robust_se_from_per_gt_if, _nan_gt_entry, ) from diff_diff.staggered import ( @@ -72,6 +77,9 @@ validate_pscore_trim, ) +if TYPE_CHECKING: + from diff_diff.survey import ResolvedSurveyDesign, SurveyDesign + __all__ = ["DMLDiD"] # Magnitude bound for time/first_treat labels (with anticipation headroom): @@ -108,6 +116,38 @@ def _validate_learner_spec(spec: Any, *, kind: str, param_name: str) -> None: validate_learner(spec, kind=kind, param_name=param_name) +def _validate_learner_sample_weight_support(spec: Any, param_name: str) -> None: + """Reject a user learner whose ``fit`` cannot take ``sample_weight``. + + Declared-survey fits pass ``sample_weight`` into ``cross_fit_predict``, + which forwards it BY KEYWORD (``fit_kwargs = {"sample_weight": w_fit}``) + and deliberately propagates the learner's ``TypeError`` — so a learner + whose ``fit`` has neither a keyword-addressable ``sample_weight`` + parameter (POSITIONAL_OR_KEYWORD or KEYWORD_ONLY; POSITIONAL_ONLY does + not qualify) nor ``**kwargs`` would hard-crash mid-fit. Raises + ``TypeError`` up front instead (the ``validate_learner`` convention for + object-capability failures). + """ + try: + sig = inspect.signature(spec.fit) + except (TypeError, ValueError): # pragma: no cover - exotic callables + return # cannot introspect; let cross_fit_predict surface any error + for param in sig.parameters.values(): + if param.kind is inspect.Parameter.VAR_KEYWORD: + return + if param.name == "sample_weight" and param.kind in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ): + return + raise TypeError( + f"survey_design= requires learners whose fit() accepts sample_weight " + f"by keyword; the {param_name} object {type(spec).__name__!r} does not. " + "Add a sample_weight parameter (or **kwargs) to its fit(), or use a " + "library-native learner name." + ) + + def _raw_label_is_infinite(value: Any) -> bool: """True iff the RAW label element is itself infinite (either sign). @@ -235,9 +275,14 @@ class DMLDiD(CallawaySantAnnaBootstrapMixin, CallawaySantAnnaAggregationMixin, B D-stratified and the outcome nuisance is the outcome-change regression ``E[dY | X, D=0]``; on RCS fits the folds are D x T stratified and the nuisance is the level regression ``E[(T - lam)Y | X, D=0]`` with the - lambda-corrected variance. Aggregation (event study, group, simple; - plus total on panel fits — RCS fits fail ``total`` closed) is POST-FIT - via ``results.aggregate()``. + lambda-corrected variance. (Under a survey/cluster design whose PSU is + strictly coarser than the sampling unit, folds are PSU-COHESIVE + instead of stratified, and the per-cell SE is design-based rather than + the plug-in ``sqrt(mean(psi_bar**2)/n)``; declared designs also weight + the moments — REGISTRY DMLDiD survey Notes.) Aggregation (event study, + group, simple; plus total on panel non-survey fits — RCS and + declared-survey fits fail ``total`` closed) is POST-FIT via + ``results.aggregate()``. Parameters ---------- @@ -265,7 +310,10 @@ class DMLDiD(CallawaySantAnnaBootstrapMixin, CallawaySantAnnaAggregationMixin, B Root seed for the per-cell fold draws and the bootstrap. With ``seed=None``, POINT ESTIMATES vary across fits — cross-fitting draws random folds; set ``seed`` for reproducible results with the - library's deterministic built-in learners. A user-supplied + library's deterministic built-in learners. The same seed yields + DIFFERENT folds with vs without a coarser-than-unit PSU design + (PSU-cohesive folds consume the RNG differently than stratified + folds — a config change, not a reproducibility break). A user-supplied STOCHASTIC learner object must additionally be seeded by the user (e.g. sklearn ``random_state``) — ``cross_fit_predict`` deep-copies the learner template where copyable but never seeds its internal @@ -296,9 +344,19 @@ class DMLDiD(CallawaySantAnnaBootstrapMixin, CallawaySantAnnaAggregationMixin, B and treatment effects are expected, not violations), which is not data-checkable — ``fit()`` warns. ``aggregate('total')`` is unavailable on RCS fits (fails closed, the library-wide RC convention). + cluster : str, optional + Column name for COARSER-than-unit clustering. The column is + synthesized into a ``SurveyDesign(psu=cluster)`` (or injected as + the PSU of a declared PSU-less design): it drives PSU-cohesive + cross-fitting folds (when strictly coarser than the sampling + unit), the per-cell cluster-robust SE, the bootstrap draw + structure, and ``df_inference`` — while the moment kernels stay + UNWEIGHTED and no ``sample_weight`` reaches the learners. When + ``survey_design=`` also supplies a PSU, the design's PSU wins + (warning on differing partitions). """ - _BOOTSTRAP_LABEL: ClassVar[str] = "DMLDiD" + _BOOTSTRAP_LABEL: str = "DMLDiD" def __init__( self, @@ -315,6 +373,7 @@ def __init__( cband: bool = True, pscore_trim: float = 0.01, panel: bool = True, + cluster: Optional[str] = None, ) -> None: # Raw assignment, then ONE shared validator (also re-run at the top # of fit() as the direct-mutation defense) validates and normalizes @@ -332,6 +391,7 @@ def __init__( self.cband = cband self.pscore_trim = pscore_trim self.panel = panel + self.cluster = cluster self._revalidate_config() # Fitted-state lifecycle (house convention; not inherited under this MRO). @@ -397,6 +457,11 @@ def _revalidate_config(self) -> None: "would silently select the panel lane" ) self.panel = bool(self.panel) + if self.cluster is not None and not isinstance(self.cluster, str): + raise ValueError( + f"cluster must be a column name (str) or None, got " + f"{self.cluster!r} (type {type(self.cluster).__name__})" + ) # ------------------------------------------------------------------ # Input validation @@ -866,6 +931,7 @@ def _precompute( time: str, first_treat: str, covariates: List[str], + resolved_survey: Optional["ResolvedSurveyDesign"] = None, ) -> Dict[str, Any]: # CS groupby form — NOT sorted(unique): sorted() raises a bare # TypeError on mixed-type unit labels that CS accepts. @@ -895,9 +961,12 @@ def _precompute( n_units = len(all_units) # obs_per_unit OMITTED deliberately: only the RC path sets it, and a # non-None value divides the aggregation WIF (would silently shrink - # SEs on a panel). Survey keys likewise omitted (mixin reads are - # .get with panel fallbacks). - return { + # SEs on a panel). Survey keys (CS names: survey_weights / + # resolved_survey / resolved_survey_unit / df_survey) are set from + # resolved_survey when a design (declared or cluster-synthesized) is + # in play; absent otherwise (mixin reads are .get with panel + # fallbacks). + out: Dict[str, Any] = { "all_units": all_units, "unit_to_idx": unit_to_idx, "unit_cohorts": unit_cohorts, @@ -913,6 +982,24 @@ def _precompute( "canonical_size": n_units, "n_units": n_units, } + if resolved_survey is not None: + from diff_diff.survey import collapse_survey_to_unit_level + + survey_weights_unit = ( + pd.Series(resolved_survey.weights, index=df.index) + .groupby(df[unit]) + .first() + .reindex(all_units) + .to_numpy() + ) + resolved_survey_unit = collapse_survey_to_unit_level( + resolved_survey, df, unit, all_units + ) + out["survey_weights"] = survey_weights_unit + out["resolved_survey"] = resolved_survey + out["resolved_survey_unit"] = resolved_survey_unit + out["df_survey"] = resolved_survey_unit.df_survey + return out def _precompute_rcs( self, @@ -922,6 +1009,7 @@ def _precompute_rcs( time: str, first_treat: str, covariates: List[str], + resolved_survey: Optional["ResolvedSurveyDesign"] = None, ) -> Dict[str, Any]: """Declared-RCS bookkeeping: rows ARE the sampling units. @@ -941,15 +1029,24 @@ def _precompute_rcs( Supplying the key would be a numeric no-op that routes lookups through float()-keyed dict access, where distinct int64 cohorts above 2**53 (admissible through 2**62 by the label pipeline) - collide. ``obs_per_unit`` and survey keys likewise omitted (true RCS - is one row per unit — a non-None obs_per_unit would divide the WIF). + collide. That rationale HOLDS under survey too: the survey branches + of ``fixed_cohort_agg_weights`` (native-key dict) and + ``_get_agg_cache`` (bincount) are collision-free, so survey cohort + masses flow without ``agg_cohort_masses``. ``obs_per_unit`` stays + omitted (true RCS is one row per unit — a non-None obs_per_unit + would divide the WIF). Survey keys: on RCS the rows ARE the + sampling units, so ``resolved_survey_unit`` is the per-obs design + itself (CS RCS precedent) and ``survey_weights`` are the per-obs + weights; ``rcs_cohort_masses`` stays unweighted (the per-cell + ``agg_weight`` becomes an inert fallback — survey masses win in + the aggregation cache). """ n_obs = len(df) unit_cohorts = df[first_treat].to_numpy() treatment_groups = sorted(g for g in df[first_treat].unique() if g > 0) time_periods = sorted(df[time].unique()) period_to_col = {t: i for i, t in enumerate(time_periods)} - return { + out: Dict[str, Any] = { "all_units": np.arange(n_obs), "unit_to_idx": None, "unit_cohorts": unit_cohorts, @@ -970,6 +1067,12 @@ def _precompute_rcs( g: int(np.count_nonzero(unit_cohorts == g)) for g in treatment_groups }, } + if resolved_survey is not None: + out["survey_weights"] = resolved_survey.weights.copy() + out["resolved_survey"] = resolved_survey + out["resolved_survey_unit"] = resolved_survey + out["df_survey"] = resolved_survey.df_survey + return out def _sanitize_learner_error(self, exc: BaseException) -> str: """Persisted error text for cross_fit_diagnostics / to_dict exports. @@ -1060,16 +1163,50 @@ def _compute_dml_gt( D_cell = treated_valid[cell_idx].astype(np.float64) dY_cell = dY[cell_idx] X_cell = X_base[cell_idx] - p_hat = n_treated / n_cell + + # Survey state (all .get with panel-safe fallbacks; see fit()). + weighted_moments = bool(precomputed.get("weighted_moments", False)) + use_psu_folds = bool(precomputed.get("use_psu_folds", False)) + resolved_survey_unit = precomputed.get("resolved_survey_unit") + survey_weights = precomputed.get("survey_weights") + df_survey = precomputed.get("df_survey") + + w_cell: Optional[np.ndarray] = None + if weighted_moments and survey_weights is not None: + w_cell = np.asarray(survey_weights, dtype=np.float64)[cell_idx] + w_treated_mass = float(np.sum(w_cell[D_cell == 1.0])) + w_control_mass = float(np.sum(w_cell[D_cell == 0.0])) + if w_treated_mass <= 0.0 or w_control_mass <= 0.0: + # A group with rows but zero survey mass: p_hat would leave + # (0, 1) and mislabel the cell as non_finite_score. + return ( + _nan_gt_entry( + n_treated=n_treated, + n_control=n_control, + skip_reason="zero_weight_mass", + ), + None, + None, + ) + p_hat = w_treated_mass / float(np.sum(w_cell)) + else: + p_hat = n_treated / n_cell # Empirical treated-share overlap diagnostic: every Chang bound # carries powers of 1/p_0 (paper review, few-treated edge case), and # the fitted-propensity check below cannot see a sparse EMPIRICAL # share when a learner returns non-extreme predictions. if min(p_hat, 1.0 - p_hat) < self.pscore_trim: + if w_cell is not None: + _share_detail = ( + f"weighted p_hat={p_hat:.4f} (treated mass " + f"{float(np.sum(w_cell[D_cell == 1.0])):.4g} / control mass " + f"{float(np.sum(w_cell[D_cell == 0.0])):.4g})" + ) + else: + _share_detail = f"p_hat={p_hat:.4f} ({n_treated} treated / {n_control} control)" warnings.warn( f"DMLDiD cell (g={g}, t={t}): empirical treated share " - f"p_hat={p_hat:.4f} ({n_treated} treated / {n_control} " - f"control) is extreme (min(p, 1-p) < pscore_trim=" + f"{_share_detail} is extreme (min(p, 1-p) < pscore_trim=" f"{self.pscore_trim}); the Chang score and variance scale " "with powers of 1/p_hat, so this cell's estimate may be " "unstable.", @@ -1080,7 +1217,11 @@ def _compute_dml_gt( # Per-cell fold draw: seeded from (root, (g_idx, t_idx)) — positions # in the sorted cohort/period rosters, invariant to other cells' # estimability. D-stratified (DoubleML parity; documented REGISTRY - # deviation from Chang's plain random partition). + # deviation from Chang's plain random partition) — except under a + # coarser-than-unit PSU design, where folds are PSU-cohesive + # (cluster-cohesive splitting replaces the stratification MECHANISM; + # the composition requirement is one period only on this lane and + # stays guarded by cross_fit_predict's degenerate checks). seed_seq = np.random.SeedSequence(entropy=root_entropy, spawn_key=(g_idx, t_idx)) rng = np.random.default_rng(seed_seq) diagnostics: Dict[str, Any] = { @@ -1089,10 +1230,19 @@ def _compute_dml_gt( "p_hat": float(p_hat), "n_clipped_ps": None, "fold_seed": {"entropy": int(root_entropy), "spawn_key": [int(g_idx), int(t_idx)]}, + "psu_folds": use_psu_folds, } try: - folds = assign_folds(n_cell, self.n_folds, rng=rng, stratify=D_cell) + if use_psu_folds and resolved_survey_unit is not None: + folds = assign_folds( + n_cell, + int(precomputed.get("effective_n_folds", self.n_folds)), + rng=rng, + cluster_ids=np.asarray(resolved_survey_unit.psu)[cell_idx], + ) + else: + folds = assign_folds(n_cell, self.n_folds, rng=rng, stratify=D_cell) except ValueError as exc: # Cell smaller than n_folds, or a singleton D-stratum (one # treated/control unit cannot be cross-fitted). @@ -1118,6 +1268,7 @@ def _compute_dml_gt( folds, predict_method="predict_proba", context_label=f"{context} propensity", + sample_weight=w_cell, ) or_res = cross_fit_predict( make_learner(self.outcome_learner, kind="regressor"), @@ -1127,6 +1278,7 @@ def _compute_dml_gt( predict_method="predict", fit_mask=(D_cell == 0.0), context_label=f"{context} outcome", + sample_weight=w_cell, ) except DegenerateFoldError as exc: diagnostics["skip_reason"] = "cross_fit_degenerate" @@ -1160,9 +1312,11 @@ def _compute_dml_gt( try: with np.errstate(over="ignore", invalid="ignore"): summand = chang_panel_score(dY_cell, D_cell, m_hat, ps, p_hat) - theta = float(np.mean(summand)) + if w_cell is not None: + theta = float(np.average(summand, weights=w_cell)) + else: + theta = float(np.mean(summand)) psi_bar = chang_panel_score_augmented(summand, D_cell, theta, p_hat) - se = float(np.sqrt(np.mean(psi_bar**2) / n_cell)) except ValueError as exc: diagnostics["skip_reason"] = "non_finite_score" diagnostics["error"] = self._sanitize_learner_error(exc) @@ -1175,7 +1329,7 @@ def _compute_dml_gt( None, diagnostics, ) - if not (np.isfinite(theta) and np.isfinite(se) and np.all(np.isfinite(psi_bar))): + if not (np.isfinite(theta) and np.all(np.isfinite(psi_bar))): diagnostics["skip_reason"] = "non_finite_score" return ( _nan_gt_entry( @@ -1187,24 +1341,21 @@ def _compute_dml_gt( diagnostics, ) - t_stat, p_value, conf_int = safe_inference(theta, se, alpha=self.alpha) - gt_entry = { - "effect": theta, - "se": se, - "t_stat": t_stat, - "p_value": p_value, - "conf_int": conf_int, - "n_treated": n_treated, - "n_control": n_control, - "skip_reason": None, - } - - # Payload: per-unit entries psi_bar_i / n_cell, so sqrt(sum(if^2)) - # IS the cell SE (the CS influence_func_info contract). + # Payload BEFORE the SE/inference block: the design-based per-cell SE + # consumes it. No-survey payload: per-unit entries psi_bar_i / + # n_cell, so sqrt(sum(if^2)) IS the cell SE (the CS + # influence_func_info contract). Weighted payload: w_i * psi_bar_i / + # sum(w) — the Hajek analogue (reduces to psi_bar/n at w == 1). n_units = precomputed["n_units"] inf_full = np.zeros(n_units) - inf_full[cell_idx] = psi_bar / n_cell + if w_cell is not None: + inf_full[cell_idx] = w_cell * psi_bar / float(np.sum(w_cell)) + else: + inf_full[cell_idx] = psi_bar / n_cell treated_idx = np.flatnonzero(treated_valid).astype(np.int64) + # Nonzero-derivation unchanged under weighting: a zero-weight control + # dropping from the payload is inert — the per-cell CR1 helper + # rebuilds the full-length psi vector over the complete design. control_idx = np.flatnonzero((inf_full != 0.0) & ~treated_valid).astype(np.int64) if_entry = { "treated_idx": treated_idx, @@ -1212,6 +1363,58 @@ def _compute_dml_gt( "treated_inf": inf_full[treated_idx], "control_inf": inf_full[control_idx], } + + # Per-cell SE. PSU designs (declared OR bare cluster=) route through + # the CS per-cell CR1 helper (3-valued contract: float = use it, + # NaN = unidentified clustered variance and MUST propagate, + # None = malformed -> fall back). Non-PSU survey designs use the + # weighted sqrt-sum (CS mirror: the full design enters aggregate SEs + # only); the no-survey branch is verbatim. + se: float + with np.errstate(over="ignore", invalid="ignore"): + if ( + resolved_survey_unit is not None + and getattr(resolved_survey_unit, "psu", None) is not None + ): + se_cr1 = _cluster_robust_se_from_per_gt_if(if_entry, resolved_survey_unit) + if se_cr1 is None: + se = float(np.sqrt(np.sum(inf_full**2))) + else: + se = float(se_cr1) + elif w_cell is not None: + se = float(np.sqrt(np.sum(inf_full**2))) + else: + se = float(np.sqrt(np.mean(psi_bar**2) / n_cell)) + # NOTE: `se` is deliberately OUTSIDE the non_finite_score gate on the + # design-based branches — a NaN from the CR1 helper is the + # unidentified-variance signal and must flow to safe_inference as a + # NaN-consistent inference tuple on a RETAINED cell. + if resolved_survey_unit is None and not np.isfinite(se): + diagnostics["skip_reason"] = "non_finite_score" + return ( + _nan_gt_entry( + n_treated=n_treated, + n_control=n_control, + skip_reason="non_finite_score", + ), + None, + diagnostics, + ) + + t_stat, p_value, conf_int = safe_inference(theta, se, alpha=self.alpha, df=df_survey) + gt_entry = { + "effect": theta, + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_int": conf_int, + "n_treated": n_treated, + "n_control": n_control, + "skip_reason": None, + } + if w_cell is not None: + # CS cell-payload convention: the TREATED survey mass. + gt_entry["survey_weight_sum"] = float(np.sum(w_cell[D_cell == 1.0])) return gt_entry, if_entry, diagnostics def _compute_dml_rcs_gt( @@ -1296,16 +1499,62 @@ def _compute_dml_rcs_gt( T_cell = at_t[cell_idx].astype(np.float64) y_cell = y_obs[cell_idx] X_cell = X_obs[cell_idx] + + # Survey state (all .get with panel-safe fallbacks; see fit()). + weighted_moments = bool(precomputed.get("weighted_moments", False)) + use_psu_folds = bool(precomputed.get("use_psu_folds", False)) + resolved_survey_unit = precomputed.get("resolved_survey_unit") + survey_weights = precomputed.get("survey_weights") + df_survey = precomputed.get("df_survey") + + w_cell: Optional[np.ndarray] = None + if weighted_moments and survey_weights is not None: + w_cell = np.asarray(survey_weights, dtype=np.float64)[cell_idx] + # FOUR-group WEIGHTED-mass guard: a group with rows but zero + # survey mass would push the weighted p_hat/lam_hat out of + # (0, 1) and mislabel the cell as non_finite_score. + _group_masses = ( + float(np.sum(w_cell[(D_cell == 1.0) & (T_cell == 1.0)])), + float(np.sum(w_cell[(D_cell == 1.0) & (T_cell == 0.0)])), + float(np.sum(w_cell[(D_cell == 0.0) & (T_cell == 1.0)])), + float(np.sum(w_cell[(D_cell == 0.0) & (T_cell == 0.0)])), + ) + if min(_group_masses) <= 0.0: + return ( + _nan_gt_entry( + n_treated=n_treated, + n_control=n_control, + skip_reason="zero_weight_mass", + ), + None, + None, + ) # Global-within-cell shares (REGISTRY convention: mirrors the Case 1 # global p_hat and DoubleML's d.mean()/t.mean()); strictly interior - # BY the four-group guard. - p_hat = float(D_cell.mean()) - lam_hat = float(T_cell.mean()) + # BY the four-group guard (weighted analogue under a declared + # design: Hajek shares, interior by the weighted-mass guard). + if w_cell is not None: + _w_total = float(np.sum(w_cell)) + p_hat = float(np.sum(w_cell * D_cell) / _w_total) + lam_hat = float(np.sum(w_cell * T_cell) / _w_total) + else: + p_hat = float(D_cell.mean()) + lam_hat = float(T_cell.mean()) if min(p_hat, 1.0 - p_hat) < self.pscore_trim: + if w_cell is not None: + _share_detail = ( + f"weighted p_hat={p_hat:.4f} over the pooled two-period rows " + f"(treated mass {float(np.sum(w_cell * D_cell)):.4g} / control " + f"mass {float(np.sum(w_cell * (1.0 - D_cell))):.4g})" + ) + else: + _share_detail = ( + f"p_hat={p_hat:.4f} over the pooled two-period rows " + f"({n_treated} treated / {n_control} control)" + ) warnings.warn( f"DMLDiD cell (g={g}, t={t}): empirical treated share " - f"p_hat={p_hat:.4f} over the pooled two-period rows " - f"({n_treated} treated / {n_control} control) is extreme " + f"{_share_detail} is extreme " f"(min(p, 1-p) < pscore_trim={self.pscore_trim}); the Chang " "score and variance scale with powers of 1/p_hat, so this " "cell's estimate may be unstable.", @@ -1327,7 +1576,10 @@ def _compute_dml_rcs_gt( # the FOUR D x T classes (DoubleML's d + 2t encoding — REGISTRY # deviation Note): every training complement then carries control # rows in both periods by construction, Chang's fold-composition - # requirement for fitting l_2 on I_kz^c. + # requirement for fitting l_2 on I_kz^c. Under a coarser-than-unit + # PSU design the folds are PSU-cohesive instead (the stratification + # MECHANISM cannot be cluster-constant), and the composition + # REQUIREMENT is preserved by the explicit guard below. seed_seq = np.random.SeedSequence(entropy=root_entropy, spawn_key=(g_idx, t_idx)) rng = np.random.default_rng(seed_seq) diagnostics: Dict[str, Any] = { @@ -1337,10 +1589,19 @@ def _compute_dml_rcs_gt( "lam_hat": float(lam_hat), "n_clipped_ps": None, "fold_seed": {"entropy": int(root_entropy), "spawn_key": [int(g_idx), int(t_idx)]}, + "psu_folds": use_psu_folds, } try: - folds = assign_folds(n_cell, self.n_folds, rng=rng, stratify=D_cell + 2.0 * T_cell) + if use_psu_folds and resolved_survey_unit is not None: + folds = assign_folds( + n_cell, + int(precomputed.get("effective_n_folds", self.n_folds)), + rng=rng, + cluster_ids=np.asarray(resolved_survey_unit.psu)[cell_idx], + ) + else: + folds = assign_folds(n_cell, self.n_folds, rng=rng, stratify=D_cell + 2.0 * T_cell) except ValueError as exc: # Cell smaller than n_folds, or a singleton D x T stratum (e.g. # ONE treated row in the base period cannot be cross-fitted). @@ -1358,6 +1619,29 @@ def _compute_dml_rcs_gt( context = f"DMLDiD (g={g}, t={t})" try: + # Chang's I_kz^c fold-composition requirement, checked EXPLICITLY + # whenever the D x T stratification no longer guarantees it (PSU + # folds) or per-row zero weights could hollow a period out even + # inside stratified folds (declared survey): every training + # complement needs positive-(weight-)mass control rows in BOTH + # periods, else the l_2 regression on r = (T - lam)y is + # finite-but-invalid. Raised here (inside this try) so the + # existing DegenerateFoldError handler skips the cell. + if use_psu_folds or weighted_moments: + _guard_w = w_cell if w_cell is not None else np.ones(n_cell) + for _k in range(folds.n_folds): + _complement = folds.fold_ids != _k + _ctrl = _complement & (D_cell == 0.0) + if ( + float(np.sum(_guard_w[_ctrl & (T_cell == 1.0)])) <= 0.0 + or float(np.sum(_guard_w[_ctrl & (T_cell == 0.0)])) <= 0.0 + ): + raise DegenerateFoldError( + f"{context} outcome: fold {_k}'s training " + "complement lacks positive-weight control rows in " + "both periods (Chang's I_kz^c fold-composition " + "requirement)" + ) with np.errstate(over="ignore", invalid="ignore"): ps_res = cross_fit_predict( make_learner(self.propensity_learner, kind="classifier"), @@ -1366,6 +1650,7 @@ def _compute_dml_rcs_gt( folds, predict_method="predict_proba", context_label=f"{context} propensity", + sample_weight=w_cell, ) r_cell = (T_cell - lam_hat) * y_cell or_res = cross_fit_predict( @@ -1376,6 +1661,7 @@ def _compute_dml_rcs_gt( predict_method="predict", fit_mask=(D_cell == 0.0), context_label=f"{context} outcome", + sample_weight=w_cell, ) except DegenerateFoldError as exc: diagnostics["skip_reason"] = "cross_fit_degenerate" @@ -1409,11 +1695,22 @@ def _compute_dml_rcs_gt( try: with np.errstate(over="ignore", invalid="ignore"): summand = chang_rcs_score(y_cell, D_cell, T_cell, m2_hat, ps, p_hat, lam_hat) - theta = float(np.mean(summand)) + if w_cell is not None: + theta = float(np.average(summand, weights=w_cell)) + else: + theta = float(np.mean(summand)) psi_bar, g2_lambda = _chang_rcs_score_augmented_with_slope( - summand, D_cell, T_cell, y_cell, m2_hat, ps, theta, p_hat, lam_hat + summand, + D_cell, + T_cell, + y_cell, + m2_hat, + ps, + theta, + p_hat, + lam_hat, + weights=w_cell, ) - se = float(np.sqrt(np.mean(psi_bar**2) / n_cell)) except ValueError as exc: diagnostics["skip_reason"] = "non_finite_score" diagnostics["error"] = self._sanitize_learner_error(exc) @@ -1426,12 +1723,7 @@ def _compute_dml_rcs_gt( None, diagnostics, ) - if not ( - np.isfinite(theta) - and np.isfinite(se) - and np.isfinite(g2_lambda) - and np.all(np.isfinite(psi_bar)) - ): + if not (np.isfinite(theta) and np.isfinite(g2_lambda) and np.all(np.isfinite(psi_bar))): diagnostics["skip_reason"] = "non_finite_score" return ( _nan_gt_entry( @@ -1444,7 +1736,61 @@ def _compute_dml_rcs_gt( ) diagnostics["g2_lambda"] = float(g2_lambda) - t_stat, p_value, conf_int = safe_inference(theta, se, alpha=self.alpha) + # Payload BEFORE the SE/inference block (the design-based per-cell + # SE consumes it): per-OBS entries psi_bar_i / n_cell over BOTH + # periods' rows, so sqrt(sum(if^2)) IS the cell SE on the no-survey + # path; weighted payload is the Hajek analogue w_i * psi_bar_i / + # sum(w). cell_idx is a flatnonzero of one mask => strictly + # increasing and duplicate-free, and the D-partition keeps + # treated_idx/control_idx disjoint (the fancy-+= scatter contract in + # the aggregation layer). + n_units = precomputed["n_units"] + inf_full = np.zeros(n_units) + if w_cell is not None: + inf_full[cell_idx] = w_cell * psi_bar / float(np.sum(w_cell)) + else: + inf_full[cell_idx] = psi_bar / n_cell + treated_idx = cell_idx[D_cell == 1.0].astype(np.int64) + control_idx = cell_idx[D_cell == 0.0].astype(np.int64) + if_entry = { + "treated_idx": treated_idx, + "control_idx": control_idx, + "treated_inf": inf_full[treated_idx], + "control_inf": inf_full[control_idx], + } + + # Per-cell SE (same dispatch as the panel cell; see the comment + # there): PSU designs -> CS per-cell CR1 helper (NaN propagates as + # the deliberate unidentified-variance signal on a RETAINED cell); + # non-PSU survey -> weighted sqrt-sum; no-survey verbatim. + se: float + with np.errstate(over="ignore", invalid="ignore"): + if ( + resolved_survey_unit is not None + and getattr(resolved_survey_unit, "psu", None) is not None + ): + se_cr1 = _cluster_robust_se_from_per_gt_if(if_entry, resolved_survey_unit) + if se_cr1 is None: + se = float(np.sqrt(np.sum(inf_full**2))) + else: + se = float(se_cr1) + elif w_cell is not None: + se = float(np.sqrt(np.sum(inf_full**2))) + else: + se = float(np.sqrt(np.mean(psi_bar**2) / n_cell)) + if resolved_survey_unit is None and not np.isfinite(se): + diagnostics["skip_reason"] = "non_finite_score" + return ( + _nan_gt_entry( + n_treated=n_treated, + n_control=n_control, + skip_reason="non_finite_score", + ), + None, + diagnostics, + ) + + t_stat, p_value, conf_int = safe_inference(theta, se, alpha=self.alpha, df=df_survey) gt_entry = { "effect": theta, "se": se, @@ -1453,29 +1799,17 @@ def _compute_dml_rcs_gt( "conf_int": conf_int, # DISPLAY counts: pooled two-period valid rows. Aggregation # weights come from agg_weight below (fixed cohort row mass, the - # CS-RCS convention) — never from these counts. + # CS-RCS convention; an inert fallback under survey — the + # aggregation cache's survey cohort masses win) — never from + # these counts. "n_treated": n_treated, "n_control": n_control, "skip_reason": None, "agg_weight": precomputed["rcs_cohort_masses"][g], } - - # Payload: per-OBS entries psi_bar_i / n_cell over BOTH periods' - # rows, so sqrt(sum(if^2)) IS the cell SE. cell_idx is a flatnonzero - # of one mask => strictly increasing and duplicate-free, and the - # D-partition keeps treated_idx/control_idx disjoint (the fancy-+= - # scatter contract in the aggregation layer). - n_units = precomputed["n_units"] - inf_full = np.zeros(n_units) - inf_full[cell_idx] = psi_bar / n_cell - treated_idx = cell_idx[D_cell == 1.0].astype(np.int64) - control_idx = cell_idx[D_cell == 0.0].astype(np.int64) - if_entry = { - "treated_idx": treated_idx, - "control_idx": control_idx, - "treated_inf": inf_full[treated_idx], - "control_inf": inf_full[control_idx], - } + if w_cell is not None: + # CS RCS cell-payload convention: the treated-at-period-t mass. + gt_entry["survey_weight_sum"] = float(np.sum(w_cell[(D_cell == 1.0) & (T_cell == 1.0)])) return gt_entry, if_entry, diagnostics # ------------------------------------------------------------------ @@ -1490,17 +1824,170 @@ def fit( time: str, first_treat: str, covariates: Optional[Iterable[str]] = None, + survey_design: Optional["SurveyDesign"] = None, ) -> DMLDiDResults: - """Estimate staggered ATT(g,t) via per-cell cross-fitted Chang scores.""" + """Estimate staggered ATT(g,t) via per-cell cross-fitted Chang scores. + + Parameters + ---------- + survey_design : SurveyDesign, optional + Complex survey design (pweight-only; full-design TSL — + weights/strata/PSU/FPC). Declared designs weight the moment + kernels (Hajek p-hat/lambda-hat/theta), pass ``sample_weight`` + into the nuisance learners (user learner objects must accept + ``sample_weight`` by keyword — a learner without it is rejected + up front), switch cross-fitting to PSU-cohesive folds when the + PSU is strictly coarser than the sampling unit, and route the + per-cell and aggregate variances through the design-based + kernels with ``df = n_PSU - n_strata`` t-inference. Survey + support is a documented library extension of Chang (2020), + which assumes i.i.d. sampling — Theorem 2's coverage claim + does not carry over (REGISTRY DMLDiD Notes). Replicate-weight + designs are not supported yet (fail closed; TODO.md). + """ df, covariates = self._validate_and_prepare( data, outcome, unit, time, first_treat, covariates ) + + # --- Survey/cluster resolution (CS transliteration, staggered.py) --- + from diff_diff.survey import ( + SurveyDesign, + _inject_cluster_as_psu, + _resolve_effective_cluster, + _resolve_survey_for_fit, + _validate_unit_constant_survey, + compute_survey_metadata, + ) + + ( + resolved_survey, + _survey_weights_raw, + survey_weight_type, + survey_metadata, + ) = _resolve_survey_for_fit(survey_design, data, "analytical") + + # Replicate designs fail closed FIRST (None-guarded): no replicate + # variance path exists for the cross-fitted scores yet. + if resolved_survey is not None and resolved_survey.uses_replicate_variance: + raise NotImplementedError( + "DMLDiD does not support replicate-weight survey designs yet " + "(tracked in TODO.md); use a full-design SurveyDesign " + "(weights/strata/psu/fpc) instead." + ) + + effective_survey_design = survey_design + cluster_ids_for_check: Optional[np.ndarray] = None + if self.cluster is not None: + if self.cluster not in data.columns: + raise ValueError(f"cluster column '{self.cluster}' not found in data") + _cluster_col = data[self.cluster] + if _cluster_col.isna().any(): + raise ValueError( + f"cluster column '{self.cluster}' contains missing values; " + "drop or impute them before fitting" + ) + cluster_ids_for_check = _cluster_col.to_numpy() + if resolved_survey is None: + # Bare cluster=: synthesize a PSU-only design. survey_metadata + # stays None DELIBERATELY (it is the declared-survey marker: + # the aggregate('total') gate and the weighted-kernel gate + # both key on it); df carried on Results.df_inference instead. + effective_survey_design = SurveyDesign(psu=self.cluster, weight_type="pweight") + ( + resolved_survey, + _survey_weights_raw, + survey_weight_type, + _synth_metadata, + ) = _resolve_survey_for_fit(effective_survey_design, data, "analytical") + elif resolved_survey.psu is None: + # Declared design without PSU: inject the cluster as the PSU. + # (resolved_survey non-None implies survey_design non-None on + # this branch - the bare-cluster synthesize path is above.) + assert survey_design is not None + from dataclasses import replace as _dc_replace + + effective_survey_design = _dc_replace(survey_design, psu=self.cluster) + resolved_survey = _inject_cluster_as_psu(resolved_survey, cluster_ids_for_check) + survey_metadata = compute_survey_metadata(resolved_survey, resolved_survey.weights) + else: + # Both supplied: the design's PSU wins (warn on differing + # partitions); the return value is intentionally unused. + _resolve_effective_cluster(resolved_survey, cluster_ids_for_check, self.cluster) + + if resolved_survey is not None: + if self.panel: + _validate_unit_constant_survey(data, unit, effective_survey_design) + if resolved_survey.weight_type != "pweight": + raise ValueError( + f"DMLDiD supports weight_type='pweight' survey designs only, " + f"got '{resolved_survey.weight_type}'" + ) + + weighted_moments = survey_design is not None + if weighted_moments: + # Learner capability gate: cross_fit_predict passes sample_weight + # BY KEYWORD, so a user learner without a keyword-addressable + # sample_weight (or **kwargs) would raise a raw TypeError mid-fit. + for spec, pname in ( + (self.propensity_learner, "propensity_learner"), + (self.outcome_learner, "outcome_learner"), + ): + if isinstance(spec, str): + continue # native learners all accept sample_weight + _validate_learner_sample_weight_support(spec, pname) + if self.panel: - precomputed = self._precompute(df, outcome, unit, time, first_treat, covariates) + precomputed = self._precompute( + df, outcome, unit, time, first_treat, covariates, resolved_survey=resolved_survey + ) cell_fn = self._compute_dml_gt else: - precomputed = self._precompute_rcs(df, outcome, unit, time, first_treat, covariates) + precomputed = self._precompute_rcs( + df, outcome, unit, time, first_treat, covariates, resolved_survey=resolved_survey + ) cell_fn = self._compute_dml_rcs_gt + + # Survey metadata reflects the estimation index space (units on the + # panel lane; the RCS lane's per-obs design is its own unit level). + resolved_survey_unit = precomputed.get("resolved_survey_unit") + if survey_metadata is not None and resolved_survey_unit is not None: + survey_metadata = compute_survey_metadata( + resolved_survey_unit, resolved_survey_unit.weights + ) + df_survey = precomputed.get("df_survey") + + # PSU-cohesive folds whenever the PSU is strictly coarser than the + # sampling unit and cluster-cohesive splitting is POSSIBLE + # (>= 2 PSUs). With 2 <= n_psu < n_folds the effective fold count is + # REDUCED to n_psu (warned) — never silently reverting to unit folds + # there, because with >= 2 PSUs the clustered variance is identified + # and would legitimize nuisances trained with within-PSU leakage. + # Only the single-PSU design falls back to stratified unit folds: + # cohesive splitting is impossible and the clustered variance is + # NaN either way (the <2-PSU contract; REGISTRY Note). + use_psu_folds = False + effective_n_folds = self.n_folds + if resolved_survey_unit is not None and resolved_survey_unit.psu is not None: + _n_psu_global = int(np.unique(resolved_survey_unit.psu).size) + if _n_psu_global < len(resolved_survey_unit.psu) and _n_psu_global >= 2: + use_psu_folds = True + if _n_psu_global < self.n_folds: + effective_n_folds = _n_psu_global + warnings.warn( + f"DMLDiD survey/cluster design has only {_n_psu_global} " + f"PSU(s), fewer than n_folds={self.n_folds}; " + f"cross-fitting uses {_n_psu_global} PSU-cohesive folds " + "instead (fold count reduced to preserve cluster " + "cohesion — unit-level folds would leak information " + "within PSUs while the clustered variance is " + "identified).", + UserWarning, + stacklevel=2, + ) + precomputed["use_psu_folds"] = use_psu_folds + precomputed["effective_n_folds"] = effective_n_folds + precomputed["weighted_moments"] = weighted_moments + treatment_groups = precomputed["treatment_groups"] time_periods = precomputed["time_periods"] observed_sorted = precomputed["observed_sorted"] @@ -1569,26 +2056,38 @@ def fit( # the sole user-visible trace). Accumulated per cell during the # estimation loop — no second O(n_units x n_cells) sweep. if dropped_units: + _has_design = precomputed.get("resolved_survey_unit") is not None if self.panel: + _weighting_note = ( + "(survey/cluster fits: point and aggregation weights both " + "use cohort masses — see REGISTRY.md)" + if _has_design + else "(point weights use per-cell valid counts; aggregation " + "cohort masses use full cohorts — see REGISTRY.md)" + ) warnings.warn( f"{len(dropped_units)} unit(s) were excluded from at least one " "(group, time) cell they would otherwise join, due to a " "missing or NON-FINITE outcome, a non-finite covariate at the " "cell's base period, or an outcome difference overflowing to " - "non-finite. DMLDiD estimates each cell on its complete cases " - "(point weights use per-cell valid counts; aggregation cohort " - "masses use full cohorts — see REGISTRY.md).", + f"non-finite. DMLDiD estimates each cell on its complete cases " + f"{_weighting_note}.", UserWarning, stacklevel=2, ) else: + _weighting_note = ( + "(survey fits: aggregation weights use survey cohort masses " + "— see REGISTRY.md)" + if _has_design + else "(aggregation weights use fixed cohort row masses — " "see REGISTRY.md)" + ) warnings.warn( f"{len(dropped_units)} observation(s) were excluded from a " "(group, time) cell they would otherwise join, due to a " "missing or NON-FINITE outcome or a non-finite covariate " "on the row. DMLDiD estimates each cell on its complete " - "cases (aggregation weights use fixed cohort row masses — " - "see REGISTRY.md).", + f"cases {_weighting_note}.", UserWarning, stacklevel=2, ) @@ -1607,6 +2106,20 @@ def fit( cohort_mass = float(np.count_nonzero(unit_cohorts == g)) if cohort_mass <= 0: continue + # Declared-survey fits gate materialization on the WEIGHTED + # cohort mass (CS precedent): a cohort with rows but zero + # survey mass must not materialize a reference cell. + _ref_weighted_mass: Optional[float] = None + if weighted_moments and precomputed.get("survey_weights") is not None: + _ref_weighted_mass = float( + np.sum( + np.asarray(precomputed["survey_weights"], dtype=np.float64)[ + unit_cohorts == g + ] + ) + ) + if _ref_weighted_mass <= 0.0: + continue ref_entry: Dict[str, Any] = { "effect": 0.0, "se": np.nan, @@ -1618,6 +2131,8 @@ def fit( "skip_reason": None, "is_reference": True, } + if _ref_weighted_mass is not None: + ref_entry["survey_weight_sum"] = _ref_weighted_mass if not self.panel: # Keep the RCS pg basis uniform: reference cells carry the # same fixed cohort row mass as estimated cells. @@ -1642,18 +2157,27 @@ def fit( reference_event_times = tuple(sorted(ref_event_times)) # Overall ATT (simple aggregation over post-treatment finite cells). - overall_att, overall_se, _ = self._aggregate_simple( + # overall_effective_df is non-None only when replicate variance + # dropped replicates — unreachable while replicate designs are + # rejected, but kept for structural parity with CS. + overall_att, overall_se, overall_effective_df = self._aggregate_simple( group_time_effects, influence_func_info, df, unit, precomputed ) + if overall_effective_df is not None: + df_survey = overall_effective_df overall_t_stat, overall_p_value, overall_conf_int = safe_inference( - overall_att, overall_se, alpha=self.alpha + overall_att, overall_se, alpha=self.alpha, df=df_survey ) # Optional multiplier bootstrap (keyword form; aggregate=None is the # supported skip-the-event-study-prep input — post-fit aggregate() # supplies its own level on replay). Adapted engine override block: # the fit-time event-study sub-block and the dead cband lines are - # dropped; every df=df_survey becomes df=None (plain normal theory). + # dropped; bootstrap overrides use df=None (plain normal theory) + # even on survey fits — the CS convention for percentile-bootstrap + # inference (staggered.py), NOT an oversight. The PSU-level survey + # bootstrap and the <2-PSU NaN contract activate inside the mixin + # via precomputed["resolved_survey_unit"]. bootstrap_results = None if self.n_bootstrap > 0: bootstrap_results = self._run_multiplier_bootstrap( @@ -1698,6 +2222,23 @@ def fit( n_treated_units = int(np.sum(unit_cohorts > 0)) n_control_units = int(np.sum(unit_cohorts == 0)) + # Survey/cluster results provenance (CS conventions): cluster_name + # follows PSU precedence (declared design PSU column > cluster=); + # df_inference is populated ONLY on the bare-cluster path (declared + # designs carry their df on survey_metadata.df_survey). + cluster_name: Optional[str] = None + n_clusters: Optional[int] = None + df_inference: Optional[float] = None + if resolved_survey_unit is not None: + if survey_design is not None and survey_design.psu is not None: + cluster_name = survey_design.psu + elif self.cluster is not None: + cluster_name = self.cluster + if resolved_survey_unit.psu is not None: + n_clusters = int(np.unique(resolved_survey_unit.psu).size) + if survey_metadata is None and df_survey is not None: + df_inference = float(df_survey) + results = DMLDiDResults( group_time_effects=group_time_effects, overall_att=overall_att, @@ -1728,18 +2269,27 @@ def fit( propensity_learner=_learner_spec_label(self.propensity_learner), outcome_learner=_learner_spec_label(self.outcome_learner), n_folds=self.n_folds, + effective_n_folds=(effective_n_folds if effective_n_folds != self.n_folds else None), cross_fit_diagnostics=cross_fit_diagnostics, seed=self.seed, n_bootstrap=self.n_bootstrap, bootstrap_weights=self.bootstrap_weights, cband=self.cband, panel=self.panel, + survey_metadata=survey_metadata, + cluster_name=cluster_name, + n_clusters=n_clusters, + df_inference=df_inference, ) results._aggregation_kit = _build_aggregation_kit( cast(Any, self), # duck-typed host contract (alpha/anticipation/cband) precomputed, influence_func_info, group_time_effects, + # Declared-survey marker: gates aggregate('total') closed on + # survey fits (bare cluster= keeps survey_metadata None and + # stays admitted). + is_survey_fit=survey_metadata is not None, bootstrap_results=bootstrap_results, ) self.results_ = results diff --git a/diff_diff/dml_did_results.py b/diff_diff/dml_did_results.py index 3707f8d79..3ed80785e 100644 --- a/diff_diff/dml_did_results.py +++ b/diff_diff/dml_did_results.py @@ -13,8 +13,11 @@ ``cluster=None`` IS the per-sampling-unit influence-function variance by definition (REGISTRY.md "IF-based variance estimators..." — the default), and DMLDiD's augmented-score SE ``sqrt(mean(psi_bar**2)/n)`` is exactly -that — per UNIT on panel fits, per OBSERVATION on repeated-cross-section -fits (rows are the sampling units there). +that on NO-DESIGN fits — per UNIT on panel fits, per OBSERVATION on +repeated-cross-section fits (rows are the sampling units there). Under a +``survey_design=``/``cluster=`` the per-cell SE is the design-based CR1 / +weighted-IF variance instead (the CS clustered-``hc1`` convention: +``SurveyDesign(psu=...)`` routed through the shared stratified-PSU meat). """ from dataclasses import dataclass, field @@ -59,12 +62,15 @@ class DMLDiDResults(CallawaySantAnnaResults): Inherits the full Callaway-Sant'Anna results surface — ``att``/``se`` aliases, ``to_dataframe``, post-fit ``aggregate()`` (simple / event_study - / group, plus total on panel fits; repeated-cross-section fits fail - ``total`` closed) with bootstrap replay — and adds the DML provenance - fields below. Every inherited CS-only field that DMLDiD never populates - (``epv_*``, ``pscore_fallback``, ``allow_unbalanced_panel``, - ``used_rc_on_unbalanced_panel``, ``cluster_name``, ``n_clusters``, - ``df_inference``, ``survey_metadata``, ``influence_functions``, + / group, plus total on panel non-survey fits; repeated-cross-section + AND declared-survey fits fail ``total`` closed) with bootstrap replay — + and adds the DML provenance fields below. ``cluster_name``/ + ``n_clusters``/``df_inference``/``survey_metadata`` are populated on + survey/``cluster=`` fits (CS conventions: ``survey_metadata`` marks a + DECLARED design; ``df_inference`` carries the bare-cluster df). Every + inherited CS-only field that DMLDiD never populates (``epv_*``, + ``pscore_fallback``, ``allow_unbalanced_panel``, + ``used_rc_on_unbalanced_panel``, ``influence_functions``, ``event_study_effects``/``event_study_vcov``/``event_study_vcov_index``/ ``event_study_df``) stays at its inherited default and is inert. ``group_effects`` stays ``None`` permanently — ``aggregate("group")`` @@ -84,7 +90,12 @@ class DMLDiDResults(CallawaySantAnnaResults): the object itself: result pickles must not retain arbitrary user objects. n_folds : int - Cross-fitting fold count K. + Cross-fitting fold count K as REQUESTED (the configuration value). + effective_n_folds : int, optional + The REALIZED fold count when it differs from ``n_folds`` — set only + when a coarse survey/cluster PSU design reduced K to the global PSU + count to preserve cluster cohesion (warned at fit); ``None`` + otherwise. cross_fit_diagnostics : dict, optional Per-``(g, t)`` cross-fit diagnostics: per-stage fold losses and fit counts, ``p_hat``, propensity clip count, and the fold-seed @@ -100,6 +111,7 @@ class DMLDiDResults(CallawaySantAnnaResults): propensity_learner: Any = "logit" outcome_learner: Any = "linear" n_folds: int = 5 + effective_n_folds: Optional[int] = None cross_fit_diagnostics: Optional[Dict[Any, Dict[str, Any]]] = field(default=None, repr=False) seed: Optional[int] = None n_bootstrap: int = 0 @@ -107,9 +119,12 @@ class DMLDiDResults(CallawaySantAnnaResults): cband: bool = True def __repr__(self) -> str: + folds = f"n_folds={self.n_folds}" + if self.effective_n_folds is not None: + folds += f" (effective {self.effective_n_folds})" return ( f"DMLDiDResults(att={self.overall_att:.6g}, se={self.overall_se:.6g}, " - f"n_cells={len(self.group_time_effects)}, n_folds={self.n_folds}, " + f"n_cells={len(self.group_time_effects)}, {folds}, " f"propensity_learner={self.propensity_learner!r}, " f"outcome_learner={self.outcome_learner!r})" ) @@ -157,9 +172,12 @@ def summary(self, alpha: Optional[float] = None) -> str: # statistic column is effect/SE (a standardized ratio, not the # p-value's source), so label the p column honestly. base = base.replace("t-stat", "z-stat").replace("P>|t|", "Boot. p") - else: - # Analytical DMLDiD inference is normal-theory throughout + elif self.survey_metadata is None and self.df_inference is None: + # Analytical NO-DESIGN inference is normal-theory throughout # (safe_inference with df=None): relabel the parent's t columns. + # Survey/bare-cluster fits use finite-df t inference + # (df=df_survey / df_inference) — the parent's t labels are + # CORRECT there and must not be relabeled. base = base.replace("t-stat", "z-stat").replace("P>|t|", "P>|z|") lines = base.split("\n") for i, line in enumerate(lines): @@ -183,6 +201,10 @@ def summary(self, alpha: Optional[float] = None) -> str: f"{'Propensity learner:':<30} {self.propensity_learner!r:>10}", f"{'Outcome learner:':<30} {self.outcome_learner!r:>10}", f"{'Cross-fitting folds (K):':<30} {self.n_folds:>10}", + ] + if self.effective_n_folds is not None: + header.append(f"{'Effective folds (PSU-reduced):':<30} {self.effective_n_folds:>10}") + header += [ # The seed line renders UNCONDITIONALLY: the fold draw moves # point estimates on every fit, bootstrap or not. With # seed=None the OS-drawn fold entropy (recorded per cell in @@ -215,6 +237,9 @@ def to_dict(self) -> Dict[str, Any]: result["propensity_learner"] = str(self.propensity_learner) result["outcome_learner"] = str(self.outcome_learner) result["n_folds"] = int(self.n_folds) + result["effective_n_folds"] = ( + None if self.effective_n_folds is None else int(self.effective_n_folds) + ) result["seed"] = None if self.seed is None else int(self.seed) result["n_bootstrap"] = int(self.n_bootstrap) result["bootstrap_weights"] = self.bootstrap_weights diff --git a/diff_diff/guides/llms-autonomous.txt b/diff_diff/guides/llms-autonomous.txt index e7cc3c58f..863870f7a 100644 --- a/diff_diff/guides/llms-autonomous.txt +++ b/diff_diff/guides/llms-autonomous.txt @@ -349,7 +349,7 @@ supported / out of scope; `warn` supported but with documented caveats; | `StackedDiD` | ✓ | ✓ | ✗ | ✗ | ✗ | ✓ | ✗ | ✗ | ✓ | | `WooldridgeDiD` (ETWFE) | ✓ | ✓ | ✗ | ✗ | ✗ | ✓ | ✗ | ✗ | ✓ | | `LWDiD` | ✓ | ✓ | ✗ | ✗ | ✓ | warn | partial | ✗ | ✓ | -| `DMLDiD` | ✓ | ✓ | ✗ | ✗ | partial | ✓ (REQUIRED) | ✗ | ✗ | ✗ (per-sampling-unit IF variance IS unit clustering; no coarser cluster=; panel=False = declared RCS) | +| `DMLDiD` | ✓ | ✓ | ✗ | ✗ | partial | ✓ (REQUIRED) | ✗ | ✗ | ✓ (pweight full-design TSL + coarser cluster=; PSU-cohesive folds; replicate designs rejected; panel=False = declared RCS) | | `EfficientDiD` | ✓ | ✓ | ✗ | ✗ | partial | ✓ | ✗ | ✗ | ✓ | | `SyntheticDiD` | ✓ | ✗ | ✗ | ✗ | ✓ | ✓ | ✓ | ✗ | partial | | `TROP` | ✓ | ✓ | ✗ | ✗ | ✗ | ✗ | ✓ | ✗ | partial | @@ -765,12 +765,12 @@ Explicit RCS support in this library: REGISTRY.md §CallawaySantAnna; use this variant on RCS data. - `DMLDiD(panel=False)` - Chang (2020) Case 2 declared-RCS mode (row-unique unit IDs; level-outcome orthogonal scores with the - lambda-corrected variance; REGISTRY.md §DMLDiD). CAVEATS: it carries - NO survey weights even though RCS data is typically survey data - (weighted RCS -> CallawaySantAnna(panel=False, survey_design=...)), - and NO cluster= - the generic cluster-on-the-unit-proxy instruction - below CANNOT be followed on DMLDiD (its per-observation IF variance - is the only inference surface); aggregate('total') fails closed. + lambda-corrected variance; REGISTRY.md §DMLDiD). Survey/cluster + support on both designs: `survey_design=` (pweight full-design TSL; + the weighted-lambda plug-in is a library extension - Theorem 2's + coverage claim does not carry over) and coarser-than-unit `cluster=` + (variance/folds only); replicate designs rejected; + aggregate('total') fails closed on RCS and declared-survey fits. - `TripleDifference` - DDD cross-sectional use cases are documented in `docs/choosing_estimator.rst`; the two-period DDD estimator does not require within-unit tracking when the third comparison axis diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index 5acaa51d7..88a14d784 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -1414,7 +1414,8 @@ DMLDiD( cband=True, # sup-t bands on the event-study replay pscore_trim=0.01, # clip fitted propensities (never drop) panel=True, # False = declared repeated cross sections (Chang Case 2) -).fit(data, outcome, unit, time, first_treat, covariates) -> DMLDiDResults + cluster=None, # coarser-than-unit clustering (variance/folds only) +).fit(data, outcome, unit, time, first_treat, covariates, survey_design=None) -> DMLDiDResults ``` Key contracts: @@ -1422,7 +1423,7 @@ Key contracts: - HonestDiD / PreTrendsPower consume the aggregate('event_study') container (admitted source); the native `compute_honest_did(results)` route raises with the container instruction. Varying-base containers warn; use base_period='universal' for clean Rambachan-Roth interpretation. - `SieveLearner(k_max=None, criterion="bic")` is the exported configurable learner (adaptive polynomial degree by IC); any sklearn-style estimator object also plugs in (seed stochastic learners yourself — the library seed pins folds, not learner internals). - Per-cell complete cases (one consolidated unbalanced-input warning); degenerate cells (fewer members than folds, singleton treated stratum, fail-closed learner) become NaN cells with machine-readable skip_reason and a consolidated warning; surviving cells still aggregate. -- No survey_design=/cluster= on EITHER design (Chang assumes i.i.d.; the per-sampling-unit IF variance IS unit-level clustering — per unit on panel, per observation on RCS). panel=False requires row-unique unit IDs and assumes stationary cross-sectional sampling (Assumption 2.3, warned, not data-checkable) — note RCS data is typically survey data (BRFSS/ACS/CPS) and DMLDiD carries NO survey weights; use CallawaySantAnna(panel=False, survey_design=...) for weighted RCS. Reproducibility: set seed (fold draws move point estimates). +- Survey/cluster support on BOTH designs. `survey_design=` (pweight-only, full-design TSL — weights/strata/PSU/FPC): weighted moment kernels (Hajek p-hat/lambda-hat/theta), sample_weight into the nuisance learners (user learner objects must accept sample_weight by keyword — rejected up front otherwise), PSU-cohesive cross-fitting folds when the PSU is strictly coarser than the sampling unit, design-based per-cell/aggregate variance with df = n_PSU - n_strata t-inference. A documented library EXTENSION of Chang's i.i.d. theory (Theorem 2's coverage claim does not carry over on the weighted-lambda RCS lane — REGISTRY DMLDiD Notes). Bare cluster= keeps the kernels unweighted (variance/folds/df only). Replicate-weight designs are rejected (tracked in TODO.md); aggregate('total') fails closed on declared-survey fits. panel=False requires row-unique unit IDs and assumes stationary cross-sectional sampling (Assumption 2.3, warned, not data-checkable). Reproducibility: set seed (fold draws move point estimates; PSU folds consume the RNG differently than stratified folds). ### TROP diff --git a/diff_diff/guides/llms-practitioner.txt b/diff_diff/guides/llms-practitioner.txt index 34048fca6..f715db9dc 100644 --- a/diff_diff/guides/llms-practitioner.txt +++ b/diff_diff/guides/llms-practitioner.txt @@ -229,8 +229,10 @@ Is treatment adoption staggered (multiple cohorts, different timing)? | | post-fit, set seed= for reproducible | | fold draws; panel=False = declared | | repeated cross sections (Case 2, -| | lambda-corrected variance; no survey -| | weights) +| | lambda-corrected variance); +| | survey_design=/cluster= supported on +| | both designs (pweight TSL; replicate +| | designs rejected) | \-- WooldridgeDiD (ETWFE) -- nonlinear outcomes (logit/Poisson) or saturated OLS | |-- NO, simple 2x2 design: diff --git a/diff_diff/guides/llms.txt b/diff_diff/guides/llms.txt index 9443ddea6..c621344b9 100644 --- a/diff_diff/guides/llms.txt +++ b/diff_diff/guides/llms.txt @@ -81,7 +81,7 @@ The site is organized into 5 sections, each with a landing page: - [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`. - [QDiD](https://diff-diff.readthedocs.io/en/stable/api/changes_in_changes.html): **Deprecated 3.9, removed 4.0 - use `ChangesInChanges(method="qdid")`.** Athey & Imbens (2006) quantile DiD comparison estimator (additive quantile-by-quantile DiD, matching R `qte::QDiD()` including its covariate branch via `covariates=`); same bootstrap machinery as ChangesInChanges. The paper recommends CiC over QDiD (scale-dependent model with testable restrictions; a non-monotonicity warning fires when violated - unconditional fits only, the covariate-path counterfactual quantile curve is monotone by construction). - [LWDiD](https://diff-diff.readthedocs.io/en/stable/api/lwdid.html): Lee & Wooldridge (2025, 2026) rolling-transformation DiD — unit-specific demean/detrend converts panel to cross-section, supports staggered adoption with flexible control groups. Signature: `LWDiD(rolling='demean', estimation_method='reg', vcov_type='hc1', cluster=None, control_group='not_yet_treated', alpha=0.05, n_bootstrap=0, seed=None, pscore_trim=0.01, n_neighbors=1, caliper=None, with_replacement=True, n_jobs=1).fit(data, outcome, unit, time, treatment, first_treat=None, covariates=None)`. `estimation_method` values: `reg` (papers' RA), `ipw`, `dr` (papers' IPWRA, doubly robust), `psm`; `vcov_type` values: `classical`/`hc1`/`hc2`/`hc3` for `reg`; `ipw`/`dr` accept `hc1` only (influence-function variance); `psm` accepts `hc1` as configuration only - PSM inference is unavailable (NaN) pending an Abadie-Imbens matching variance; cluster-robust inference via the constructor's `cluster=` column (hc1/CR1 only, not a `vcov_type` value; rejected for `psm`). Per-period effects: post-fit `results.aggregate('event_study')`. -- [DMLDiD](https://diff-diff.readthedocs.io/en/stable/api/dml_did.html): Chang (2020) double/debiased machine learning DiD — staggered ATT(g,t) with cross-fitted ML nuisances (DML2) and Neyman-orthogonal scores; covariates REQUIRED (conditional parallel trends). Signature: `DMLDiD(propensity_learner='logit', outcome_learner='linear', n_folds=5, control_group='never_treated', anticipation=0, alpha=0.05, n_bootstrap=0, bootstrap_weights=None, seed=None, base_period='varying', cband=True, pscore_trim=0.01, panel=True).fit(data, outcome, unit, time, first_treat, covariates)`. `panel=False` = declared repeated cross sections (Chang Case 2: level outcomes, row-unique unit IDs, lambda-corrected variance; no survey weights — weighted RCS belongs to CallawaySantAnna(panel=False, survey_design=...)). Learners: string names (`linear`/`ridge`/`sieve` regressors, `logit` classifier) or any object with fit/predict(_proba) (sklearn-compatible); `SieveLearner(k_max, criterion)` is exported for adaptive polynomial nuisances. Aggregation is POST-FIT: `results.aggregate('event_study'/'group'/'simple')`, plus `'total'` on panel fits (RCS fits fail 'total' closed); sup-t bands via bootstrap replay. With seed=None point estimates vary across fits (random folds); set seed for reproducibility. +- [DMLDiD](https://diff-diff.readthedocs.io/en/stable/api/dml_did.html): Chang (2020) double/debiased machine learning DiD — staggered ATT(g,t) with cross-fitted ML nuisances (DML2) and Neyman-orthogonal scores; covariates REQUIRED (conditional parallel trends). Signature: `DMLDiD(propensity_learner='logit', outcome_learner='linear', n_folds=5, control_group='never_treated', anticipation=0, alpha=0.05, n_bootstrap=0, bootstrap_weights=None, seed=None, base_period='varying', cband=True, pscore_trim=0.01, panel=True, cluster=None).fit(data, outcome, unit, time, first_treat, covariates, survey_design=None)`. `panel=False` = declared repeated cross sections (Chang Case 2: level outcomes, row-unique unit IDs, lambda-corrected variance). Survey/cluster support on BOTH designs: `survey_design=` (pweight full-design TSL — weighted moments, PSU-cohesive folds, design-based variance with t-inference; a library extension of Chang's i.i.d. theory) and coarser-than-unit `cluster=` (variance/folds only, kernels stay unweighted); replicate-weight designs are not supported yet. Learners: string names (`linear`/`ridge`/`sieve` regressors, `logit` classifier) or any object with fit/predict(_proba) (sklearn-compatible); `SieveLearner(k_max, criterion)` is exported for adaptive polynomial nuisances. Aggregation is POST-FIT: `results.aggregate('event_study'/'group'/'simple')`, plus `'total'` on panel non-survey fits (RCS and declared-survey fits fail 'total' closed); sup-t bands via bootstrap replay. With seed=None point estimates vary across fits (random folds); set seed for reproducibility. - [BaconDecomposition](https://diff-diff.readthedocs.io/en/stable/api/bacon.html): Goodman-Bacon (2021) decomposition for diagnosing TWFE bias in staggered settings ## Diagnostics and Sensitivity Analysis diff --git a/diff_diff/mmm.py b/diff_diff/mmm.py index 38c39a59a..2050f160f 100644 --- a/diff_diff/mmm.py +++ b/diff_diff/mmm.py @@ -124,7 +124,10 @@ "results.aggregate('total') instead - its single row is the " "estimator-owned total incremental outcome over the per-cell " "complete-case treated units and needs no scale. " - "PANEL fits (the default) always support the total route (bootstrapped fits replay); repeated-cross-section fits " + "PANEL fits support the total route (bootstrapped fits replay) except " + "declared-survey_design fits and bare-cluster fits whose kept cells " + "have incomplete treated support (both raise with the reason - the " + "same CS exclusions); repeated-cross-section fits " "(panel=False) fail aggregate('total') closed on EVERY adopter " "(CallawaySantAnna's RCS fits included) - estimator-owned totals need " "per-unit tracking; on RCS fits pass a caller-defined numeric scale, " diff --git a/diff_diff/practitioner.py b/diff_diff/practitioner.py index c0ae286fa..b0d018e4c 100644 --- a/diff_diff/practitioner.py +++ b/diff_diff/practitioner.py @@ -539,12 +539,25 @@ def _handle_dml_did(results: Any): "default linear specification." ), code=( - "# Refit with alternative nuisance learners:\n" + "# Refit with alternative nuisance learners (carry the fit's\n" + "# design forward so the comparison isolates the learner):\n" "alt = DMLDiD(outcome_learner='sieve', seed=0" + (", panel=False" if getattr(results, "panel", True) is False else "") + + ( + ", cluster=..." + if getattr(results, "survey_metadata", None) is None + and getattr(results, "df_inference", None) is not None + else "" + ) + ").fit(\n" " df, outcome=..., unit=..., time=..., first_treat=...,\n" - " covariates=[...])\n" + " covariates=[...]" + + ( + ",\n survey_design=... # the original fit's SurveyDesign" + if getattr(results, "survey_metadata", None) is not None + else "" + ) + + ")\n" "print(alt.att, results.att) # should be close" ), step_name="learner_sensitivity", diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index c67176a45..94ddd0f1f 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, ClassVar, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd @@ -583,7 +583,7 @@ class CallawaySantAnna( # 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" + _BOOTSTRAP_LABEL: str = "CallawaySantAnna" def __init__( self, @@ -5156,6 +5156,11 @@ def _build_aggregation_kit( for (g, t), data in group_time_effects.items() ) bookkeeping["is_survey_fit"] = bool(is_survey_fit) + # Producer identity for replayed-bootstrap warning attribution: the + # kit-replay host brands its <2-PSU warning with this label so a DMLDiD + # (or DDD) survey fit does not warn as "CallawaySantAnna" on post-fit + # aggregate(). Legacy kits without the key default at the read site. + bookkeeping["bootstrap_label"] = getattr(estimator, "_BOOTSTRAP_LABEL", "CallawaySantAnna") # Data minimization: the results object is picklable and users share # result artifacts, so the kit must not turn it into a carrier for raw diff --git a/diff_diff/staggered_bootstrap.py b/diff_diff/staggered_bootstrap.py index 62675c1a9..138c06fef 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, ClassVar, Dict, Iterator, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple import numpy as np @@ -145,13 +145,14 @@ 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] + # `self._BOOTSTRAP_LABEL` would otherwise be [attr-defined]. A PLAIN + # annotation (not ClassVar): estimator hosts assign it as class-level + # constant data (valid for an instance-variable declaration), while the + # kit-replay host (`staggered_results._KitBootstrapAggregator`) sets it + # per-instance from the kit's recorded label — a ClassVar here would make + # that instance assignment a "cannot assign to class variable via + # instance" [misc] error. + _BOOTSTRAP_LABEL: str n_bootstrap: int bootstrap_weights: str alpha: float diff --git a/diff_diff/staggered_results.py b/diff_diff/staggered_results.py index 3e8fa7103..709a7820d 100644 --- a/diff_diff/staggered_results.py +++ b/diff_diff/staggered_results.py @@ -69,16 +69,22 @@ class _KitBootstrapAggregator(CallawaySantAnnaBootstrapMixin, CallawaySantAnnaAg replayed bootstrap. """ - _BOOTSTRAP_LABEL = "CallawaySantAnna" - def __init__( - self, alpha: float, anticipation: int, n_bootstrap: int, bootstrap_weights: str + self, + alpha: float, + anticipation: int, + n_bootstrap: int, + bootstrap_weights: str, + bootstrap_label: str = "CallawaySantAnna", ) -> None: self.alpha = alpha self.anticipation = anticipation self.n_bootstrap = n_bootstrap self.bootstrap_weights = bootstrap_weights self.seed = None # unused — the replay injects the captured state + # Producer identity from the kit (legacy kits default to CS), so the + # replayed <2-PSU warning is branded with the fitting estimator. + self._BOOTSTRAP_LABEL = bootstrap_label @dataclass @@ -418,7 +424,11 @@ def _aggregate_compute( "extension install)." ) host = _KitBootstrapAggregator( - kit.alpha, kit.anticipation, spec.n_bootstrap, spec.weight_type + kit.alpha, + kit.anticipation, + spec.n_bootstrap, + spec.weight_type, + bootstrap_label=kit.bookkeeping.get("bootstrap_label", "CallawaySantAnna"), ) boot_replay = host._run_multiplier_bootstrap( group_time_effects=self.group_time_effects, diff --git a/diff_diff/staggered_triple_diff.py b/diff_diff/staggered_triple_diff.py index 977baafc6..34638e9d2 100644 --- a/diff_diff/staggered_triple_diff.py +++ b/diff_diff/staggered_triple_diff.py @@ -16,7 +16,7 @@ """ import warnings -from typing import TYPE_CHECKING, ClassVar, List, Optional +from typing import TYPE_CHECKING, List, Optional from diff_diff._base import BaseEstimator from diff_diff._staggered_triple_diff_engine import _StaggeredTripleDiffEngineMixin @@ -121,7 +121,7 @@ class StaggeredTripleDifference( # 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" + _BOOTSTRAP_LABEL: str = "StaggeredTripleDifference" def __init__( self, diff --git a/diff_diff/triple_diff.py b/diff_diff/triple_diff.py index 0ee82c674..67d76e32f 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, ClassVar, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import numpy as np import pandas as pd @@ -596,7 +596,7 @@ class TripleDifference( # 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" + _BOOTSTRAP_LABEL: str = "TripleDifference" _PARAM_ATTR_ALIASES = {"robust": "_robust_arg"} _DERIVED_CONFIG_ATTRS = ("robust",) diff --git a/docs/api/dml_did.rst b/docs/api/dml_did.rst index ace59d885..69d1ca8a0 100644 --- a/docs/api/dml_did.rst +++ b/docs/api/dml_did.rst @@ -131,9 +131,12 @@ Learner configuration (library defaults: ``"logit"``; ``"linear"``, ``"ridge"``, ``"sieve"``) or ANY object satisfying the duck-typed learner protocol — ``fit(X, y)`` plus ``predict`` (regressor) or ``predict_proba`` -(classifier). ``DMLDiD``'s cross-fitting is unweighted, so a -``sample_weight`` parameter on ``fit`` is not required (it is passed only -by consumers that actually use weighted cross-fitting). A user-constructed or scikit-learn +(classifier). On no-design and bare-``cluster=`` fits the cross-fitting +is unweighted, so a ``sample_weight`` parameter on ``fit`` is not +required there; a DECLARED ``survey_design=`` fit passes +``sample_weight`` into both nuisances, and a user learner whose ``fit`` +cannot take ``sample_weight`` by keyword is rejected up front with a +``TypeError``. A user-constructed or scikit-learn estimator object plugs in directly; string names select library defaults only. :class:`~diff_diff.SieveLearner` is the exported configurable learner (``DMLDiD(outcome_learner=SieveLearner(k_max=3))``). @@ -174,16 +177,25 @@ Restrictions while outcomes are the period-specific potential outcomes; warned at fit, not data-checkable). ``aggregate('total')`` is unavailable on RCS fits (fails closed). - NOTE: repeated-cross-section data is typically SURVEY data (BRFSS / ACS / - CPS), and ``DMLDiD`` carries NO survey weights — use - :class:`~diff_diff.CallawaySantAnna` ``(panel=False, survey_design=...)`` - for weighted RCS designs. -- **No survey/cluster support** — Chang (2020) assumes i.i.d. sampling; - ``fit()`` accepts no ``survey_design=`` or ``cluster=``. The - per-sampling-unit influence-function variance — per unit on panel fits, - per observation on repeated-cross-section fits — IS unit-level - clustering; coarser - clustering is a tracked follow-up (``DEFERRED.md``). + Repeated-cross-section data is typically SURVEY data (BRFSS / ACS / + CPS) — pass a pweight :class:`~diff_diff.SurveyDesign` via + ``survey_design=`` for weighted RCS. +- **Survey/cluster support (both lanes)** — ``survey_design=`` + (pweight-only, full-design TSL: weights/strata/PSU/FPC) weights the + moment kernels, passes ``sample_weight`` into the nuisance learners, + switches to PSU-cohesive cross-fitting folds when the PSU is strictly + coarser than the sampling unit, and uses design-based variance with + ``df = n_PSU - n_strata`` t-inference: PSU designs get the + cluster-robust survey kernel per cell, strata/FPC-only designs use the + weighted influence-function per-cell SE with the full design entering + the AGGREGATE variances (the CallawaySantAnna convention). This is a documented LIBRARY + EXTENSION of Chang (2020), which assumes i.i.d. sampling — on the + weighted-λ RCS lane Theorem 2's coverage claim does not carry over + (REGISTRY DMLDiD Notes). Bare ``cluster=`` (constructor) keeps the + kernels unweighted and affects folds, variance and df only. + Replicate-weight designs are not supported yet (fail closed; + ``TODO.md``); ``aggregate('total')`` also fails closed on + declared-survey fits. - **Propensity clipping, never dropping** — fitted propensities are clipped to ``[pscore_trim, 1 - pscore_trim]`` after an extremeness warning (the paper gives no trimming rule). diff --git a/docs/choosing_estimator.rst b/docs/choosing_estimator.rst index 4bedb1391..6fff4ecf8 100644 --- a/docs/choosing_estimator.rst +++ b/docs/choosing_estimator.rst @@ -687,10 +687,11 @@ insensitive to the nuisance learners' regularization bias. **vs Callaway-Sant'Anna**: same cell architecture and aggregation surface; DMLDiD replaces CS's parametric nuisances with cross-fitted ML learners — prefer it when the covariate relationship is nonlinear/high-dimensional, -prefer CS otherwise (fewer moving parts, survey/cluster support). Both -handle declared repeated cross sections via ``panel=False``, but only CS -carries survey weights there — RCS data is typically survey data, so -weighted RCS belongs to CS. +prefer CS otherwise (fewer moving parts, replicate-weight support). Both +handle declared repeated cross sections via ``panel=False``, and both +carry pweight survey designs and ``cluster=`` there (DMLDiD's survey +lane is a documented library extension of Chang's i.i.d. theory; +replicate-weight designs stay CS-only). **Example**:: @@ -881,7 +882,7 @@ Survey Design Support Most estimators support an optional ``survey_design`` parameter in ``fit()`` (``SyntheticControl`` accepts the parameter but raises ``NotImplementedError``; -``LWDiD`` and ``DMLDiD`` accept no ``survey_design`` parameter at all — +``LWDiD`` accepts no ``survey_design`` parameter at all — passing it raises ``TypeError``). Pass a :class:`~diff_diff.SurveyDesign` object to get design-based variance estimation. The depth of support varies by estimator and variance method: @@ -1014,10 +1015,10 @@ estimation. The depth of support varies by estimator and variance method: - -- - -- * - ``DMLDiD`` - - N/A (no survey support) - - N/A - - -- + - Full (pweight only) + - ``n_PSU - n_strata`` - -- + - Multiplier (PSU) * - ``ChangesInChanges`` / ``QDiD`` - -- - -- diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index ba5eb0179..b4f64a023 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -979,6 +979,12 @@ sources: - path: docs/methodology/REGISTRY.md section: "DMLDiD" type: methodology + - path: docs/methodology/survey-theory.md + section: "IF-based TSL roster / variance-dispatch table / PSU-multiplier-bootstrap roster" + type: methodology + - path: docs/methodology/variance-conventions.md + section: "dml_did / dml_did_rcs inventory rows (CI-pinned to tests/test_variance_conventions.py)" + type: methodology - path: docs/index.rst section: "Supported Estimators (one-line catalog row)" type: user_guide diff --git a/docs/index.rst b/docs/index.rst index dc6744f54..5df12e3b8 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -186,7 +186,7 @@ Supported Estimators * - :class:`~diff_diff.LWDiD` - Lee & Wooldridge (2025, 2026) rolling-transformation DiD; ``rolling='detrend'`` handles heterogeneous linear trends * - :class:`~diff_diff.DMLDiD` - - Chang (2020) double/debiased ML DiD; staggered ATT(g,t) with cross-fitted nuisance learners (panel or declared repeated cross sections) + - Chang (2020) double/debiased ML DiD; staggered ATT(g,t) with cross-fitted nuisance learners (panel or declared repeated cross sections; survey/cluster support) * - :class:`~diff_diff.QDiD` - Quantile DiD comparison estimator applying DiD quantile-by-quantile (deprecated 3.9 - use :class:`~diff_diff.ChangesInChanges` with ``method="qdid"``) * - :class:`~diff_diff.RegressionDiscontinuity` diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index fb7deaab1..fbb3bf532 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -2878,7 +2878,11 @@ the finite-dimensional `p_0` is handled by the variance correction below. treated (or single control) unit is a singleton stratum → `assign_folds` raises → the cell is recorded as a `cross_fit_degenerate` NaN cell where CS would have estimated it. Rationale: one treated unit cannot be - cross-fitted; CS is the right tool there. + cross-fitted; CS is the right tool there. Under a survey/cluster design + whose PSU is strictly coarser than the sampling unit the folds are + PSU-COHESIVE instead (the stratification cannot be cluster-constant; see + the survey Note below) — the degenerate guards in `cross_fit_predict` + keep covering the panel lane's composition requirement there. - **Note:** Pooled θ̂ AND pooled Σ̂ — both the point estimate and the variance are computed as POOLED means over all cell members rather than the paper's `1/K` equal-fold averages. With fold sizes differing by ±1 the pooled forms @@ -2940,8 +2944,11 @@ the finite-dimensional `p_0` is handled by the variance correction below. covariate, or a ΔY overflow from two finite outcomes — the errstate boundaries make the warning the only user-visible trace of the latter). - **Note:** Skip-reason vocabulary — DMLDiD cells carry `skip_reason ∈` - {`missing_period`, `zero_treated_control`, `cross_fit_degenerate`, - `non_finite_score`}: the first two are the CS meanings; `cross_fit_degenerate` + {`missing_period`, `zero_treated_control`, `zero_weight_mass`, + `cross_fit_degenerate`, `non_finite_score`}: the first two are the CS + meanings; `zero_weight_mass` (declared-survey fits only, the CS meaning) = + a required group has rows but zero survey mass, so the weighted p̂/λ̂ would + leave (0, 1); `cross_fit_degenerate` = fold assignment or a fail-closed learner made the cell un-cross-fittable (chained learner message quoted in the consolidated skip warning); `non_finite_score` = the score/variance computation produced or received @@ -2959,10 +2966,11 @@ the finite-dimensional `p_0` is handled by the variance correction below. influence-function variance IS unit-level clustering, so the 4.0 auto-cluster DEFAULT flip is inert for it — mirroring how CallawaySantAnna is absent from M-080's `code_refs` and how M-010 pre-empted M-080 for the - TWFE event-study mode. CS additionally exposes `cluster=` for - COARSER-than-unit clustering, a surface DMLDiD ships without: the paper - assumes i.i.d. sampling, and the survey/cluster follow-up is the tracked - `DEFERRED.md` row. + TWFE event-study mode. COARSER-than-unit clustering ships via the CS-style + `cluster=` constructor parameter (3.11): the synthesized/injected + `SurveyDesign(psu=...)` drives PSU-cohesive folds, the per-cell CR1 SE, + the bootstrap draw structure and `df_inference`, while the moment kernels + stay unweighted (see the survey Notes below). - **Note:** Global λ̂ convention (Case 2) — `λ̂` is the FULL-SAMPLE-within-cell post-period sampling share `mean(T)` over the pooled two-period rows, mirroring the global-p̂ convention above and DoubleML's `t.mean()` (the @@ -2984,7 +2992,13 @@ the finite-dimensional `p_0` is handled by the variance correction below. guarantees control rows in BOTH periods in every training complement (Chang's `I_kz^c` fold-composition requirement). Consequence: any singleton D×T stratum (e.g. ONE treated row in the base period) dies as - `cross_fit_degenerate`. + `cross_fit_degenerate`. Under PSU-cohesive folds (coarser-than-unit + survey/cluster designs) the stratification MECHANISM is unavailable, and + under a declared survey per-row zero weights can hollow a period out even + inside stratified folds — in both cases the REQUIREMENT is preserved by an + explicit per-complement guard (positive-weight control rows in BOTH + periods, else `DegenerateFoldError` → `cross_fit_degenerate` skip); the + no-design stratified path never evaluates the guard (bit-identity). - **Note:** Case 2 four-group guard — a cell needs treated AND control rows in BOTH periods; any empty group skips as `zero_treated_control` (the vocabulary is reused, not widened). A new λ̂-extremeness warning mirrors @@ -3000,14 +3014,82 @@ the finite-dimensional `p_0` is handled by the variance correction below. through `float()`-keyed dict access, where distinct int64 cohorts above 2^53 (admissible through 2^62 by the label pipeline) collide. - **Note:** `aggregate('total')` fails closed on RCS fits (the library-wide - repeated-cross-section convention; `is_panel: False` kit bookkeeping); - panel fits keep the total route. Per-observation IF entries make - `vcov_type="hc1"` the per-SAMPLING-UNIT variance on RCS (rows are the - sampling units). + repeated-cross-section convention; `is_panel: False` kit bookkeeping) AND + on declared-`survey_design=` fits (`is_survey_fit` kit gate, the CS/ + EfficientDiD/Imputation/TwoStage convention); panel no-design and + bare-`cluster=` fits keep the total route (the latter subject to the + shared `_cs_total_mass` complete-treated-support coincidence guard). + Per-observation IF entries make `vcov_type="hc1"` the per-SAMPLING-UNIT + variance on RCS (rows are the sampling units). - **Note:** Case-2-only moment conditions — Assumption 3.2's level-outcome bounds (`E[Y²|X] ≤ C`, `|E[YU]| ≤ C`) belong to the repeated-cross-section case ONLY and are NOT imposed on the panel path; stationary sampling (Assumption 2.3) is warned at fit (not data-checkable). +- **Note:** Survey support (3.11) is a LIBRARY EXTENSION — Chang (2020) + assumes i.i.d. sampling and never discusses clustering or weighting. + Declared `survey_design=` (pweight-only, full-design TSL: + weights/strata/PSU/FPC, resolved and validated by the shared CS + machinery) enters the ESTIMATOR, not just the variance: p̂ and (Case 2) + λ̂ become Hájek weighted shares, θ̂ the weighted score mean, Ĝ₂λ the + weighted slope mean, the nuisance learners receive `sample_weight` + (user learner objects must accept it by keyword — a `TypeError` fires + up front otherwise), and the IF payload is the weighted analogue + `w_i·ψ̄_i/Σw`. The per-cell SE is PSU-gated exactly like CS: designs + with a PSU route through `_cluster_robust_se_from_per_gt_if` + (`compute_survey_if_variance`; a NaN return is the deliberate + unidentified-variance signal and flows to `safe_inference` on a + RETAINED cell), strata/FPC-only designs use the weighted + `sqrt(sum(if²))` with the full design entering aggregate SEs via + `_se_from_psi`. Analytical inference uses `df = df_survey` + (`n_PSU − n_strata`) t-statistics; bootstrap overrides keep normal + theory (the CS convention). On the weighted-λ̂ RCS lane Theorem 2's + coverage claim does NOT carry over — the weighted plug-in is validated + by the survey invariant battery (`tests/test_survey_dml.py`), not by + the paper. Replicate-weight designs fail closed (TODO.md row). +- **Note:** PSU-cohesive cross-fitting (survey/cluster designs) — when the + effective design's PSU is strictly coarser than the sampling unit AND + there are at least `n_folds` PSUs globally, fold assignment switches + from D (panel) / D×T (RCS) stratification to cluster-cohesive folds + (`assign_folds(cluster_ids=psu)`): unit-level folds with within-PSU + dependence would leak information across the train/test split, + undermining the cross-fitting argument (the DoubleML clustered-DML + prescription). With `2 <= n_psu < n_folds` the EFFECTIVE fold count is + REDUCED to `n_psu` under a `UserWarning` — PSU cohesion is preserved, + never silently traded for the requested fold count, because with >= 2 + PSUs the clustered variance is IDENTIFIED and finite inference from + leaky unit folds would be reported as clustered. Only the SINGLE-PSU + design falls back to stratified unit folds: cluster-cohesive splitting + is impossible there and the variance layer NaNs out either way + (per-cell CR1 returns NaN below 2 PSUs; the bootstrap emits the <2-PSU + warning and NaNs every bootstrap surface — remediation is a coarser + design or more PSUs). Under a globally-eligible design a + CELL with fewer in-cell PSUs than folds skips as `cross_fit_degenerate`; + if EVERY cell does, the existing all-degenerate check aborts with the + generic "Could not estimate any group-time effects" ValueError (raised + before the consolidated skip warning; no skip detail in the message) — + remediation there is fewer folds. + The same `seed` yields DIFFERENT folds with vs without a coarse PSU + design (different RNG consumption — a config change, not a + reproducibility break); `diagnostics["psu_folds"]` records the mode. + Note the shared mixin's <2-PSU bootstrap warning suggests + "n_bootstrap=0 (analytical IF variance)" — pre-existing CS-shared + wording; on a one-PSU design the analytical path also yields NaN + inference (the design is unidentified either way). +- **Note:** Bare `cluster=` (constructor) — synthesized into + `SurveyDesign(psu=cluster, weight_type="pweight")` (or injected as the + PSU of a declared PSU-less design; a design-supplied PSU wins with a + warning on differing partitions). The moment KERNELS stay unweighted + and no `sample_weight` reaches the learners — `cluster=` is not a + weighting request, and `(X, y)`-only learners keep working. Point + estimates can still move through two accepted channels: PSU-cohesive + folds (inherent to clustered cross-fitting) and, on incomplete panels, + the synthesized all-ones `survey_weights` switching aggregation masses + from per-cell complete-case `n_treated` to full cohort mass (the same + documented CS bare-`cluster=` divergence). On a fully complete panel + with an identity PSU both channels are inert and the fit is + bit-identical to unclustered. `survey_metadata` stays `None` (it is the + DECLARED-design marker); the bare-cluster df (`n_PSU − 1`) is carried + on `df_inference`. *Edge cases:* - Propensity near 0/1: clipped per the trimming Note (error bounds blow up as diff --git a/docs/methodology/papers/chang-2020-review.md b/docs/methodology/papers/chang-2020-review.md index b53b81976..a554dbc4a 100644 --- a/docs/methodology/papers/chang-2020-review.md +++ b/docs/methodology/papers/chang-2020-review.md @@ -181,15 +181,15 @@ the same variance estimators remain consistent under kernel first stages. per-item conventions/deviations are the REGISTRY "DMLDiD" Notes; Case 3 items remain open, tracked in DEFERRED.md): - [x] Neyman-orthogonal Case 1 score (3.1) implemented exactly (Abadie score + mean-zero adjustment); [x] score (3.2) — SHIPPED as `chang_rcs_score` (`DMLDiD(panel=False)`); [ ] score (3.3) — Case 3 open -- [x] DML2 cross-fitting: K-fold partition (D-stratified — documented deviation), nuisances fit on fold complements, never on the evaluation fold +- [x] DML2 cross-fitting: K-fold partition (D-stratified — documented deviation; PSU-cohesive instead under a coarser-than-unit survey/cluster design, a further documented library extension), nuisances fit on fold complements, never on the evaluation fold - [x] Outcome nuisance `ℓ̂` fit on the UNTREATED subsample of the auxiliary fold only (`I_kz^c`) - [x] Scalar nuisance p̂: the global (full-sample-within-cell) convention adopted and documented (the I_k vs I_k^c printing contradiction is thereby sidestepped — see Gaps); [x] `λ̂` — same global convention (`mean(T)` within cell; REGISTRY Note) - [x] Final estimator: pooled mean (equals the paper's `1/K` average at equal fold sizes — documented deviation) - [x] Variance from the AUGMENTED score: `Ĝ_1p = -θ̃/p̂` folded in; [x] `Ĝ_2λ (T - λ̂)` — explicit term in `chang_rcs_score_augmented` (`Ĝ_2λ` = sample mean of the closed-form `∂λψ₂`, `chang_rcs_lambda_slope`; the paper prints no estimator — REGISTRY Note) - [x] Strict-overlap enforcement: fitted propensities clipped to `[trim, 1-trim]` (documented deviation; paper gives no rule) -- [x] Per-fold/per-cell degenerate guards (zero treated/control, cell < K, singleton stratum) — closed skip vocabulary; [x] Case 2 pre/post-share guards — four-group guard + `λ̂` extremeness warning + D×T-stratified folds (control rows in both periods per training complement by construction) +- [x] Per-fold/per-cell degenerate guards (zero treated/control, cell < K, singleton stratum) — closed skip vocabulary; [x] Case 2 pre/post-share guards — four-group guard + `λ̂` extremeness warning + D×T-stratified folds (control rows in both periods per training complement — by construction under stratified folds; preserved by an explicit per-complement composition guard under PSU folds / weighted fits, REGISTRY Note) - [x] Auxiliary-sample feasibility: an empty untreated training complement raises `DegenerateFoldError` (targeted, before the learner) → `cross_fit_degenerate` cell -- [x] Normal-approximation inference via `safe_inference()` +- [x] Normal-approximation inference via `safe_inference()` (no-design fits; survey/bare-cluster fits use finite-df t inference via `df=df_survey` — library extension, not from the paper) - [ ] Multilevel treatment (Case 3): open — DEFERRED - [ ] Case 3 guards: open — DEFERRED (see the Case 3 caution above) - [x] Validation (Case 1): `doubleml.DoubleMLDID` (2-period) + `DoubleMLDIDBinary` (staggered per-cell, end-to-end public fit) parity spikes, doubleml==0.11.4 pinned, golden literals in-repo; [x] Case 2 equation-level fixtures — SHIPPED (closed-form/oracle fixtures, derivative-identity checks, DR both directions, `DoubleMLDIDCSBinary` characterization spike — no parity oracle exists). CAVEAT: the paper's own §4 RCS simulation DGPs (pp. 17-21) are NOT replicated — the shipped recovery/coverage tests use a library-authored RCS design; replication is a tracked TODO.md row (needs the paper PDF); [ ] Case 3 fixtures — open diff --git a/docs/methodology/survey-theory.md b/docs/methodology/survey-theory.md index c5d642c01..89f9cbcde 100644 --- a/docs/methodology/survey-theory.md +++ b/docs/methodology/survey-theory.md @@ -481,7 +481,7 @@ The total meat is sum_h V_h, computed by `_compute_stratified_psu_meat()` in ### IF-based TSL variance For scalar IF-based estimators (CallawaySantAnna, ImputationDiD, TwoStageDiD, -TripleDifference, StaggeredTripleDifference, EfficientDiD), the variance is +TripleDifference, StaggeredTripleDifference, EfficientDiD, DMLDiD), the variance is computed directly from per-unit influence function values without the bread matrix: @@ -692,6 +692,7 @@ Each estimator uses one of three variance strategies under survey designs: | TwoWayFixedEffects | TSL sandwich | OLS-based, all weight types | | MultiPeriodDiD | TSL sandwich | OLS-based, all weight types | | CallawaySantAnna | TSL on IFs | pweight only | +| DMLDiD | TSL on IFs | pweight only; cross-fitted augmented scores; library extension of Chang (2020) i.i.d. theory | | SunAbraham | TSL sandwich | OLS-based, all weight types | | TripleDifference | TSL on IFs | pweight only | | StaggeredTripleDifference | TSL on IFs | pweight only | @@ -728,8 +729,9 @@ For IF-based estimators, the variance computation proceeds as: Two bootstrap strategies interact with survey designs: -- **Multiplier bootstrap at PSU level** (CallawaySantAnna, ImputationDiD, - TwoStageDiD, ContinuousDiD, EfficientDiD, StaggeredTripleDifference): +- **Multiplier bootstrap at PSU level** (CallawaySantAnna, DMLDiD, + ImputationDiD, TwoStageDiD, ContinuousDiD, EfficientDiD, + StaggeredTripleDifference): Generates multiplier weights at the PSU level within strata, with FPC scaling. Each bootstrap draw reweights the IF values. diff --git a/docs/methodology/variance-conventions.md b/docs/methodology/variance-conventions.md index e9b017e00..ca3c39e24 100644 --- a/docs/methodology/variance-conventions.md +++ b/docs/methodology/variance-conventions.md @@ -38,8 +38,8 @@ produced wrong figures three separate times while this inventory was drafted). | `imputation_pretrends_event_study` | unpinned | 157, 157, 157, None, None, None, None, None | **legitimate** | pretrends lead regression: CR1 k = K_reference (D2 fixed); tail df converged in 3.9 (M-127): leads use t(residual df) via the df_convention knob (was silent normal theory); the None calls are the knob-independent BJS aggregates (L3) | | `two_stage_default` | — | None | **legitimate** | L3: Gardner two-stage variance, not the shared CR1 sandwich | | `callaway_santanna_default` | — | None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None | **legitimate** | L3: influence-function variance anchored to Stata csdid | -| `dml_did` | — | None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None | **legitimate** | L3: Chang (2020) Thm 2 augmented-score plug-in variance (per-unit influence function; normal-theory safe_inference throughout, no cluster surface) | -| `dml_did_rcs` | — | None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None | **legitimate** | L3: Chang (2020) Thm 2 lambda-corrected augmented-score plug-in variance (Case 2; per-observation influence function; normal-theory safe_inference throughout, no cluster surface) | +| `dml_did` | — | None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None | **legitimate** | L3: Chang (2020) Thm 2 augmented-score plug-in variance (per-unit influence function; normal-theory safe_inference on no-design fits; survey_design=/cluster= fits use the design-based CR1/weighted-IF per-cell variance with df=df_survey t-inference outside the shared CR1 sandwich) | +| `dml_did_rcs` | — | None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None | **legitimate** | L3: Chang (2020) Thm 2 lambda-corrected augmented-score plug-in variance (Case 2; per-observation influence function; normal-theory safe_inference on no-design fits; survey_design=/cluster= fits use the design-based CR1/weighted-IF per-cell variance with df=df_survey t-inference outside the shared CR1 sandwich) | cr1_k is the sorted multiset of K_reference counts reaching the shared clustered CR1 denominator — visible columns + the signed cluster_k_adjustment (linalg._compute_robust_vcov_numpy with diff --git a/docs/migration-4.0.md b/docs/migration-4.0.md index 54dcda9f3..58a0f1f73 100644 --- a/docs/migration-4.0.md +++ b/docs/migration-4.0.md @@ -164,7 +164,7 @@ backend fail closed with a refit message. On `ImputationDiD`, `TwoStageDiD` and `ContinuousDiD`, the recompute levels still raise `NotImplementedError` when the fit used `n_bootstrap > 0` — keep the fit-time call there for now and track their open `TODO.md` rows. -`aggregate("simple")` — and, on its five adopters, `aggregate("total")` (3.10; DMLDiD joined in 3.11 — panel fits only, RCS fits fail `total` closed) — does relay, +`aggregate("simple")` — and, on its five adopters, `aggregate("total")` (3.10; DMLDiD joined in 3.11 — panel non-survey fits only: RCS fits, declared-`survey_design=` fits, and bare-`cluster=` fits with incomplete treated support fail `total` closed) — does relay, and `StackedDiD`, `ChaisemartinDHaultfoeuille` and `HeterogeneousAdoptionDiD` are unaffected — their `aggregate()` is a pure view over stored fields. ``` diff --git a/docs/practitioner_decision_tree.rst b/docs/practitioner_decision_tree.rst index ed9b80c44..aa43c72f7 100644 --- a/docs/practitioner_decision_tree.rst +++ b/docs/practitioner_decision_tree.rst @@ -438,10 +438,11 @@ satisfaction, NPS, or similar. The survey uses stratified sampling, clustering ( by geography), or probability weights. **Answer:** Use a survey-capable method above, combined with -:class:`~diff_diff.SurveyDesign`. (:class:`~diff_diff.LWDiD` and -:class:`~diff_diff.DMLDiD` are the exceptions: they accept no -``survey_design`` parameter at all — see the -:ref:`survey-design-support` matrix.) +:class:`~diff_diff.SurveyDesign`. (:class:`~diff_diff.LWDiD` is the +exception: it accepts no ``survey_design`` parameter at all — see the +:ref:`survey-design-support` matrix. :class:`~diff_diff.DMLDiD` supports +pweight full-design TSL and ``cluster=`` on both lanes; replicate-weight +designs are not supported yet.) Ignoring survey weights and clustering makes your confidence intervals too narrow - you will be overconfident about the result. Passing a ``SurveyDesign`` to ``fit()`` @@ -573,9 +574,10 @@ The six scenarios above cover the most common business use cases. (``results.aggregate('event_study')`` feeds HonestDiD/PreTrendsPower); set ``seed=`` for reproducible fold draws. Panel data by default; ``panel=False`` runs declared repeated cross sections (Chang Case 2, - λ-corrected variance). No survey/cluster support on either lane - (per-sampling-unit influence-function variance) — weighted RCS belongs - to ``CallawaySantAnna(panel=False, survey_design=...)``. + λ-corrected variance). Survey/cluster support on both lanes: + ``survey_design=`` (pweight full-design TSL — a documented library + extension of Chang's i.i.d. theory) and coarser-than-unit + ``cluster=``; replicate-weight designs are not supported yet. For the full academic decision tree with all estimators, see :doc:`choosing_estimator`. diff --git a/docs/survey-roadmap.md b/docs/survey-roadmap.md index 3b6f67df1..f7a3eca42 100644 --- a/docs/survey-roadmap.md +++ b/docs/survey-roadmap.md @@ -271,7 +271,7 @@ the limitation and suggested alternative. | Estimator | Limitation | Alternative | |-----------|-----------|-------------| | LWDiD | Any `survey_design` / sampling weights | No weight argument exists on any path, so the failure mode is a bare `TypeError: unexpected keyword argument` rather than a descriptive error (the exception to the preamble above). The LW papers derive the transformation and exact-inference layer for unweighted panels; a weighted counterpart is DEFERRED pending user demand. Use `CallawaySantAnna` (or another survey-capable staggered estimator) when design-based variance is required. | -| DMLDiD | Any `survey_design` / sampling weights / `cluster=` | No weight or cluster argument exists on any path (bare `TypeError` on the kwarg) — on the panel lane AND the declared-RCS lane (`panel=False`). Chang (2020) assumes i.i.d. sampling; the per-sampling-unit influence-function variance IS unit-level clustering, and the survey/cluster extension is the tracked DEFERRED.md row (pending user demand). NOTE the RCS lane sharpens this gap: repeated-cross-section data is typically survey data (BRFSS/ACS/CPS) — use `CallawaySantAnna(panel=False, survey_design=...)` for weighted RCS, and `CallawaySantAnna` generally when design-based variance or coarser clustering is required. | +| DMLDiD | Replicate weights | Full-design pweight TSL (weights/strata/PSU/FPC) + survey multiplier bootstrap + `cluster=` shipped in 3.11 on BOTH lanes (a documented library extension of Chang 2020's i.i.d. theory; REGISTRY DMLDiD Notes). Replicate-weight designs remain rejected (`NotImplementedError`; the IF-reweighting route is tracked as a TODO.md Actionable row) — use a full-design SurveyDesign, or `CallawaySantAnna` when replicate weights are required. | | SyntheticDiD | Replicate weights | Pre-existing limitation: no replicate-weight survey support on SDID. All three variance methods (bootstrap, placebo, jackknife) now support pweight-only and strata/PSU/FPC designs; replicate-weight designs remain rejected. | | TROP | Replicate weights | Use strata/PSU/FPC design with Rao-Wu rescaled bootstrap | | BaconDecomposition | Replicate weights | Diagnostic only, no inference | @@ -282,7 +282,7 @@ the limitation and suggested alternative. | WooldridgeDiD | Unsupported-period filtering + `survey_design` | Restrict the frame to the supported periods explicitly and re-fit. Deleting rows in-place is naive subsetting: it removes their PSUs and strata from the TSL meat and from `df_survey = n_PSU - n_strata`. Exact only if every PSU and stratum survives the restriction (true on a balanced panel; NOT in general — an unbalanced frame can hold a PSU observed only at unsupported periods). Verify before relying on it. | | WooldridgeDiD | Unidentified-cohort exclusion + `survey_design` | Same reason (ledger `M-123`). Drop the cohort from the frame yourself, or supply a panel where every cohort has a pre-treatment period. | | All bootstrap estimators | Bootstrap + replicate weights | These are alternative variance methods; pick one | -| CS, EfficientDiD, ImputationDiD, TwoStageDiD | `aggregate('total')` on fits declaring a `survey_design` | The estimator-owned total (3.10) is panel non-survey only: the realized-mass relay omits the survey mass-uncertainty (att*dC) variance term and design-aware population-scale totals are not implemented (retained weight scale differs by design family). Pass a caller-derived numeric `scale=` to the MMM exporters instead, or use `cluster=` (without `survey_design`) for an unweighted clustered fit. Tracked in DEFERRED.md (Paper-gated). | +| CS, DMLDiD, EfficientDiD, ImputationDiD, TwoStageDiD | `aggregate('total')` on fits declaring a `survey_design` | The estimator-owned total (3.10) is panel non-survey only: the realized-mass relay omits the survey mass-uncertainty (att*dC) variance term and design-aware population-scale totals are not implemented (retained weight scale differs by design family). Pass a caller-derived numeric `scale=` to the MMM exporters instead, or use `cluster=` (without `survey_design`) for an unweighted clustered fit. Tracked in DEFERRED.md (Paper-gated). | **Warning/fallback (no error):** MultiPeriodDiD with `wild_bootstrap` + `survey_design` warns and falls back to analytical inference. diff --git a/docs/v4-design.md b/docs/v4-design.md index 510bde1b7..30a71cf13 100644 --- a/docs/v4-design.md +++ b/docs/v4-design.md @@ -530,14 +530,16 @@ exact relay `C x overall` CONDITIONAL on the realized aggregation mass, where `C` is the adopter's finite-masked complete-case aggregation support (CallawaySantAnna: the kept post-anticipation finite cells' mass, replayed from a fit-time kit snapshot; DMLDiD: the same kit-snapshot replay via the -inherited CS channel over its per-cell complete-case `n_treated`; +inherited CS channel — per-cell complete-case `n_treated` on no-design +fits, cohort masses once a survey/cluster design populates the kit's +`survey_weights` (3.11 survey support); ImputationDiD: the finite-tau support - fixing the documented raw-|Omega_1| overcount of the MMM `scale="auto"` route for total exports; TwoStageDiD: the post-filter treatment-indicator support; EfficientDiD: the kept cells' integer `n_treated` sum). Adopters: those five, panel non-survey fits only - repeated-cross-section-routed -fits, fits declaring a `survey_design=`, and CS bare-`cluster=` fits whose -cohort-mass weighting diverges from the complete-case count all raise +fits, fits declaring a `survey_design=`, and CS/DMLDiD bare-`cluster=` +fits whose cohort-mass weighting diverges from the complete-case count all raise `NotImplementedError` naming the reason (DEFERRED tracks the survey/RC remainder and its att*dC mass-uncertainty variance term). StackedDiD is staged out (estimand ill-defined under `weighting=` variants; DEFERRED). diff --git a/tests/test_dml_did.py b/tests/test_dml_did.py index a7fc4a93c..6c54a63d3 100644 --- a/tests/test_dml_did.py +++ b/tests/test_dml_did.py @@ -1016,9 +1016,12 @@ def test_anticipation_aggregation_includes_t_g_minus_1(self): class TestUnweightedLearnerSignature: def test_fit_x_y_only_learners_through_public_fit(self, data): - # The advertised duck-typed contract is fit/predict(_proba); DMLDiD - # never passes sample weights, so learners whose fit signature is - # only (X, y) must work end to end through the public fit(). + # The advertised duck-typed contract is fit/predict(_proba); on + # no-survey and bare-cluster fits DMLDiD passes no sample weights, + # so learners whose fit signature is only (X, y) must work end to + # end through the public fit(). (Declared survey_design= fits DO + # pass sample_weight and gate on it up front — + # tests/test_survey_dml.py.) class XYOnlyRegressor: def fit(self, X, y): import numpy as _np @@ -1064,6 +1067,7 @@ class TestFullConfigMutationDefense: ("outcome_learner", "logit"), ("pscore_trim", 0.9), ("n_folds", 1), + ("cluster", 123), ], ) def test_mutated_config_raises_before_any_cell(self, data, attr, bad): diff --git a/tests/test_dr_scores.py b/tests/test_dr_scores.py index 4ab912419..c0e366f0d 100644 --- a/tests/test_dr_scores.py +++ b/tests/test_dr_scores.py @@ -341,3 +341,122 @@ def test_internal_with_slope_variant_matches_public_pair(self): psi_bar, chang_rcs_score_augmented(summand, D, T, y, m2, ps, theta, p_hat, lam_hat) ) assert g2_lambda == chang_rcs_lambda_slope(y, D, T, m2, ps, p_hat, lam_hat) + + +class TestChangRCSSlopeWeights: + """Survey weights on the lambda-slope family (DMLDiD survey support).""" + + def _inputs(self, n=24): + rng = np.random.default_rng(5) + y = rng.normal(size=n) + D = (rng.uniform(size=n) < 0.5).astype(float) + T = (rng.uniform(size=n) < 0.5).astype(float) + m2_hat = rng.normal(size=n) + ps = np.clip(rng.uniform(size=n), 0.1, 0.9) + return y, D, T, m2_hat, ps + + def test_uniform_weights_match_unweighted_exactly(self): + from diff_diff._dr_scores import chang_rcs_lambda_slope + + y, D, T, m2, ps = self._inputs() + base = chang_rcs_lambda_slope(y, D, T, m2, ps, 0.5, 0.45) + weighted = chang_rcs_lambda_slope(y, D, T, m2, ps, 0.5, 0.45, weights=np.ones(len(y))) + # np.average with uniform weights is not guaranteed bit-identical to + # np.mean, but must agree to float64 roundoff. + np.testing.assert_allclose(weighted, base, rtol=1e-14) + + def test_weighted_slope_equals_weighted_mean_of_terms(self): + from diff_diff._dr_scores import chang_rcs_lambda_slope + + y, D, T, m2, ps = self._inputs() + rng = np.random.default_rng(9) + w = rng.uniform(0.5, 2.0, size=len(y)) + p_hat, lam_hat = 0.5, 0.45 + odds = (D - ps) / (1.0 - ps) + term1 = ( + -((1.0 - 2.0 * lam_hat) / (lam_hat**2 * (1.0 - lam_hat) ** 2)) + * (odds / p_hat) + * ((T - lam_hat) * y - m2) + ) + term2 = -(y / (p_hat * lam_hat * (1.0 - lam_hat))) * odds + expected = float(np.average(term1 + term2, weights=w)) + got = chang_rcs_lambda_slope(y, D, T, m2, ps, p_hat, lam_hat, weights=w) + np.testing.assert_allclose(got, expected, rtol=1e-14) + + def test_weighted_slope_finite_difference_identity(self): + # The slope is d/d(lambda) of the WEIGHTED mean of psi_2 evaluated at + # the plug-ins: central finite difference on the weighted score mean. + from diff_diff._dr_scores import chang_rcs_lambda_slope, chang_rcs_score + + y, D, T, m2, ps = self._inputs() + rng = np.random.default_rng(11) + w = rng.uniform(0.5, 2.0, size=len(y)) + p_hat, lam_hat, eps = 0.5, 0.45, 1e-6 + up = np.average(chang_rcs_score(y, D, T, m2, ps, p_hat, lam_hat + eps), weights=w) + dn = np.average(chang_rcs_score(y, D, T, m2, ps, p_hat, lam_hat - eps), weights=w) + fd = (up - dn) / (2 * eps) + got = chang_rcs_lambda_slope(y, D, T, m2, ps, p_hat, lam_hat, weights=w) + np.testing.assert_allclose(got, fd, rtol=1e-5) + + @pytest.mark.parametrize( + "bad_weights, match", + [ + (np.ones((4, 6)), "1-dimensional"), + (np.ones(5), "length"), + (np.array([1.0, np.nan] + [1.0] * 22), "non-finite"), + (np.array([-1.0] + [1.0] * 23), "non-negative"), + (np.zeros(24), "positive sum"), + ], + ) + def test_weight_validation_rejections_public(self, bad_weights, match): + from diff_diff._dr_scores import chang_rcs_lambda_slope + + y, D, T, m2, ps = self._inputs() + with pytest.raises(ValueError, match=match): + chang_rcs_lambda_slope(y, D, T, m2, ps, 0.5, 0.45, weights=bad_weights) + + @pytest.mark.parametrize( + "bad_weights, match", + [ + (np.ones((4, 6)), "1-dimensional"), + (np.ones(5), "length"), + (np.array([1.0, np.inf] + [1.0] * 22), "non-finite"), + (np.array([-1.0] + [1.0] * 23), "non-negative"), + (np.zeros(24), "positive sum"), + ], + ) + def test_weight_validation_rejections_internal_entry(self, bad_weights, match): + # The DMLDiD RCS cell calls the internal _with_slope entry directly, + # so its weight validation must hold there too, not only on the + # public wrappers. + from diff_diff._dr_scores import _chang_rcs_score_augmented_with_slope + + y, D, T, m2, ps = self._inputs() + summand = np.zeros(len(y)) + with pytest.raises(ValueError, match=match): + _chang_rcs_score_augmented_with_slope( + summand, D, T, y, m2, ps, 0.0, 0.5, 0.45, weights=bad_weights + ) + + def test_weighted_augmented_matches_public_pair(self): + from diff_diff._dr_scores import ( + _chang_rcs_score_augmented_with_slope, + chang_rcs_lambda_slope, + chang_rcs_score, + chang_rcs_score_augmented, + ) + + y, D, T, m2, ps = self._inputs() + rng = np.random.default_rng(13) + w = rng.uniform(0.5, 2.0, size=len(y)) + p_hat, lam_hat = 0.5, 0.45 + summand = chang_rcs_score(y, D, T, m2, ps, p_hat, lam_hat) + theta = float(np.average(summand, weights=w)) + psi_bar, g2 = _chang_rcs_score_augmented_with_slope( + summand, D, T, y, m2, ps, theta, p_hat, lam_hat, weights=w + ) + np.testing.assert_array_equal( + psi_bar, + chang_rcs_score_augmented(summand, D, T, y, m2, ps, theta, p_hat, lam_hat, weights=w), + ) + assert g2 == chang_rcs_lambda_slope(y, D, T, m2, ps, p_hat, lam_hat, weights=w) diff --git a/tests/test_guides.py b/tests/test_guides.py index ed2e8ab81..ef39c2920 100644 --- a/tests/test_guides.py +++ b/tests/test_guides.py @@ -806,10 +806,14 @@ def test_llms_full_dml_fit_documents_covariates_required(self): from diff_diff import DMLDiD assert "covariates" in inspect.signature(DMLDiD.fit).parameters + assert "survey_design" in inspect.signature(DMLDiD.fit).parameters section = self._dml_section() fit_start = section.index(").fit(") fit_block = section[fit_start : section.index("DMLDiDResults", fit_start)] assert "covariates" in fit_block + # The new public fit parameter must be documented in the guide's + # .fit( block too (constructor parity is pinned separately). + assert "survey_design" in fit_block assert "REQUIRED" in section # the covariates-required contract def test_llms_full_dml_documents_post_fit_aggregation_and_seed(self): diff --git a/tests/test_survey_dml.py b/tests/test_survey_dml.py new file mode 100644 index 000000000..eaa6fb12b --- /dev/null +++ b/tests/test_survey_dml.py @@ -0,0 +1,781 @@ +"""DMLDiD survey-design support (both lanes): TSL + survey bootstrap + cluster=. + +Survey support is a documented library extension of Chang (2020), which +assumes i.i.d. sampling (Assumption 2.3) — no external oracle exists +(DoubleML has no survey support; R ``did::`` is survey-naive), so the +evidence is the library's standard survey invariant battery (mirroring +``tests/test_survey_phase4.py``'s CS coverage) plus direct kernel +cross-checks against ``compute_survey_if_variance``. +""" + +import warnings + +import numpy as np +import pandas as pd +import pytest + +from diff_diff import DMLDiD +from diff_diff.staggered_aggregation import fixed_cohort_agg_weights +from diff_diff.survey import SurveyDesign, compute_survey_if_variance +from diff_diff.utils import safe_inference +from tests.conftest import assert_nan_inference + +FIT_KW = dict(outcome="y", unit="unit", time="time", first_treat="g") + + +# --------------------------------------------------------------------------- +# DGPs +# --------------------------------------------------------------------------- + + +def _make_panel(n_units=120, n_periods=4, seed=7, weight_fn=None, drop_one=False): + rng = np.random.RandomState(seed) + units = np.arange(n_units) + psu = units // 6 # 20 PSUs, strictly coarser than the unit + stratum = psu % 2 + g = rng.choice([0, 3, 4], size=n_units, p=[0.5, 0.25, 0.25]) + w = weight_fn(rng, n_units) if weight_fn is not None else rng.uniform(0.5, 2.0, n_units) + rows = [] + for u in units: + for t in range(1, n_periods + 1): + y = 1.0 + 0.4 * t + 0.2 * rng.randn() + (1.5 if g[u] > 0 and t >= g[u] else 0.0) + rows.append((u, t, y, g[u], rng.randn(), psu[u], stratum[u], w[u])) + df = pd.DataFrame(rows, columns=["unit", "time", "y", "g", "x1", "psu", "stratum", "w"]) + if drop_one: + # One treated unit's post outcome goes non-finite -> incomplete panel. + victim = df.index[(df["g"] == 3) & (df["time"] == 3)][0] + df.loc[victim, "y"] = np.nan + return df + + +def _make_rcs(n_obs=900, seed=11, weight_fn=None): + rng = np.random.RandomState(seed) + obs = np.arange(n_obs) + psu = obs // 45 # 20 PSUs + stratum = psu % 2 + g = rng.choice([0, 3, 4], size=n_obs, p=[0.5, 0.25, 0.25]) + t = rng.choice([1, 2, 3, 4], size=n_obs) + w = weight_fn(rng, n_obs) if weight_fn is not None else rng.uniform(0.5, 2.0, n_obs) + y = 1.0 + 0.4 * t + 0.2 * rng.randn(n_obs) + np.where((g > 0) & (t >= g), 1.5, 0.0) + return pd.DataFrame( + { + "unit": obs, + "time": t, + "y": y, + "g": g, + "x1": rng.randn(n_obs), + "psu": psu, + "stratum": stratum, + "w": w, + } + ) + + +def _fit(df, *, panel=True, survey=None, cluster=None, seed=42, ignore_warnings=True, **kw): + est = DMLDiD(seed=seed, panel=panel, cluster=cluster, **kw) + if not ignore_warnings: + return est.fit(df, covariates=["x1"], survey_design=survey, **FIT_KW) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return est.fit(df, covariates=["x1"], survey_design=survey, **FIT_KW) + + +_DESIGN = SurveyDesign(weights="w", strata="stratum", psu="psu") + + +@pytest.fixture(scope="module") +def panel_df(): + return _make_panel() + + +@pytest.fixture(scope="module") +def rcs_df(): + return _make_rcs() + + +@pytest.fixture(scope="module") +def panel_plain(panel_df): + return _fit(panel_df) + + +@pytest.fixture(scope="module") +def panel_survey(panel_df): + return _fit(panel_df, survey=_DESIGN) + + +@pytest.fixture(scope="module") +def rcs_plain(rcs_df): + return _fit(rcs_df, panel=False) + + +@pytest.fixture(scope="module") +def rcs_survey(rcs_df): + return _fit(rcs_df, panel=False, survey=_DESIGN) + + +# --------------------------------------------------------------------------- +# Duck learners for the capability-gate and weight-threading pins +# --------------------------------------------------------------------------- + + +class _XYOnlyRegressor: + """(X, y)-only fit: valid on no-design fits, rejected under survey.""" + + def fit(self, X, y): + self.mean_ = float(np.mean(y)) + return self + + def predict(self, X): + return np.full(X.shape[0], self.mean_) + + +class _WeightedRegressor: + def fit(self, X, y, sample_weight=None): + if sample_weight is None: + self.mean_ = float(np.mean(y)) + else: + self.mean_ = float(np.average(y, weights=np.where(sample_weight > 0, sample_weight, 0))) + return self + + def predict(self, X): + return np.full(X.shape[0], self.mean_) + + +class _KwargsRegressor: + def fit(self, X, y, **kwargs): + self.mean_ = float(np.mean(y)) + return self + + def predict(self, X): + return np.full(X.shape[0], self.mean_) + + +class _SpyRegressor: + """Records every sample_weight its fit() receives. + + The log lives on the CLASS: cross_fit_predict deep-copies the learner + template per fold, and class attributes stay shared across copies. + """ + + log: list = [] + + def fit(self, X, y, sample_weight=None): + type(self).log.append(None if sample_weight is None else np.asarray(sample_weight).copy()) + self.mean_ = float(np.mean(y)) + return self + + def predict(self, X): + return np.full(X.shape[0], self.mean_) + + +class _SpyClassifier: + log: list = [] + + def fit(self, X, y, sample_weight=None): + type(self).log.append(None if sample_weight is None else np.asarray(sample_weight).copy()) + self.p_ = float(np.mean(y)) + return self + + def predict_proba(self, X): + p = min(max(self.p_, 0.05), 0.95) + return np.column_stack([np.full(X.shape[0], 1 - p), np.full(X.shape[0], p)]) + + +# --------------------------------------------------------------------------- +# 1. Invariants +# --------------------------------------------------------------------------- + + +class TestInvariants: + @pytest.mark.parametrize("panel", [True, False]) + def test_uniform_weights_match_unweighted( + self, panel, panel_df, rcs_df, panel_plain, rcs_plain + ): + df = (panel_df if panel else rcs_df).copy() + df["w1"] = 1.0 + # Weights-only design: no PSU, no strata -> stratified folds retained, + # weighted kernels with w == 1 (1e-12, not bit-level: the + # sample_weight path through the native solvers is a different + # numerical route than the unweighted one). + res_w = _fit(df, panel=panel, survey=SurveyDesign(weights="w1")) + base = panel_plain if panel else rcs_plain + np.testing.assert_allclose(res_w.overall_att, base.overall_att, rtol=0, atol=1e-12) + np.testing.assert_allclose(res_w.overall_se, base.overall_se, rtol=0, atol=1e-12) + + @pytest.mark.parametrize("panel", [True, False]) + def test_weight_scale_invariance(self, panel, panel_df, rcs_df): + df = (panel_df if panel else rcs_df).copy() + df["w3"] = df["w"] * 3.0 + r1 = _fit(df, panel=panel, survey=_DESIGN) + r3 = _fit(df, panel=panel, survey=SurveyDesign(weights="w3", strata="stratum", psu="psu")) + np.testing.assert_allclose(r3.overall_att, r1.overall_att, rtol=1e-10) + np.testing.assert_allclose(r3.overall_se, r1.overall_se, rtol=1e-10) + + @pytest.mark.parametrize("panel", [True, False]) + def test_weights_move_the_point_estimate( + self, panel, panel_survey, rcs_survey, panel_plain, rcs_plain + ): + res, base = (panel_survey, panel_plain) if panel else (rcs_survey, rcs_plain) + assert np.isfinite(res.overall_att) + assert res.overall_att != base.overall_att + + @pytest.mark.parametrize("panel", [True, False]) + def test_full_design_smoke(self, panel, panel_survey, rcs_survey): + res = panel_survey if panel else rcs_survey + assert np.isfinite(res.overall_att) and np.isfinite(res.overall_se) + assert res.survey_metadata is not None + assert res.survey_metadata.n_strata == 2 + assert res.survey_metadata.n_psu == 20 + assert res.survey_metadata.df_survey == 18 # n_psu - n_strata + assert res.cluster_name == "psu" + assert res.n_clusters == 20 + + def test_fpc_shrinks_se(self, panel_df): + df = panel_df.copy() + # Census-like FPC: each stratum's PSU population barely exceeds the + # sample -> (1 - f_h) shrinks every variance contribution. + df["fpc"] = 11 + r_nofpc = _fit(df, survey=_DESIGN) + r_fpc = _fit(df, survey=SurveyDesign(weights="w", strata="stratum", psu="psu", fpc="fpc")) + assert r_fpc.overall_se < r_nofpc.overall_se + + +# --------------------------------------------------------------------------- +# 2. Per-cell SE +# --------------------------------------------------------------------------- + + +def _reconstruct_cell_psi(kit, gt_key): + info = kit.influence[gt_key] + resolved = kit.bookkeeping["resolved_survey_unit"] + n = len(resolved.weights) + psi = np.zeros(n) + psi[info["treated_idx"]] = info["treated_inf"] + psi[info["control_idx"]] = info["control_inf"] + return psi, resolved + + +class TestPerCellSE: + @pytest.mark.parametrize("fixture", ["panel_survey", "rcs_survey"]) + def test_psu_design_matches_compute_survey_if_variance(self, fixture, request): + res = request.getfixturevalue(fixture) + kit = res._aggregation_kit + checked = 0 + for gt_key, data in res.group_time_effects.items(): + if data.get("skip_reason") is not None or data.get("is_reference"): + continue + if gt_key not in kit.influence: + continue + psi, resolved = _reconstruct_cell_psi(kit, gt_key) + expected = compute_survey_if_variance(psi, resolved) + np.testing.assert_allclose(data["se"], np.sqrt(expected), rtol=1e-12) + checked += 1 + assert checked > 0 + + def test_strata_only_design_uses_weighted_sqrt_sum(self, panel_df): + res = _fit(panel_df, survey=SurveyDesign(weights="w", strata="stratum")) + kit = res._aggregation_kit + checked = 0 + for gt_key, data in res.group_time_effects.items(): + if data.get("skip_reason") is not None or data.get("is_reference"): + continue + info = kit.influence.get(gt_key) + if info is None: + continue + ssq = float(np.sum(info["treated_inf"] ** 2) + np.sum(info["control_inf"] ** 2)) + np.testing.assert_allclose(data["se"], np.sqrt(ssq), rtol=1e-12) + checked += 1 + assert checked > 0 + + def test_single_psu_retained_cells_nan_inference(self, panel_df): + # ONE global PSU: cluster-cohesive folds impossible -> stratified + # fallback fits the points; the clustered variance is unidentified + # -> retained cells with NaN-consistent inference, NOT skips. + df = panel_df.copy() + df["one_psu"] = 0 + with pytest.warns(UserWarning, match="PSU"): + res = _fit(df, survey=SurveyDesign(weights="w", psu="one_psu"), ignore_warnings=False) + assert np.isfinite(res.overall_att) + found = False + for data in res.group_time_effects.values(): + if data.get("is_reference") or data.get("skip_reason") is not None: + continue + assert np.isfinite(data["effect"]) + assert_nan_inference(data) + found = True + assert found + + def test_all_lonely_psu_remove_nan_inference(self, panel_df): + # >= 2 PSUs but every stratum is a singleton under lonely_psu + # "remove": the meat collapses -> same retained-cell NaN contract. + df = panel_df.copy() + df["lone_stratum"] = df["unit"] % 4 + df["lone_psu"] = df["unit"] % 4 + with pytest.warns(UserWarning, match="lonely_psu|PSU"): + res = _fit( + df, + survey=SurveyDesign( + weights="w", strata="lone_stratum", psu="lone_psu", lonely_psu="remove" + ), + ignore_warnings=False, + ) + for data in res.group_time_effects.values(): + if data.get("is_reference") or data.get("skip_reason") is not None: + continue + assert np.isfinite(data["effect"]) + assert_nan_inference(data) + + +# --------------------------------------------------------------------------- +# 3. cluster= wiring +# --------------------------------------------------------------------------- + + +class TestClusterWiring: + def test_bare_cluster_synthesizes_design(self, panel_df): + res = _fit(panel_df, cluster="psu") + assert res.survey_metadata is None # declared-survey marker stays off + assert res.cluster_name == "psu" + assert res.n_clusters == 20 + assert res.df_inference == 19.0 # n_psu - 1 + + def test_bare_cluster_matches_explicit_psu_design(self, panel_df): + r_cluster = _fit(panel_df, cluster="psu") + r_design = _fit(panel_df, survey=SurveyDesign(psu="psu")) + np.testing.assert_allclose(r_cluster.overall_att, r_design.overall_att, rtol=1e-12) + np.testing.assert_allclose(r_cluster.overall_se, r_design.overall_se, rtol=1e-12) + + def test_identity_psu_bare_cluster_bit_identical_on_complete_panel(self, panel_df, panel_plain): + # Identity PSU (cluster == unit) on a fully complete panel: folds + # stay stratified (predicate's coarser-than-unit conjunct), kernels + # stay unweighted, aggregation masses coincide -> bit-identical. + res = _fit(panel_df, cluster="unit") + assert res.overall_att == panel_plain.overall_att + # The SE deliberately moves: identity-PSU CR1 carries the + # per-PSU centering + Bessel factor the plain sqrt(sum(if^2)) + # does not (that IS the clustering request). + assert res.overall_se != panel_plain.overall_se + assert np.isfinite(res.overall_se) + + def test_incomplete_panel_bare_cluster_moves_overall_att(self): + # CS-parity divergence (staggered_results.py bare-cluster note): the + # synthesized all-ones survey_weights switch aggregation masses from + # per-cell complete-case n_treated to full cohort mass, so on an + # INCOMPLETE panel the overall ATT moves. Expected, documented. + df = _make_panel(drop_one=True) + base = _fit(df) + res = _fit(df, cluster="unit") + assert np.isfinite(res.overall_att) + assert res.overall_att != base.overall_att + + def test_xy_only_learner_works_under_bare_cluster(self, panel_df): + res = _fit(panel_df, cluster="psu", outcome_learner=_XYOnlyRegressor()) + assert np.isfinite(res.overall_att) + + def test_design_psu_wins_over_cluster_with_warning(self, panel_df): + df = panel_df.copy() + df["other"] = df["unit"] % 7 + est = DMLDiD(seed=42, cluster="other") + with pytest.warns(UserWarning, match="PSU|cluster"): + with warnings.catch_warnings(): + warnings.simplefilter("always") + res = est.fit(df, covariates=["x1"], survey_design=_DESIGN, **FIT_KW) + assert res.cluster_name == "psu" # design PSU wins + + def test_design_without_psu_injects_cluster(self, panel_df): + res = _fit(panel_df, survey=SurveyDesign(weights="w", strata="stratum"), cluster="psu") + assert res.survey_metadata is not None + assert res.survey_metadata.n_psu == 20 + assert res.cluster_name == "psu" + + def test_missing_cluster_column_raises(self, panel_df): + with pytest.raises(ValueError, match="cluster column"): + _fit(panel_df, cluster="nope") + + def test_nan_cluster_raises(self, panel_df): + df = panel_df.copy() + df.loc[df.index[0], "psu"] = np.nan + with pytest.raises(ValueError, match="missing values"): + _fit(df, cluster="psu") + + def test_cluster_constructor_rejects_non_str(self): + with pytest.raises(ValueError, match="cluster must be"): + DMLDiD(cluster=123) + + def test_cluster_set_params_transactional(self): + est = DMLDiD() + with pytest.raises(ValueError, match="cluster must be"): + est.set_params(cluster=3.5) + assert est.cluster is None # probe re-init validated before mutating + + def test_cluster_direct_mutation_revalidated_at_fit(self, panel_df): + est = DMLDiD(seed=42) + est.cluster = 123 # bypasses __init__ validation + with pytest.raises(ValueError, match="cluster must be"): + est.fit(panel_df, covariates=["x1"], **FIT_KW) + + +# --------------------------------------------------------------------------- +# 4. Folds +# --------------------------------------------------------------------------- + + +class TestFolds: + def test_psu_folds_diagnostic_true_under_coarse_design(self, panel_survey): + diags = [d for d in panel_survey.cross_fit_diagnostics.values() if "psu_folds" in d] + assert diags and all(d["psu_folds"] is True for d in diags) + + def test_identity_psu_keeps_stratified_folds(self, panel_df): + res = _fit(panel_df, cluster="unit") + diags = [d for d in res.cross_fit_diagnostics.values() if "psu_folds" in d] + assert diags and all(d["psu_folds"] is False for d in diags) + + @pytest.mark.parametrize("panel", [True, False]) + def test_few_psus_reduce_fold_count_keep_cohesion(self, panel, panel_df, rcs_df): + # 2 <= n_psu_global < n_folds: with >= 2 PSUs the clustered variance + # is IDENTIFIED, so silently reverting to unit folds would legitimize + # nuisances trained with within-PSU leakage (review R1 P1). Instead + # the effective fold count is reduced to n_psu with a warning, + # preserving PSU cohesion — on BOTH lanes. + df = (panel_df if panel else rcs_df).copy() + df["psu3"] = df["unit"] % 3 + with pytest.warns(UserWarning, match="fold count reduced to preserve cluster"): + res = _fit( + df, + panel=panel, + survey=SurveyDesign(weights="w", psu="psu3"), + n_folds=5, + ignore_warnings=False, + ) + assert np.isfinite(res.overall_att) + assert np.isfinite(res.overall_se) + diags = [d for d in res.cross_fit_diagnostics.values() if d.get("propensity")] + assert diags and all(d["psu_folds"] is True for d in diags) + # Effective fold count == n_psu (3), visible in the per-fold counts. + assert all(len(d["propensity"]["n_fit_per_fold"]) == 3 for d in diags) + # Provenance: requested vs realized fold counts both serialized. + assert res.n_folds == 5 + assert res.effective_n_folds == 3 + d = res.to_dict() + assert d["n_folds"] == 5 and d["effective_n_folds"] == 3 + assert "Effective folds (PSU-reduced):" in res.summary() + assert "n_folds=5 (effective 3)" in repr(res) + + @pytest.mark.parametrize("panel", [True, False]) + def test_exact_psu_fold_match_uses_psu_folds_unreduced(self, panel, panel_df, rcs_df): + # n_psu == n_folds: PSU folds at the requested count, no reduction. + df = (panel_df if panel else rcs_df).copy() + df["psu3"] = df["unit"] % 3 + res = _fit(df, panel=panel, survey=SurveyDesign(weights="w", psu="psu3"), n_folds=3) + assert np.isfinite(res.overall_att) + assert np.isfinite(res.overall_se) + diags = [d for d in res.cross_fit_diagnostics.values() if d.get("propensity")] + assert diags and all(d["psu_folds"] is True for d in diags) + assert all(len(d["propensity"]["n_fit_per_fold"]) == 3 for d in diags) + assert res.effective_n_folds is None + assert res.to_dict()["effective_n_folds"] is None + assert "Effective folds (PSU-reduced):" not in res.summary() + assert "effective" not in repr(res) + + def test_rcs_composition_guard_skips_cell(self): + # Coarse PSU folds on the RCS lane where whole periods live inside + # single PSUs: some training complement loses one period's controls + # entirely -> the 6a guard converts the cell to a + # cross_fit_degenerate skip instead of a finite-but-invalid fit. + rng = np.random.RandomState(3) + n = 240 + t = np.repeat([1, 2, 3, 4], n // 4) + # PSU == period block: any fold built from whole PSUs drops periods. + psu = t.copy() + g = rng.choice([0, 3], size=n, p=[0.6, 0.4]) + df = pd.DataFrame( + { + "unit": np.arange(n), + "time": t, + "y": 1.0 + 0.3 * t + rng.randn(n) * 0.2 + np.where((g > 0) & (t >= g), 1.0, 0.0), + "g": g, + "x1": rng.randn(n), + "psu": psu, + "w": np.ones(n), + } + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + est = DMLDiD(seed=0, panel=False, n_folds=2) + try: + res = est.fit( + df, + covariates=["x1"], + survey_design=SurveyDesign(weights="w", psu="psu"), + **FIT_KW, + ) + except ValueError as exc: + # Every cell degenerate is also an acceptable outcome here. + assert "Could not estimate any group-time effects" in str(exc) + return + skips = { + d.get("skip_reason") for d in res.cross_fit_diagnostics.values() if d.get("skip_reason") + } + assert "cross_fit_degenerate" in skips + + def test_spy_learner_receives_weights_on_both_nuisances(self, panel_df): + _SpyClassifier.log = [] + _SpyRegressor.log = [] + _fit( + panel_df, + survey=_DESIGN, + propensity_learner=_SpyClassifier(), + outcome_learner=_SpyRegressor(), + ) + ps_log, or_log = _SpyClassifier.log, _SpyRegressor.log + assert ps_log and or_log + assert all(w is not None for w in ps_log) + assert all(w is not None for w in or_log) + # Every received vector is a per-fold slice of the RESOLVED design + # weights (resolve() normalizes pweights to sum(w) == n_obs; the + # panel collapse then takes each unit's first row). + raw = panel_df["w"].to_numpy() + norm = raw * (len(raw) / raw.sum()) + unit_w = pd.Series(norm, index=panel_df.index).groupby(panel_df["unit"]).first() + w_all = set(np.round(unit_w.to_numpy(), 9)) + for rec in ps_log + or_log: + assert set(np.round(rec, 9)).issubset(w_all) + + def test_spy_learner_receives_no_weights_without_design(self, panel_df): + _SpyClassifier.log = [] + _SpyRegressor.log = [] + _fit( + panel_df, + propensity_learner=_SpyClassifier(), + outcome_learner=_SpyRegressor(), + ) + ps_log, or_log = _SpyClassifier.log, _SpyRegressor.log + assert ps_log and or_log + assert all(w is None for w in ps_log + or_log) + + +# --------------------------------------------------------------------------- +# 5. Skips +# --------------------------------------------------------------------------- + + +class TestZeroWeightMass: + def test_panel_zero_treated_mass_skips(self, panel_df): + df = panel_df.copy() + df.loc[df["g"] == 3, "w"] = 0.0 # cohort present, zero survey mass + with pytest.warns(UserWarning, match="zero_weight_mass"): + res = _fit(df, survey=_DESIGN, ignore_warnings=False) + reasons = { + d.get("skip_reason") for d in res.group_time_effects.values() if d.get("skip_reason") + } + assert "zero_weight_mass" in reasons + + def test_rcs_zero_group_mass_skips(self, rcs_df): + df = rcs_df.copy() + df.loc[(df["g"] == 3) & (df["time"] >= 3), "w"] = 0.0 + with pytest.warns(UserWarning, match="zero_weight_mass"): + res = _fit(df, panel=False, survey=_DESIGN, ignore_warnings=False) + reasons = { + d.get("skip_reason") for d in res.group_time_effects.values() if d.get("skip_reason") + } + assert "zero_weight_mass" in reasons + + +# --------------------------------------------------------------------------- +# 6. Bootstrap +# --------------------------------------------------------------------------- + + +class TestSurveyBootstrap: + @pytest.mark.parametrize("panel", [True, False]) + def test_survey_bootstrap_smoke(self, panel, panel_df, rcs_df, ci_params): + df = panel_df if panel else rcs_df + res = _fit(df, panel=panel, survey=_DESIGN, n_bootstrap=ci_params.bootstrap(49)) + assert np.isfinite(res.overall_att) + assert np.isfinite(res.overall_se) + + def test_single_psu_bootstrap_nan_with_dml_label(self, panel_df, ci_params): + df = panel_df.copy() + df["one_psu"] = 0 + with pytest.warns(UserWarning, match="DMLDiD bootstrap with survey/cluster design"): + res = _fit( + df, + survey=SurveyDesign(weights="w", psu="one_psu"), + n_bootstrap=ci_params.bootstrap(29), + ignore_warnings=False, + ) + assert np.isnan(res.overall_se) + + def test_replay_label_and_legacy_kit_fallback(self, panel_df, ci_params): + df = panel_df.copy() + df["one_psu"] = 0 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = _fit( + df, + survey=SurveyDesign(weights="w", psu="one_psu"), + n_bootstrap=ci_params.bootstrap(29), + ) + # Post-fit replay re-emits the <2-PSU warning branded DMLDiD (the + # kit label fix; the replay host used to hardcode CallawaySantAnna). + with pytest.warns(UserWarning, match="DMLDiD bootstrap with survey/cluster design"): + res.aggregate("event_study") + # Legacy kits without the key fall back to the CS default, no crash. + res._aggregation_kit.bookkeeping.pop("bootstrap_label") + with pytest.warns(UserWarning, match="CallawaySantAnna bootstrap with survey/cluster"): + res.aggregate("event_study") + + def test_survey_bootstrap_replay_matches_fit(self, panel_df, ci_params): + res = _fit(panel_df, survey=_DESIGN, n_bootstrap=ci_params.bootstrap(49)) + agg = res.aggregate("simple") + np.testing.assert_allclose(float(agg.att[0]), res.overall_att, rtol=1e-12) + + +# --------------------------------------------------------------------------- +# 7. Rejections +# --------------------------------------------------------------------------- + + +class TestRejections: + def test_replicate_design_fails_closed(self, panel_df): + df = panel_df.copy() + df["rw1"] = df["w"] * 1.1 + df["rw2"] = df["w"] * 0.9 + design = SurveyDesign(weights="w", replicate_weights=["rw1", "rw2"], replicate_method="JK1") + with pytest.raises(NotImplementedError, match="replicate"): + _fit(panel_df.assign(rw1=df["rw1"], rw2=df["rw2"]), survey=design) + + def test_cluster_plus_replicate_hits_blanket_replicate_message(self, panel_df): + df = panel_df.copy() + df["rw1"] = df["w"] * 1.1 + df["rw2"] = df["w"] * 0.9 + design = SurveyDesign(weights="w", replicate_weights=["rw1", "rw2"], replicate_method="JK1") + with pytest.raises(NotImplementedError, match="replicate"): + _fit(df, survey=design, cluster="psu") + + def test_non_pweight_rejected(self, panel_df): + with pytest.raises(ValueError, match="pweight"): + _fit(panel_df, survey=SurveyDesign(weights="w", weight_type="fweight")) + + def test_total_fails_closed_on_survey_fit(self, panel_survey): + with pytest.raises(NotImplementedError): + panel_survey.aggregate("total") + + def test_total_admitted_on_complete_panel_bare_cluster(self, panel_df): + res = _fit(panel_df, cluster="psu") + total = res.aggregate("total") + assert np.isfinite(float(total.att[0])) + + def test_panel_mover_survey_column_rejected(self, panel_df): + df = panel_df.copy() + # weight varies WITHIN a unit -> panel-lane unit-constancy rejection + df.loc[df["time"] == 1, "w"] = df.loc[df["time"] == 1, "w"] * 2.0 + with pytest.raises(ValueError): + _fit(df, survey=_DESIGN) + + def test_xy_only_learner_rejected_up_front_under_survey(self, panel_df): + with pytest.raises(TypeError, match="sample_weight"): + _fit(panel_df, survey=_DESIGN, outcome_learner=_XYOnlyRegressor()) + + def test_weighted_and_kwargs_learners_accepted_under_survey(self, panel_df): + r1 = _fit(panel_df, survey=_DESIGN, outcome_learner=_WeightedRegressor()) + assert np.isfinite(r1.overall_att) + r2 = _fit(panel_df, survey=_DESIGN, outcome_learner=_KwargsRegressor()) + assert np.isfinite(r2.overall_att) + + +# --------------------------------------------------------------------------- +# 8. df threading + rendering +# --------------------------------------------------------------------------- + + +class TestDfThreading: + def test_per_cell_inference_uses_survey_df(self, panel_survey): + df_survey = panel_survey.survey_metadata.df_survey + for data in panel_survey.group_time_effects.values(): + if data.get("skip_reason") is not None or data.get("is_reference"): + continue + t_ref, p_ref, ci_ref = safe_inference( + data["effect"], data["se"], alpha=0.05, df=df_survey + ) + np.testing.assert_allclose(data["p_value"], p_ref, rtol=1e-12) + np.testing.assert_allclose(data["conf_int"], ci_ref, rtol=1e-12) + + def test_summary_keeps_t_labels_under_survey(self, panel_survey): + text = panel_survey.summary() + assert "t-stat" in text and "z-stat" not in text + assert "Survey Design" in text or "survey" in text.lower() + + def test_summary_keeps_t_labels_under_bare_cluster(self, panel_df): + res = _fit(panel_df, cluster="psu") + text = res.summary() + assert "t-stat" in text and "z-stat" not in text + + def test_summary_keeps_z_labels_without_design(self, panel_plain): + text = panel_plain.summary() + assert "z-stat" in text + + def test_to_dict_carries_cluster_fields(self, panel_survey): + d = panel_survey.to_dict() + assert d.get("cluster_name") == "psu" + assert d.get("n_clusters") == 20 + + def test_event_study_df_threads_kit_df_survey(self, panel_survey): + es = panel_survey.aggregate("event_study") + assert es.df is not None + df_vals = np.asarray(es.df, dtype=float) + finite = df_vals[np.isfinite(df_vals)] + assert finite.size > 0 + assert np.all(finite == panel_survey.survey_metadata.df_survey) + + def test_describe_target_parameter_names_cohort_masses(self, panel_survey, panel_plain): + from diff_diff._reporting_helpers import describe_target_parameter + + surv = describe_target_parameter(panel_survey) + assert "cohort-mass-weighted" in surv["name"] + plain = describe_target_parameter(panel_plain) + assert "valid-treated-count-weighted" in plain["name"] + + def test_design_effect_diagnostic_live_on_survey_fit(self, panel_survey): + from diff_diff.diagnostic_report import DiagnosticReport + + report = DiagnosticReport(panel_survey).to_dict() + block = report.get("design_effect") + assert block is not None + assert "not applicable" not in str(block).lower() + + +# --------------------------------------------------------------------------- +# 9. Aggregation masses +# --------------------------------------------------------------------------- + + +class TestAggregationMasses: + @pytest.mark.parametrize("fixture", ["panel_survey", "rcs_survey"]) + def test_survey_cohort_masses_drive_aggregations(self, fixture, request): + res = request.getfixturevalue(fixture) + for level in ("simple", "event_study", "group"): + agg = res.aggregate(level) + assert np.all(np.isfinite(np.asarray(agg.att, dtype=float))) + + def test_rcs_masses_equal_hand_computed_weighted_masses(self, rcs_df, rcs_survey): + bk = rcs_survey._aggregation_kit.bookkeeping + masses = fixed_cohort_agg_weights(bk) + raw = rcs_df["w"].to_numpy() + norm_factor = len(raw) / raw.sum() # resolve() normalizes to sum == n + for g in (3, 4): + expected = float(rcs_df.loc[rcs_df["g"] == g, "w"].sum()) * norm_factor + np.testing.assert_allclose(masses[g], expected, rtol=1e-10) + + def test_survey_mass_dict_no_int64_collision(self): + # Regression for the >2**53 float-key collision the RCS precompute + # deliberately avoids: the survey mass branch keys by NATIVE cohort + # values, so two int64 cohorts colliding as float64 stay distinct. + g1 = 2**53 + g2 = 2**53 + 1 + assert float(g1) == float(g2) # they WOULD collide as float keys + unit_cohorts = np.array([g1, g1, g2, 0], dtype=np.int64) + sw = np.array([1.0, 2.0, 5.0, 7.0]) + masses = fixed_cohort_agg_weights({"survey_weights": sw, "unit_cohorts": unit_cohorts}) + assert masses[g1] == 3.0 + assert masses[g2] == 5.0 diff --git a/tests/test_variance_conventions.py b/tests/test_variance_conventions.py index 28ab4ac66..c9b1a1e18 100644 --- a/tests/test_variance_conventions.py +++ b/tests/test_variance_conventions.py @@ -450,8 +450,10 @@ def snapshot(self): status="legitimate", reason=( "L3: Chang (2020) Thm 2 augmented-score plug-in variance " - "(per-unit influence function; normal-theory safe_inference " - "throughout, no cluster surface)" + "(per-unit influence function; normal-theory safe_inference on " + "no-design fits; survey_design=/cluster= fits use the design-based " + "CR1/weighted-IF per-cell variance with df=df_survey t-inference " + "outside the shared CR1 sandwich)" ), ), dict( @@ -476,7 +478,10 @@ def snapshot(self): reason=( "L3: Chang (2020) Thm 2 lambda-corrected augmented-score " "plug-in variance (Case 2; per-observation influence function; " - "normal-theory safe_inference throughout, no cluster surface)" + "normal-theory safe_inference on no-design fits; survey_design=/" + "cluster= fits use the design-based CR1/weighted-IF per-cell " + "variance with df=df_survey t-inference outside the shared CR1 " + "sandwich)" ), ), ] From f870fb7b590cc56683b571675c43a68463d2d258 Mon Sep 17 00:00:00 2001 From: igerber Date: Fri, 28 Aug 2026 07:37:00 -0400 Subject: [PATCH 2/4] fix(dml): bare-cluster provenance in reporting + practitioner refit snippet MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address PR #796 review feedback: - describe_target_parameter: split the declared-survey marker (survey_metadata) from bare-cluster provenance (df_inference) so an unweighted panel=False cluster= fit is described with the fixed cohort-row-mass RCS wording, not "SURVEY cohort masses" (its synthesized all-ones weights reduce exactly to the fixed row masses). - RCS dropped-observations warning keys its weighting note on survey_metadata for the same reason. - practitioner learner-sensitivity snippet: include the fit's actual cluster= (from cluster_name) whenever populated, so a survey_design-without-PSU + cluster= fit's refit keeps PSU-cohesive folds and clustered inference instead of silently dropping them. - docs survey matrix: DMLDiD Strata/PSU/FPC cell now states the support level — "Full (TSL; df = n_PSU - n_strata)". - regression tests for all of the above (TestClusterReportingProvenance). --- diff_diff/_reporting_helpers.py | 18 +++++----- diff_diff/dml_did.py | 5 ++- diff_diff/practitioner.py | 10 ++++-- docs/choosing_estimator.rst | 2 +- tests/test_survey_dml.py | 60 +++++++++++++++++++++++++++++++++ 5 files changed, 82 insertions(+), 13 deletions(-) diff --git a/diff_diff/_reporting_helpers.py b/diff_diff/_reporting_helpers.py index 9c79d64a7..e0f102d1c 100644 --- a/diff_diff/_reporting_helpers.py +++ b/diff_diff/_reporting_helpers.py @@ -152,14 +152,14 @@ def describe_target_parameter(results: Any) -> Dict[str, Any]: if name == "DMLDiDResults": # Survey/cluster fits switch the aggregation weight source to cohort # masses (survey masses on declared designs; all-ones cohort masses - # on bare cluster=). getattr-defaulted reads only: this function is - # exercised on attribute-less stubs by the exhaustiveness guard. - _has_design = ( - getattr(results, "survey_metadata", None) is not None - or getattr(results, "df_inference", None) is not None - ) + # on bare cluster=). survey_metadata is the declared-survey marker; + # bare cluster= keeps it None and carries its df on df_inference + # instead. getattr-defaulted reads only: this function is exercised + # on attribute-less stubs by the exhaustiveness guard. + _is_survey = getattr(results, "survey_metadata", None) is not None + _is_clustered = getattr(results, "df_inference", None) is not None if getattr(results, "panel", True) is False: - if _has_design: + if _is_survey: weight_name = "survey-cohort-mass-weighted" weight_clause = ( "cell weights are SURVEY cohort masses (per-observation " @@ -168,13 +168,15 @@ def describe_target_parameter(results: Any) -> Dict[str, Any]: "DMLDiD survey Note)" ) else: + # Plain AND bare-cluster RCS: a bare cluster='s synthesized + # all-ones weights reduce exactly to the fixed row masses. weight_name = "cohort-mass-weighted" weight_clause = ( "cell weights are FIXED cohort row masses (the CS-RCS " "convention, WIF-consistent; the per-cell complete-case " "``n_treated`` is display-only; REGISTRY DMLDiD RCS Note)" ) - elif _has_design: + elif _is_survey or _is_clustered: weight_name = "cohort-mass-weighted" weight_clause = ( "cell weights are cohort masses (survey masses on declared " diff --git a/diff_diff/dml_did.py b/diff_diff/dml_did.py index 7c3fe164d..61915d3f5 100644 --- a/diff_diff/dml_did.py +++ b/diff_diff/dml_did.py @@ -2076,10 +2076,13 @@ def fit( stacklevel=2, ) else: + # Declared-survey marker, NOT resolved_survey_unit: a bare + # cluster='s synthesized all-ones weights reduce exactly to + # the fixed row masses, so it takes the no-design wording. _weighting_note = ( "(survey fits: aggregation weights use survey cohort masses " "— see REGISTRY.md)" - if _has_design + if survey_metadata is not None else "(aggregation weights use fixed cohort row masses — " "see REGISTRY.md)" ) warnings.warn( diff --git a/diff_diff/practitioner.py b/diff_diff/practitioner.py index b0d018e4c..78b406986 100644 --- a/diff_diff/practitioner.py +++ b/diff_diff/practitioner.py @@ -544,9 +544,13 @@ def _handle_dml_did(results: Any): "alt = DMLDiD(outcome_learner='sieve', seed=0" + (", panel=False" if getattr(results, "panel", True) is False else "") + ( - ", cluster=..." - if getattr(results, "survey_metadata", None) is None - and getattr(results, "df_inference", None) is not None + # cluster_name covers bare cluster=, a design-injected + # cluster (survey_design without PSU + cluster=), and a + # design-owned PSU (where re-passing the same column is a + # silent no-op) — conservative so the refit never drops + # PSU-cohesive folds or clustered inference. + f", cluster={getattr(results, 'cluster_name', None)!r}" + if getattr(results, "cluster_name", None) is not None else "" ) + ").fit(\n" diff --git a/docs/choosing_estimator.rst b/docs/choosing_estimator.rst index 6fff4ecf8..8506fea08 100644 --- a/docs/choosing_estimator.rst +++ b/docs/choosing_estimator.rst @@ -1016,7 +1016,7 @@ estimation. The depth of support varies by estimator and variance method: - -- * - ``DMLDiD`` - Full (pweight only) - - ``n_PSU - n_strata`` + - Full (TSL; df = ``n_PSU - n_strata``) - -- - Multiplier (PSU) * - ``ChangesInChanges`` / ``QDiD`` diff --git a/tests/test_survey_dml.py b/tests/test_survey_dml.py index eaa6fb12b..c0d95c17e 100644 --- a/tests/test_survey_dml.py +++ b/tests/test_survey_dml.py @@ -113,6 +113,11 @@ def rcs_survey(rcs_df): return _fit(rcs_df, panel=False, survey=_DESIGN) +@pytest.fixture(scope="module") +def rcs_cluster(rcs_df): + return _fit(rcs_df, panel=False, cluster="psu") + + # --------------------------------------------------------------------------- # Duck learners for the capability-gate and weight-threading pins # --------------------------------------------------------------------------- @@ -415,6 +420,61 @@ def test_cluster_direct_mutation_revalidated_at_fit(self, panel_df): est.fit(panel_df, covariates=["x1"], **FIT_KW) +class TestClusterReportingProvenance: + """Bare cluster= carries df_inference but is NOT a survey fit — the + reporting and practitioner surfaces must not describe it as one.""" + + def test_bare_cluster_rcs_headline_uses_fixed_row_masses(self, rcs_cluster, rcs_survey): + from diff_diff._reporting_helpers import describe_target_parameter + + assert rcs_cluster.survey_metadata is None + assert rcs_cluster.df_inference is not None + bare = describe_target_parameter(rcs_cluster) + assert "survey" not in bare["name"] + assert "SURVEY" not in bare["definition"] + assert "FIXED cohort row masses" in bare["definition"] + surv = describe_target_parameter(rcs_survey) + assert "survey-cohort-mass-weighted" in surv["name"] + + def test_bare_cluster_panel_headline_keeps_cohort_mass_wording(self, panel_df): + # Panel bare cluster= DOES switch to cohort masses (all-ones survey + # weights replace per-cell complete-case n_treated) — pin that the + # provenance split keeps this branch intact. + from diff_diff._reporting_helpers import describe_target_parameter + + res = _fit(panel_df, cluster="psu") + desc = describe_target_parameter(res) + assert "cohort-mass-weighted" in desc["name"] + assert "bare ``cluster=``" in desc["definition"] + + def test_rcs_dropped_warning_wording_by_provenance(self, rcs_df): + df = rcs_df.copy() + victim = df.index[(df["g"] == 3) & (df["time"] == 3)][0] + df.loc[victim, "y"] = np.nan + with pytest.warns(UserWarning, match="fixed cohort row masses"): + _fit(df, panel=False, cluster="psu", ignore_warnings=False) + with pytest.warns(UserWarning, match="survey cohort masses"): + _fit(df, panel=False, survey=_DESIGN, ignore_warnings=False) + + def test_practitioner_refit_snippet_preserves_cluster(self, panel_df, rcs_cluster): + from diff_diff import practitioner_next_steps + + # Injected cluster on a PSU-less design must survive into the + # learner-sensitivity refit (it carries the PSU-cohesive folds and + # clustered inference that isolate the learner comparison). + res = _fit(panel_df, survey=SurveyDesign(weights="w", strata="stratum"), cluster="psu") + text = str(practitioner_next_steps(res)) + assert "survey_design=" in text + assert "cluster='psu'" in text + # Bare cluster= keeps its cluster too. + assert "cluster='psu'" in str(practitioner_next_steps(rcs_cluster)) + + def test_practitioner_refit_snippet_no_cluster_on_plain(self, panel_plain): + from diff_diff import practitioner_next_steps + + assert "cluster=" not in str(practitioner_next_steps(panel_plain)) + + # --------------------------------------------------------------------------- # 4. Folds # --------------------------------------------------------------------------- From f0c5f8b7ae7986d1ad149d3f6c89434d8ce961fa Mon Sep 17 00:00:00 2001 From: igerber Date: Fri, 28 Aug 2026 09:57:18 -0400 Subject: [PATCH 3/4] test(rdd): loosen headstart p3 golden rtol to 5e-8 (OpenBLAS CI leg breached 1e-8 by 0.5%) --- tests/test_rddensity.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_rddensity.py b/tests/test_rddensity.py index bbf4cc16f..fcf6eaaf1 100644 --- a/tests/test_rddensity.py +++ b/tests/test_rddensity.py @@ -60,12 +60,19 @@ # pilot quantities carry ~1e-8 relative float noise (floors_0_0 and # floors_50_0 are numerically the same selector run where the nLocalMin # gate does not bind). +# - the headstart p=3 configs: the real-data selector chain at order p+2=5 +# leaves the jackknife SEs with ~6e-9 relative noise under Accelerate and +# ~1e-8 under CI's OpenBLAS (se_right breached the 1e-8 bar by 0.5% on +# the ubuntu py3.14 pure-Python leg); 5e-8 keeps ~5x headroom over the +# worst observed BLAS while staying an order under the floor overrides. RTOL_OVERRIDES = { ("dgp_normal", "p7"): 2e-3, ("dgp_masspoints", "floors_0_0"): 1e-7, ("dgp_masspoints", "floors_50_0"): 1e-7, ("dgp_masspoints_small", "floors_0_0"): 1e-7, ("dgp_masspoints_small", "floors_50_0"): 1e-7, + ("headstart", "p3_each"): 5e-8, + ("headstart", "p3_common"): 5e-8, } From eafd6270b56a809f30a2286eb37c0b08eb7c91d8 Mon Sep 17 00:00:00 2001 From: igerber Date: Fri, 28 Aug 2026 10:07:47 -0400 Subject: [PATCH 4/4] fix(dml): recompute survey_metadata from raw weights, not normalized Address the P2 from the PR #796 review rerun: resolve() rescales pweights to mean 1, but both DMLDiD survey_metadata recompute sites (the cluster-inject path and the unit-level recompute) passed the resolved weights as compute_survey_metadata's raw_weights, so sum_weights and weight_range reported normalized-scale values (scale-invariant fields - effective_n, design_effect, df - were unaffected; estimates/inference untouched). Fix: capture the raw per-obs design weights once at fit, collapse them per unit (groupby-first, all_units order) on the panel lane, and pass raw to both sites. Rename the misnamed unused _survey_weights_raw unpack (it holds resolved/normalized weights) to _survey_weights_resolved. Raw-scale metadata pins added for panel, RCS, and injected-cluster fits (the module fixtures' uniform(0.5, 2.0) weights are non-unit-scale). The same normalized-weights pattern pre-exists in CS (staggered.py:2412), the staggered DDD engine (_staggered_triple_diff_engine.py:299), and ContinuousDiD (continuous_did.py:1107) - tracked as a TODO.md Actionable Backlog row rather than expanding this PR across three more estimator families mid-review. --- TODO.md | 1 + diff_diff/dml_did.py | 41 +++++++++++++++++++++++++++++++++------- tests/test_survey_dml.py | 30 +++++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 7 deletions(-) diff --git a/TODO.md b/TODO.md index 0a89c6d02..c4124e13a 100644 --- a/TODO.md +++ b/TODO.md @@ -37,6 +37,7 @@ Related tracking surfaces: | `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 | | Opt-in singleton-group pruning for TwoWayFixedEffects (static + event-study mode; reghdfe parity): singleton units/periods are currently RETAINED class-wide - the within-demeaned row is zero so points are unchanged, but N/G/residual-df count it and CR1/finite-sample SEs shift (~0.41019 -> 0.40962 measured; REGISTRY "Deviation from R" Note, R5 review) - reghdfe iteratively drops singletons by default while fixest retains them (diff-diff matches fixest); an opt-in knob needs iterative unit+period pruning with consistent cluster/survey/replicate/Conley array subsetting and a default-flip decision protocol (moves published SEs) | `diff_diff/twfe.py`, `diff_diff/estimators.py`, `diff_diff/utils.py` | 3(a) R5 | Mid | Low | | Cohort-timing validation input for the simultaneous-adoption event-study family (TWFE `event_study=True` + MultiPeriodDiD through 3.9): an optional `first_treat=`/`cohort=` column so simultaneous adoption becomes checkable under the contract-valid time-invariant `D_i` indicator - today the staggered-adoption advisory derives timing from within-unit 0->1 transitions, so it can only fire on off-contract time-varying `D_it` input, and with valid `D_i` adoption timing is not observable in the inputs at all (REGISTRY "staggered-adoption detection limit" Notes, both sections); design questions: validate-only vs steering error, and interplay with the M-011 removal | `diff_diff/twfe.py`, `diff_diff/estimators.py` | 3(a) R2 | Mid | Medium | +| Unit-level `survey_metadata` recompute passes NORMALIZED weights as `compute_survey_metadata`'s `raw_weights` in three pre-existing families — CS (`staggered.py:2412`), the staggered DDD engine (`_staggered_triple_diff_engine.py:299`), and ContinuousDiD (`continuous_did.py:1107`, where the variable is even named `raw_w_unit`) — so `sum_weights`/`weight_range` report the mean-1 rescaled values (scale-invariant fields are unaffected; estimates/inference untouched). Same class as the DMLDiD fix in the DML survey PR (raw per-obs weights collapsed per unit, groupby-first): apply the identical pattern per family with raw-scale metadata pins on non-unit-scale weights | `diff_diff/staggered.py`, `diff_diff/_staggered_triple_diff_engine.py`, `diff_diff/continuous_did.py` | DML survey PR review | Quick | Medium | | Bootstrapped ES REPLAY containers (CS, DMLDiD and EfficientDiD) publish the analytical `survey_metadata.df_survey` provenance scalar beside percentile inference on survey fits — a cross-estimator provenance residual (shipped CS behaves identically; the per-row df/inference channels are correctly NaN, the scalar metadata field is the residual; DMLDiD joined with its 3.11 survey support). Evaluate clearing/gating it on all three estimators together (cross-surface twins) | `diff_diff/staggered_results.py`, `diff_diff/dml_did_results.py`, `diff_diff/efficient_did_results.py`, `diff_diff/results_base.py` | EDiD-replay review | Quick | Low | | DMLDiD replicate-weight survey designs (currently fail closed with `NotImplementedError` pointing here): the IF-reweighting route (`compute_replicate_if_variance` on the augmented scores) is BETTER justified for DMLDiD than for CS — Neyman orthogonality makes the score first-order insensitive to nuisance perturbation — but needs per-cell replicate SE plumbing, the aggregate `_se_from_psi` replicate branch already activates on kit keys, and its own rejection/df test matrix (QR-rank df, `n_valid - 1` tightening) | `diff_diff/dml_did.py` | DML survey PR | Mid | Low | | ImputationDiD/TwoStageDiD `aggregate()` recompute levels on bootstrapped fits fail closed ('simple' relays since the M-027 per-level convergence; M-021/M-022); ImputationDiD's per-target psi machinery makes seeded replay tractable (the panel-backed kit retains everything the psi precompute reads), TwoStageDiD's per-level GMM scores are function-locals and would need retention | `diff_diff/imputation_results.py`, `diff_diff/two_stage_results.py`, `diff_diff/aggregation.py` | 2(b) PR-3b | Mid | Low | diff --git a/diff_diff/dml_did.py b/diff_diff/dml_did.py index 61915d3f5..40ff713e9 100644 --- a/diff_diff/dml_did.py +++ b/diff_diff/dml_did.py @@ -1861,7 +1861,7 @@ def fit( ( resolved_survey, - _survey_weights_raw, + _survey_weights_resolved, survey_weight_type, survey_metadata, ) = _resolve_survey_for_fit(survey_design, data, "analytical") @@ -1875,6 +1875,21 @@ def fit( "(weights/strata/psu/fpc) instead." ) + # Raw (pre-normalization) per-obs design weights, for metadata + # provenance only: compute_survey_metadata expects the ORIGINAL + # scale (resolve() rescales pweights to mean 1, so the resolved + # .weights would misreport sum_weights/weight_range; the + # scale-invariant fields — effective_n, design_effect, df — are + # unaffected either way). + raw_obs_weights: Optional[np.ndarray] = None + if resolved_survey is not None: + assert survey_design is not None + raw_obs_weights = ( + data[survey_design.weights].values.astype(np.float64) + if survey_design.weights + else np.ones(len(data), dtype=np.float64) + ) + effective_survey_design = survey_design cluster_ids_for_check: Optional[np.ndarray] = None if self.cluster is not None: @@ -1895,7 +1910,7 @@ def fit( effective_survey_design = SurveyDesign(psu=self.cluster, weight_type="pweight") ( resolved_survey, - _survey_weights_raw, + _survey_weights_resolved, survey_weight_type, _synth_metadata, ) = _resolve_survey_for_fit(effective_survey_design, data, "analytical") @@ -1908,7 +1923,8 @@ def fit( effective_survey_design = _dc_replace(survey_design, psu=self.cluster) resolved_survey = _inject_cluster_as_psu(resolved_survey, cluster_ids_for_check) - survey_metadata = compute_survey_metadata(resolved_survey, resolved_survey.weights) + assert raw_obs_weights is not None + survey_metadata = compute_survey_metadata(resolved_survey, raw_obs_weights) else: # Both supplied: the design's PSU wins (warn on differing # partitions); the return value is intentionally unused. @@ -1948,12 +1964,23 @@ def fit( cell_fn = self._compute_dml_rcs_gt # Survey metadata reflects the estimation index space (units on the - # panel lane; the RCS lane's per-obs design is its own unit level). + # panel lane; the RCS lane's per-obs design is its own unit level), + # recomputed with RAW weights collapsed the same way (the resolved + # unit weights are normalized — see raw_obs_weights above). resolved_survey_unit = precomputed.get("resolved_survey_unit") if survey_metadata is not None and resolved_survey_unit is not None: - survey_metadata = compute_survey_metadata( - resolved_survey_unit, resolved_survey_unit.weights - ) + assert raw_obs_weights is not None + if self.panel: + raw_unit_weights = ( + pd.Series(raw_obs_weights, index=df.index) + .groupby(df[unit]) + .first() + .reindex(precomputed["all_units"]) + .to_numpy(dtype=np.float64) + ) + else: + raw_unit_weights = raw_obs_weights + survey_metadata = compute_survey_metadata(resolved_survey_unit, raw_unit_weights) df_survey = precomputed.get("df_survey") # PSU-cohesive folds whenever the PSU is strictly coarser than the diff --git a/tests/test_survey_dml.py b/tests/test_survey_dml.py index c0d95c17e..d36369972 100644 --- a/tests/test_survey_dml.py +++ b/tests/test_survey_dml.py @@ -475,6 +475,36 @@ def test_practitioner_refit_snippet_no_cluster_on_plain(self, panel_plain): assert "cluster=" not in str(practitioner_next_steps(panel_plain)) +class TestSurveyMetadataRawScale: + """survey_metadata provenance must be on the RAW weight scale: resolve() + rescales pweights to mean 1, so recomputing at the estimation index + level with the resolved weights would misreport sum_weights and + weight_range (the module fixtures' uniform(0.5, 2.0) weights are + deliberately non-unit-scale).""" + + def test_panel_metadata_uses_raw_unit_weights(self, panel_df, panel_survey): + md = panel_survey.survey_metadata + unit_w = panel_df.groupby("unit")["w"].first() + np.testing.assert_allclose(md.sum_weights, unit_w.sum(), rtol=1e-12) + np.testing.assert_allclose(md.weight_range, (unit_w.min(), unit_w.max()), rtol=1e-12) + + def test_rcs_metadata_uses_raw_obs_weights(self, rcs_df, rcs_survey): + md = rcs_survey.survey_metadata + w = rcs_df["w"] + np.testing.assert_allclose(md.sum_weights, w.sum(), rtol=1e-12) + np.testing.assert_allclose(md.weight_range, (w.min(), w.max()), rtol=1e-12) + + def test_injected_cluster_metadata_uses_raw_unit_weights(self, panel_df): + # PSU-less design + cluster=: metadata is recomputed on the inject + # path AND again at unit level — the final values must still be the + # raw unit-level scale. + res = _fit(panel_df, survey=SurveyDesign(weights="w", strata="stratum"), cluster="psu") + md = res.survey_metadata + unit_w = panel_df.groupby("unit")["w"].first() + np.testing.assert_allclose(md.sum_weights, unit_w.sum(), rtol=1e-12) + np.testing.assert_allclose(md.weight_range, (unit_w.min(), unit_w.max()), rtol=1e-12) + + # --------------------------------------------------------------------------- # 4. Folds # ---------------------------------------------------------------------------