Fix OLS survey edge cases and harden solve_poisson (round 3)

- P1: Compute survey TSL vcov on kept columns only when solve_ols drops
  rank-deficient columns; expand back with NaN. Prevents singular bread
  matrix on all-eventually-treated ETWFE designs.
- P1: Guard against zero-weight unit/time groups before within_transform;
  raise targeted ValueError instead of letting NaN propagate.
- P2: Add weight validation (shape, NaN, Inf, non-negative, positive sum)
  to solve_poisson(weights=...) matching solve_logit pattern.
- P2: Add regression tests for rank-deficient survey OLS and zero-weight
  unit rejection.
- P3: Add pweight-only note to REGISTRY.md and survey-roadmap.md.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: igerber

diff_diff/linalg.py

@@ -2401,6 +2401,20 @@ def solve_poisson(
     """
     n, k_orig = X.shape
 
+    # Validate weights (mirrors solve_logit validation)
+    if weights is not None:
+        weights = np.asarray(weights, dtype=np.float64)
+        if weights.shape != (n,):
+            raise ValueError(f"weights must have shape ({n},), got {weights.shape}")
+        if np.any(np.isnan(weights)):
+            raise ValueError("weights contain NaN values")
+        if np.any(~np.isfinite(weights)):
+            raise ValueError("weights contain Inf values")
+        if np.any(weights < 0):
+            raise ValueError("weights must be non-negative")
+        if np.sum(weights) <= 0:
+            raise ValueError("weights sum to zero — no observations have positive weight")
+
     # Validate rank_deficient_action (same as solve_logit/solve_ols)
     valid_actions = ("warn", "error", "silent")
     if rank_deficient_action not in valid_actions:

diff_diff/wooldridge.py

@@ -647,6 +647,23 @@ def _fit_ols(
         # weighted FWL projection — all columns (treatment interactions +
         # covariates) are demeaned together.
         wt_weights = survey_weights if survey_weights is not None else np.ones(len(tmp))
+
+        # Guard: zero-weight unit/time groups cause 0/0 in within_transform
+        if survey_weights is not None and np.any(survey_weights == 0):
+            for grp_col, grp_label in [(unit, "unit"), (time, "time period")]:
+                grp_sums = sample.groupby(grp_col).apply(
+                    lambda g: survey_weights[g.index].sum(),
+                    include_groups=False,
+                )
+                zero_grps = grp_sums[grp_sums == 0].index.tolist()
+                if zero_grps:
+                    raise ValueError(
+                        f"Survey weights sum to zero for {grp_label}(s) "
+                        f"{zero_grps[:3]}. Cannot compute weighted "
+                        f"within-transformation. Remove zero-weight "
+                        f"{grp_label}s or use non-zero weights."
+                    )
+
         transformed = within_transform(
             tmp, all_vars, unit=unit, time=time, suffix="_demeaned",
             weights=wt_weights,
@@ -671,7 +688,15 @@ def _fit_ols(
         # Survey TSL vcov replaces cluster-robust vcov
         if resolved is not None:
             from diff_diff.survey import compute_survey_vcov
-            vcov = compute_survey_vcov(X, resids, resolved)
+            nan_mask_ols = np.isnan(coefs)
+            if np.any(nan_mask_ols):
+                kept = ~nan_mask_ols
+                vcov_kept = compute_survey_vcov(X[:, kept], resids, resolved)
+                vcov = np.full((len(coefs), len(coefs)), np.nan)
+                kept_idx = np.where(kept)[0]
+                vcov[np.ix_(kept_idx, kept_idx)] = vcov_kept
+            else:
+                vcov = compute_survey_vcov(X, resids, resolved)
 
         # 7. Extract β_{g,t} and build gt_effects dict
         gt_effects: Dict[Tuple, Dict] = {}

docs/methodology/REGISTRY.md

@@ -1179,6 +1179,7 @@ where `g(·)` is the link inverse (logistic or exp), `η_i` is the individual li
 **Survey design notes:**
 - **OLS path:** Survey-weighted within-transformation + WLS via `solve_ols(weights=...)` + TSL vcov via `compute_survey_vcov()`.
 - **Logit/Poisson paths:** Survey-weighted IRLS via `solve_logit(weights=...)`/`solve_poisson(weights=...)` + X_tilde linearization trick for TSL vcov: `X_tilde = X * sqrt(V)`, `r_tilde = (y - mu) / sqrt(V)`, then `compute_survey_vcov(X_tilde, r_tilde, resolved)` gives correct QMLE sandwich. ASF means and gradients use survey-weighted averaging.
+- **Note:** Only `pweight` (probability weights) are supported; `fweight`/`aweight` raise `ValueError` because the composed survey/QMLE weighting changes their semantics.
 - **Note:** Replicate-weight variance is not yet supported (`NotImplementedError`). Use TSL (strata/PSU/FPC) instead.
 - **Note:** Bootstrap inference (`n_bootstrap > 0`) cannot be combined with `survey_design` — no survey-aware bootstrap variant is implemented.
 

docs/survey-roadmap.md

@@ -209,10 +209,11 @@ co-sign.
 ### 10f. WooldridgeDiD Survey Support — SHIPPED
 
 WooldridgeDiD (ETWFE) now supports `survey_design` for all three methods
-(OLS, logit, Poisson). OLS uses survey-weighted within-transformation +
-WLS + TSL vcov. Logit/Poisson use survey-weighted IRLS + X_tilde
-linearization for TSL vcov. Replicate-weight designs raise
-`NotImplementedError`; bootstrap + survey is rejected.
+(OLS, logit, Poisson) with `pweight` only (`fweight`/`aweight` rejected).
+OLS uses survey-weighted within-transformation + WLS + TSL vcov.
+Logit/Poisson use survey-weighted IRLS + X_tilde linearization for TSL
+vcov. Replicate-weight designs raise `NotImplementedError`; bootstrap +
+survey is rejected.
 
 ### 10g. Practitioner Guidance (LOW priority)
 

tests/test_wooldridge.py

@@ -1521,3 +1521,29 @@ def test_poisson_zero_weight_cell(self, survey_panel):
         )
         assert np.isfinite(r.overall_att)
         assert np.isfinite(r.overall_se)
+
+    def test_ols_survey_rank_deficient(self, survey_panel):
+        """Survey OLS handles rank-deficient all-eventually-treated designs."""
+        from diff_diff.survey import SurveyDesign
+        # Remove never-treated (cohort=0) to create rank-deficient design
+        df = survey_panel[survey_panel["cohort"] > 0].copy()
+        sd = SurveyDesign(weights="weight", strata="stratum", psu="unit")
+        r = WooldridgeDiD(control_group="not_yet_treated").fit(
+            df, outcome="y", unit="unit", time="time",
+            cohort="cohort", survey_design=sd,
+        )
+        assert np.isfinite(r.overall_att)
+        assert np.isfinite(r.overall_se)
+
+    def test_ols_survey_zero_weight_unit_rejected(self, survey_panel):
+        """Zero-weight unit raises ValueError before within_transform."""
+        from diff_diff.survey import SurveyDesign
+        df = survey_panel.copy()
+        # Zero out all weights for unit 0
+        df.loc[df["unit"] == 0, "weight"] = 0.0
+        sd = SurveyDesign(weights="weight", strata="stratum", psu="unit")
+        with pytest.raises(ValueError, match="Survey weights sum to zero for unit"):
+            WooldridgeDiD().fit(
+                df, outcome="y", unit="unit", time="time",
+                cohort="cohort", survey_design=sd,
+            )