Address AI review round 3: ICC covariate variance, DGP REGISTRY notes

- Include covariate variance (0.2725) in ICC calibration when add_covariates=True
- Document informative_sampling structural Y(0) ranking in docstring and REGISTRY.md
- Add cross-section default-weights test and ICC + covariates regression test

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

Author: igerber

diff_diff/prep_dgp.py

@@ -1271,6 +1271,8 @@ def generate_survey_did_data(
         repeated cross-sections, ranking is refreshed each period. Within
         each stratum, rank-based weights are scaled to preserve the
         stratum's baseline weight level from ``weight_variation``.
+        Ranking is based on structural Y(0) (unit FE + PSU effects),
+        excluding covariates from ``add_covariates``.
     heterogeneous_te_by_strata : bool, default=False
         If True, treatment effect varies by stratum:
         ``TE_h = TE * (1 + 0.5 * (h - mean) / std)``. Creates a gap
@@ -1399,8 +1401,11 @@ def generate_survey_did_data(
 
     # --- ICC -> psu_re_sd resolution ---
     if icc is not None:
+        # Include covariate variance: Var(0.5*x1) + Var(0.3*x2)
+        # where x1 ~ N(0,1), x2 ~ Bernoulli(0.5)
+        cov_var = (0.25 + 0.09 * 0.25) if add_covariates else 0.0
         psu_re_sd = np.sqrt(
-            icc * (unit_fe_sd**2 + noise_sd**2)
+            icc * (unit_fe_sd**2 + noise_sd**2 + cov_var)
             / ((1 - icc) * (1 + psu_period_factor**2))
         )
 

docs/methodology/REGISTRY.md

@@ -2515,6 +2515,17 @@ The 8-step workflow in `docs/llms-practitioner.txt` is adapted from Baker et al.
   covariates). Paper's Step 8 is "Keep learning." The mandatory with/without covariate
   comparison is a diff-diff convention.
 
+### Survey DGP (`generate_survey_did_data`)
+
+- **Note:** The `icc` parameter calibrates `psu_re_sd` using the variance
+  decomposition `Var(Y) = sigma²_psu * (1 + psu_period_factor²) + sigma²_unit +
+  sigma²_noise + sigma²_cov`. When `add_covariates=True`, the covariate variance
+  `Var(0.5*x1) + Var(0.3*x2) = 0.2725` is included in the calibration.
+- **Note:** Defensive enhancement: `informative_sampling` ranks units on structural
+  Y(0) (unit FE + PSU effects + time trend), excluding covariate contributions from
+  `add_covariates`. This models selection on structural characteristics (geography,
+  demographics) rather than residual variation, matching real survey sampling frames.
+
 ---
 
 # Version History

tests/test_prep.py

@@ -1639,6 +1639,45 @@ def test_informative_sampling_cross_section(self):
         corr = np.corrcoef(p1["weight"], p1["outcome"])[0, 1]
         assert corr > 0.1
 
+    def test_informative_sampling_cross_section_default_weights(self):
+        """Cross-section informative sampling with default weight_variation."""
+        from diff_diff.prep_dgp import generate_survey_did_data
+
+        df = generate_survey_did_data(
+            n_units=1000,
+            informative_sampling=True,
+            panel=False,
+            seed=42,
+        )
+        p1 = df[df["period"] == 1]
+        for s in range(5):
+            expected_mean = 1.0 + 1.0 * (s / 4)
+            stratum_weights = p1.loc[p1["stratum"] == s, "weight"]
+            assert abs(stratum_weights.mean() - expected_mean) < 0.15
+            assert stratum_weights.std() > 0.01
+
+    def test_icc_with_covariates(self):
+        """ICC calibration should account for covariate variance."""
+        from diff_diff.prep_dgp import generate_survey_did_data
+
+        target_icc = 0.3
+        df = generate_survey_did_data(
+            n_units=1000, icc=target_icc, add_covariates=True, seed=42
+        )
+        # ANOVA-based ICC on period 1
+        p1 = df[df["period"] == 1]
+        groups = p1.groupby("psu")["outcome"]
+        grand_mean = p1["outcome"].mean()
+        n_total = len(p1)
+        n_groups = groups.ngroups
+        n_bar = n_total / n_groups
+        ssb = (groups.size() * (groups.mean() - grand_mean) ** 2).sum()
+        msb = ssb / (n_groups - 1)
+        ssw = groups.apply(lambda x: ((x - x.mean()) ** 2).sum()).sum()
+        msw = ssw / (n_total - n_groups)
+        realized_icc = (msb - msw) / (msb + (n_bar - 1) * msw)
+        assert abs(realized_icc - target_icc) / target_icc < 0.50
+
     def test_heterogeneous_te_by_strata(self):
         """Unweighted mean TE should differ from population ATT."""
         from diff_diff.prep_dgp import generate_survey_did_data