Fix CI review R9: align controls aggregation to effective sample

R8's controls-block fix scoped NaN/Inf validation to the positive-weight
subset (shorter data_controls) but then assigned those shorter arrays
into an x_agg_input built from the full-length frame, causing a
length-mismatch on any SurveyDesign.subpopulation() / zero-weight
excluded row before covariate aggregation could run.

Root-caused fix: derive both the validation frame AND the aggregation
frame from the same positive-weight effective sample (data_eff,
survey_weights_eff). Zero-weight rows are genuinely out-of-sample
throughout the DID^X path now. Non-survey fits unchanged.

Added TestZeroWeightSubpopulation.test_zero_weight_row_with_nan_control
pinning the subpopulation contract for the DID^X path — injects a
zero-weight row with NaN control value and asserts fit() succeeds.

All 263 targeted tests pass.

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

Author: igerber

diff_diff/chaisemartin_dhaultfoeuille.py

@@ -730,16 +730,18 @@ def fit(
                     f"data. Available columns: {list(data.columns)}"
                 )
             # SurveyDesign.subpopulation() contract: zero-weight rows are
-            # out-of-sample. Scope NaN/Inf validation to positive-weight
-            # rows so that excluded obs with missing covariates do not
-            # abort the fit. The downstream weighted aggregation
-            # (sum(w*x)/sum(w)) handles zero-weight rows correctly on
-            # its own.
+            # out-of-sample. Scope BOTH validation and aggregation to the
+            # positive-weight subset so excluded rows with missing/invalid
+            # covariates do not abort the fit and cell aggregation aligns
+            # with the effective sample used by _validate_and_aggregate_to_cells.
             if survey_weights is not None:
                 pos_mask_ctrl = np.asarray(survey_weights) > 0
-                data_controls = data.loc[pos_mask_ctrl, controls].copy()
+                data_eff = data.loc[pos_mask_ctrl]
+                survey_weights_eff = np.asarray(survey_weights)[pos_mask_ctrl]
             else:
-                data_controls = data[controls].copy()
+                data_eff = data
+                survey_weights_eff = None
+            data_controls = data_eff[controls].copy()
             for c in controls:
                 try:
                     data_controls[c] = pd.to_numeric(data_controls[c])
@@ -760,14 +762,15 @@ def fit(
                         "Remove or replace non-finite covariates before fitting."
                     )
             # Aggregate covariates to cell means (same groupby as treatment/outcome).
-            # Use the coerced copy joined with group/time from original data.
-            x_agg_input = data[[group, time]].copy()
+            # Build x_agg_input from the same effective-sample frame so rows
+            # align with data_controls.
+            x_agg_input = data_eff[[group, time]].copy()
             x_agg_input[controls] = data_controls[controls].values
-            if survey_weights is not None:
+            if survey_weights_eff is not None:
                 # Survey-weighted covariate cell means: sum(w*x)/sum(w)
-                x_agg_input["_w_"] = survey_weights
+                x_agg_input["_w_"] = survey_weights_eff
                 for c in controls:
-                    x_agg_input[f"_wx_{c}"] = survey_weights * x_agg_input[c].values
+                    x_agg_input[f"_wx_{c}"] = survey_weights_eff * x_agg_input[c].values
                 wx_cols = [f"_wx_{c}" for c in controls]
                 g_agg = x_agg_input.groupby([group, time], as_index=False).agg(
                     {**{wc: "sum" for wc in wx_cols}, "_w_": "sum"}

tests/test_survey_dcdh.py

@@ -861,6 +861,28 @@ def test_zero_weight_row_with_nan_outcome(self, base_data):
         )
         assert np.isfinite(result.overall_att)
 
+    def test_zero_weight_row_with_nan_control(self, base_data):
+        """A zero-weight row with NaN in a control column must not abort
+        the DID^X path, and the covariate cell aggregation must use only
+        positive-weight rows (no length-mismatch error)."""
+        rng = np.random.default_rng(13)
+        df_ = base_data.copy()
+        df_["pw"] = 1.0
+        df_["x"] = rng.normal(0, 1, size=len(df_))
+        # Inject a zero-weight row with NaN control value
+        sample = df_.iloc[0].copy()
+        sample["x"] = np.nan
+        sample["pw"] = 0.0
+        df_ = pd.concat([df_, pd.DataFrame([sample])], ignore_index=True)
+        sd = SurveyDesign(weights="pw")
+        result = ChaisemartinDHaultfoeuille(seed=1).fit(
+            df_,
+            outcome="outcome", group="group",
+            time="period", treatment="treatment",
+            L_max=1, controls=["x"], survey_design=sd,
+        )
+        assert np.isfinite(result.overall_att)
+
     def test_zero_weight_row_with_nan_heterogeneity(self, base_data):
         """A zero-weight row with NaN in the heterogeneity column must
         not trip the heterogeneity time-invariance validator."""