diff-diff 2.8.4__tar.gz → 2.9.0__tar.gz

{diff_diff-2.8.4 → diff_diff-2.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diff-diff
-Version: 2.8.4
+Version: 2.9.0
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
@@ -134,7 +134,7 @@ Detailed guide: [`docs/llms-practitioner.txt`](docs/llms-practitioner.txt)
 - **Wild cluster bootstrap**: Valid inference with few clusters (<50) using Rademacher, Webb, or Mammen weights
 - **Panel data support**: Two-way fixed effects estimator for panel designs
 - **Multi-period analysis**: Event-study style DiD with period-specific treatment effects
-- **Staggered adoption**: Callaway-Sant'Anna (2021), Sun-Abraham (2021), Borusyak-Jaravel-Spiess (2024) imputation, Two-Stage DiD (Gardner 2022), Stacked DiD (Wing, Freedman & Hollingsworth 2024), and Efficient DiD (Chen, Sant'Anna & Xie 2025) estimators for heterogeneous treatment timing
+- **Staggered adoption**: Callaway-Sant'Anna (2021), Sun-Abraham (2021), Borusyak-Jaravel-Spiess (2024) imputation, Two-Stage DiD (Gardner 2022), Stacked DiD (Wing, Freedman & Hollingsworth 2024), Efficient DiD (Chen, Sant'Anna & Xie 2025), and Wooldridge ETWFE (2021/2023) estimators for heterogeneous treatment timing
 - **Triple Difference (DDD)**: Ortiz-Villavicencio & Sant'Anna (2025) estimators with proper covariate handling
 - **Synthetic DiD**: Combined DiD with synthetic control for improved robustness
 - **Triply Robust Panel (TROP)**: Factor-adjusted DiD with synthetic weights (Athey et al. 2025)
@@ -167,6 +167,7 @@ All estimators have short aliases for convenience:
 | `Stacked` | `StackedDiD` | Stacked DiD |
 | `Bacon` | `BaconDecomposition` | Goodman-Bacon decomposition |
 | `EDiD` | `EfficientDiD` | Efficient DiD |
+| `ETWFE` | `WooldridgeDiD` | Wooldridge ETWFE (2021/2023) |
 
 `TROP` already uses its short canonical name and needs no alias.
 

{diff_diff-2.8.4 → diff_diff-2.9.0}/README.md

@@ -84,7 +84,7 @@ Detailed guide: [`docs/llms-practitioner.txt`](docs/llms-practitioner.txt)
 - **Wild cluster bootstrap**: Valid inference with few clusters (<50) using Rademacher, Webb, or Mammen weights
 - **Panel data support**: Two-way fixed effects estimator for panel designs
 - **Multi-period analysis**: Event-study style DiD with period-specific treatment effects
-- **Staggered adoption**: Callaway-Sant'Anna (2021), Sun-Abraham (2021), Borusyak-Jaravel-Spiess (2024) imputation, Two-Stage DiD (Gardner 2022), Stacked DiD (Wing, Freedman & Hollingsworth 2024), and Efficient DiD (Chen, Sant'Anna & Xie 2025) estimators for heterogeneous treatment timing
+- **Staggered adoption**: Callaway-Sant'Anna (2021), Sun-Abraham (2021), Borusyak-Jaravel-Spiess (2024) imputation, Two-Stage DiD (Gardner 2022), Stacked DiD (Wing, Freedman & Hollingsworth 2024), Efficient DiD (Chen, Sant'Anna & Xie 2025), and Wooldridge ETWFE (2021/2023) estimators for heterogeneous treatment timing
 - **Triple Difference (DDD)**: Ortiz-Villavicencio & Sant'Anna (2025) estimators with proper covariate handling
 - **Synthetic DiD**: Combined DiD with synthetic control for improved robustness
 - **Triply Robust Panel (TROP)**: Factor-adjusted DiD with synthetic weights (Athey et al. 2025)
@@ -117,6 +117,7 @@ All estimators have short aliases for convenience:
 | `Stacked` | `StackedDiD` | Stacked DiD |
 | `Bacon` | `BaconDecomposition` | Goodman-Bacon decomposition |
 | `EDiD` | `EfficientDiD` | Efficient DiD |
+| `ETWFE` | `WooldridgeDiD` | Wooldridge ETWFE (2021/2023) |
 
 `TROP` already uses its short canonical name and needs no alias.
 

{diff_diff-2.8.4 → diff_diff-2.9.0}/diff_diff/__init__.py

@@ -164,6 +164,8 @@ from diff_diff.trop import (
     TROPResults,
     trop,
 )
+from diff_diff.wooldridge import WooldridgeDiD
+from diff_diff.wooldridge_results import WooldridgeDiDResults
 from diff_diff.utils import (
     WildBootstrapResults,
     check_parallel_trends,
@@ -210,8 +212,9 @@ SDDD = StaggeredTripleDifference
 Stacked = StackedDiD
 Bacon = BaconDecomposition
 EDiD = EfficientDiD
+ETWFE = WooldridgeDiD
 
-__version__ = "2.8.4"
+__version__ = "2.9.0"
 __all__ = [
     # Estimators
     "DifferenceInDifferences",
@@ -276,6 +279,10 @@ __all__ = [
     "EfficientDiDResults",
     "EDiDBootstrapResults",
     "EDiD",
+    # WooldridgeDiD (ETWFE)
+    "WooldridgeDiD",
+    "WooldridgeDiDResults",
+    "ETWFE",
     # Visualization
     "plot_bacon",
     "plot_event_study",

{diff_diff-2.8.4 → diff_diff-2.9.0}/diff_diff/efficient_did.py

@@ -347,8 +347,6 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
         ValueError
             Missing columns, unbalanced panel, non-absorbing treatment,
             or PT-Post without a never-treated group.
-        NotImplementedError
-            If ``covariates`` and ``survey_design`` are both set.
         """
         self._validate_params()
 
@@ -381,16 +379,6 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
 
         # Bootstrap + survey supported via PSU-level multiplier bootstrap.
 
-        # Guard covariates + survey (DR path does not yet thread survey weights)
-        if covariates is not None and len(covariates) > 0 and resolved_survey is not None:
-            raise NotImplementedError(
-                "Survey weights with covariates are not yet supported for "
-                "EfficientDiD. The doubly robust covariate path does not "
-                "thread survey weights through nuisance estimation. "
-                "Use covariates=None with survey_design, or drop survey_design "
-                "when using covariates."
-            )
-
         # Normalize empty covariates list to None (use nocov path)
         if covariates is not None and len(covariates) == 0:
             covariates = None
@@ -583,6 +571,7 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
             # Use the resolved survey's weights (already normalized per weight_type)
             # subset to unit level via _unit_first_panel_row (aligned to all_units)
             unit_level_weights = self._unit_resolved_survey.weights
+        self._unit_level_weights = unit_level_weights
 
         cohort_fractions: Dict[float, float] = {}
         if unit_level_weights is not None:
@@ -617,6 +606,15 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
                     stacklevel=2,
                 )
 
+        # Guard: never-treated with zero survey weight → no valid comparisons
+        # Applies to both covariates (DR nuisance) and nocov (weighted means) paths
+        if cohort_fractions.get(np.inf, 0.0) <= 0 and unit_level_weights is not None:
+            raise ValueError(
+                "Never-treated group has zero survey weight. EfficientDiD "
+                "requires a never-treated control group with positive "
+                "survey weight for estimation."
+            )
+
         # ----- Covariate preparation (if provided) -----
         covariate_matrix: Optional[np.ndarray] = None
         m_hat_cache: Dict[Tuple, np.ndarray] = {}
@@ -686,6 +684,15 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
             else:
                 effective_p1_col = period_1_col
 
+            # Guard: skip cohorts with zero survey weight (all units zero-weighted)
+            if cohort_fractions[g] <= 0:
+                warnings.warn(
+                    f"Cohort {g} has zero survey weight; skipping.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                continue
+
             # Estimate all (g, t) cells including pre-treatment. Under PT-Post,
             # pre-treatment cells serve as placebo/pre-trend diagnostics, matching
             # the CallawaySantAnna implementation. Users filter to t >= g for
@@ -707,6 +714,15 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
                     anticipation=self.anticipation,
                 )
 
+                # Filter out comparison pairs with zero survey weight
+                if unit_level_weights is not None and pairs:
+                    pairs = [
+                        (gp, tpre) for gp, tpre in pairs
+                        if np.sum(unit_level_weights[
+                            never_treated_mask if np.isinf(gp) else cohort_masks[gp]
+                        ]) > 0
+                    ]
+
                 if not pairs:
                     warnings.warn(
                         f"No valid comparison pairs for (g={g}, t={t}). " "ATT will be NaN.",
@@ -742,6 +758,7 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
                                 never_treated_mask,
                                 t_col_val,
                                 tpre_col_val,
+                                unit_weights=unit_level_weights,
                             )
                         # m_{g', tpre, 1}(X)
                         key_gp_tpre = (gp, tpre_col_val, effective_p1_col)
@@ -755,6 +772,7 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
                                 gp_mask_for_reg,
                                 tpre_col_val,
                                 effective_p1_col,
+                                unit_weights=unit_level_weights,
                             )
                         # r_{g, inf}(X) and r_{g, g'}(X) via sieve (Eq 4.1-4.2)
                         for comp in {np.inf, gp}:
@@ -770,6 +788,7 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
                                     k_max=self.sieve_k_max,
                                     criterion=self.sieve_criterion,
                                     ratio_clip=self.ratio_clip,
+                                    unit_weights=unit_level_weights,
                                 )
 
                     # Per-unit DR generated outcomes: shape (n_units, H)
@@ -801,6 +820,7 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
                                 group_mask_s,
                                 k_max=self.sieve_k_max,
                                 criterion=self.sieve_criterion,
+                                unit_weights=unit_level_weights,
                             )
 
                     # Conditional Omega*(X) with per-unit propensities (Eq 3.12)
@@ -817,14 +837,19 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
                         covariate_matrix=covariate_matrix,
                         s_hat_cache=s_hat_cache,
                         bandwidth=self.kernel_bandwidth,
+                        unit_weights=unit_level_weights,
                     )
 
                     # Per-unit weights: (n_units, H)
                     per_unit_w = compute_per_unit_weights(omega_cond)
 
-                    # ATT = mean_i( w(X_i) @ gen_out[i] )
+                    # ATT = (survey-)weighted mean of per-unit DR scores
                     if per_unit_w.shape[1] > 0:
-                        att_gt = float(np.mean(np.sum(per_unit_w * gen_out, axis=1)))
+                        per_unit_scores = np.sum(per_unit_w * gen_out, axis=1)
+                        if unit_level_weights is not None:
+                            att_gt = float(np.average(per_unit_scores, weights=unit_level_weights))
+                        else:
+                            att_gt = float(np.mean(per_unit_scores))
                     else:
                         att_gt = np.nan
 
@@ -979,6 +1004,7 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
                 cluster_indices=unit_cluster_indices,
                 n_clusters=n_clusters,
                 resolved_survey=self._unit_resolved_survey,
+                unit_level_weights=self._unit_level_weights,
             )
             # Update estimates with bootstrap inference
             overall_se = bootstrap_results.overall_att_se
@@ -1140,6 +1166,7 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
         unit_cohorts: np.ndarray,
         cohort_fractions: Dict[float, float],
         n_units: int,
+        unit_weights: Optional[np.ndarray] = None,
     ) -> np.ndarray:
         """Compute weight influence function correction (O(1) scale, matching EIF).
 
@@ -1159,6 +1186,9 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
             ``{cohort: n_cohort / n}`` for each cohort.
         n_units : int
             Total number of units.
+        unit_weights : ndarray, shape (n_units,), optional
+            Survey weights at the unit level.  When provided, uses the
+            survey-weighted WIF formula: IF_i(p_g) = (w_i * 1{G_i=g} - pg_k).
 
         Returns
         -------
@@ -1172,10 +1202,19 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
             return np.zeros(n_units)
 
         indicator = (unit_cohorts[:, None] == groups_for_keepers[None, :]).astype(float)
-        indicator_sum = np.sum(indicator - pg_keepers, axis=1)
+
+        if unit_weights is not None:
+            # Survey-weighted WIF (matches staggered_aggregation.py:392-401):
+            # IF_i(p_g) = (w_i * 1{G_i=g} - pg_k), NOT (1{G_i=g} - pg_k)
+            weighted_indicator = indicator * unit_weights[:, None]
+            indicator_diff = weighted_indicator - pg_keepers
+            indicator_sum = np.sum(indicator_diff, axis=1)
+        else:
+            indicator_diff = indicator - pg_keepers
+            indicator_sum = np.sum(indicator_diff, axis=1)
 
         with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
-            if1 = (indicator - pg_keepers) / sum_pg
+            if1 = indicator_diff / sum_pg
             if2 = np.outer(indicator_sum, pg_keepers) / sum_pg**2
             wif_matrix = if1 - if2
             wif_contrib = wif_matrix @ effects
@@ -1232,13 +1271,34 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
 
         # WIF correction: accounts for uncertainty in cohort-size weights
         wif = self._compute_wif_contribution(
-            keepers, effects, unit_cohorts, cohort_fractions, n_units
+            keepers, effects, unit_cohorts, cohort_fractions, n_units,
+            unit_weights=self._unit_level_weights,
         )
-        agg_eif_total = agg_eif + wif  # both O(1) scale
+        # Compute SE: survey path uses score-level psi to avoid double-weighting
+        # (compute_survey_vcov applies w_i internally, which would double-weight
+        # the survey-weighted WIF term). Dispatch replicate vs TSL.
+        if self._unit_resolved_survey is not None:
+            uw = self._unit_level_weights
+            total_w = float(np.sum(uw))
+            psi_total = uw * agg_eif / total_w + wif / total_w
+
+            if (hasattr(self._unit_resolved_survey, 'uses_replicate_variance')
+                    and self._unit_resolved_survey.uses_replicate_variance):
+                from diff_diff.survey import compute_replicate_if_variance
+
+                variance, _ = compute_replicate_if_variance(
+                    psi_total, self._unit_resolved_survey
+                )
+            else:
+                from diff_diff.survey import compute_survey_if_variance
 
-        # SE = sqrt(mean(EIF^2) / n) — standard IF-based SE
-        # (dispatches to survey TSL or cluster-robust when active)
-        se = self._eif_se(agg_eif_total, n_units, cluster_indices, n_clusters)
+                variance = compute_survey_if_variance(
+                    psi_total, self._unit_resolved_survey
+                )
+            se = float(np.sqrt(max(variance, 0.0))) if np.isfinite(variance) else np.nan
+        else:
+            agg_eif_total = agg_eif + wif
+            se = self._eif_se(agg_eif_total, n_units, cluster_indices, n_clusters)
 
         return overall_att, se
 
@@ -1324,15 +1384,37 @@ class EfficientDiD(EfficientDiDBootstrapMixin):
                 agg_eif += w[k] * eif_by_gt[gt]
 
             # WIF correction for event-study aggregation
+            wif_e = np.zeros(n_units)
             if unit_cohorts is not None:
                 es_keepers = [(g, t) for (g, t) in gt_pairs]
                 es_effects = effs
-                wif = self._compute_wif_contribution(
-                    es_keepers, es_effects, unit_cohorts, cohort_fractions, n_units
+                wif_e = self._compute_wif_contribution(
+                    es_keepers, es_effects, unit_cohorts, cohort_fractions, n_units,
+                    unit_weights=self._unit_level_weights,
                 )
-                agg_eif = agg_eif + wif
 
-            agg_se = self._eif_se(agg_eif, n_units, cluster_indices, n_clusters)
+            if self._unit_resolved_survey is not None:
+                uw = self._unit_level_weights
+                total_w = float(np.sum(uw))
+                psi_total = uw * agg_eif / total_w + wif_e / total_w
+
+                if (hasattr(self._unit_resolved_survey, 'uses_replicate_variance')
+                        and self._unit_resolved_survey.uses_replicate_variance):
+                    from diff_diff.survey import compute_replicate_if_variance
+
+                    variance, _ = compute_replicate_if_variance(
+                        psi_total, self._unit_resolved_survey
+                    )
+                else:
+                    from diff_diff.survey import compute_survey_if_variance
+
+                    variance = compute_survey_if_variance(
+                        psi_total, self._unit_resolved_survey
+                    )
+                agg_se = float(np.sqrt(max(variance, 0.0))) if np.isfinite(variance) else np.nan
+            else:
+                agg_eif = agg_eif + wif_e
+                agg_se = self._eif_se(agg_eif, n_units, cluster_indices, n_clusters)
 
             t_stat, p_val, ci = safe_inference(
                 agg_eff, agg_se, alpha=self.alpha, df=self._survey_df

{diff_diff-2.8.4 → diff_diff-2.9.0}/diff_diff/efficient_did_bootstrap.py

@@ -63,6 +63,7 @@ class EfficientDiDBootstrapMixin:
         cluster_indices: Optional[np.ndarray] = None,
         n_clusters: Optional[int] = None,
         resolved_survey: object = None,
+        unit_level_weights: Optional[np.ndarray] = None,
     ) -> EDiDBootstrapResults:
         """Run multiplier bootstrap on stored EIF values.
 
@@ -136,11 +137,18 @@ class EfficientDiDBootstrapMixin:
         original_atts = np.array([group_time_effects[gt]["effect"] for gt in gt_pairs])
 
         # Perturbed ATTs: (n_bootstrap, n_gt)
+        # Under survey design, perturb survey-score object w_i * eif_i / sum(w)
+        # to match the analytical variance convention (compute_survey_if_variance).
         bootstrap_atts = np.zeros((self.n_bootstrap, n_gt))
         for j, gt in enumerate(gt_pairs):
             eif_gt = eif_by_gt[gt]  # shape (n_units,)
             with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
-                perturbation = (all_weights @ eif_gt) / n_units
+                if unit_level_weights is not None:
+                    total_w = float(np.sum(unit_level_weights))
+                    eif_scaled = unit_level_weights * eif_gt / total_w
+                    perturbation = all_weights @ eif_scaled
+                else:
+                    perturbation = (all_weights @ eif_gt) / n_units
             bootstrap_atts[:, j] = original_atts[j] + perturbation
 
         # Post-treatment mask — also exclude NaN effects

{diff_diff-2.8.4 → diff_diff-2.9.0}/diff_diff/efficient_did_covariates.py

@@ -37,6 +37,7 @@ def estimate_outcome_regression(
     group_mask: np.ndarray,
     t_col: int,
     tpre_col: int,
+    unit_weights: Optional[np.ndarray] = None,
 ) -> np.ndarray:
     """Estimate conditional mean outcome change m_hat(X) for a comparison group.
 
@@ -56,6 +57,9 @@ def estimate_outcome_regression(
         Mask selecting units in the comparison group.
     t_col, tpre_col : int
         Column indices in ``outcome_wide`` for the two time periods.
+    unit_weights : ndarray, shape (n_units,), optional
+        Survey weights at the unit level.  When provided, uses WLS
+        instead of OLS for the within-group regression.
 
     Returns
     -------
@@ -68,9 +72,13 @@ def estimate_outcome_regression(
     X_group = covariate_matrix[group_mask]
     X_design = np.column_stack([np.ones(len(X_group)), X_group])
 
+    w_group = unit_weights[group_mask] if unit_weights is not None else None
+
     coef, _, _ = solve_ols(
         X_design,
         delta_y,
+        weights=w_group,
+        weight_type="pweight" if w_group is not None else None,
         return_vcov=False,
         rank_deficient_action="warn",
     )
@@ -121,7 +129,9 @@ def _polynomial_sieve_basis(X: np.ndarray, degree: int) -> np.ndarray:
     """
     n, d = X.shape
 
-    # Standardize for numerical stability
+    # Standardize for numerical stability (unweighted mean/std intentional —
+    # this is only for conditioning, not for the statistical estimand; with
+    # survey weights the sieve basis is the same, only the objective changes)
     X_mean = X.mean(axis=0)
     X_std = X.std(axis=0)
     X_std[X_std < 1e-10] = 1.0  # avoid division by zero for constant columns
@@ -146,6 +156,7 @@ def estimate_propensity_ratio_sieve(
     k_max: Optional[int] = None,
     criterion: str = "bic",
     ratio_clip: float = 20.0,
+    unit_weights: Optional[np.ndarray] = None,
 ) -> np.ndarray:
     r"""Estimate propensity ratio via sieve convex minimization (Eq 4.1-4.2).
 
@@ -176,6 +187,9 @@ def estimate_propensity_ratio_sieve(
         ``"aic"`` or ``"bic"``.
     ratio_clip : float
         Clip ratios to ``[1/ratio_clip, ratio_clip]``.
+    unit_weights : ndarray, shape (n_units,), optional
+        Survey weights at the unit level.  When provided, uses weighted
+        normal equations for the sieve estimation.
 
     Returns
     -------
@@ -197,9 +211,20 @@ def estimate_propensity_ratio_sieve(
     k_max = max(k_max, 1)
 
     # Penalty multiplier for IC
+    # BIC penalty uses observation count (not weighted) — complexity vs distinct obs
     n_total = int(np.sum(mask_g)) + n_gp
     c_n = 2.0 if criterion == "aic" else np.log(max(n_total, 2))
 
+    # Weighted totals for loss normalization (raw probability weights)
+    if unit_weights is not None:
+        w_g = unit_weights[mask_g]
+        w_gp = unit_weights[mask_gp]
+        n_total_w = float(np.sum(w_g)) + float(np.sum(w_gp))
+    else:
+        w_g = None
+        w_gp = None
+        n_total_w = float(n_total)
+
     best_ic = np.inf
     best_ratio = np.ones(n_units)  # fallback: constant ratio 1
 
@@ -214,9 +239,15 @@ def estimate_propensity_ratio_sieve(
         Psi_gp = basis_all[mask_gp]  # (n_gp, n_basis)
         Psi_g = basis_all[mask_g]  # (n_g, n_basis)
 
-        # Normal equations: (Psi_gp' Psi_gp) beta = Psi_g.sum(axis=0)
-        A = Psi_gp.T @ Psi_gp
-        b = Psi_g.sum(axis=0)
+        # Normal equations (weighted when survey weights present):
+        # Unweighted: (Psi_gp' Psi_gp) beta = Psi_g.sum(axis=0)
+        # Weighted:   (Psi_gp' W_gp Psi_gp) beta = (w_g * Psi_g).sum(axis=0)
+        if w_gp is not None:
+            A = Psi_gp.T @ (w_gp[:, None] * Psi_gp)
+            b = (w_g[:, None] * Psi_g).sum(axis=0)
+        else:
+            A = Psi_gp.T @ Psi_gp
+            b = Psi_g.sum(axis=0)
 
         try:
             beta = np.linalg.solve(A, b)
@@ -230,11 +261,12 @@ def estimate_propensity_ratio_sieve(
         # Predicted ratio for all units
         r_hat = basis_all @ beta
 
-        # IC selection: loss at optimum = -(1/n) * b'beta
-        # Derivation: L(beta) = (1/n)(beta'A*beta - 2*b'beta).
+        # IC selection: loss at optimum = -(1/n_w) * b'beta
+        # Derivation: L(beta) = (1/n_w)(beta'A*beta - 2*b'beta).
         # At optimum A*beta = b, so beta'A*beta = b'beta.
-        # Therefore L = (1/n)(b'beta - 2*b'beta) = -(1/n)*b'beta.
-        loss = -float(b @ beta) / n_total
+        # Therefore L = (1/n_w)(b'beta - 2*b'beta) = -(1/n_w)*b'beta.
+        # Loss uses weighted totals; BIC penalty uses observation count.
+        loss = -float(b @ beta) / n_total_w
         ic_val = 2.0 * loss + c_n * n_basis / n_total
 
         if ic_val < best_ic:
@@ -280,6 +312,7 @@ def estimate_inverse_propensity_sieve(
     group_mask: np.ndarray,
     k_max: Optional[int] = None,
     criterion: str = "bic",
+    unit_weights: Optional[np.ndarray] = None,
 ) -> np.ndarray:
     r"""Estimate s_{g'}(X) = 1/p_{g'}(X) via sieve convex minimization.
 
@@ -305,6 +338,9 @@ def estimate_inverse_propensity_sieve(
         Maximum polynomial degree. None = auto.
     criterion : str
         ``"aic"`` or ``"bic"``.
+    unit_weights : ndarray, shape (n_units,), optional
+        Survey weights at the unit level.  When provided, uses weighted
+        normal equations for the sieve estimation.
 
     Returns
     -------
@@ -322,10 +358,25 @@ def estimate_inverse_propensity_sieve(
         k_max = min(int(n_group**0.2), 5)
     k_max = max(k_max, 1)
 
+    # BIC penalty uses observation count (not weighted)
     c_n = 2.0 if criterion == "aic" else np.log(max(n_units, 2))
 
+    # Weighted loss normalization and fallback
+    if unit_weights is not None:
+        w_group = unit_weights[group_mask]
+        sum_w_group = float(np.sum(w_group))
+        if sum_w_group <= 0:
+            # Zero survey weight for this group — return unconditional fallback
+            return np.ones(n_units)
+        n_units_w = float(np.sum(unit_weights))
+        fallback_ratio = n_units_w / sum_w_group
+    else:
+        w_group = None
+        n_units_w = float(n_units)
+        fallback_ratio = n_units / n_group
+
     best_ic = np.inf
-    best_s = np.full(n_units, n_units / n_group)  # fallback: unconditional
+    best_s = np.full(n_units, fallback_ratio)  # fallback: unconditional
 
     for K in range(1, k_max + 1):
         n_basis = comb(K + d, d)
@@ -335,9 +386,16 @@ def estimate_inverse_propensity_sieve(
         basis_all = _polynomial_sieve_basis(covariate_matrix, K)
         Psi_gp = basis_all[group_mask]
 
-        A = Psi_gp.T @ Psi_gp
-        # RHS: sum of basis over ALL units (not just one group)
-        b = basis_all.sum(axis=0)
+        # Normal equations (weighted when survey weights present):
+        # Unweighted: (Psi_gp' Psi_gp) beta = Psi_all.sum(axis=0)
+        # Weighted:   (Psi_gp' W_group Psi_gp) beta = (w_all * Psi_all).sum(axis=0)
+        if w_group is not None:
+            A = Psi_gp.T @ (w_group[:, None] * Psi_gp)
+            b = (unit_weights[:, None] * basis_all).sum(axis=0)
+        else:
+            A = Psi_gp.T @ Psi_gp
+            # RHS: sum of basis over ALL units (not just one group)
+            b = basis_all.sum(axis=0)
 
         try:
             beta = np.linalg.solve(A, b)
@@ -348,8 +406,9 @@ def estimate_inverse_propensity_sieve(
 
         s_hat = basis_all @ beta
 
-        # IC: loss = -(1/n) * b'beta (same derivation as ratio estimator)
-        loss = -float(b @ beta) / n_units
+        # IC: loss = -(1/n_w) * b'beta (same derivation as ratio estimator)
+        # Loss uses weighted totals; BIC penalty uses observation count.
+        loss = -float(b @ beta) / n_units_w
         ic_val = 2.0 * loss + c_n * n_basis / n_units
 
         if ic_val < best_ic:
@@ -433,6 +492,10 @@ def compute_generated_outcomes_cov(
     g_mask = cohort_masks[target_g]
     pi_g = cohort_fractions[target_g]
 
+    # Guard: zero survey weight for the target cohort → no DR estimation possible
+    if pi_g <= 0:
+        return np.zeros((n_units, H))
+
     gen_out = np.zeros((n_units, H))
 
     for j, (gp, tpre) in enumerate(valid_pairs):
@@ -496,6 +559,7 @@ def _kernel_weights_matrix(
     X_all: np.ndarray,
     X_group: np.ndarray,
     bandwidth: float,
+    group_weights: Optional[np.ndarray] = None,
 ) -> np.ndarray:
     """Gaussian kernel weight matrix.
 
@@ -503,11 +567,21 @@ def _kernel_weights_matrix(
     normalized kernel weight ``K_h(X_group[j], X_all[i])``.
 
     Each row sums to 1 (Nadaraya-Watson normalization).
+
+    Parameters
+    ----------
+    group_weights : ndarray, shape (n_group,), optional
+        Survey weights for the group units.  When provided, kernel
+        weights are multiplied by survey weights before row-normalization,
+        making the Nadaraya-Watson estimator survey-weighted.
     """
     # Squared distances: (n_all, n_group)
     dist_sq = cdist(X_all, X_group, metric="sqeuclidean")
     # Gaussian kernel
     raw = np.exp(-dist_sq / (2.0 * bandwidth**2))
+    # Survey-weight: each group unit j contributes ∝ w_j * K_h(X_i, X_j)
+    if group_weights is not None:
+        raw = raw * group_weights[np.newaxis, :]
     # Normalize each row
     row_sums = raw.sum(axis=1, keepdims=True)
     row_sums[row_sums < 1e-15] = 1.0  # avoid division by zero
@@ -559,6 +633,7 @@ def compute_omega_star_conditional(
     covariate_matrix: np.ndarray,
     s_hat_cache: Dict[float, np.ndarray],
     bandwidth: Optional[float] = None,
+    unit_weights: Optional[np.ndarray] = None,
     never_treated_val: float = np.inf,
 ) -> np.ndarray:
     r"""Kernel-smoothed conditional Omega\*(X_i) for each unit (Eq 3.12).
@@ -583,6 +658,9 @@ def compute_omega_star_conditional(
         value is shape ``(n_units,)``. Keyed by group identifier.
     bandwidth : float or None
         Kernel bandwidth. None = Silverman's rule.
+    unit_weights : ndarray, shape (n_units,), optional
+        Survey weights at the unit level.  When provided, kernel-smoothed
+        covariances use survey-weighted Nadaraya-Watson regression.
     never_treated_val : float
 
     Returns
@@ -622,13 +700,17 @@ def compute_omega_star_conditional(
             stacklevel=2,
         )
 
+    # Per-group survey weights for kernel smoothing
+    w_g = unit_weights[g_mask] if unit_weights is not None else None
+    w_inf = unit_weights[never_treated_mask] if unit_weights is not None else None
+
     # Pre-compute kernel weight matrices per group
     Y_g = outcome_wide[g_mask]
     X_g = covariate_matrix[g_mask]
     Yg_t_minus_1 = Y_g[:, t_col] - Y_g[:, y1_col]
 
-    W_g = _kernel_weights_matrix(covariate_matrix, X_g, bandwidth)
-    W_inf = _kernel_weights_matrix(covariate_matrix, X_inf, bandwidth)
+    W_g = _kernel_weights_matrix(covariate_matrix, X_g, bandwidth, group_weights=w_g)
+    W_inf = _kernel_weights_matrix(covariate_matrix, X_inf, bandwidth, group_weights=w_inf)
 
     inf_t_minus_tpre = {}
     for _, tpre in valid_pairs:
@@ -683,7 +765,10 @@ def compute_omega_star_conditional(
                     )
                     if gp_j not in W_gp_cache:
                         X_gp = covariate_matrix[cohort_masks[gp_j]]
-                        W_gp_cache[gp_j] = _kernel_weights_matrix(covariate_matrix, X_gp, bandwidth)
+                        w_gp_j = unit_weights[cohort_masks[gp_j]] if unit_weights is not None else None
+                        W_gp_cache[gp_j] = _kernel_weights_matrix(
+                            covariate_matrix, X_gp, bandwidth, group_weights=w_gp_j
+                        )
                         gp_outcomes_cache[gp_j] = outcome_wide[cohort_masks[gp_j]]
                     W_gp = W_gp_cache[gp_j]
                     Y_gp = gp_outcomes_cache[gp_j]