diff --git a/TODO.md b/TODO.md index aafe0f4e9..0b60cb7c2 100644 --- a/TODO.md +++ b/TODO.md @@ -69,6 +69,7 @@ Deferred items from PR reviews that were not addressed before merge. | StaggeredTripleDifference R parity: benchmark only tests no-covariate path (xformla=~1). Add covariate-adjusted scenarios and aggregation SE parity assertions. | `benchmarks/R/benchmark_staggered_triplediff.R` | #245 | Medium | | StaggeredTripleDifference: per-cohort group-effect SEs include WIF (conservative vs R's wif=NULL). Documented in REGISTRY. Could override mixin for exact R match. | `staggered_triple_diff.py` | #245 | Low | | HonestDiD Delta^RM: uses naive FLCI instead of paper's ARP conditional/hybrid confidence sets (Sections 3.2.1-3.2.2). ARP infrastructure exists but moment inequality transformation needs calibration. CIs are conservative (wider, valid coverage). | `honest_did.py` | #248 | Medium | +| Replicate weight tests use Fay-like BRR perturbations (0.5/1.5), not true half-sample BRR. Add true BRR regressions per estimator family. Existing `test_survey_phase6.py` covers true BRR at the helper level. | `tests/test_replicate_weight_expansion.py` | #253 | Low | #### Performance diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index 1b8a2ac07..61695d04a 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -240,14 +240,14 @@ def fit( resolved_survey, survey_weights, survey_weight_type, survey_metadata = ( _resolve_survey_for_fit(survey_design, data, self.inference) ) - # Reject replicate-weight designs — base DiD uses compute_survey_vcov - # (TSL) directly, not LinearRegression's replicate dispatch. - if resolved_survey is not None and resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "DifferenceInDifferences does not yet support replicate-weight " - "survey designs. Use CallawaySantAnna, EfficientDiD, " - "ContinuousDiD, or TripleDifference for replicate-weight " - "inference, or use a TSL-based survey design (strata/psu/fpc)." + _uses_replicate = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate and self.inference == "wild_bootstrap": + raise ValueError( + "Cannot use inference='wild_bootstrap' with replicate-weight " + "survey designs. Replicate weights provide their own variance " + "estimation." ) # Handle absorbed fixed effects (within-transformation) @@ -358,6 +358,13 @@ def fit( ) survey_metadata = compute_survey_metadata(resolved_survey, raw_w) + # When absorb + replicate: pass survey_design=None to prevent + # LinearRegression from computing replicate vcov on already-demeaned + # data (demeaning depends on weights, so replicate refits must re-demean). + _lr_survey = resolved_survey + if _uses_replicate and absorbed_vars: + _lr_survey = None + reg = LinearRegression( include_intercept=False, # Intercept already in X robust=self.robust, @@ -366,7 +373,7 @@ def fit( rank_deficient_action=self.rank_deficient_action, weights=survey_weights, weight_type=survey_weight_type, - survey_design=resolved_survey, + survey_design=_lr_survey, ).fit(X, y, df_adjustment=n_absorbed_effects) coefficients = reg.coefficients_ @@ -375,14 +382,69 @@ def fit( assert coefficients is not None att = coefficients[att_idx] - # Get inference - either from bootstrap or analytical - if self.inference == "wild_bootstrap" and self.cluster is not None: + # Get inference - replicate absorb override, bootstrap, or analytical + if _uses_replicate and absorbed_vars: + # Estimator-level replicate variance: re-demean + re-solve per replicate + from diff_diff.survey import compute_replicate_refit_variance + from diff_diff.utils import safe_inference + + _absorb_list = list(absorbed_vars) # capture for closure + + # Handle rank-deficient nuisance: refit only identified columns + _id_mask = ~np.isnan(coefficients) + _id_cols = np.where(_id_mask)[0] + _att_idx_reduced = int(np.searchsorted(_id_cols, att_idx)) + + def _refit_did_absorb(w_r): + nz = w_r > 0 + wd = data[nz].copy() + w_nz = w_r[nz] + wd["_treat_time"] = ( + wd[treatment].values.astype(float) * wd[time].values.astype(float) + ) + vars_dm = [outcome, treatment, time, "_treat_time"] + (covariates or []) + for ab_var in _absorb_list: + wd, _ = demean_by_group(wd, vars_dm, ab_var, inplace=True, weights=w_nz) + y_r = wd[outcome].values.astype(float) + d_r = wd[treatment].values.astype(float) + t_r = wd[time].values.astype(float) + dt_r = wd["_treat_time"].values.astype(float) + X_r = np.column_stack([np.ones(len(y_r)), d_r, t_r, dt_r]) + if covariates: + for cov in covariates: + X_r = np.column_stack([X_r, wd[cov].values.astype(float)]) + coef_r, _, _ = solve_ols( + X_r[:, _id_cols], y_r, + weights=w_nz, weight_type=survey_weight_type, + rank_deficient_action="silent", return_vcov=False, + ) + return coef_r + + vcov_reduced, _n_valid_rep = compute_replicate_refit_variance( + _refit_did_absorb, coefficients[_id_mask], resolved_survey + ) + vcov = _expand_vcov_with_nan(vcov_reduced, len(coefficients), _id_cols) + se = float(np.sqrt(max(vcov[att_idx, att_idx], 0.0))) + _df_rep = ( + survey_metadata.df_survey + if survey_metadata and survey_metadata.df_survey + else 0 # rank-deficient replicate → NaN inference + ) + if _n_valid_rep < resolved_survey.n_replicates: + _df_rep = _n_valid_rep - 1 if _n_valid_rep > 1 else 0 + if survey_metadata is not None: + survey_metadata.df_survey = _df_rep if _df_rep > 0 else None + t_stat, p_value, conf_int = safe_inference( + att, se, alpha=self.alpha, df=_df_rep + ) + elif self.inference == "wild_bootstrap" and self.cluster is not None: # Override with wild cluster bootstrap inference se, p_value, conf_int, t_stat, vcov, _ = self._run_wild_bootstrap_inference( X, y, residuals, cluster_ids, att_idx ) else: # Use analytical inference from LinearRegression + # (handles replicate vcov for no-absorb path automatically) vcov = reg.vcov_ inference = reg.get_inference(att_idx) se = inference.se @@ -1017,14 +1079,14 @@ def fit( # type: ignore[override] resolved_survey, survey_weights, survey_weight_type, survey_metadata = ( _resolve_survey_for_fit(survey_design, data, effective_inference) ) - # Reject replicate-weight designs — MultiPeriodDiD uses - # compute_survey_vcov (TSL) directly without replicate dispatch. - if resolved_survey is not None and resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "MultiPeriodDiD does not yet support replicate-weight survey " - "designs. Use CallawaySantAnna for staggered adoption with " - "replicate weights, or use a TSL-based survey design " - "(strata/psu/fpc)." + _uses_replicate_mp = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate_mp and effective_inference == "wild_bootstrap": + raise ValueError( + "Cannot use inference='wild_bootstrap' with replicate-weight " + "survey designs. Replicate weights provide their own variance " + "estimation." ) # Handle absorbed fixed effects (within-transformation) @@ -1177,7 +1239,80 @@ def fit( # type: ignore[override] ) # Compute survey vcov if applicable - if _use_survey_vcov: + _n_valid_rep_mp = None + if _use_survey_vcov and _uses_replicate_mp and absorb: + # Absorb + replicate: estimator-level refit (demeaning depends on weights) + from diff_diff.survey import compute_replicate_refit_variance + + _absorb_list_mp = list(absorb) + # Handle rank-deficient nuisance: refit only identified columns + _id_mask_mp = ~np.isnan(coefficients) + _id_cols_mp = np.where(_id_mask_mp)[0] + + def _refit_mp_absorb(w_r): + nz = w_r > 0 + wd = data[nz].copy() + w_nz = w_r[nz] + d_raw_ = wd[treatment].values.astype(float) + t_raw_ = wd[time].values + wd["_did_treatment"] = d_raw_ + for period_ in non_ref_periods: + wd[f"_did_period_{period_}"] = (t_raw_ == period_).astype(float) + wd[f"_did_interact_{period_}"] = d_raw_ * (t_raw_ == period_).astype(float) + vars_dm_ = ( + [outcome, "_did_treatment"] + + [f"_did_period_{p}" for p in non_ref_periods] + + [f"_did_interact_{p}" for p in non_ref_periods] + + (covariates or []) + ) + for ab_var_ in _absorb_list_mp: + wd, _ = demean_by_group(wd, vars_dm_, ab_var_, inplace=True, weights=w_nz) + y_r = wd[outcome].values.astype(float) + d_r = wd["_did_treatment"].values.astype(float) + X_r = np.column_stack([np.ones(len(y_r)), d_r]) + for period_ in non_ref_periods: + X_r = np.column_stack( + [X_r, wd[f"_did_period_{period_}"].values.astype(float)] + ) + for period_ in non_ref_periods: + X_r = np.column_stack( + [X_r, wd[f"_did_interact_{period_}"].values.astype(float)] + ) + if covariates: + for cov_ in covariates: + X_r = np.column_stack([X_r, wd[cov_].values.astype(float)]) + coef_r, _, _ = solve_ols( + X_r[:, _id_cols_mp], y_r, + weights=w_nz, weight_type=survey_weight_type, + rank_deficient_action="silent", return_vcov=False, + ) + return coef_r + + vcov_reduced_mp, _n_valid_rep_mp = compute_replicate_refit_variance( + _refit_mp_absorb, coefficients[_id_mask_mp], resolved_survey + ) + vcov = _expand_vcov_with_nan(vcov_reduced_mp, len(coefficients), _id_cols_mp) + elif _use_survey_vcov and _uses_replicate_mp: + # No absorb + replicate: X is fixed, use compute_replicate_vcov directly + from diff_diff.survey import compute_replicate_vcov + + nan_mask = np.isnan(coefficients) + if np.any(nan_mask): + kept_cols = np.where(~nan_mask)[0] + if len(kept_cols) > 0: + vcov_reduced, _n_valid_rep_mp = compute_replicate_vcov( + X[:, kept_cols], y, coefficients[kept_cols], resolved_survey, + weight_type=survey_weight_type, + ) + vcov = _expand_vcov_with_nan(vcov_reduced, X.shape[1], kept_cols) + else: + vcov = np.full((X.shape[1], X.shape[1]), np.nan) + _n_valid_rep_mp = 0 + else: + vcov, _n_valid_rep_mp = compute_replicate_vcov( + X, y, coefficients, resolved_survey, weight_type=survey_weight_type, + ) + elif _use_survey_vcov: from diff_diff.survey import compute_survey_vcov nan_mask = np.isnan(coefficients) @@ -1201,9 +1336,18 @@ def fit( # type: ignore[override] df = n_eff_df - k_effective - n_absorbed_effects if resolved_survey is not None and resolved_survey.df_survey is not None: df = resolved_survey.df_survey + # Replicate df: rank-deficient → NaN inference; dropped replicates → n_valid-1 + if _uses_replicate_mp: + if resolved_survey.df_survey is None: + df = 0 # rank-deficient replicate → NaN inference + if _n_valid_rep_mp is not None and _n_valid_rep_mp < resolved_survey.n_replicates: + df = _n_valid_rep_mp - 1 if _n_valid_rep_mp > 1 else 0 + if survey_metadata is not None: + survey_metadata.df_survey = df if df > 0 else None # Guard: fall back to normal distribution if df is non-positive - if df is not None and df <= 0: + # Skip for replicate designs — df=0 is intentional for NaN inference + if df is not None and df <= 0 and not _uses_replicate_mp: warnings.warn( f"Degrees of freedom is non-positive (df={df}). " "Using normal distribution instead of t-distribution for inference.", diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 4f32a6638..70fc112e1 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -244,13 +244,16 @@ def fit( survey_design, data, "analytical" ) + _uses_replicate_imp = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate_imp and self.n_bootstrap > 0: + raise ValueError( + "Cannot use n_bootstrap > 0 with replicate-weight survey designs. " + "Replicate weights provide their own variance estimation." + ) # Validate within-unit constancy for panel survey designs if resolved_survey is not None: - if resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "ImputationDiD does not yet support replicate-weight survey " - "designs. Use a TSL-based survey design (strata/psu/fpc)." - ) _validate_unit_constant_survey(data, unit, survey_design) if resolved_survey.weight_type != "pweight": raise ValueError( @@ -463,95 +466,196 @@ def fit( else: overall_att = float(np.mean(valid_tau)) - # ---- Conservative variance (Theorem 3) ---- - # Build weights matching the ATT: proportional to survey weights for - # finite tau_hat, uniform when no survey - overall_weights = np.zeros(n_omega_1) - n_valid = int(finite_mask.sum()) - if n_valid > 0: - if survey_weights is not None: - treated_sw = survey_weights[omega_1_mask.values] - sw_finite = treated_sw[finite_mask] - overall_weights[finite_mask] = sw_finite / sw_finite.sum() - else: - overall_weights[finite_mask] = 1.0 / n_valid + # ---- Variance ---- + _n_valid_rep_imp = None + _vcov_rep_imp = None + overall_se = np.nan # placeholder; overridden by replicate or conservative path + + if not _uses_replicate_imp: + # Conservative variance (Theorem 3) + overall_weights = np.zeros(n_omega_1) + n_valid = int(finite_mask.sum()) + if n_valid > 0: + if survey_weights is not None: + treated_sw = survey_weights[omega_1_mask.values] + sw_finite = treated_sw[finite_mask] + overall_weights[finite_mask] = sw_finite / sw_finite.sum() + else: + overall_weights[finite_mask] = 1.0 / n_valid - if n_valid == 0: - overall_se = np.nan - else: - overall_se = self._compute_conservative_variance( - df=df, - outcome=outcome, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - weights=overall_weights, - cluster_var=cluster_var, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, - ) + if n_valid == 0: + overall_se = np.nan + else: + overall_se = self._compute_conservative_variance( + df=df, + outcome=outcome, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=unit_fe, + time_fe=time_fe, + grand_mean=grand_mean, + delta_hat=delta_hat, + weights=overall_weights, + cluster_var=cluster_var, + kept_cov_mask=kept_cov_mask, + survey_weights=survey_weights, + ) # Survey degrees of freedom for t-distribution inference _survey_df = resolved_survey.df_survey if resolved_survey is not None else None + # Replicate df: rank-deficient → NaN inference; dropped replicates → n_valid-1 + if _uses_replicate_imp and _survey_df is None: + _survey_df = 0 # rank-deficient replicate → NaN inference + # Compute overall inference (may be overridden by replicate below) overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df ) - # Event study and group aggregation + # Event study and group aggregation (full-sample, for point estimates) event_study_effects = None group_effects = None if aggregate in ("event_study", "all"): event_study_effects = self._aggregate_event_study( - df=df, - outcome=outcome, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - cluster_var=cluster_var, - treatment_groups=treatment_groups, - balance_e=balance_e, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, + df=df, outcome=outcome, unit=unit, time=time, + first_treat=first_treat, covariates=covariates, + omega_0_mask=omega_0_mask, omega_1_mask=omega_1_mask, + unit_fe=unit_fe, time_fe=time_fe, grand_mean=grand_mean, + delta_hat=delta_hat, cluster_var=cluster_var, + treatment_groups=treatment_groups, balance_e=balance_e, + kept_cov_mask=kept_cov_mask, survey_weights=survey_weights, survey_df=_survey_df, ) if aggregate in ("group", "all"): group_effects = self._aggregate_group( - df=df, - outcome=outcome, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - cluster_var=cluster_var, + df=df, outcome=outcome, unit=unit, time=time, + first_treat=first_treat, covariates=covariates, + omega_0_mask=omega_0_mask, omega_1_mask=omega_1_mask, + unit_fe=unit_fe, time_fe=time_fe, grand_mean=grand_mean, + delta_hat=delta_hat, cluster_var=cluster_var, treatment_groups=treatment_groups, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, + kept_cov_mask=kept_cov_mask, survey_weights=survey_weights, survey_df=_survey_df, ) + # Replicate variance: derive keys from actual outputs (after filtering) + if _uses_replicate_imp: + from diff_diff.survey import compute_replicate_refit_variance + + _rel_times_treated = df.loc[omega_1_mask, "_rel_time"].values + _cohorts_treated = df.loc[omega_1_mask, first_treat].values + + # Derive keys from actual outputs (excludes filtered/Prop5/ref) + _sorted_rel_times = sorted( + e for e in (event_study_effects or {}).keys() + if np.isfinite(event_study_effects[e]["effect"]) + and event_study_effects[e].get("n_obs", 1) > 0 + ) + _sorted_groups = sorted( + g for g in (group_effects or {}).keys() + if np.isfinite(group_effects[g]["effect"]) + ) + _n_es = len(_sorted_rel_times) + + # Pre-compute balanced cohort mask for balance_e + _balanced_mask_treated = None + if balance_e is not None and _sorted_rel_times: + df_1 = df.loc[omega_1_mask] + rel_times_all = df_1["_rel_time"].values + all_horizons_full = sorted( + set(int(h) for h in rel_times_all if np.isfinite(h)) + ) + if self.horizon_max is not None: + all_horizons_full = [ + h for h in all_horizons_full if abs(h) <= self.horizon_max + ] + cohort_rel_times = self._build_cohort_rel_times(df, first_treat) + _balanced_mask_treated = self._compute_balanced_cohort_mask( + df_1, first_treat, all_horizons_full, balance_e, cohort_rel_times + ) + + # Single vectorized refit: [overall, es_e0..., grp_g0...] + def _refit_imp(w_r): + ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( + df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, + ) + tau_r, _ = self._impute_treatment_effects( + df, outcome, unit, time, covariates, omega_1_mask, + ufe_r, tfe_r, gm_r, delta_r, + ) + fin = np.isfinite(tau_r) + treated_w = w_r[omega_1_mask.values] + results = [] + # [0] Overall ATT + tw_fin = treated_w[fin] + tw_sum = np.sum(tw_fin) + results.append( + float(np.sum(tau_r[fin] * tw_fin) / tw_sum) if tw_sum > 0 else np.nan + ) + # [1..n_es] Event-study (identified only) + for e in _sorted_rel_times: + mask_e = fin & (_rel_times_treated == e) + if _balanced_mask_treated is not None: + mask_e = mask_e & _balanced_mask_treated + tw_e = treated_w[mask_e] + s = np.sum(tw_e) + results.append(float(np.sum(tau_r[mask_e] * tw_e) / s) if s > 0 else np.nan) + # [n_es+1..] Group (identified only) + for g in _sorted_groups: + mask_g = fin & (_cohorts_treated == g) + tw_g = treated_w[mask_g] + s = np.sum(tw_g) + results.append(float(np.sum(tau_r[mask_g] * tw_g) / s) if s > 0 else np.nan) + return np.array(results) + + # Build full-sample estimate from actual effects + _full_est = [overall_att] + _full_est.extend([event_study_effects[e]["effect"] for e in _sorted_rel_times]) + _full_est.extend([group_effects[g]["effect"] for g in _sorted_groups]) + + _vcov_rep_imp, _n_valid_rep_imp = compute_replicate_refit_variance( + _refit_imp, np.array(_full_est), resolved_survey + ) + overall_se = float(np.sqrt(max(_vcov_rep_imp[0, 0], 0.0))) + + # Override df if replicates were dropped + if _n_valid_rep_imp < resolved_survey.n_replicates: + _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 + if survey_metadata is not None: + survey_metadata.df_survey = _survey_df if _survey_df and _survey_df > 0 else None + + overall_t, overall_p, overall_ci = safe_inference( + overall_att, overall_se, alpha=self.alpha, df=_survey_df + ) + + # Override event-study SEs from vcov diagonal + for i, e in enumerate(_sorted_rel_times): + if event_study_effects is not None and e in event_study_effects: + se_e = float(np.sqrt(max(_vcov_rep_imp[1 + i, 1 + i], 0.0))) + eff_e = event_study_effects[e]["effect"] + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) + event_study_effects[e]["se"] = se_e + event_study_effects[e]["t_stat"] = t_e + event_study_effects[e]["p_value"] = p_e + event_study_effects[e]["conf_int"] = ci_e + + # Override group SEs from vcov diagonal + for j, g in enumerate(_sorted_groups): + if group_effects is not None and g in group_effects: + se_g = float(np.sqrt(max(_vcov_rep_imp[1 + _n_es + j, 1 + _n_es + j], 0.0))) + eff_g = group_effects[g]["effect"] + t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) + group_effects[g]["se"] = se_g + group_effects[g]["t_stat"] = t_g + group_effects[g]["p_value"] = p_g + group_effects[g]["conf_int"] = ci_g + # Build treatment effects dataframe treated_df = df.loc[omega_1_mask, [unit, time, "_tau_hat", "_rel_time"]].copy() treated_df = treated_df.rename(columns={"_tau_hat": "tau_hat", "_rel_time": "rel_time"}) diff --git a/diff_diff/stacked_did.py b/diff_diff/stacked_did.py index d923c7071..7c610b5c0 100644 --- a/diff_diff/stacked_did.py +++ b/diff_diff/stacked_did.py @@ -242,15 +242,9 @@ def fit( resolved_survey, survey_weights, survey_weight_type, survey_metadata = ( _resolve_survey_for_fit(survey_design, data, "analytical") ) - # Reject replicate-weight designs — StackedDiD uses - # compute_survey_vcov (TSL) directly without replicate dispatch. - if resolved_survey is not None and resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "StackedDiD does not yet support replicate-weight survey " - "designs. Use CallawaySantAnna for staggered adoption with " - "replicate weights, or use a TSL-based survey design " - "(strata/psu/fpc)." - ) + _uses_replicate_sd = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) # Reject fweight and aweight — Q-weight composition is ratio-valued # and breaks both frequency-weight (integer) and analytic-weight @@ -273,6 +267,9 @@ def fit( col_name = getattr(survey_design, attr, None) if col_name is not None: survey_cols.append(col_name) + # Propagate replicate weight columns through stacked dataset + if survey_design.replicate_weights is not None: + survey_cols.extend(survey_design.replicate_weights) df = data.copy() df[time] = pd.to_numeric(df[time]) @@ -428,8 +425,42 @@ def fit( ) assert vcov is not None - # ---- Survey VCV override (TSL variance) ---- - if resolved_survey is not None: + # ---- Survey VCV override ---- + _n_valid_rep_sd = None + resolved_stacked = None + if resolved_survey is not None and _uses_replicate_sd: + # Replicate variance: re-run WLS per replicate with composed weights + from diff_diff.survey import compute_replicate_refit_variance, compute_survey_metadata + + resolved_stacked = survey_design.resolve(stacked_df) + + # Refit closure: compose Q-weights with replicate survey weights + def _refit_stacked(w_r): + composed_r = Q_weights * w_r + w_sum = np.sum(composed_r) + if w_sum > 0: + composed_r = composed_r * (n_stacked / w_sum) + sqrt_w_r = np.sqrt(composed_r) + coef_r, _, _ = solve_ols( + X * sqrt_w_r[:, np.newaxis], Y * sqrt_w_r, + cluster_ids=cluster_ids, + rank_deficient_action="silent", return_vcov=False, + ) + return coef_r + + # Full-sample cohort effect vector + vcov, _n_valid_rep_sd = compute_replicate_refit_variance( + _refit_stacked, coef, resolved_stacked + ) + + # Compute survey metadata + raw_w_stacked = ( + stacked_df[survey_design.weights].values.astype(np.float64) + if survey_design.weights is not None + else np.ones(n_stacked, dtype=np.float64) + ) + survey_metadata = compute_survey_metadata(resolved_stacked, raw_w_stacked) + elif resolved_survey is not None: from diff_diff.survey import ( _inject_cluster_as_psu, _resolve_effective_cluster, @@ -485,8 +516,14 @@ def fit( _survey_df = ( max(survey_metadata.df_survey, 1) if survey_metadata is not None and survey_metadata.df_survey is not None - else None + else (0 if _uses_replicate_sd else None) ) + # Override df when replicate replicates were dropped + if _n_valid_rep_sd is not None and resolved_stacked is not None: + if _n_valid_rep_sd < resolved_stacked.n_replicates: + _survey_df = _n_valid_rep_sd - 1 if _n_valid_rep_sd > 1 else 0 + if survey_metadata is not None: + survey_metadata.df_survey = _survey_df if _survey_df > 0 else None t_stat, p_value, conf_int = safe_inference( effect, se, alpha=self.alpha, df=_survey_df ) @@ -522,8 +559,13 @@ def fit( _survey_df_overall = ( max(survey_metadata.df_survey, 1) if survey_metadata is not None and survey_metadata.df_survey is not None - else None + else (0 if _uses_replicate_sd else None) ) + if _n_valid_rep_sd is not None and resolved_stacked is not None: + if _n_valid_rep_sd < resolved_stacked.n_replicates: + _survey_df_overall = _n_valid_rep_sd - 1 if _n_valid_rep_sd > 1 else 0 + if survey_metadata is not None: + survey_metadata.df_survey = _survey_df_overall if _survey_df_overall > 0 else None overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df_overall ) diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index 8ab8a9162..c93544ab8 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -501,16 +501,13 @@ def fit( if resolved_survey is not None: _validate_unit_constant_survey(data, unit, survey_design) - # Reject replicate-weight designs — SunAbraham's weighted - # within-transformation bakes survey weights into X and y, so - # replicate refits on the already-transformed design are incorrect. - # Full estimator-level replicate refits are not yet implemented. - if resolved_survey is not None and resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "SunAbraham does not yet support replicate-weight survey designs. " - "The weighted within-transformation must be recomputed for each " - "replicate, which requires estimator-level replicate refits. " - "Use a TSL-based survey design (strata/psu/fpc) instead." + _uses_replicate_sa = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate_sa and self.n_bootstrap > 0: + raise ValueError( + "Cannot use n_bootstrap > 0 with replicate-weight survey designs. " + "Replicate weights provide their own variance estimation." ) # Bootstrap + survey supported via Rao-Wu rescaled bootstrap. @@ -630,9 +627,52 @@ def fit( cluster_var, survey_weights=survey_weights, survey_weight_type=survey_weight_type, - resolved_survey=resolved_survey, + # For replicate designs: pass None to prevent LinearRegression from + # computing bogus replicate vcov on already-demeaned data. We + # override vcov_cohort below with the correct estimator-level refit. + resolved_survey=None if _uses_replicate_sa else resolved_survey, ) + # Replicate variance override: fully refit the IW estimator per + # replicate, including recomputing cohort-share aggregation weights + # from w_r, so replicate SEs reflect the complete estimator. + _n_valid_rep_sa = None + if _uses_replicate_sa: + from diff_diff.survey import compute_replicate_refit_variance + + # The refit returns [overall_att, es_e0, es_e1, ...] after + # full re-aggregation with replicate-weighted cohort shares. + _sa_rel_periods = list(rel_periods_to_estimate) + + def _refit_sa(w_r): + # Drop zero-weight obs for within-transform safety + nz = w_r > 0 + df_reg_nz = df_reg[nz] if not np.all(nz) else df_reg + w_nz = w_r[nz] if not np.all(nz) else w_r + ce_r, _, vcov_r, cim_r = self._fit_saturated_regression( + df_reg_nz, outcome, unit, time, first_treat, + treatment_groups, _sa_rel_periods, covariates, + cluster_var, survey_weights=w_nz, + survey_weight_type=survey_weight_type, + resolved_survey=None, + ) + # Create temp weight column for IW aggregation with w_r + # Use full w_r (including zeros) for correct mass computation + _wt_col = "_rep_wt" + df[_wt_col] = w_r + es_r, _ = self._compute_iw_effects( + df, unit, first_treat, treatment_groups, _sa_rel_periods, + ce_r, {}, vcov_r, cim_r, survey_weight_col=_wt_col, + ) + att_r, _ = self._compute_overall_att( + df, first_treat, es_r, ce_r, _, vcov_r, cim_r, + survey_weight_col=_wt_col, + ) + results = [att_r] + for e in _sa_rel_periods: + 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 survey_weight_col = ( survey_design.weights @@ -648,6 +688,10 @@ def fit( if survey_metadata is not None and survey_metadata.df_survey is not None else None ) + # Replicate df: rank-deficient → NaN inference (dropped-replicate + # override happens after replicate refit below) + if _uses_replicate_sa and _sa_survey_df is None: + _sa_survey_df = 0 # rank-deficient replicate → NaN inference # Compute interaction-weighted event study effects event_study_effects, cohort_weights = self._compute_iw_effects( @@ -680,6 +724,66 @@ def fit( overall_att, overall_se, alpha=self.alpha, df=_sa_survey_df ) + # Replicate variance override: refit fully re-aggregated estimates + if _uses_replicate_sa: + # Build full-sample estimate vector from actual outputs + _full_est_sa = [overall_att] + for e in _sa_rel_periods: + _full_est_sa.append( + event_study_effects[e]["effect"] if e in event_study_effects else np.nan + ) + + _vcov_sa, _n_valid_rep_sa = compute_replicate_refit_variance( + _refit_sa, np.array(_full_est_sa), resolved_survey + ) + + # Override df if replicates dropped + if _n_valid_rep_sa < resolved_survey.n_replicates: + _sa_survey_df = _n_valid_rep_sa - 1 if _n_valid_rep_sa > 1 else 0 + if survey_metadata is not None: + survey_metadata.df_survey = _sa_survey_df if _sa_survey_df and _sa_survey_df > 0 else None + + # Override overall ATT SE + overall_se = float(np.sqrt(max(_vcov_sa[0, 0], 0.0))) + overall_t, overall_p, overall_ci = safe_inference( + overall_att, overall_se, alpha=self.alpha, df=_sa_survey_df + ) + + # Override event-study SEs + for i, e in enumerate(_sa_rel_periods): + if e in event_study_effects and np.isfinite(event_study_effects[e]["effect"]): + se_e = float(np.sqrt(max(_vcov_sa[1 + i, 1 + i], 0.0))) + eff_e = event_study_effects[e]["effect"] + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_sa_survey_df) + event_study_effects[e]["se"] = se_e + event_study_effects[e]["t_stat"] = t_e + event_study_effects[e]["p_value"] = p_e + event_study_effects[e]["conf_int"] = ci_e + + # Cohort-level replicate SEs: second refit for raw (g,e) coefficients + _keys_ordered = sorted(coef_index_map.keys(), key=lambda k: coef_index_map[k]) + _full_cohort_vec = np.array([cohort_effects.get(k, np.nan) for k in _keys_ordered]) + + def _refit_sa_cohort(w_r): + nz = w_r > 0 + df_reg_nz = df_reg[nz] if not np.all(nz) else df_reg + w_nz = w_r[nz] if not np.all(nz) else w_r + ce_r, _, _, _ = self._fit_saturated_regression( + df_reg_nz, outcome, unit, time, first_treat, + treatment_groups, _sa_rel_periods, covariates, + cluster_var, survey_weights=w_nz, + survey_weight_type=survey_weight_type, + resolved_survey=None, + ) + return np.array([ce_r.get(k, np.nan) for k in _keys_ordered]) + + _vcov_cohort_rep, _ = compute_replicate_refit_variance( + _refit_sa_cohort, _full_cohort_vec, resolved_survey + ) + for key in _keys_ordered: + idx = coef_index_map[key] + cohort_ses[key] = float(np.sqrt(max(_vcov_cohort_rep[idx, idx], 0.0))) + # Run bootstrap if requested bootstrap_results = None if self.n_bootstrap > 0: diff --git a/diff_diff/survey.py b/diff_diff/survey.py index 9fa2be457..bd18700e2 100644 --- a/diff_diff/survey.py +++ b/diff_diff/survey.py @@ -16,7 +16,7 @@ import warnings from dataclasses import dataclass, field, replace -from typing import List, Optional, Tuple +from typing import Callable, List, Optional, Tuple import numpy as np import pandas as pd @@ -1703,6 +1703,151 @@ def compute_replicate_if_variance( raise ValueError(f"Unknown replicate method: {method}") +def compute_replicate_refit_variance( + refit_fn: Callable[[np.ndarray], np.ndarray], + full_sample_estimate: np.ndarray, + resolved: "ResolvedSurveyDesign", +) -> Tuple[np.ndarray, int]: + """Compute replicate variance by re-running an arbitrary estimation function. + + For each replicate weight column, calls ``refit_fn(w_r)`` and collects + the resulting estimate vector. Variance is computed from the distribution + of replicate estimates using method-specific scaling. + + This generalises :func:`compute_replicate_vcov` (which hard-codes + ``solve_ols`` as the refit) for estimators whose estimation procedure + is more complex than a single OLS call (e.g. within-transformation, + two-stage imputation, stacked regression). + + Parameters + ---------- + refit_fn : callable + ``(n,) weight array -> (k,) estimate array``. Must return the same + length *k* on every call. Should return all-NaN when the estimation + fails for that replicate. + full_sample_estimate : np.ndarray + Estimate vector from the full-sample weights, shape ``(k,)``. + resolved : ResolvedSurveyDesign + Must have ``uses_replicate_variance == True``. + + Returns + ------- + tuple of (np.ndarray, int) + ``(vcov, n_valid)`` where *vcov* has shape ``(k, k)`` and *n_valid* + is the number of replicates that produced finite estimates. + """ + full_sample_estimate = np.asarray(full_sample_estimate, dtype=np.float64).ravel() + k = len(full_sample_estimate) + rep_weights = resolved.replicate_weights + method = resolved.replicate_method + R = resolved.n_replicates + + # Collect replicate estimate vectors + est_reps = np.full((R, k), np.nan) + for r in range(R): + w_r = rep_weights[:, r].copy() + if not resolved.combined_weights: + w_r = w_r * resolved.weights + if np.sum(w_r) == 0: + continue + try: + est_r = refit_fn(w_r) + est_r = np.asarray(est_r, dtype=np.float64).ravel() + if len(est_r) == k: + est_reps[r] = est_r + except (np.linalg.LinAlgError, ValueError, RuntimeError): + pass # NaN row for failed replicate + + # Remove replicates with NaN estimates + valid = np.all(np.isfinite(est_reps), axis=1) + n_invalid = int(R - np.sum(valid)) + if n_invalid > 0: + warnings.warn( + f"{n_invalid} of {R} replicate refits failed. " + f"Variance computed from {int(np.sum(valid))} valid replicates.", + UserWarning, + stacklevel=2, + ) + n_valid = int(np.sum(valid)) + if n_valid < 2: + if n_valid == 0: + warnings.warn( + "All replicate refits failed. Returning NaN variance.", + UserWarning, + stacklevel=2, + ) + else: + warnings.warn( + f"Only {n_valid} valid replicate(s) — variance is not estimable " + f"with fewer than 2. Returning NaN.", + UserWarning, + stacklevel=2, + ) + return np.full((k, k), np.nan), n_valid + + est_valid = est_reps[valid] + c = full_sample_estimate + + # --- Centering (mse flag) --- + if resolved.mse: + center = c + else: + if resolved.replicate_rscales is not None: + pos_scale = resolved.replicate_rscales[valid] > 0 + if np.any(pos_scale): + center = np.mean(est_valid[pos_scale], axis=0) + else: + center = np.mean(est_valid, axis=0) + else: + center = np.mean(est_valid, axis=0) + diffs = est_valid - center[np.newaxis, :] + + outer_sum = diffs.T @ diffs # (k, k) + + # --- Method-specific scaling --- + # BRR/Fay: fixed scaling, ignore user-supplied scale/rscales + if method in ("BRR", "Fay"): + if resolved.replicate_scale is not None or resolved.replicate_rscales is not None: + warnings.warn( + f"Custom replicate_scale/replicate_rscales ignored for {method} " + f"(BRR/Fay use fixed scaling).", + UserWarning, + stacklevel=2, + ) + factor = _replicate_variance_factor(method, R, resolved.fay_rho) + return factor * outer_sum, n_valid + + # JK1/JKn: apply scale * rscales multiplicatively + scale = resolved.replicate_scale if resolved.replicate_scale is not None else 1.0 + + if resolved.replicate_rscales is not None: + valid_rscales = resolved.replicate_rscales[valid] + V = np.zeros((k, k)) + for i in range(len(diffs)): + V += valid_rscales[i] * np.outer(diffs[i], diffs[i]) + return scale * V, n_valid + + if method == "JK1": + factor = _replicate_variance_factor(method, R, resolved.fay_rho) + return scale * factor * outer_sum, n_valid + elif method == "JKn": + rep_strata = resolved.replicate_strata + if rep_strata is None: + raise ValueError("JKn requires replicate_strata") + valid_strata = rep_strata[valid] + V = np.zeros((k, k)) + for h in np.unique(rep_strata): + n_h_original = int(np.sum(rep_strata == h)) + mask_h = valid_strata == h + if not np.any(mask_h): + continue + diffs_h = diffs[mask_h] + V += ((n_h_original - 1.0) / n_h_original) * (diffs_h.T @ diffs_h) + return scale * V, n_valid + else: + raise ValueError(f"Unknown replicate method: {method}") + + def aggregate_to_psu( values: np.ndarray, resolved: "ResolvedSurveyDesign", diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index 0d5815895..aa420a930 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -127,15 +127,14 @@ def fit( # type: ignore[override] resolved_survey, survey_weights, survey_weight_type, survey_metadata = ( _resolve_survey_for_fit(survey_design, data, self.inference) ) - # Reject replicate-weight designs — TWFE within-transformation must - # be recomputed per replicate (same reason as SunAbraham rejection) - if resolved_survey is not None and resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "TwoWayFixedEffects does not yet support replicate-weight " - "survey designs. The weighted within-transformation must be " - "recomputed for each replicate. Use CallawaySantAnna or " - "TripleDifference for replicate-weight inference, or use a " - "TSL-based survey design (strata/psu/fpc)." + _uses_replicate_twfe = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate_twfe and self.inference == "wild_bootstrap": + raise ValueError( + "Cannot use inference='wild_bootstrap' with replicate-weight " + "survey designs. Replicate weights provide their own variance " + "estimation." ) # Use unit-level clustering if not specified (use local variable to avoid mutation) @@ -210,6 +209,10 @@ def fit( # type: ignore[override] # If "error", let LinearRegression raise immediately # If "warn" or "silent", suppress generic warning and use TWFE's context-specific # error/warning messages (more informative for panel data) + # For replicate designs: pass survey_design=None to prevent LinearRegression + # from computing replicate vcov on already-demeaned data (demeaning depends + # on weights, so replicate refits must re-demean at the estimator level). + _lr_survey_twfe = None if _uses_replicate_twfe else resolved_survey if self.rank_deficient_action == "error": reg = LinearRegression( include_intercept=False, @@ -219,7 +222,7 @@ def fit( # type: ignore[override] rank_deficient_action="error", weights=survey_weights, weight_type=survey_weight_type, - survey_design=resolved_survey, + survey_design=_lr_survey_twfe, ).fit(X, y, df_adjustment=df_adjustment) else: # Suppress generic warning, TWFE provides context-specific messages below @@ -235,7 +238,7 @@ def fit( # type: ignore[override] rank_deficient_action="silent", weights=survey_weights, weight_type=survey_weight_type, - survey_design=resolved_survey, + survey_design=_lr_survey_twfe, ).fit(X, y, df_adjustment=df_adjustment) coefficients = reg.coefficients_ @@ -281,8 +284,57 @@ def fit( # type: ignore[override] stacklevel=2, ) - # Get inference - either from bootstrap or analytical - if self.inference == "wild_bootstrap": + # Get inference - replicate, bootstrap, or analytical + if _uses_replicate_twfe: + # Estimator-level replicate variance: re-do within-transform per replicate + from diff_diff.linalg import solve_ols + from diff_diff.survey import compute_replicate_refit_variance + from diff_diff.utils import safe_inference as _safe_inf + + _all_vars_twfe = list(all_vars) + _covariates_twfe = list(covariates) if covariates else [] + # Handle rank-deficient nuisance: refit only identified columns + _id_mask_twfe = ~np.isnan(coefficients) + _id_cols_twfe = np.where(_id_mask_twfe)[0] + + def _refit_twfe(w_r): + # Drop zero-weight obs to prevent zero-sum demeaning + # (JK1/BRR half-samples zero entire clusters) + nz = w_r > 0 + data_nz = data[nz].copy() + w_nz = w_r[nz] + data_dem_r = _within_transform_util( + data_nz, _all_vars_twfe, unit, time, suffix="_demeaned", weights=w_nz, + ) + y_r = data_dem_r[f"{outcome}_demeaned"].values + X_list_r = [data_dem_r["_treatment_post_demeaned"].values] + for cov_ in _covariates_twfe: + X_list_r.append(data_dem_r[f"{cov_}_demeaned"].values) + X_r = np.column_stack([np.ones(len(y_r))] + X_list_r) + coef_r, _, _ = solve_ols( + X_r[:, _id_cols_twfe], y_r, + weights=w_nz, weight_type=survey_weight_type, + rank_deficient_action="silent", return_vcov=False, + ) + return coef_r + + from diff_diff.linalg import _expand_vcov_with_nan as _expand_twfe + vcov_reduced, _n_valid_rep_twfe = compute_replicate_refit_variance( + _refit_twfe, coefficients[_id_mask_twfe], resolved_survey + ) + vcov = _expand_twfe(vcov_reduced, len(coefficients), _id_cols_twfe) + se = float(np.sqrt(max(vcov[att_idx, att_idx], 0.0))) + _df_rep = ( + survey_metadata.df_survey + if survey_metadata and survey_metadata.df_survey + else 0 # rank-deficient replicate → NaN inference + ) + if _n_valid_rep_twfe < resolved_survey.n_replicates: + _df_rep = _n_valid_rep_twfe - 1 if _n_valid_rep_twfe > 1 else 0 + if survey_metadata is not None: + survey_metadata.df_survey = _df_rep if _df_rep > 0 else None + t_stat, p_value, conf_int = _safe_inf(att, se, alpha=self.alpha, df=_df_rep) + elif self.inference == "wild_bootstrap": # Override with wild cluster bootstrap inference se, p_value, conf_int, t_stat, vcov, _ = self._run_wild_bootstrap_inference( X, y, residuals, cluster_ids, att_idx diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index cd4baa4e1..28f7eb3cf 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -240,13 +240,16 @@ def fit( _resolve_survey_for_fit(survey_design, data, "analytical") ) + _uses_replicate_ts = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate_ts and self.n_bootstrap > 0: + raise ValueError( + "Cannot use n_bootstrap > 0 with replicate-weight survey designs. " + "Replicate weights provide their own variance estimation." + ) # Validate within-unit constancy for panel survey designs if resolved_survey is not None: - if resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "TwoStageDiD does not yet support replicate-weight survey " - "designs. Use a TSL-based survey design (strata/psu/fpc)." - ) _validate_unit_constant_survey(data, unit, survey_design) if resolved_survey.weight_type != "pweight": raise ValueError( @@ -329,6 +332,11 @@ def fit( if resolved_survey.fpc is not None else None ), + replicate_weights=( + resolved_survey.replicate_weights[keep_mask.values] + if resolved_survey.replicate_weights is not None + else None + ), ) # Recompute n_psu/n_strata after subsetting new_n_psu = ( @@ -475,6 +483,9 @@ def fit( # Survey degrees of freedom for t-distribution inference _survey_df = resolved_survey.df_survey if resolved_survey is not None else None + # Replicate df: rank-deficient → NaN inference + if _uses_replicate_ts and _survey_df is None: + _survey_df = 0 # Always compute overall ATT (static specification) overall_att, overall_se = self._stage2_static( @@ -495,58 +506,152 @@ def fit( survey_weight_type=survey_weight_type, ) + # Compute overall ATT inference (may be overridden by replicate below) overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df ) - # Event study and group aggregation + # Event study and group aggregation (full-sample, for point estimates) event_study_effects = None group_effects = None if aggregate in ("event_study", "all"): event_study_effects = self._stage2_event_study( - df=df, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - cluster_var=cluster_var, - treatment_groups=treatment_groups, - ref_period=ref_period, - balance_e=balance_e, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, - survey_weight_type=survey_weight_type, - survey_df=_survey_df, + df=df, unit=unit, time=time, first_treat=first_treat, + covariates=covariates, omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, unit_fe=unit_fe, time_fe=time_fe, + grand_mean=grand_mean, delta_hat=delta_hat, + cluster_var=cluster_var, treatment_groups=treatment_groups, + ref_period=ref_period, balance_e=balance_e, + kept_cov_mask=kept_cov_mask, survey_weights=survey_weights, + survey_weight_type=survey_weight_type, survey_df=_survey_df, ) if aggregate in ("group", "all"): group_effects = self._stage2_group( - df=df, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - cluster_var=cluster_var, - treatment_groups=treatment_groups, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, - survey_weight_type=survey_weight_type, - survey_df=_survey_df, + df=df, unit=unit, time=time, first_treat=first_treat, + covariates=covariates, omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, unit_fe=unit_fe, time_fe=time_fe, + grand_mean=grand_mean, delta_hat=delta_hat, + cluster_var=cluster_var, treatment_groups=treatment_groups, + kept_cov_mask=kept_cov_mask, survey_weights=survey_weights, + survey_weight_type=survey_weight_type, survey_df=_survey_df, ) + # Replicate variance override: derive keys from actual outputs, then refit + _n_valid_rep_ts = None + _vcov_rep_ts = None + if _uses_replicate_ts: + from diff_diff.survey import compute_replicate_refit_variance + + # Derive keys from actual outputs (excludes filtered/Prop5 horizons) + _sorted_es_periods_ts = sorted( + e for e in (event_study_effects or {}).keys() + if np.isfinite(event_study_effects[e]["effect"]) + ) + _sorted_groups_ts = sorted( + g for g in (group_effects or {}).keys() + if np.isfinite(group_effects[g]["effect"]) + ) + _n_es_ts = len(_sorted_es_periods_ts) + _n_grp_ts = len(_sorted_groups_ts) + + # Build full-sample estimate from actual outputs + _full_est_ts = [overall_att] + _full_est_ts.extend([event_study_effects[e]["effect"] for e in _sorted_es_periods_ts]) + _full_est_ts.extend([group_effects[g]["effect"] for g in _sorted_groups_ts]) + + def _refit_ts(w_r): + ufe_r, tfe_r, gm_r, delta_r, kcm_r = self._fit_untreated_model( + df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, + ) + y_tilde_r = self._residualize( + df, outcome, unit, time, covariates, + ufe_r, tfe_r, gm_r, delta_r, + ) + df_tmp = df.copy() + df_tmp["_y_tilde"] = y_tilde_r + results = [] + + att_r, _ = self._stage2_static( + df=df_tmp, unit=unit, time=time, first_treat=first_treat, + covariates=covariates, omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, unit_fe=ufe_r, time_fe=tfe_r, + grand_mean=gm_r, delta_hat=delta_r, cluster_var=cluster_var, + kept_cov_mask=kcm_r, survey_weights=w_r, + survey_weight_type="pweight", + ) + results.append(att_r) + + if _sorted_es_periods_ts: + es_r = self._stage2_event_study( + df=df_tmp, unit=unit, time=time, first_treat=first_treat, + covariates=covariates, omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, unit_fe=ufe_r, time_fe=tfe_r, + grand_mean=gm_r, delta_hat=delta_r, + cluster_var=cluster_var, treatment_groups=treatment_groups, + ref_period=ref_period, balance_e=balance_e, + kept_cov_mask=kcm_r, survey_weights=w_r, + survey_weight_type="pweight", survey_df=None, + ) + for e in _sorted_es_periods_ts: + results.append(es_r[e]["effect"] if e in es_r else np.nan) + + if _sorted_groups_ts: + grp_r = self._stage2_group( + df=df_tmp, unit=unit, time=time, first_treat=first_treat, + covariates=covariates, omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, unit_fe=ufe_r, time_fe=tfe_r, + grand_mean=gm_r, delta_hat=delta_r, + cluster_var=cluster_var, treatment_groups=treatment_groups, + kept_cov_mask=kcm_r, survey_weights=w_r, + survey_weight_type="pweight", survey_df=None, + ) + for g in _sorted_groups_ts: + results.append(grp_r[g]["effect"] if g in grp_r else np.nan) + + return np.array(results) + + _vcov_rep_ts, _n_valid_rep_ts = compute_replicate_refit_variance( + _refit_ts, np.array(_full_est_ts), resolved_survey + ) + overall_se = float(np.sqrt(max(_vcov_rep_ts[0, 0], 0.0))) + + # Override df if replicates were dropped + if _n_valid_rep_ts < resolved_survey.n_replicates: + _survey_df = _n_valid_rep_ts - 1 if _n_valid_rep_ts > 1 else 0 + if survey_metadata is not None: + survey_metadata.df_survey = _survey_df if _survey_df and _survey_df > 0 else None + + # Recompute overall inference with replicate SE/df + overall_t, overall_p, overall_ci = safe_inference( + overall_att, overall_se, alpha=self.alpha, df=_survey_df + ) + + # Override event-study SEs (only for identified effects) + for i, e in enumerate(_sorted_es_periods_ts): + if event_study_effects is not None and e in event_study_effects: + se_e = float(np.sqrt(max(_vcov_rep_ts[1 + i, 1 + i], 0.0))) + eff_e = event_study_effects[e]["effect"] + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) + event_study_effects[e]["se"] = se_e + event_study_effects[e]["t_stat"] = t_e + event_study_effects[e]["p_value"] = p_e + event_study_effects[e]["conf_int"] = ci_e + + # Override group SEs (only for identified effects) + for j, g in enumerate(_sorted_groups_ts): + if group_effects is not None and g in group_effects: + se_g = float(np.sqrt(max( + _vcov_rep_ts[1 + _n_es_ts + j, 1 + _n_es_ts + j], 0.0 + ))) + eff_g = group_effects[g]["effect"] + t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) + group_effects[g]["se"] = se_g + group_effects[g]["t_stat"] = t_g + group_effects[g]["p_value"] = p_g + group_effects[g]["conf_int"] = ci_g + # Build treatment effects DataFrame treated_df = df.loc[omega_1_mask, [unit, time, "_y_tilde", "_rel_time"]].copy() treated_df = treated_df.rename(columns={"_y_tilde": "tau_hat", "_rel_time": "rel_time"}) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index a39822395..bc6aa5fb4 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -763,7 +763,7 @@ where weights ŵ_{g,e} = n_{g,e} / Σ_g n_{g,e} (sample share of cohort g at eve - [x] R comparison: ATT matches within machine precision (<1e-11) - [x] R comparison: SE matches within 0.3% (well within 1% threshold) - [x] R comparison: Event study effects match perfectly (correlation 1.0) -- [x] Survey design support (Phase 3): weighted within-transform, survey weights in LinearRegression with TSL vcov; bootstrap+survey supported (Phase 6) via Rao-Wu rescaled bootstrap +- [x] Survey design support (Phase 3): weighted within-transform, survey weights in LinearRegression with TSL vcov; bootstrap+survey supported (Phase 6) via Rao-Wu rescaled bootstrap. Replicate weights supported via estimator-level refit (see Replicate Weight Variance section); replicate+bootstrap rejected. --- @@ -1036,7 +1036,7 @@ The paper text states a stricter bound (T_min + 1) but the R code by the co-auth - [x] Overall ATT as average of post-treatment delta_h with delta-method SE - [x] Anticipation parameter support - [x] Never-treated encoding (0 and inf) -- [x] Survey design support (Phase 3): Q-weights compose multiplicatively with survey weights; TSL vcov on composed weights; survey design columns propagated through sub-experiments +- [x] Survey design support (Phase 3): Q-weights compose multiplicatively with survey weights; TSL vcov on composed weights; survey design columns propagated through sub-experiments. Replicate weights supported via estimator-level refit with Q-weight composition (see Replicate Weight Variance section). - **Note:** Survey weights compose multiplicatively with Q-weights for StackedDiD; only `weight_type="pweight"` (default) is supported — `fweight` and `aweight` are rejected because Q-weight composition changes weight semantics (non-integer for fweight, non-inverse-variance for aweight) --- @@ -2245,16 +2245,18 @@ variance from the distribution of replicate estimates. t-based inference. - **Note:** Replicate-weight support matrix: - **Supported**: CallawaySantAnna (reg/ipw/dr without covariates, no - bootstrap), ContinuousDiD - (no bootstrap), EfficientDiD (no bootstrap), TripleDifference (all - methods), LinearRegression (OLS path) - - **Rejected with NotImplementedError**: SunAbraham, TwoWayFixedEffects - (within-transformation must be recomputed per replicate), - DifferenceInDifferences, MultiPeriodDiD, StackedDiD (use - compute_survey_vcov directly), ImputationDiD, TwoStageDiD (custom - variance), SyntheticDiD, TROP (bootstrap-based variance), - BaconDecomposition (diagnostic only) - - CS/ContinuousDiD/EfficientDiD reject replicate + `n_bootstrap > 0` + bootstrap), ContinuousDiD (no bootstrap), EfficientDiD (no bootstrap), + TripleDifference (all methods), LinearRegression (OLS path), + DifferenceInDifferences (no-absorb via LinearRegression dispatch, + absorb via estimator-level refit), MultiPeriodDiD (no-absorb via + `compute_replicate_vcov`, absorb via estimator-level refit), + TwoWayFixedEffects (estimator-level refit with within-transformation), + SunAbraham (estimator-level refit, replaces `vcov_cohort`), + StackedDiD (estimator-level refit with Q-weight composition), + ImputationDiD (two-stage refit), TwoStageDiD (two-stage refit) + - **Rejected with NotImplementedError**: SyntheticDiD, TROP + (bootstrap-based variance), BaconDecomposition (diagnostic only) + - Estimators with replicate support reject replicate + bootstrap (replicate weights provide analytical variance) - **Note:** When invalid replicates are dropped in `compute_replicate_vcov` (OLS path), `n_valid` is returned and used for `df_survey = n_valid - 1` @@ -2300,6 +2302,13 @@ Domain estimation preserving full design structure. - **Note:** For replicate-weight designs, `subpopulation()` zeros out both full-sample and replicate weight columns for excluded observations, preserving all replicate metadata. +- **Note:** Estimator-level replicate refits (TWFE, SunAbraham, DiD/MultiPeriodDiD + with `absorb`) drop zero-weight observations before weighted demeaning to + prevent division-by-zero in within-transformation group means. This matches + R's `survey::withReplicates()` convention where zero-weight units are excluded + from per-replicate estimation. Replicates that fail despite this (e.g., + rank-deficient after unit deletion) are counted as invalid and excluded from + variance computation. - **Note:** Defensive enhancement: ContinuousDiD and TripleDifference validate the positive-weight effective sample size before WLS cell fits. After `subpopulation()` zeroes weights, raw row counts may exceed the diff --git a/tests/test_replicate_weight_expansion.py b/tests/test_replicate_weight_expansion.py new file mode 100644 index 000000000..3f6fa4109 --- /dev/null +++ b/tests/test_replicate_weight_expansion.py @@ -0,0 +1,423 @@ +"""Tests for replicate weight support expansion to 7 additional estimators.""" + +import numpy as np +import pandas as pd +import pytest + +from diff_diff import ( + DifferenceInDifferences, + ImputationDiD, + MultiPeriodDiD, + StackedDiD, + SunAbraham, + TwoStageDiD, + TwoWayFixedEffects, +) +from diff_diff.survey import SurveyDesign + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +def _make_simple_panel(): + """2-period panel for DiD/MultiPeriodDiD (treatment/post binary columns).""" + np.random.seed(123) + n_units = 40 + rows = [] + for i in range(n_units): + treated = 1 if i < 20 else 0 + wt = 1.0 + 0.2 * (i % 5) + for t in [0, 1]: + y = 5.0 + 0.5 * treated + 1.0 * t + if treated and t == 1: + y += 2.0 # ATT = 2 + y += np.random.normal(0, 0.3) + rows.append({ + "unit": i, "time": t, "treated": treated, "post": t, + "outcome": y, "weight": wt, + }) + data = pd.DataFrame(rows) + return data + + +def _make_staggered_panel(): + """Multi-period staggered panel for TWFE/SA/Stacked/Imputation/TwoStage.""" + np.random.seed(456) + n_units, n_periods = 50, 8 + rows = [] + for i in range(n_units): + if i < 15: + ft = 4 # cohort 1 + elif i < 30: + ft = 6 # cohort 2 + else: + ft = 0 # never-treated + wt = 1.0 + 0.3 * (i % 5) + for t in range(1, n_periods + 1): + y = 10.0 + i * 0.03 + t * 0.2 + if ft > 0 and t >= ft: + y += 2.0 + y += np.random.normal(0, 0.4) + rows.append({ + "unit": i, "time": t, "first_treat": ft, + "outcome": y, "weight": wt, + "treated": 1 if ft > 0 else 0, + "post": 1 if ft > 0 and t >= ft else 0, + }) + data = pd.DataFrame(rows) + return data + + +def _add_jk1_replicates(data, n_rep=15, unit_col="unit"): + """Add JK1 (delete-cluster jackknife) replicate weight columns.""" + units = sorted(data[unit_col].unique()) + cluster_size = max(1, len(units) // n_rep) + rep_cols = [] + for r in range(n_rep): + start = r * cluster_size + end = min((r + 1) * cluster_size, len(units)) + deleted_units = set(units[start:end]) + w_r = data["weight"].values.copy() + mask = data[unit_col].isin(deleted_units).values + w_r[mask] = 0.0 + w_r[~mask] *= n_rep / (n_rep - 1) + col = f"rep_{r}" + data[col] = w_r + rep_cols.append(col) + return rep_cols + + +def _add_brr_replicates(data, n_rep=16, unit_col="unit"): + """Add BRR replicate weight columns (random sign perturbation).""" + rng = np.random.RandomState(789) + units = sorted(data[unit_col].unique()) + rep_cols = [] + for r in range(n_rep): + signs = rng.choice([-1, 1], size=len(units)) + sign_map = dict(zip(units, signs)) + perturbation = data[unit_col].map(sign_map).values.astype(float) + # BRR: w_r = w * (1 + epsilon) / 2, where epsilon in {-1, 1} + # Simplified: w_r = w * (1 + perturbation) (combined_weights=True style) + w_r = data["weight"].values * (1.0 + 0.5 * perturbation) + w_r = np.maximum(w_r, 0.0) + col = f"brr_{r}" + data[col] = w_r + rep_cols.append(col) + return rep_cols + + +# --------------------------------------------------------------------------- +# Smoke tests — each estimator × {JK1, BRR} +# --------------------------------------------------------------------------- + +class TestDiDReplicate: + """DifferenceInDifferences with replicate weights.""" + + def test_did_jk1(self): + data = _make_simple_panel() + rep_cols = _add_jk1_replicates(data, n_rep=10) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = DifferenceInDifferences().fit( + data, "outcome", "treated", "post", survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_did_brr(self): + data = _make_simple_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = DifferenceInDifferences().fit( + data, "outcome", "treated", "post", survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 + + def test_did_wild_bootstrap_rejected(self): + """Wild bootstrap + survey is rejected before replicate check.""" + data = _make_simple_panel() + rep_cols = _add_jk1_replicates(data, n_rep=10) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + with pytest.raises((ValueError, NotImplementedError)): + DifferenceInDifferences(inference="wild_bootstrap", cluster="unit").fit( + data, "outcome", "treated", "post", survey_design=sd, + ) + + +class TestDiDAbsorbReplicate: + """DiD absorb path with replicate weights.""" + + def test_did_absorb_brr(self): + """Absorb path should produce finite replicate SE.""" + data = _make_simple_panel() + # Add a group column for absorb + data["group"] = (data["unit"] % 4).astype(str) + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = DifferenceInDifferences().fit( + data, "outcome", "treated", "post", absorb=["group"], survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 + + +class TestMultiPeriodDiDReplicate: + """MultiPeriodDiD with replicate weights.""" + + def test_multiperiod_jk1(self): + data = _make_simple_panel() + rep_cols = _add_jk1_replicates(data, n_rep=10) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = MultiPeriodDiD().fit( + data, "outcome", "treated", "time", post_periods=[1], survey_design=sd, + ) + assert np.isfinite(result.avg_att) + assert np.isfinite(result.avg_se) and result.avg_se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_multiperiod_brr(self): + data = _make_simple_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = MultiPeriodDiD().fit( + data, "outcome", "treated", "time", post_periods=[1], survey_design=sd, + ) + assert np.isfinite(result.avg_att) + assert np.isfinite(result.avg_se) and result.avg_se > 0 + + +class TestTWFEReplicate: + """TwoWayFixedEffects with replicate weights.""" + + @staticmethod + def _make_twfe_panel(): + """Balanced 2-period panel with variation in treatment timing.""" + np.random.seed(321) + n_units = 40 + rows = [] + for i in range(n_units): + treated = 1 if i < 20 else 0 + wt = 1.0 + 0.2 * (i % 5) + for t in [0, 1]: + y = 5.0 + i * 0.05 + t * 1.0 + if treated and t == 1: + y += 2.0 + y += np.random.normal(0, 0.3) + rows.append({ + "unit": i, "time": t, "treated": treated, + "post": t, "outcome": y, "weight": wt, + }) + return pd.DataFrame(rows) + + def test_twfe_brr(self): + """BRR works well with TWFE: perturbation doesn't zero out units.""" + data = self._make_twfe_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = TwoWayFixedEffects().fit( + data, "outcome", "treated", "post", "unit", survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_twfe_brr_larger(self): + """Second BRR test with different seed.""" + data = self._make_twfe_panel() + rep_cols = _add_brr_replicates(data, n_rep=20) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = TwoWayFixedEffects().fit( + data, "outcome", "treated", "post", "unit", survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 + + +class TestSunAbrahamReplicate: + """SunAbraham with replicate weights.""" + + def test_sun_abraham_brr(self): + """BRR replicates are less aggressive than JK1 for SunAbraham.""" + data = _make_staggered_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = SunAbraham(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_sun_abraham_bootstrap_rejected(self): + data = _make_staggered_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + with pytest.raises(ValueError, match="n_bootstrap"): + SunAbraham(n_bootstrap=100).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + + +class TestStackedDiDReplicate: + """StackedDiD with replicate weights.""" + + def test_stacked_jk1(self): + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = StackedDiD().fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_stacked_brr(self): + data = _make_staggered_panel() + rep_cols = _add_brr_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = StackedDiD().fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + + +class TestImputationDiDReplicate: + """ImputationDiD with replicate weights.""" + + def test_imputation_jk1(self): + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = ImputationDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_imputation_event_study_replicate(self): + """Event-study with replicate weights: overall ATT SE must be finite, + and at least some per-period SEs should be finite.""" + data = _make_staggered_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = ImputationDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", + aggregate="event_study", survey_design=sd, + ) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + assert result.event_study_effects is not None + # At least some identified periods should have finite SE + finite_ses = [ + e for e, eff in result.event_study_effects.items() + if np.isfinite(eff["effect"]) and np.isfinite(eff["se"]) and eff["se"] > 0 + ] + assert len(finite_ses) > 0, "No event-study periods have finite replicate SE" + + def test_imputation_group_replicate(self): + """Group SEs should use replicate variance.""" + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = ImputationDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", + aggregate="group", survey_design=sd, + ) + assert result.group_effects is not None + for g, eff in result.group_effects.items(): + assert np.isfinite(eff["se"]) and eff["se"] > 0, f"group {g}: SE not finite" + + def test_imputation_bootstrap_rejected(self): + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + with pytest.raises(ValueError, match="n_bootstrap"): + ImputationDiD(n_bootstrap=100).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + + +class TestTwoStageDiDReplicate: + """TwoStageDiD with replicate weights.""" + + def test_two_stage_jk1(self): + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = TwoStageDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_two_stage_event_study_replicate(self): + """Event-study SEs should use replicate variance, not GMM SE.""" + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = TwoStageDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", + aggregate="event_study", survey_design=sd, + ) + assert result.event_study_effects is not None + non_ref = {e: eff for e, eff in result.event_study_effects.items() if eff["effect"] != 0.0} + assert len(non_ref) > 0, "No non-reference event-study effects" + for e, eff in non_ref.items(): + assert np.isfinite(eff["se"]) and eff["se"] > 0, f"period {e}: SE not finite" + + def test_two_stage_always_treated(self): + """Replicate weights should be subsetted when always-treated units are excluded.""" + data = _make_staggered_panel() + # Add always-treated units (first_treat <= min time) + for i in range(50, 55): + for t in range(1, 9): + data = pd.concat([data, pd.DataFrame([{ + "unit": i, "time": t, "first_treat": 1, + "outcome": 12.0 + np.random.normal(0, 0.3), + "weight": 1.5, "treated": 1, "post": 1, + }])], ignore_index=True) + rep_cols = _add_jk1_replicates(data, n_rep=10, unit_col="unit") + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + # Should not crash despite always-treated unit exclusion + result = TwoStageDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + + def test_two_stage_bootstrap_rejected(self): + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + with pytest.raises(ValueError, match="n_bootstrap"): + TwoStageDiD(n_bootstrap=100).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + + +class TestSunAbrahamCohortSEs: + """SunAbraham cohort-level SEs should be consistent with replicate vcov.""" + + def test_cohort_ses_finite(self): + """Cohort SEs should be recomputed from replicate vcov, not stale.""" + data = _make_staggered_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = SunAbraham(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert result.cohort_effects is not None + for key, eff in result.cohort_effects.items(): + se = eff["se"] + assert np.isfinite(se), f"cohort {key}: SE is {se}" diff --git a/tests/test_survey_phase6.py b/tests/test_survey_phase6.py index bb54ada36..1ebe589d5 100644 --- a/tests/test_survey_phase6.py +++ b/tests/test_survey_phase6.py @@ -498,8 +498,8 @@ def test_replicate_metadata(self, replicate_data): assert sm.n_replicates == len(rep_cols) assert sm.df_survey == len(rep_cols) - 1 - def test_replicate_rejected_by_base_did(self, replicate_data): - """DifferenceInDifferences rejects replicate-weight designs.""" + def test_replicate_accepted_by_base_did(self, replicate_data): + """DifferenceInDifferences now supports replicate-weight designs.""" data, rep_cols = replicate_data n = len(data) data["treated"] = (np.arange(n) < n // 2).astype(int) @@ -510,11 +510,12 @@ def test_replicate_rejected_by_base_did(self, replicate_data): weights="weight", replicate_weights=rep_cols, replicate_method="JK1", ) - with pytest.raises(NotImplementedError, match="DifferenceInDifferences"): - DifferenceInDifferences().fit( - data, outcome="outcome", treatment="treated", time="post", - survey_design=sd, - ) + result = DifferenceInDifferences().fit( + data, outcome="outcome", treatment="treated", time="post", + survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 def test_replicate_if_variance(self, replicate_data): """IF-based replicate variance produces finite results.""" @@ -1284,8 +1285,8 @@ def test_triple_diff_replicate_all_methods(self, est_method): assert np.isfinite(result.se) assert result.survey_metadata is not None - def test_sun_abraham_replicate_rejected(self): - """SunAbraham should reject replicate-weight survey designs.""" + def test_sun_abraham_replicate_accepted(self): + """SunAbraham now supports replicate-weight survey designs.""" from diff_diff import SunAbraham data, rep_cols = self._make_staggered_replicate_data() @@ -1293,11 +1294,11 @@ def test_sun_abraham_replicate_rejected(self): weights="weight", replicate_weights=rep_cols, replicate_method="JK1", ) - with pytest.raises(NotImplementedError, match="SunAbraham"): - SunAbraham(n_bootstrap=0).fit( - data, "outcome", "unit", "time", "first_treat", - survey_design=sd, - ) + result = SunAbraham(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", + survey_design=sd, + ) + assert np.isfinite(result.overall_att) class TestSubpopulationMaskValidation: @@ -1357,23 +1358,49 @@ def test_callaway_santanna_replicate_bootstrap_rejected(self): survey_design=sd, ) - def test_twfe_replicate_rejected(self): - """TwoWayFixedEffects should reject replicate-weight designs.""" + def test_twfe_replicate_accepted(self): + """TwoWayFixedEffects now supports replicate-weight designs.""" from diff_diff.twfe import TwoWayFixedEffects - data, rep_cols = TestEstimatorReplicateWeights._make_staggered_replicate_data() + # Create a simple 2-period panel that works with within-transformation + np.random.seed(77) + rows = [] + for i in range(40): + treated = 1 if i < 20 else 0 + wt = 1.0 + 0.2 * (i % 5) + for t in [0, 1]: + y = 5.0 + i * 0.05 + t * 1.0 + if treated and t == 1: + y += 2.0 + y += np.random.normal(0, 0.3) + rows.append({ + "unit": i, "time": t, "treated": treated, + "post": t, "outcome": y, "weight": wt, + }) + data = pd.DataFrame(rows) + # Use BRR to avoid zero-weight units breaking within-transform + rng = np.random.RandomState(42) + units = sorted(data["unit"].unique()) + rep_cols = [] + for r in range(16): + signs = rng.choice([-1, 1], size=len(units)) + sm = dict(zip(units, signs)) + p = data["unit"].map(sm).values.astype(float) + data[f"brr_{r}"] = np.maximum(data["weight"].values * (1.0 + 0.5 * p), 0.0) + rep_cols.append(f"brr_{r}") sd = SurveyDesign( weights="weight", replicate_weights=rep_cols, - replicate_method="JK1", + replicate_method="BRR", ) - with pytest.raises(NotImplementedError, match="TwoWayFixedEffects"): - TwoWayFixedEffects().fit( - data, outcome="outcome", treatment="first_treat", - unit="unit", time="time", survey_design=sd, - ) + result = TwoWayFixedEffects().fit( + data, outcome="outcome", treatment="treated", + time="post", unit="unit", survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 - def test_stacked_did_replicate_rejected(self): - """StackedDiD should reject replicate-weight designs.""" + def test_stacked_did_replicate_accepted(self): + """StackedDiD now supports replicate-weight designs.""" from diff_diff import StackedDiD data, rep_cols = TestEstimatorReplicateWeights._make_staggered_replicate_data() @@ -1381,11 +1408,11 @@ def test_stacked_did_replicate_rejected(self): weights="weight", replicate_weights=rep_cols, replicate_method="JK1", ) - with pytest.raises(NotImplementedError, match="StackedDiD"): - StackedDiD().fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", survey_design=sd, - ) + result = StackedDiD().fit( + data, outcome="outcome", unit="unit", time="time", + first_treat="first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) def test_invalid_replicate_scale_rejected(self): """Negative or zero replicate_scale should be rejected.""" @@ -1416,37 +1443,37 @@ def _replicate_sd_and_data(self): ) return data, sd - def test_multi_period_did_replicate_rejected(self): - """MultiPeriodDiD rejects replicate-weight designs.""" + def test_multi_period_did_replicate_accepted(self): + """MultiPeriodDiD now supports replicate-weight designs.""" from diff_diff.estimators import MultiPeriodDiD data, sd = self._replicate_sd_and_data() data["treated"] = (data["first_treat"] > 0).astype(int) data["post"] = (data["time"] >= 3).astype(int) - with pytest.raises(NotImplementedError): - MultiPeriodDiD().fit( - data, outcome="outcome", treatment="treated", - time="post", survey_design=sd, - ) + result = MultiPeriodDiD().fit( + data, outcome="outcome", treatment="treated", + time="post", survey_design=sd, + ) + assert np.isfinite(result.avg_att) - def test_imputation_did_replicate_rejected(self): - """ImputationDiD rejects replicate-weight designs.""" + def test_imputation_did_replicate_accepted(self): + """ImputationDiD now supports replicate-weight designs.""" from diff_diff.imputation import ImputationDiD data, sd = self._replicate_sd_and_data() - with pytest.raises(NotImplementedError): - ImputationDiD().fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", survey_design=sd, - ) + result = ImputationDiD(n_bootstrap=0).fit( + data, outcome="outcome", unit="unit", time="time", + first_treat="first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) - def test_two_stage_did_replicate_rejected(self): - """TwoStageDiD rejects replicate-weight designs.""" + def test_two_stage_did_replicate_accepted(self): + """TwoStageDiD now supports replicate-weight designs.""" from diff_diff.two_stage import TwoStageDiD data, sd = self._replicate_sd_and_data() - with pytest.raises(NotImplementedError): - TwoStageDiD().fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", survey_design=sd, - ) + result = TwoStageDiD(n_bootstrap=0).fit( + data, outcome="outcome", unit="unit", time="time", + first_treat="first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) def test_bacon_replicate_rejected(self): """BaconDecomposition rejects replicate-weight designs."""