From 0d310af657bfbfd94279253675cb194372d4826a Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 29 Aug 2026 10:30:26 -0400 Subject: [PATCH 1/4] fix(survey): unit-level survey_metadata recompute reports raw weight scale CS (panel+RC), the staggered DDD engine, ContinuousDiD (analytic branch), and EfficientDiD passed resolve()'s mean-1 normalized weights as compute_survey_metadata's raw_weights, misreporting sum_weights and weight_range (scale-invariant fields and all estimates/inference unaffected). Capture the raw weight column per family (the DMLDiD #796 pattern), snapshotting before working-frame mutations so a weight column aliasing a mutable role column (e.g. weights == dose in ContinuousDiD) still surfaces the user's original values. Metadata provenance only: estimates, SEs, p-values, CIs, df_survey, n_strata, n_psu byte-identical; effective_n/design_effect unchanged within float round-off. --- CHANGELOG.md | 14 + TODO.md | 2 +- diff_diff/_staggered_triple_diff_engine.py | 16 +- diff_diff/continuous_did.py | 27 +- diff_diff/efficient_did.py | 30 +- diff_diff/staggered.py | 27 +- docs/methodology/REGISTRY.md | 2 +- tests/test_survey_metadata_raw_scale.py | 442 +++++++++++++++++++++ 8 files changed, 548 insertions(+), 12 deletions(-) create mode 100644 tests/test_survey_metadata_raw_scale.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e4991162e..9e76d8edf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- **`survey_metadata` raw-scale provenance on the unit-level recompute** + (CallawaySantAnna panel + repeated-cross-section lanes, + `TripleDifference`/`StaggeredTripleDifference` staggered engine, + `ContinuousDiD` analytical branch, `EfficientDiD`): the recompute passed + the RESOLVED (mean-1 rescaled) weights as `compute_survey_metadata`'s + raw weights, so `sum_weights`/`weight_range` reported the normalized + scale instead of the user's original weight scale. They now report the + raw scale, matching every other estimator (DMLDiD got the pattern in + its survey PR). Metadata-provenance only: estimates, SEs, p-values, + CIs, `df_survey`, `n_strata`, `n_psu` are byte-identical, and + `effective_n`/`design_effect` are scale-invariant (unchanged within + floating-point round-off). + ## [3.11.0] - 2026-08-29 ### Added diff --git a/TODO.md b/TODO.md index c4124e13a..879a71bfe 100644 --- a/TODO.md +++ b/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 | -| 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 | +| `ContinuousDiD` `survey_metadata` granularity diverges across inference branches: the analytical branch recomputes at UNIT level (raw unit-weight sum/range since the raw-scale provenance fix) while `n_bootstrap>0` keeps the OBS-level metadata from `_resolve_survey_for_fit` (raw per-obs sum over the dose-filtered frame) — both raw-scale, but `sum_weights` differs by panel length between branches on the same data (characterization pin: `TestContinuousDiDRawScale::test_bootstrap_metadata_stays_obs_level_raw`). Unify on one granularity (unit-level, matching the analytic branch and CS/EfficientDiD convention) | `diff_diff/continuous_did.py` | survey-meta-raw | 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 | diff --git a/diff_diff/_staggered_triple_diff_engine.py b/diff_diff/_staggered_triple_diff_engine.py index 886990547..1f24ebe42 100644 --- a/diff_diff/_staggered_triple_diff_engine.py +++ b/diff_diff/_staggered_triple_diff_engine.py @@ -295,8 +295,20 @@ def _fit_staggered_core( if resolved_survey is not None and survey_metadata is not None: resolved_survey_unit = precomputed.get("resolved_survey_unit") if resolved_survey_unit is not None: - unit_w = resolved_survey_unit.weights - survey_metadata = compute_survey_metadata(resolved_survey_unit, unit_w) + from diff_diff.survey import _extract_unit_survey_weights + + # Raw (pre-normalization) unit weights for metadata + # provenance: compute_survey_metadata expects the ORIGINAL + # scale (resolve() rescales pweights to mean 1, so the + # resolved weights would misreport sum_weights/ + # weight_range; scale-invariant fields are unaffected + # either way). Read from ``data``: ``df`` has first_treat + # overwritten by this point. + assert survey_design is not None + raw_unit_w = _extract_unit_survey_weights( + data, unit, survey_design, precomputed["all_units"] + ) + survey_metadata = compute_survey_metadata(resolved_survey_unit, raw_unit_w) # Survey df for t-distribution critical values df_survey = precomputed.get("df_survey") diff --git a/diff_diff/continuous_did.py b/diff_diff/continuous_did.py index dd1ce123d..ee1488677 100644 --- a/diff_diff/continuous_did.py +++ b/diff_diff/continuous_did.py @@ -506,6 +506,16 @@ def fit( if col not in df.columns: raise ValueError(f"Column '{col}' not found in data.") + # Snapshot the raw survey-weight column BEFORE any df mutation + # (never-treated dose zeroing, first_treat inf->0, to_numeric): a + # design whose weight column aliases a mutable role column (e.g. + # weights == dose) must still surface the user's ORIGINAL values in + # survey_metadata. Per-unit via groupby-first so the later + # dose-filter (which drops whole units) cannot desync alignment. + raw_unit_w_meta: Optional[pd.Series] = None + if survey_design is not None and survey_design.weights: + raw_unit_w_meta = data.groupby(unit)[survey_design.weights].first() + # Covariate-path guards (conditional parallel trends). if cov_cols: if survey_design is not None: @@ -1103,7 +1113,22 @@ def fit( if _unit_resolved is not None: from diff_diff.survey import compute_survey_metadata - raw_w_unit = _unit_resolved.weights + # Raw (pre-normalization) unit weights for metadata + # provenance: compute_survey_metadata expects the + # ORIGINAL scale (resolve() rescales pweights to mean 1, + # so _unit_resolved.weights would misreport sum_weights/ + # weight_range; scale-invariant fields are unaffected + # either way). ``raw_unit_w_meta`` was snapshotted from + # pristine ``data`` before the df mutations; reindexing + # to ``all_units`` (the dose-filtered unit order that + # built unit_resolved) keeps alignment — survey weights + # are unit-constant (validated at resolve time). + assert survey_design is not None + raw_w_unit = ( + raw_unit_w_meta.reindex(precomp["all_units"]).to_numpy(dtype=np.float64) + if raw_unit_w_meta is not None + else np.ones(precomp["n_units"], dtype=np.float64) + ) survey_metadata = compute_survey_metadata(_unit_resolved, raw_w_unit) # Propagate replicate df override to survey_metadata for display diff --git a/diff_diff/efficient_did.py b/diff_diff/efficient_did.py index 4df2051d8..b28844b8f 100644 --- a/diff_diff/efficient_did.py +++ b/diff_diff/efficient_did.py @@ -1397,6 +1397,20 @@ def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any apply_bootstrap_group_overrides(group_effects, bootstrap_results, self.alpha) # ----- Build results ----- + # Raw (pre-normalization) unit weights for the metadata recompute: + # compute_survey_metadata expects the ORIGINAL scale (resolve() + # rescales pweights to mean 1, so the resolved unit weights would + # misreport sum_weights/weight_range; scale-invariant fields are + # unaffected either way). + raw_unit_w_meta: Optional[np.ndarray] = None + if self._unit_resolved_survey is not None: + assert survey_design is not None + raw_obs_w_meta = ( + data[survey_design.weights].values.astype(np.float64) + if survey_design.weights + else np.ones(len(data), dtype=np.float64) + ) + raw_unit_w_meta = raw_obs_w_meta[self._unit_first_panel_row] self.results_ = EfficientDiDResults( group_time_effects=group_time_effects, overall_att=overall_att, @@ -1452,7 +1466,7 @@ def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any kernel_bandwidth=self.kernel_bandwidth, omega_ridge=self.omega_ridge, survey_metadata=( - self._recompute_unit_survey_metadata(survey_metadata) + self._recompute_unit_survey_metadata(survey_metadata, raw_unit_w_meta) if survey_metadata is not None else None ), @@ -1483,14 +1497,22 @@ def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any self.is_fitted_ = True return self.results_ - def _recompute_unit_survey_metadata(self, panel_metadata): - """Recompute survey metadata from unit-level design if available.""" + def _recompute_unit_survey_metadata(self, panel_metadata, raw_unit_weights=None): + """Recompute survey metadata from unit-level design if available. + + ``raw_unit_weights`` carries the ORIGINAL-scale (pre-normalization) + unit weights and MUST be passed whenever ``_unit_resolved_survey`` + is set — never fall back to ``_unit_resolved_survey.weights``, which + resolve() rescaled to mean 1 and would misreport sum_weights/ + weight_range (scale-invariant fields are unaffected either way). + """ if self._unit_resolved_survey is not None: from diff_diff.survey import compute_survey_metadata + assert raw_unit_weights is not None meta = compute_survey_metadata( self._unit_resolved_survey, - self._unit_resolved_survey.weights, + raw_unit_weights, ) # Propagate effective replicate df if available # (but not the df=0 sentinel — keep metadata as None for undefined df) diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 94ddd0f1f..1389f5a67 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -2406,10 +2406,31 @@ def fit( if resolved_survey is not None and survey_metadata is not None: resolved_survey_unit = precomputed.get("resolved_survey_unit") if resolved_survey_unit is not None: - from diff_diff.survey import compute_survey_metadata + from diff_diff.survey import ( + _extract_unit_survey_weights, + compute_survey_metadata, + ) - unit_w = resolved_survey_unit.weights - survey_metadata = compute_survey_metadata(resolved_survey_unit, unit_w) + # Raw (pre-normalization) weights for metadata provenance: + # 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). Read from ``data``: ``df`` + # has first_treat/_never_treated overwritten by this point. + assert survey_design is not None + if precomputed.get("is_panel", True): + raw_unit_w = _extract_unit_survey_weights( + data, unit, survey_design, precomputed["all_units"] + ) + else: + # RC lane: resolved_survey_unit is per-observation. + raw_unit_w = ( + data[survey_design.weights].values.astype(np.float64) + if survey_design.weights + else np.ones(len(data), dtype=np.float64) + ) + survey_metadata = compute_survey_metadata(resolved_survey_unit, raw_unit_w) # Survey df for safe_inference calls — use the unit-level resolved # survey df computed in _precompute_structures for consistency. diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index fbb3bf532..20b26ccb0 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1188,7 +1188,7 @@ The multiplier bootstrap uses random weights w_i with E[w]=0 and Var(w)=1: - **Note:** Non-survey DR path also includes nuisance IF corrections (PS + OR), matching the survey path structure (Phase 7a). Previously used plug-in IF only. As of v3.7 the non-survey reg and ipw paths carry their corrections too (OR estimation-effect / PS score), so the nuisance-IF treatment is method-uniform. - **Note (post-fit aggregate() - rows M-020/M-117):** `fit(aggregate=)` is deprecated (3.9; removed 4.0) in favor of post-fit `results.aggregate(type, balance_e=)` on the fit-retained aggregation kit ('simple' relays the stored overall inference; 'event_study'/'group' recompute from the kit — on BOOTSTRAPPED fits they REPLAY the fit-time multiplier bootstrap from the kit's `BootstrapReplaySpec` RNG-state snapshot: percentile se/CI/cband matching a fit-time aggregation to BLAS reassociation (~1 ULP, `assert_allclose`, never bit-identity; the discrete percentile p-value is a count statistic outside that claim), vcov/df cleared (no analytical provenance beside percentile inference), backend-stamped and failing closed on a Rust-vs-NumPy weight-backend mismatch or a pre-replay legacy pickle — the M-027 per-level policy, with the former blanket fail-closed retired). `DiagnosticReport` now derives the event-study container internally on plain fits (raw `event_study_effects` absent), so its `parallel_trends`, `pretrends_power`, and `sensitivity` checks run without the deprecated fit-time kwarg — CS containers are M-093-admitted into `compute_pretrends_power` / `HonestDiD.sensitivity_analysis` with pinned raw-route parity, and the raw field, when present (the requested-but-empty `{}` included), always takes precedence. `heterogeneity` is unaffected (it reads `group_time_effects` on plain fits). Bootstrapped fits now DERIVE successfully via the percentile-bootstrap replay: the derived container carries `vcov=None`, so `parallel_trends` runs on the Bonferroni per-period fallback and `pretrends_power`/`sensitivity` ride the diagonal-covariance fallback; the replay's re-emitted fit-time warnings are recorded and republished per section. The fail-closed skip remains for kit-less/legacy pickles, backend-mismatched replays, and the sibling estimators' bootstrap gates. The recompute levels' replay RE-EMITS the fit-time bootstrap warnings by design (the relay levels' no-re-warn convention is scoped to 'simple'/'total'). -- **Note (post-fit `aggregate('total')` - the estimator-owned total incremental outcome, 3.10):** (a) **Estimand:** `total = C x overall_att` with `se = C x overall_se`, CI scaled by `C`, `t`/`p`/`df` inherited unchanged - an exact relay CONDITIONAL ON THE REALIZED AGGREGATION MASS (never call it unqualified "exact"): `C` is the design-fixed complete-case treated-observation count on every admitted routing, the inherited SE prices the overall's full inference (WIF/estimated-share terms included) but treats `C` as fixed at its realized value, and the random-mass `att*dC` variance term is deliberately out of scope (the DEFERRED survey/RC remainder). Every SE branch is homogeneous of degree 1 in the scaling, so the relay is exact under analytical AND percentile-bootstrap inference. (b) **Fail-closed routing gates**, in order: (1) repeated-cross-section-routed fits (`panel=False`, or genuinely-unbalanced `allow_unbalanced_panel=True` fits - a balanced panel with the flag stays panel-routed and admitted) raise `NotImplementedError` - the RC cohort mass counts UNITS and weights every kept post period by it, a ~T-fold overcount of treated observations on a T-wave RCS (execution-verified); (2) fits declaring a `survey_design=` raise - the realized-mass relay omits the survey mass-uncertainty term and design-aware population-scale totals are not implemented (retained weight scale differs by design family: analytic pweight resolves normalized to sum = n and CS retains no raw record - CS accepts pweight only - while replicate designs retain raw scale); the gate reads the fit-time `is_survey_fit` KIT snapshot, so bare-`cluster=` fits are ADMITTED (their synthesized all-ones internal design is not a survey) and a post-fit edit of the public `survey_metadata` cannot bypass it; the DECLARATION BOUNDARY is deliberate - an explicit unweighted `SurveyDesign(psu=...)` fit is numerically identical to bare `cluster=` yet fails closed by declaration (declaring survey_design opts into the survey contract; conservative, never wrong numbers; the message points at the `cluster=` route); a combined survey+RC fit deterministically gets the RC message; (3) the coincidence check: on bare-`cluster=` fits `_aggregate_simple` weights kept cells by the synthesized cohort masses (full cohort size per kept cell), and when kept cells have INCOMPLETE treated support that realized mass diverges from the complete-case count - the same overcount class as (1) - so the fit raises rather than publishing an ambiguous mass; clean bare-`cluster=` fits (the normal case) are admitted with `C` equal to both counts. Panel-routed pre-upgrade (legacy) kits missing the fit-time snapshots raise the refit message - never a mutable-field or public-replay fallback. (c) **Filter-replay contract:** `C` replays `_aggregate_simple`'s cell selection verbatim over the immutable `agg_gt_cells` fit-time kit snapshot (EDiD snapshot discipline; the kit builder stashes `(g, t, effect, n_treated)` 4-tuples in dict order) - anticipation filter FIRST (cells with `g - anticipation <= t < g` COUNT as post, the business-facing subtlety; universal-base reference cells are excluded by this filter alone, their base period being `< g - anticipation` by construction), finite-effect mask SECOND, weight = cohort mass where `fixed_cohort_agg_weights` provides one else per-cell complete-case `n_treated` (the source's `agg_weight` fallback is deliberately absent - RC-only, gated out), pairwise `np.sum`; the overall scalars come from the same stored fields the `'simple'` relay reads, so a post-fit mutation of `group_time_effects` moves NEITHER factor and `total == C x aggregate('simple').att` cannot desynchronize (pinned by a mutation test). NEVER use `agg_total_weight` (the RC all-units WIF mass). (d) **NaN rules:** empty post set / all-NaN effects give `C = NaN` and an all-NaN row WITHOUT re-emitting the fit-time UserWarnings (the RELAY-level no-re-warn convention — the NaN row is the signal; the recompute levels' bootstrap replay re-emits by design); with finite `C` the relay is VERBATIM - inherited NaNs (a degenerate fit's `(nan, nan)` CI beside finite att/se = 0) pass through, mirroring `aggregate('simple')` per the repo's non-estimable-row convention (att/n are never blanked by NaN inference fields); the ONLY additional blanking is true float overflow (a finite overall value made non-finite BY the `x C` scaling - unreachable in practice, guarded per this code family's finiteness convention). (e) **Bootstrap:** 'total' is a RELAY level - available on bootstrapped fits with the df column NaN (the M-027 per-level policy). (f) **Container conventions:** single row, `level="total"`, `label=["total"]`, `target=["total"]` (a total incremental outcome, NOT an ATT - the single-non-att neutral "estimate" rendering applies), `n = [C]` with `n_kind="obs"`, `weight=[1.0]`, per-class df carrier (CS: `resolve_inference_df` - finite on admitted bare-`cluster=` fits, NaN on plain analytical panel fits). Per-level `n` semantics are deliberately DIFFERENT: 'total''s `n` is the treated observations entering the support, while 'simple''s `n` stays treated+control units by contract - `n_kind` is per-container and each is documented, satisfying the never-conflate rule. (g) **MMM admission:** both `diff_diff.mmm` exporters accept the total container ALONE - `scale` (numeric or `"auto"`) is rejected as double-counting; for overall-total exports this route supersedes `scale="auto"` (see the MMM section). The total remains meaningful only for outcomes additive in levels - the same unverifiable caveat every scale route carries. +- **Note (post-fit `aggregate('total')` - the estimator-owned total incremental outcome, 3.10):** (a) **Estimand:** `total = C x overall_att` with `se = C x overall_se`, CI scaled by `C`, `t`/`p`/`df` inherited unchanged - an exact relay CONDITIONAL ON THE REALIZED AGGREGATION MASS (never call it unqualified "exact"): `C` is the design-fixed complete-case treated-observation count on every admitted routing, the inherited SE prices the overall's full inference (WIF/estimated-share terms included) but treats `C` as fixed at its realized value, and the random-mass `att*dC` variance term is deliberately out of scope (the DEFERRED survey/RC remainder). Every SE branch is homogeneous of degree 1 in the scaling, so the relay is exact under analytical AND percentile-bootstrap inference. (b) **Fail-closed routing gates**, in order: (1) repeated-cross-section-routed fits (`panel=False`, or genuinely-unbalanced `allow_unbalanced_panel=True` fits - a balanced panel with the flag stays panel-routed and admitted) raise `NotImplementedError` - the RC cohort mass counts UNITS and weights every kept post period by it, a ~T-fold overcount of treated observations on a T-wave RCS (execution-verified); (2) fits declaring a `survey_design=` raise - the realized-mass relay omits the survey mass-uncertainty term and design-aware population-scale totals are not implemented (retained weight scale differs by design family: analytic pweight resolves normalized to sum = n and CS retains no PER-OBSERVATION raw-weight record - the `survey_metadata` scalars report raw-scale `sum_weights`/`weight_range` provenance, but not the per-obs raw vector a population-scale total would need; CS accepts pweight only - while replicate designs retain raw scale); the gate reads the fit-time `is_survey_fit` KIT snapshot, so bare-`cluster=` fits are ADMITTED (their synthesized all-ones internal design is not a survey) and a post-fit edit of the public `survey_metadata` cannot bypass it; the DECLARATION BOUNDARY is deliberate - an explicit unweighted `SurveyDesign(psu=...)` fit is numerically identical to bare `cluster=` yet fails closed by declaration (declaring survey_design opts into the survey contract; conservative, never wrong numbers; the message points at the `cluster=` route); a combined survey+RC fit deterministically gets the RC message; (3) the coincidence check: on bare-`cluster=` fits `_aggregate_simple` weights kept cells by the synthesized cohort masses (full cohort size per kept cell), and when kept cells have INCOMPLETE treated support that realized mass diverges from the complete-case count - the same overcount class as (1) - so the fit raises rather than publishing an ambiguous mass; clean bare-`cluster=` fits (the normal case) are admitted with `C` equal to both counts. Panel-routed pre-upgrade (legacy) kits missing the fit-time snapshots raise the refit message - never a mutable-field or public-replay fallback. (c) **Filter-replay contract:** `C` replays `_aggregate_simple`'s cell selection verbatim over the immutable `agg_gt_cells` fit-time kit snapshot (EDiD snapshot discipline; the kit builder stashes `(g, t, effect, n_treated)` 4-tuples in dict order) - anticipation filter FIRST (cells with `g - anticipation <= t < g` COUNT as post, the business-facing subtlety; universal-base reference cells are excluded by this filter alone, their base period being `< g - anticipation` by construction), finite-effect mask SECOND, weight = cohort mass where `fixed_cohort_agg_weights` provides one else per-cell complete-case `n_treated` (the source's `agg_weight` fallback is deliberately absent - RC-only, gated out), pairwise `np.sum`; the overall scalars come from the same stored fields the `'simple'` relay reads, so a post-fit mutation of `group_time_effects` moves NEITHER factor and `total == C x aggregate('simple').att` cannot desynchronize (pinned by a mutation test). NEVER use `agg_total_weight` (the RC all-units WIF mass). (d) **NaN rules:** empty post set / all-NaN effects give `C = NaN` and an all-NaN row WITHOUT re-emitting the fit-time UserWarnings (the RELAY-level no-re-warn convention — the NaN row is the signal; the recompute levels' bootstrap replay re-emits by design); with finite `C` the relay is VERBATIM - inherited NaNs (a degenerate fit's `(nan, nan)` CI beside finite att/se = 0) pass through, mirroring `aggregate('simple')` per the repo's non-estimable-row convention (att/n are never blanked by NaN inference fields); the ONLY additional blanking is true float overflow (a finite overall value made non-finite BY the `x C` scaling - unreachable in practice, guarded per this code family's finiteness convention). (e) **Bootstrap:** 'total' is a RELAY level - available on bootstrapped fits with the df column NaN (the M-027 per-level policy). (f) **Container conventions:** single row, `level="total"`, `label=["total"]`, `target=["total"]` (a total incremental outcome, NOT an ATT - the single-non-att neutral "estimate" rendering applies), `n = [C]` with `n_kind="obs"`, `weight=[1.0]`, per-class df carrier (CS: `resolve_inference_df` - finite on admitted bare-`cluster=` fits, NaN on plain analytical panel fits). Per-level `n` semantics are deliberately DIFFERENT: 'total''s `n` is the treated observations entering the support, while 'simple''s `n` stays treated+control units by contract - `n_kind` is per-container and each is documented, satisfying the never-conflate rule. (g) **MMM admission:** both `diff_diff.mmm` exporters accept the total container ALONE - `scale` (numeric or `"auto"`) is rejected as double-counting; for overall-total exports this route supersedes `scale="auto"` (see the MMM section). The total remains meaningful only for outcomes additive in levels - the same unverifiable caveat every scale route carries. **Reference implementation(s):** - R: `did::att_gt()` (Callaway & Sant'Anna's official package) diff --git a/tests/test_survey_metadata_raw_scale.py b/tests/test_survey_metadata_raw_scale.py new file mode 100644 index 000000000..53abd1258 --- /dev/null +++ b/tests/test_survey_metadata_raw_scale.py @@ -0,0 +1,442 @@ +"""survey_metadata raw-scale provenance pins for the unit-level recompute. + +``SurveyDesign.resolve()`` rescales pweights to mean 1, so any metadata +recompute that feeds the RESOLVED weights back into +``compute_survey_metadata`` misreports the two scale-dependent fields +(``sum_weights``, ``weight_range``) while leaving the scale-invariant +fields (``effective_n``, ``design_effect``, ``df_survey``, ...) and all +estimates/inference untouched. DMLDiD received the raw-capture pattern in +its survey PR (``tests/test_survey_dml.py::TestSurveyMetadataRawScale``); +these pins cover the four remaining families with a unit-level recompute: +CallawaySantAnna (panel + RC lanes), the staggered DDD engine, +ContinuousDiD (analytic branch), and EfficientDiD. + +All fixtures use deliberately non-unit-scale weights (uniform(0.5, 2.0) / +1 + exponential) so a mean-1 rescale is detectable. +""" + +import warnings + +import numpy as np +import pandas as pd +import pytest + +from diff_diff import ( + CallawaySantAnna, + ContinuousDiD, + EfficientDiD, + StaggeredTripleDifference, +) +from diff_diff.survey import SurveyDesign + +# --------------------------------------------------------------------------- +# DGPs (weights constant within unit; strata + PSU present) +# --------------------------------------------------------------------------- + + +def _make_staggered_panel(n_units=60, n_periods=8, seed=42): + rng = np.random.default_rng(seed) + weights = rng.uniform(0.5, 2.0, n_units) + strata = np.arange(n_units) // 12 + psu = np.arange(n_units) // 5 # crosses strata; designs pass nest=True + rows = [] + for u in range(n_units): + ft = 4 if u < 20 else (6 if u < 40 else 0) + for t in range(1, n_periods + 1): + y = 10 + 0.05 * u + 0.2 * t + (2.0 if ft > 0 and t >= ft else 0.0) + y += rng.normal(0, 0.5) + rows.append( + { + "unit": u, + "time": t, + "first_treat": ft, + "outcome": y, + "weight": weights[u], + "stratum": strata[u], + "psu": psu[u], + } + ) + return pd.DataFrame(rows) + + +def _make_rc_data(n=400, seed=42): + """One row per unit: CallawaySantAnna(panel=False) rejects duplicates.""" + rng = np.random.default_rng(seed) + df = pd.DataFrame( + { + "unit": np.arange(n), + "time": rng.choice([1, 2, 3, 4], n), + "weight": rng.uniform(0.5, 2.0, n), + "stratum": rng.choice(3, n), + } + ) + df["psu"] = df["stratum"] * 4 + rng.choice(4, n) + df["first_treat"] = rng.choice([0, 3], n) + df["outcome"] = rng.normal(0, 1, n) + 1.5 * ( + (df["first_treat"] > 0) & (df["time"] >= df["first_treat"]) + ) + return df + + +def _make_sddd_data(n_units=200, n_periods=6, seed=42): + rng = np.random.default_rng(seed) + cohorts = rng.choice([3, 4, 0], size=n_units, p=[0.4, 0.3, 0.3]) + elig = rng.binomial(1, 0.5, size=n_units) + weights = 1.0 + rng.exponential(0.5, size=n_units) + strata = rng.choice(4, size=n_units) + psu = strata * 2 + rng.choice(2, size=n_units) # globally unique labels + rows = [] + for i in range(n_units): + for t in range(1, n_periods + 1): + te = 2.0 if (cohorts[i] > 0 and t >= cohorts[i] and elig[i] == 1) else 0.0 + y = rng.normal(0, 1) + 0.5 * t + te + rng.normal(0, 0.5) + rows.append( + { + "unit": i, + "period": t, + "outcome": y, + "first_treat": cohorts[i], + "eligibility": elig[i], + "weight": weights[i], + "stratum": strata[i], + "psu": psu[i], + } + ) + return pd.DataFrame(rows) + + +def _make_continuous_data(n_u=80, n_t=4, seed=42): + """Units 0-4 are TREATED with dose 0: the drop_units dose filter fires, + so ``len(df) < len(data)`` inside fit and the survey re-resolve runs on + the FILTERED frame — expected metadata below must be computed on the + filtered frame (a raw capture reading the unfiltered input would fail).""" + rng = np.random.default_rng(seed) + units = np.repeat(range(n_u), n_t) + times = np.tile(range(1, n_t + 1), n_u) + ft_unit = np.where(np.arange(n_u) < 40, 3, 0) + dose_unit = np.where(np.arange(n_u) < 40, rng.uniform(0.5, 2.0, n_u), 0.0) + dose_unit[:5] = 0.0 + ft = np.repeat(ft_unit, n_t) + dose = np.repeat(dose_unit, n_t) + y = rng.normal(size=len(units)) + 0.5 * dose * (times >= ft) * (ft > 0) + w = np.repeat(rng.uniform(0.5, 2.0, n_u), n_t) + strata = np.repeat(np.where(np.arange(n_u) < 40, 1, 2), n_t) + psu_unit = np.arange(n_u) // 10 + return pd.DataFrame( + { + "unit": units, + "time": times, + "first_treat": ft, + "dose": dose, + "outcome": y, + "weight": w, + "stratum": strata, + "psu": np.repeat(psu_unit, n_t), + } + ) + + +_DESIGN = SurveyDesign(weights="weight", strata="stratum", psu="psu", nest=True) +_DESIGN_NO_W = SurveyDesign(strata="stratum", psu="psu", nest=True) + + +def _assert_raw_scale(md, raw_w): + """The two scale-dependent fields match the raw weights; the two + weight-derived scale-invariant ratios match too (raw == normalized + mathematically; rtol absorbs the constant-rescale last-ULP wiggle).""" + raw_w = np.asarray(raw_w, dtype=np.float64) + np.testing.assert_allclose(md.sum_weights, raw_w.sum(), rtol=1e-12) + np.testing.assert_allclose(md.weight_range, (raw_w.min(), raw_w.max()), rtol=1e-12) + n = len(raw_w) + sum_w, sum_w2 = raw_w.sum(), (raw_w**2).sum() + np.testing.assert_allclose(md.effective_n, sum_w**2 / sum_w2, rtol=1e-12) + np.testing.assert_allclose(md.design_effect, n * sum_w2 / sum_w**2, rtol=1e-12) + + +def _fit_quiet(est, *args, **kwargs): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return est.fit(*args, **kwargs) + + +# --------------------------------------------------------------------------- +# CallawaySantAnna +# --------------------------------------------------------------------------- + + +class TestCSRawScale: + @pytest.fixture(scope="class") + def panel_df(self): + return _make_staggered_panel() + + def test_panel_metadata_uses_raw_unit_weights(self, panel_df): + res = _fit_quiet( + CallawaySantAnna(), + panel_df, + "outcome", + "unit", + "time", + "first_treat", + survey_design=_DESIGN, + ) + _assert_raw_scale(res.survey_metadata, panel_df.groupby("unit")["weight"].first()) + + def test_rc_metadata_uses_raw_obs_weights(self): + rc = _make_rc_data() + res = _fit_quiet( + CallawaySantAnna(panel=False), + rc, + "outcome", + "unit", + "time", + "first_treat", + survey_design=_DESIGN, + ) + _assert_raw_scale(res.survey_metadata, rc["weight"]) + + def test_injected_cluster_metadata_uses_raw_unit_weights(self, panel_df): + # PSU-less design + cluster=: metadata is computed on the inject + # path and recomputed at unit level — final values must still be + # the raw unit-level scale. + res = _fit_quiet( + CallawaySantAnna(cluster="psu"), + panel_df, + "outcome", + "unit", + "time", + "first_treat", + survey_design=SurveyDesign(weights="weight", strata="stratum", nest=True), + ) + _assert_raw_scale(res.survey_metadata, panel_df.groupby("unit")["weight"].first()) + + def test_weights_none_ones_fallback(self, panel_df): + res = _fit_quiet( + CallawaySantAnna(), + panel_df, + "outcome", + "unit", + "time", + "first_treat", + survey_design=_DESIGN_NO_W, + ) + md = res.survey_metadata + n_units = panel_df["unit"].nunique() + assert md.sum_weights == n_units + assert md.weight_range == (1.0, 1.0) + # Length guard: a wrong-length ones array shifts effective_n. + assert md.effective_n == n_units + + def test_bare_cluster_never_reaches_metadata_recompute(self, panel_df): + # Bare cluster= (no survey_design) synthesizes an internal design but + # leaves survey_metadata None, so the recompute's + # `assert survey_design is not None` is unreachable on this route + # (both lanes; also pinned upstream by + # test_bare_cluster_does_not_set_survey_metadata). + res = _fit_quiet( + CallawaySantAnna(cluster="psu"), + panel_df, + "outcome", + "unit", + "time", + "first_treat", + ) + assert np.isfinite(res.overall_att) + assert res.survey_metadata is None + rc = _make_rc_data() + rc["cluster_col"] = rc["unit"] // 20 + res_rc = _fit_quiet( + CallawaySantAnna(panel=False, cluster="cluster_col"), + rc, + "outcome", + "unit", + "time", + "first_treat", + ) + assert np.isfinite(res_rc.overall_att) + assert res_rc.survey_metadata is None + + +# --------------------------------------------------------------------------- +# Staggered DDD engine +# --------------------------------------------------------------------------- + + +class TestStaggeredDDDRawScale: + @pytest.fixture(scope="class") + def sddd_df(self): + return _make_sddd_data() + + def _fit(self, df, design): + return _fit_quiet( + StaggeredTripleDifference(), + df, + "outcome", + "unit", + "period", + "first_treat", + "eligibility", + survey_design=design, + ) + + def test_metadata_uses_raw_unit_weights(self, sddd_df): + res = self._fit(sddd_df, _DESIGN) + _assert_raw_scale(res.survey_metadata, sddd_df.groupby("unit")["weight"].first()) + + def test_weights_none_ones_fallback(self, sddd_df): + res = self._fit(sddd_df, _DESIGN_NO_W) + md = res.survey_metadata + n_units = sddd_df["unit"].nunique() + assert md.sum_weights == n_units + assert md.weight_range == (1.0, 1.0) + assert md.effective_n == n_units + + def test_bare_cluster_never_reaches_metadata_recompute(self, sddd_df): + # Same unreachability pin as the CS variant, on the DDD engine. + df = sddd_df.assign(cluster_col=sddd_df["unit"] // 10) + res = _fit_quiet( + StaggeredTripleDifference(cluster="cluster_col"), + df, + "outcome", + "unit", + "period", + "first_treat", + "eligibility", + ) + assert np.isfinite(res.overall_att) + assert res.survey_metadata is None + + +# --------------------------------------------------------------------------- +# ContinuousDiD +# --------------------------------------------------------------------------- + + +class TestContinuousDiDRawScale: + @pytest.fixture(scope="class") + def cont_df(self): + return _make_continuous_data() + + @pytest.fixture(scope="class") + def cont_filtered(self, cont_df): + # Mirror the estimator's dose filter: drop treated units with dose 0. + info = cont_df.groupby("unit").first()[["first_treat", "dose"]] + drop = info[(info["first_treat"] > 0) & (info["dose"] == 0)].index + assert len(drop) > 0 # the fixture must actually exercise the filter + return cont_df[~cont_df["unit"].isin(drop)] + + def _fit(self, df, design, **kw): + return _fit_quiet( + ContinuousDiD(**kw), + df, + "outcome", + "unit", + "time", + "first_treat", + "dose", + survey_design=design, + ) + + def test_analytic_metadata_uses_raw_unit_weights_of_filtered_frame( + self, cont_df, cont_filtered + ): + res = self._fit(cont_df, _DESIGN, n_bootstrap=0) + _assert_raw_scale(res.survey_metadata, cont_filtered.groupby("unit")["weight"].first()) + + def test_bootstrap_metadata_stays_obs_level_raw(self, cont_df, cont_filtered): + # Characterization pin for the deferred analytic-vs-bootstrap + # granularity divergence (TODO.md): the bootstrap branch keeps the + # OBS-level raw metadata from the (re-)resolve on the filtered frame. + res = self._fit(cont_df, _DESIGN, n_bootstrap=20, seed=42) + w = cont_filtered["weight"] + md = res.survey_metadata + 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_weight_column_aliasing_dose_reports_original_values(self): + # weights == dose: the never-treated nonzero-dose coercion zeroes the + # dose column on the working frame BEFORE metadata construction; the + # metadata must still report the user's ORIGINAL column values + # (snapshotted pre-mutation), not the coerced ones. + rng = np.random.default_rng(3) + n_u, n_t = 60, 4 + units = np.repeat(range(n_u), n_t) + times = np.tile(range(1, n_t + 1), n_u) + ft = np.repeat(np.where(np.arange(n_u) < 30, 3, 0), n_t) + dose = np.repeat(rng.uniform(0.5, 2.0, n_u), n_t) # NT rows nonzero + y = rng.normal(size=len(units)) + 0.5 * dose * (times >= ft) * (ft > 0) + df = pd.DataFrame( + {"unit": units, "time": times, "first_treat": ft, "dose": dose, "outcome": y} + ) + res = _fit_quiet( + ContinuousDiD(n_bootstrap=0), + df, + "outcome", + "unit", + "time", + "first_treat", + "dose", + survey_design=SurveyDesign(weights="dose"), + ) + _assert_raw_scale(res.survey_metadata, df.groupby("unit")["dose"].first()) + + def test_weights_none_ones_fallback(self, cont_df, cont_filtered): + res = self._fit(cont_df, _DESIGN_NO_W, n_bootstrap=0) + md = res.survey_metadata + n_units = cont_filtered["unit"].nunique() + assert md.sum_weights == n_units + assert md.weight_range == (1.0, 1.0) + assert md.effective_n == n_units + + +# --------------------------------------------------------------------------- +# EfficientDiD +# --------------------------------------------------------------------------- + + +class TestEfficientDiDRawScale: + @pytest.fixture(scope="class") + def panel_df(self): + return _make_staggered_panel() + + def _fit(self, df, design): + return _fit_quiet( + EfficientDiD(n_bootstrap=0), + df, + "outcome", + "unit", + "time", + "first_treat", + survey_design=design, + ) + + def test_metadata_uses_raw_unit_weights(self, panel_df): + res = self._fit(panel_df, _DESIGN) + _assert_raw_scale(res.survey_metadata, panel_df.groupby("unit")["weight"].first()) + + def test_weights_none_ones_fallback(self, panel_df): + res = self._fit(panel_df, _DESIGN_NO_W) + md = res.survey_metadata + n_units = panel_df["unit"].nunique() + assert md.sum_weights == n_units + assert md.weight_range == (1.0, 1.0) + assert md.effective_n == n_units + + def test_shuffled_rows_last_cohort_alignment(self, panel_df): + # Hardening for the positional data/_unit_first_panel_row alignment: + # shuffled input rows + all-eventually-treated panel routed through + # control_group="last_cohort" (unit_info reclassification, trimmed + # period list — but no df row drops) must still surface one raw + # weight per unit. + df = panel_df.copy() + df.loc[df["first_treat"] == 0, "first_treat"] = 7 # all eventually treated + df = df.sample(frac=1.0, random_state=7).reset_index(drop=True) + res = _fit_quiet( + EfficientDiD(n_bootstrap=0, control_group="last_cohort"), + df, + "outcome", + "unit", + "time", + "first_treat", + survey_design=_DESIGN, + ) + _assert_raw_scale(res.survey_metadata, df.groupby("unit")["weight"].first()) From db1e9cd17a6870b554a56e1d451583816949ef5a Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 29 Aug 2026 10:44:45 -0400 Subject: [PATCH 2/4] fix(continuous_did): re-resolve filtered survey design from pristine data rows The dose-filter re-resolve read the mutated working frame, so a design column aliasing a mutated role column (weights == dose) zero-weighted every never-treated unit on filtered fits ('No valid (g,t) cells') and diverged from the unfiltered path's pristine-data resolution. Resolve on the pristine kept-unit rows instead (row-for-row identical to df; byte-identical on non-aliased designs) and pin the combined filter + aliasing case against an immutable duplicate column (CI review finding). --- diff_diff/continuous_did.py | 15 +++++-- tests/test_survey_metadata_raw_scale.py | 55 +++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/diff_diff/continuous_did.py b/diff_diff/continuous_did.py index ee1488677..e65003cab 100644 --- a/diff_diff/continuous_did.py +++ b/diff_diff/continuous_did.py @@ -780,11 +780,20 @@ def fit( ) lowest_dose = d_L - # Re-resolve survey design on filtered df if rows were dropped - # (survey arrays must align with df, not the original data) + # Re-resolve survey design on the filtered rows if rows were dropped + # (survey arrays must align with df, not the original data). Resolve + # from PRISTINE ``data`` rows, not the mutated working frame: the + # unfiltered path resolves from ``data``, and df's role-column + # coercions (never-treated dose zeroing, first_treat inf->0, + # to_numeric) must not leak into a design whose column aliases a + # mutated role column — resolving on df previously zero-weighted + # every never-treated unit when ``weights == dose``. The dose filter + # drops whole units and preserves row order, so the pristine + # unit-mask selection is row-for-row identical to df. if resolved_survey is not None and len(df) < len(data): + _kept_row_mask = data[unit].isin(set(df[unit].unique())).to_numpy() resolved_survey, survey_weights, survey_weight_type, survey_metadata = ( - _resolve_survey_for_fit(survey_design, df, "analytical") + _resolve_survey_for_fit(survey_design, data[_kept_row_mask], "analytical") ) # 2. Precompute structures diff --git a/tests/test_survey_metadata_raw_scale.py b/tests/test_survey_metadata_raw_scale.py index 53abd1258..e0e1e8759 100644 --- a/tests/test_survey_metadata_raw_scale.py +++ b/tests/test_survey_metadata_raw_scale.py @@ -379,6 +379,61 @@ def test_weight_column_aliasing_dose_reports_original_values(self): ) _assert_raw_scale(res.survey_metadata, df.groupby("unit")["dose"].first()) + def test_filtered_aliased_weights_match_immutable_duplicate(self): + # COMBINED interaction: zero-dose TREATED units (the drop_units + # filter fires -> the survey is re-resolved on the filtered rows) + # AND weights == dose (the never-treated nonzero-dose coercion + # mutates that column on the working frame). The re-resolve must + # read PRISTINE data rows: resolving on the mutated frame + # zero-weighted every never-treated unit (previously "No valid + # (g,t) cells"). The aliased fit must match a fit on an immutable + # duplicate of the original column on estimates, SEs, and every + # metadata field. + rng = np.random.default_rng(3) + n_u, n_t = 60, 4 + units = np.repeat(range(n_u), n_t) + times = np.tile(range(1, n_t + 1), n_u) + ft = np.repeat(np.where(np.arange(n_u) < 30, 3, 0), n_t) + dose_u = rng.uniform(0.5, 2.0, n_u) + dose_u[:5] = 0.0 # treated zero-dose: filter fires + dose = np.repeat(dose_u, n_t) # never-treated rows nonzero: coercion fires + y = rng.normal(size=len(units)) + 0.5 * dose * (times >= ft) * (ft > 0) + df = pd.DataFrame( + {"unit": units, "time": times, "first_treat": ft, "dose": dose, "outcome": y} + ) + df["w_dup"] = df["dose"] # immutable duplicate of the ORIGINAL values + + def fit_with(wcol): + return _fit_quiet( + ContinuousDiD(n_bootstrap=0), + df, + "outcome", + "unit", + "time", + "first_treat", + "dose", + survey_design=SurveyDesign(weights=wcol), + ) + + r_alias, r_dup = fit_with("dose"), fit_with("w_dup") + assert r_alias.overall_att == r_dup.overall_att + assert r_alias.overall_att_se == r_dup.overall_att_se + ma, md_ = r_alias.survey_metadata, r_dup.survey_metadata + for field in ( + "sum_weights", + "weight_range", + "effective_n", + "design_effect", + "df_survey", + "n_strata", + "n_psu", + "weight_type", + ): + assert getattr(ma, field) == getattr(md_, field), field + # And both report the raw scale of the surviving (filtered) units. + kept = df[~df["unit"].isin(range(5))] + _assert_raw_scale(ma, kept.groupby("unit")["dose"].first()) + def test_weights_none_ones_fallback(self, cont_df, cont_filtered): res = self._fit(cont_df, _DESIGN_NO_W, n_bootstrap=0) md = res.survey_metadata From b38e8bf7ce397a77574806866029304d574c4218 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 29 Aug 2026 10:53:18 -0400 Subject: [PATCH 3/4] fix(survey): use 'is not None' for weight-column guards (empty-string column names) resolve() treats any non-None string as a column name, so truthiness checks on survey_design.weights silently substituted all-ones metadata for a column literally named ''. Fix the three new raw-capture sites (CI review finding) and the same pre-existing pattern at 13 sibling sites repo-wide (_resolve_survey_for_fit included); add nonuniform weights='' pins for CS RC, EfficientDiD, and ContinuousDiD. Qualify the CHANGELOG byte-identity claim to non-aliased designs and document the alias/filter correction. --- CHANGELOG.md | 15 +++++-- diff_diff/continuous_did.py | 4 +- diff_diff/dml_did.py | 2 +- diff_diff/efficient_did.py | 4 +- diff_diff/estimators.py | 4 +- diff_diff/imputation.py | 2 +- diff_diff/lpdid.py | 2 +- diff_diff/spillover.py | 5 ++- diff_diff/staggered.py | 7 +++- diff_diff/sun_abraham.py | 2 +- diff_diff/survey.py | 2 +- diff_diff/triple_diff.py | 2 +- diff_diff/twfe.py | 2 +- diff_diff/two_stage.py | 2 +- diff_diff/wooldridge.py | 2 +- tests/test_survey_metadata_raw_scale.py | 52 +++++++++++++++++++++++++ 16 files changed, 89 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9e76d8edf..70d637a61 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,10 +16,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 raw weights, so `sum_weights`/`weight_range` reported the normalized scale instead of the user's original weight scale. They now report the raw scale, matching every other estimator (DMLDiD got the pattern in - its survey PR). Metadata-provenance only: estimates, SEs, p-values, - CIs, `df_survey`, `n_strata`, `n_psu` are byte-identical, and - `effective_n`/`design_effect` are scale-invariant (unchanged within - floating-point round-off). + its survey PR). For previously-successful fits whose survey design does + not alias a mutated role column, this is metadata-provenance only: + estimates, SEs, p-values, CIs, `df_survey`, `n_strata`, `n_psu` are + byte-identical, and `effective_n`/`design_effect` are scale-invariant + (unchanged within floating-point round-off). Additionally, + `ContinuousDiD`'s zero-dose-unit filter now re-resolves the survey + design from pristine input rows: a design column aliasing a mutated + role column (e.g. `weights` naming the dose column) previously + zero-weighted every never-treated unit on filtered fits (failing with + "No valid (g,t) cells"); such fits now estimate under the user's + original weights, consistent with the unfiltered path. ## [3.11.0] - 2026-08-29 diff --git a/diff_diff/continuous_did.py b/diff_diff/continuous_did.py index e65003cab..f1fafd085 100644 --- a/diff_diff/continuous_did.py +++ b/diff_diff/continuous_did.py @@ -513,7 +513,9 @@ def fit( # survey_metadata. Per-unit via groupby-first so the later # dose-filter (which drops whole units) cannot desync alignment. raw_unit_w_meta: Optional[pd.Series] = None - if survey_design is not None and survey_design.weights: + if survey_design is not None and survey_design.weights is not None: + # `is not None`, not truthiness: resolve() treats any non-None + # string — an empty-string column name included — as a column. raw_unit_w_meta = data.groupby(unit)[survey_design.weights].first() # Covariate-path guards (conditional parallel trends). diff --git a/diff_diff/dml_did.py b/diff_diff/dml_did.py index 40ff713e9..1ebff5ef2 100644 --- a/diff_diff/dml_did.py +++ b/diff_diff/dml_did.py @@ -1886,7 +1886,7 @@ def fit( assert survey_design is not None raw_obs_weights = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) diff --git a/diff_diff/efficient_did.py b/diff_diff/efficient_did.py index b28844b8f..4e05ad692 100644 --- a/diff_diff/efficient_did.py +++ b/diff_diff/efficient_did.py @@ -1405,9 +1405,11 @@ def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any raw_unit_w_meta: Optional[np.ndarray] = None if self._unit_resolved_survey is not None: assert survey_design is not None + # `is not None`, not truthiness: resolve() treats any non-None + # string — an empty-string column name included — as a column. raw_obs_w_meta = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) raw_unit_w_meta = raw_obs_w_meta[self._unit_first_panel_row] diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index 87f3b4a53..07d81da20 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -709,7 +709,7 @@ def fit( if resolved_survey.psu is not None and survey_metadata is not None: raw_w = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) survey_metadata = compute_survey_metadata(resolved_survey, raw_w) @@ -1889,7 +1889,7 @@ def _fit_event_study_core( if resolved_survey.psu is not None and survey_metadata is not None: raw_w = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) survey_metadata = compute_survey_metadata(resolved_survey, raw_w) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 56b441ff5..7e1296e60 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -619,7 +619,7 @@ def fit( assert survey_design is not None raw_w = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) survey_metadata = compute_survey_metadata(resolved_survey, raw_w) diff --git a/diff_diff/lpdid.py b/diff_diff/lpdid.py index dc9fc7eb0..f6b232d8f 100644 --- a/diff_diff/lpdid.py +++ b/diff_diff/lpdid.py @@ -1501,7 +1501,7 @@ def fit( resolved_panel = _inject_cluster_as_psu(resolved_panel, data[cluster].to_numpy()) raw_weights = ( data[survey_design.weights].to_numpy(dtype=float) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=float) ) survey_metadata = compute_survey_metadata(resolved_panel, raw_weights) diff --git a/diff_diff/spillover.py b/diff_diff/spillover.py index 5cf6a6713..2d736d266 100644 --- a/diff_diff/spillover.py +++ b/diff_diff/spillover.py @@ -3035,7 +3035,10 @@ def fit( # Matches the post-injection resolved_survey_fit length. raw_w_for_meta = ( np.asarray(data[survey_design.weights].values, dtype=np.float64) - if (survey_design is not None and getattr(survey_design, "weights", None)) + if ( + survey_design is not None + and getattr(survey_design, "weights", None) is not None + ) else np.ones(len(data), dtype=np.float64) ) survey_metadata = _csm(resolved_survey_fit, raw_w_for_meta) diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 1389f5a67..37a544907 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -2253,7 +2253,7 @@ def fit( assert survey_design is not None raw_w = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) survey_metadata = compute_survey_metadata(resolved_survey, raw_w) @@ -2425,9 +2425,12 @@ def fit( ) else: # RC lane: resolved_survey_unit is per-observation. + # `is not None`, not truthiness: resolve() treats any + # non-None string — an empty-string column name + # included — as a column. raw_unit_w = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) survey_metadata = compute_survey_metadata(resolved_survey_unit, raw_unit_w) diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index 37409a9ba..66690c040 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -1053,7 +1053,7 @@ def fit( assert survey_design is not None raw_w = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) survey_metadata = compute_survey_metadata(resolved_survey, raw_w) diff --git a/diff_diff/survey.py b/diff_diff/survey.py index 0f97bc77a..622564b62 100644 --- a/diff_diff/survey.py +++ b/diff_diff/survey.py @@ -1288,7 +1288,7 @@ def _resolve_survey_for_fit( resolved = survey_design.resolve(data) raw_w = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) metadata = compute_survey_metadata(resolved, raw_w) diff --git a/diff_diff/triple_diff.py b/diff_diff/triple_diff.py index 67d76e32f..34ee6f8b2 100644 --- a/diff_diff/triple_diff.py +++ b/diff_diff/triple_diff.py @@ -964,7 +964,7 @@ def fit( if resolved_survey.psu is not None and survey_metadata is not None: raw_w = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) survey_metadata = compute_survey_metadata(resolved_survey, raw_w) diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index a7e365e3c..39891450b 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -622,7 +622,7 @@ def fit( # type: ignore[override] assert survey_design is not None raw_w = ( data[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) survey_metadata = compute_survey_metadata(resolved_survey, raw_w) diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 2f1615899..76fd56bc2 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -1674,7 +1674,7 @@ def fit( assert survey_design is not None raw_w = ( np.asarray(data[survey_design.weights].values, dtype=np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(data), dtype=np.float64) ) survey_metadata = compute_survey_metadata(resolved_survey, raw_w) diff --git a/diff_diff/wooldridge.py b/diff_diff/wooldridge.py index c2fdc3483..3354507ab 100644 --- a/diff_diff/wooldridge.py +++ b/diff_diff/wooldridge.py @@ -313,7 +313,7 @@ def _resolve_survey_for_wooldridge(survey_design, sample, cluster_ids, cluster_n if resolved.psu is not None and survey_metadata is not None: raw_w = ( sample[survey_design.weights].values.astype(np.float64) - if survey_design.weights + if survey_design.weights is not None else np.ones(len(sample), dtype=np.float64) ) survey_metadata = compute_survey_metadata(resolved, raw_w) diff --git a/tests/test_survey_metadata_raw_scale.py b/tests/test_survey_metadata_raw_scale.py index e0e1e8759..aee392441 100644 --- a/tests/test_survey_metadata_raw_scale.py +++ b/tests/test_survey_metadata_raw_scale.py @@ -256,6 +256,58 @@ def test_bare_cluster_never_reaches_metadata_recompute(self, panel_df): assert res_rc.survey_metadata is None +# --------------------------------------------------------------------------- +# Empty-string weight-column name (falsy but valid: resolve() checks +# `is not None`, so "" names a real column — truthiness checks would +# silently substitute ones) +# --------------------------------------------------------------------------- + + +class TestEmptyStringWeightColumn: + def test_cs_rc_lane(self): + rc = _make_rc_data().rename(columns={"weight": ""}) + res = _fit_quiet( + CallawaySantAnna(panel=False), + rc, + "outcome", + "unit", + "time", + "first_treat", + survey_design=SurveyDesign(weights="", strata="stratum", psu="psu", nest=True), + ) + _assert_raw_scale(res.survey_metadata, rc[""]) + + def test_efficient_did(self): + panel = _make_staggered_panel().rename(columns={"weight": ""}) + res = _fit_quiet( + EfficientDiD(n_bootstrap=0), + panel, + "outcome", + "unit", + "time", + "first_treat", + survey_design=SurveyDesign(weights="", strata="stratum", psu="psu", nest=True), + ) + _assert_raw_scale(res.survey_metadata, panel.groupby("unit")[""].first()) + + def test_continuous_did_analytical(self): + cont = _make_continuous_data().rename(columns={"weight": ""}) + info = cont.groupby("unit").first()[["first_treat", "dose"]] + drop = info[(info["first_treat"] > 0) & (info["dose"] == 0)].index + kept = cont[~cont["unit"].isin(drop)] + res = _fit_quiet( + ContinuousDiD(n_bootstrap=0), + cont, + "outcome", + "unit", + "time", + "first_treat", + "dose", + survey_design=SurveyDesign(weights="", strata="stratum", psu="psu", nest=True), + ) + _assert_raw_scale(res.survey_metadata, kept.groupby("unit")[""].first()) + + # --------------------------------------------------------------------------- # Staggered DDD engine # --------------------------------------------------------------------------- From 1ef24253d12c7e3cab47fe9eeb08341377264fa4 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 29 Aug 2026 11:00:33 -0400 Subject: [PATCH 4/4] fix(survey): complete the weight-column 'is not None' sweep (SunAbraham, WooldridgeDiD) SunAbraham's cohort-aggregation weight-column resolution and WooldridgeDiD's pre-exclusion zero-weight-group validation still used truthiness, so a column literally named '' silently fell back to unweighted cohort mass (SA: moved att AND se, execution-verified) or skipped the early zero-weight rejection (Wooldridge). Both now check 'is not None' like the rest of the surface; exhaustive grep finds no remaining weight-name truthiness (HAD's isinstance(str) guard already handled ''). Adds SA ''-vs-named bit-parity (overall + event study) and a Wooldridge zero-weight-comparison-unit parity pin (CI review round 3). --- diff_diff/sun_abraham.py | 8 +-- diff_diff/wooldridge.py | 4 +- tests/test_survey_metadata_raw_scale.py | 77 +++++++++++++++++++++++++ 3 files changed, 84 insertions(+), 5 deletions(-) diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index 66690c040..a83f79a00 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -1147,12 +1147,12 @@ def _refit_sa(w_r): results.append(es_r[e]["effect"] if e in es_r else np.nan) return np.array(results) - # Resolve survey weight column name for cohort aggregation + # Resolve survey weight column name for cohort aggregation. + # `is not None`, not truthiness: resolve() treats any non-None + # string — an empty-string column name included — as a column. survey_weight_col = ( survey_design.weights - if survey_design is not None - and hasattr(survey_design, "weights") - and survey_design.weights + if survey_design is not None and getattr(survey_design, "weights", None) is not None else None ) diff --git a/diff_diff/wooldridge.py b/diff_diff/wooldridge.py index 3354507ab..fac91a168 100644 --- a/diff_diff/wooldridge.py +++ b/diff_diff/wooldridge.py @@ -1418,7 +1418,9 @@ def fit( warnings.simplefilter("ignore") survey_design.resolve(sample) - if getattr(survey_design, "weights", None): + # `is not None`, not truthiness: resolve() treats any non-None + # string — an empty-string column name included — as a column. + if getattr(survey_design, "weights", None) is not None: _cell_w = sample[survey_design.weights].to_numpy(dtype=float) # The weighted within-transform's 0/0 constraint, checked before # exclusion can delete the offending rows. OLS only: logit and diff --git a/tests/test_survey_metadata_raw_scale.py b/tests/test_survey_metadata_raw_scale.py index aee392441..f2985da91 100644 --- a/tests/test_survey_metadata_raw_scale.py +++ b/tests/test_survey_metadata_raw_scale.py @@ -290,6 +290,83 @@ def test_efficient_did(self): ) _assert_raw_scale(res.survey_metadata, panel.groupby("unit")[""].first()) + def test_sun_abraham_parity_with_named_duplicate(self): + # SunAbraham's cohort/event-time aggregation reads the weight COLUMN + # NAME; a truthiness check treated "" as no-weights and silently fell + # back to unweighted cohort mass (behavioral: moved att AND se). + # A column named "" must give bit-identical results to the same + # values under a normal name. + from diff_diff import SunAbraham + + panel = _make_staggered_panel() + panel[""] = panel["weight"] + + def fit(col): + return _fit_quiet( + SunAbraham(), + panel, + "outcome", + "unit", + "time", + "first_treat", + survey_design=SurveyDesign(weights=col), + ) + + a, b = fit(""), fit("weight") + assert a.overall_att == b.overall_att + assert a.overall_se == b.overall_se + for e in a.event_study_effects: + assert a.event_study_effects[e]["effect"] == b.event_study_effects[e]["effect"] + assert a.event_study_effects[e]["se"] == b.event_study_effects[e]["se"] + + def test_wooldridge_zero_weight_group_parity(self): + # WooldridgeDiD's pre-exclusion zero-weight-group validation reads + # the weight column name; a truthiness check skipped it for "" so a + # zero-weight comparison cell passed validation under "" but was + # rejected under a normal name. Both aliases must behave identically. + from diff_diff import WooldridgeDiD + + rng = np.random.default_rng(11) + n_u, n_t = 30, 4 + rows = [] + w = rng.uniform(0.5, 2.0, n_u) + for u in range(n_u): + ft = 3 if u < 15 else 0 + for t in range(1, n_t + 1): + rows.append( + { + "unit": u, + "time": t, + "first_treat": ft, + "outcome": rng.normal() + (1.0 if ft and t >= ft else 0.0), + "w": w[u], + } + ) + df = pd.DataFrame(rows) + # Zero out one comparison unit's weights entirely so the weighted + # within-transform's 0/0 guard has something to reject. + df.loc[df["unit"] == 20, "w"] = 0.0 + df[""] = df["w"] + + def outcome_of(col): + try: + r = _fit_quiet( + WooldridgeDiD(), + df, + "outcome", + "unit", + "time", + "first_treat", + survey_design=SurveyDesign(weights=col), + ) + return ("fit", r.overall_att, r.overall_se) + except (ValueError, NotImplementedError) as exc: + # Normalize only the quoted column name, so messages that + # embed it still compare equal across the two aliases. + return ("raise", type(exc).__name__, str(exc).replace(f"'{col}'", "''")) + + assert outcome_of("") == outcome_of("w") + def test_continuous_did_analytical(self): cont = _make_continuous_data().rename(columns={"weight": ""}) info = cont.groupby("unit").first()[["first_treat", "dose"]]