feat(dml): replicate-weight survey designs via IF-reweighting (per-cell + aggregate) (#807)

Author: igerber

TODO.md

@@ -37,7 +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 |
-| 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 |
+| CS parity with DMLDiD's replicate conventions: per-cell replicate SEs for `CallawaySantAnna` (CS leaves per-cell SEs on the weighted sqrt-sum under replicate designs; DMLDiD now computes them via `compute_replicate_if_variance`) AND flip the replace-style effective-df relays to min-cap (`min(df_survey, n_valid - 1)`) at all five sites: CS overall (`staggered.py` overall relay), the shared event-study/group aggregation sites (`staggered_aggregation.py` `df_survey_val = min(non_none_dfs)` x2), and the staggered-DDD engine twins (`_staggered_triple_diff_engine.py` overall + event-study) — the replace convention can RAISE df above the QR-rank design df (anti-conservative); latent, fires only when `n_valid < R` (all-zero/non-finite replicate column). DMLDiD's min-cap + the pinned inherited ES/group behavior are the templates (REGISTRY DMLDiD replicate Note) | `diff_diff/staggered.py`, `diff_diff/staggered_aggregation.py`, `diff_diff/_staggered_triple_diff_engine.py` | DML replicate 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 |
 | practitioner `step_name="heterogeneity"` producer-side collisions: three OTHER estimators' advice steps reuse the key with non-heterogeneity labels (`:975` ContinuousDiD dose-response, `:1022` Triple placebo-group, `:1413` LPDiD WAS arrays), so DiagnosticReport's heterogeneity completion silently drops that unrelated advice from `next_steps` via `_filter_steps` - the same latent collision fixed for StackedDiD in M-024 (renamed to `sub_experiment_balance`). Renaming these changes those estimators' report output; audit + rename with per-estimator pins. | `diff_diff/practitioner.py` | 2(b) PR-2 review R9 | Quick | Low |

changelog.d/20260830-dml-replicate-weights.md

@@ -0,0 +1,12 @@
+### Added
+- **DMLDiD replicate-weight survey designs**: `DMLDiD` now accepts
+  replicate-weight `SurveyDesign`s (BRR / Fay / JK1 / JKn / SDR) on both
+  lanes (panel and repeated cross sections), computing per-cell AND
+  aggregate variances by IF-reweighting the augmented cross-fitted scores
+  (`compute_replicate_if_variance`; nuisances are not re-estimated per
+  replicate). Inference uses `df = rank(replicate matrix) - 1` with
+  `min(df_survey, n_valid - 1)` capping; degenerate cells (zero or
+  non-finite replicate variance) fail closed to NaN inference. Replicate +
+  `cluster=` and replicate + `n_bootstrap > 0` are rejected with targeted
+  errors (previously all replicate designs failed closed with a blanket
+  `NotImplementedError`).

diff_diff/dml_did.py

@@ -1364,15 +1364,43 @@ def _compute_dml_gt(
             "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.
+        # Per-cell SE. Replicate designs use IF-reweighting on the Hajek
+        # payload (compute_replicate_if_variance; the same psi the aggregate
+        # _se_from_psi call consumes) — a zero or non-finite replicate
+        # variance is degenerate and fails closed to NaN (stricter than the
+        # shared aggregate clamp; REGISTRY DMLDiD Note). 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 (
+            if resolved_survey_unit is not None and resolved_survey_unit.uses_replicate_variance:
+                from diff_diff.survey import compute_replicate_if_variance
+
+                variance, n_valid_rep = compute_replicate_if_variance(
+                    inf_full, resolved_survey_unit
+                )
+                if not np.isfinite(variance) or variance <= 0.0:
+                    se = float("nan")
+                else:
+                    se = float(np.sqrt(variance))
+                # Per-cell df: min(design df, n_valid - 1) — the dCDH
+                # _effective_df_survey rule, inlined. n_valid is computed
+                # over the WHOLE replicate columns, so it equals R for every
+                # cell in practice (defensive; REGISTRY Note). df_survey is
+                # a CELL-LOCAL binding — never mutate
+                # precomputed["df_survey"], which feeds the post-fit
+                # aggregation kit.
+                if df_survey is not None:
+                    df_survey = min(int(df_survey), int(n_valid_rep) - 1)
+                else:
+                    # Undefined replicate df (QR rank <= 1): df=0 sentinel
+                    # -> NaN inference (CS sentinel parity).
+                    df_survey = 0
+            elif (
                 resolved_survey_unit is not None
                 and getattr(resolved_survey_unit, "psu", None) is not None
             ):
@@ -1386,9 +1414,10 @@ def _compute_dml_gt(
             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.
+        # design-based branches — a NaN from the CR1 helper (or a degenerate
+        # replicate variance) 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 (
@@ -1760,12 +1789,28 @@ def _compute_dml_rcs_gt(
         }
 
         # Per-cell SE (same dispatch as the panel cell; see the comment
-        # there): PSU designs -> CS per-cell CR1 helper (NaN propagates as
+        # there): replicate designs -> IF-reweighting with the degenerate
+        # fail-closed guard and the cell-local min(df, n_valid - 1) rule;
+        # 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 (
+            if resolved_survey_unit is not None and resolved_survey_unit.uses_replicate_variance:
+                from diff_diff.survey import compute_replicate_if_variance
+
+                variance, n_valid_rep = compute_replicate_if_variance(
+                    inf_full, resolved_survey_unit
+                )
+                if not np.isfinite(variance) or variance <= 0.0:
+                    se = float("nan")
+                else:
+                    se = float(np.sqrt(variance))
+                if df_survey is not None:
+                    df_survey = min(int(df_survey), int(n_valid_rep) - 1)
+                else:
+                    df_survey = 0
+            elif (
                 resolved_survey_unit is not None
                 and getattr(resolved_survey_unit, "psu", None) is not None
             ):
@@ -1831,19 +1876,24 @@ def fit(
         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).
+            Complex survey design (pweight-only). Declared designs weight
+            the moment kernels (Hajek p-hat/lambda-hat/theta) and pass
+            ``sample_weight`` into the nuisance learners (user learner
+            objects must accept ``sample_weight`` by keyword — a learner
+            without it is rejected up front). Two variance lanes:
+            full-design TSL (weights/strata/PSU/FPC) switches cross-fitting
+            to PSU-cohesive folds when the PSU is strictly coarser than the
+            sampling unit and routes the per-cell and aggregate variances
+            through the design-based kernels with ``df = n_PSU - n_strata``
+            t-inference; replicate-weight designs (BRR / Fay / JK1 / JKn /
+            SDR) compute per-cell AND aggregate variances by IF-reweighting
+            the cross-fitted scores with ``df = rank(replicate matrix) - 1``
+            t-inference (nuisances are not re-estimated per replicate;
+            REGISTRY DMLDiD Note). Replicate designs reject ``cluster=``
+            and ``n_bootstrap > 0`` combinations. 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).
         """
         df, covariates = self._validate_and_prepare(
             data, outcome, unit, time, first_treat, covariates
@@ -1866,13 +1916,18 @@ def fit(
             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:
+        # Replicate + bootstrap rejected FIRST (before any fit work) —
+        # replicate variance is an analytical alternative, not compatible
+        # with bootstrap (CS parity, staggered.py).
+        if (
+            self.n_bootstrap > 0
+            and 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."
+                "DMLDiD bootstrap (n_bootstrap > 0) is not supported "
+                "with replicate-weight survey designs. Replicate weights provide "
+                "analytical variance; use n_bootstrap=0 instead."
             )
 
         # Raw (pre-normalization) per-obs design weights, for metadata
@@ -1902,6 +1957,25 @@ def fit(
                     "drop or impute them before fitting"
                 )
             cluster_ids_for_check = _cluster_col.to_numpy()
+            # Reject replicate-weight + cluster= AFTER the column checks (CS
+            # ordering: a bogus cluster name raises ValueError first).
+            # Replicate IF variance is computed by replicate reweighting and
+            # ignores PSU/cluster entirely (replicate_weights are mutually
+            # exclusive with strata/psu/fpc) — honoring cluster= would
+            # silently have no effect on the variance, and the inject-as-PSU
+            # paths below would violate that mutual exclusion.
+            if resolved_survey is not None and resolved_survey.uses_replicate_variance:
+                raise NotImplementedError(
+                    f"DMLDiD(cluster={self.cluster!r}) is not "
+                    "supported with replicate-weight survey designs. "
+                    "Replicate-weight variance is computed by replicate "
+                    "reweighting (BRR / Fay / JK1 / JKn / SDR) and ignores "
+                    "PSU/cluster entirely — setting cluster= would silently "
+                    "have no effect on the variance estimate. Either omit "
+                    "cluster= (the replicate weights encode the design "
+                    "structure implicitly) or use a non-replicate survey "
+                    "design (with explicit strata/psu/fpc)."
+                )
             if resolved_survey is None:
                 # Bare cluster=: synthesize a PSU-only design. survey_metadata
                 # stays None DELIBERATELY (it is the declared-survey marker:
@@ -2188,15 +2262,31 @@ def fit(
 
         # Overall ATT (simple aggregation over post-treatment finite cells).
         # 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.
+        # dropped replicates (n_valid < R). MIN-CAP, not CS's replace: the
+        # QR-rank design df stays the ceiling (CS's replace convention can
+        # RAISE df above the design df — anti-conservative; deliberate
+        # documented divergence, REGISTRY DMLDiD Note + CS-parity TODO row).
         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
+        if overall_effective_df is not None and df_survey is not None:
+            df_survey = min(int(df_survey), int(overall_effective_df))
+            # Propagate to survey_metadata for display consistency (CS
+            # parity) — the capped value, never the sentinel below.
+            if survey_metadata is not None:
+                survey_metadata.df_survey = df_survey
+        # Replicate design with undefined df (QR rank <= 1): df=0 sentinel
+        # -> NaN inference, applied to the LOCAL df only (survey_metadata
+        # keeps None).
+        df_overall = df_survey
+        if (
+            df_survey is None
+            and resolved_survey is not None
+            and resolved_survey.uses_replicate_variance
+        ):
+            df_overall = 0
         overall_t_stat, overall_p_value, overall_conf_int = safe_inference(
-            overall_att, overall_se, alpha=self.alpha, df=df_survey
+            overall_att, overall_se, alpha=self.alpha, df=df_overall
         )
 
         # Optional multiplier bootstrap (keyword form; aggregate=None is the

diff_diff/dml_did_results.py

@@ -15,9 +15,12 @@
 and DMLDiD's augmented-score SE ``sqrt(mean(psi_bar**2)/n)`` is exactly
 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).
+``survey_design=``/``cluster=`` the per-cell SE is design-based instead:
+full-design TSL fits use the CR1 / weighted-IF variance (the CS
+clustered-``hc1`` convention: ``SurveyDesign(psu=...)`` routed through the
+shared stratified-PSU meat), while replicate-weight fits use IF-reweighting
+via ``compute_replicate_if_variance`` on the same per-cell payload
+(``df = rank(replicate matrix) - 1``; REGISTRY DMLDiD Note).
 """
 
 from dataclasses import dataclass, field