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).

Author: igerber

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

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 |

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

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 |