Fix weighted-mass IF linearization in ContinuousDiD and TripleDifference from PR #226 review (round 3)

- ContinuousDiD: store survey-weighted treated/control masses and
  weighted dpsi_bar in bootstrap_info; use weighted masses for p_1, p_0,
  n_total in IF construction so TSL linearizes the weighted estimator
- TripleDifference: use survey-weighted subgroup mass (sum(w_sub))
  instead of raw counts (n_sub) for pairwise IF combination weights
  w3, w2, w1 when survey design is active

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

Author: igerber

diff_diff/continuous_did.py

@@ -608,6 +608,10 @@ def fit(
                             control_idx = b_info["control_indices"]
                             n_t = b_info["n_treated"]
                             n_c = b_info["n_control"]
+                            # Use survey-weighted masses when available
+                            if "w_treated" in b_info:
+                                n_t = b_info["w_treated"]
+                                n_c = b_info["w_control"]
                             n_total_gt = n_t + n_c
                             p_1 = n_t / n_total_gt
                             p_0 = n_c / n_total_gt
@@ -972,11 +976,18 @@ def _compute_dose_response_gt(
         treated_indices = np.where(treated_mask)[0]
         control_indices = np.where(control_mask)[0]
 
+        # dpsi_bar: mean derivative basis vector (weighted when survey)
+        if w_treated is not None:
+            dpsi_bar = np.average(dPsi_treated, axis=0, weights=w_treated)
+        else:
+            dpsi_bar = np.mean(dPsi_treated, axis=0)
+
         bootstrap_info = {
             "bread": bread,
             "ee_treated": ee_treated,
             "ee_control": ee_control,
             "psi_bar": psi_bar,
+            "dpsi_bar": dpsi_bar,
             "beta_hat": beta_hat,
             "beta_pred": beta_pred,
             "treated_indices": treated_indices,
@@ -993,6 +1004,11 @@ def _compute_dose_response_gt(
             "acrt_glob": acrt_glob,
         }
 
+        # Store survey-weighted masses for IF linearization
+        if w_treated is not None:
+            bootstrap_info["w_treated"] = float(np.sum(w_treated))
+            bootstrap_info["w_control"] = float(np.sum(w_control))
+
         return {
             "att_d": att_d,
             "acrt_d": acrt_d,
@@ -1080,13 +1096,17 @@ def _compute_analytical_se(
             control_idx = info["control_indices"]
             n_t = info["n_treated"]
             n_c = info["n_control"]
+            # Use survey-weighted masses when available
+            if "w_treated" in info:
+                n_t = info["w_treated"]
+                n_c = info["w_control"]
             bread = info["bread"]
             ee_treated = info["ee_treated"]
             ee_control = info["ee_control"]
             psi_bar = info["psi_bar"]
+            dpsi_bar = info["dpsi_bar"]
             Psi_eval = info["Psi_eval"]
             dPsi_eval = info["dPsi_eval"]
-            dPsi_treated = info["dPsi_treated"]
             att_glob_gt = info["att_glob"]
             mu_0 = info["mu_0"]
             delta_y_treated = info["delta_y_treated"]
@@ -1119,7 +1139,6 @@ def _compute_analytical_se(
                 if_acrt_d[idx] += w * (dPsi_eval @ beta_pert)
 
             # ACRT_glob IF: (1/n_t) sum_j dpsi(D_j)' @ beta_pert
-            dpsi_bar = np.mean(dPsi_treated, axis=0)
             for k, idx in enumerate(treated_idx):
                 beta_pert = bread @ ee_treated[k] / n_t
                 if_acrt_glob[idx] += w * float(dpsi_bar @ beta_pert)

diff_diff/triple_diff.py

@@ -1025,12 +1025,19 @@ def _estimate_ddd_decomposition(
         )
 
         # Influence function weights (matching R's att_dr_rc)
-        n3 = np.sum((subgroup == 3) | (subgroup == 4))
-        n2 = np.sum((subgroup == 2) | (subgroup == 4))
-        n1 = np.sum((subgroup == 1) | (subgroup == 4))
-        w3 = n / n3
-        w2 = n / n2
-        w1 = n / n1
+        if survey_weights is not None:
+            n3 = np.sum(survey_weights[(subgroup == 3) | (subgroup == 4)])
+            n2 = np.sum(survey_weights[(subgroup == 2) | (subgroup == 4)])
+            n1 = np.sum(survey_weights[(subgroup == 1) | (subgroup == 4)])
+            n_total = np.sum(survey_weights)
+        else:
+            n3 = np.sum((subgroup == 3) | (subgroup == 4))
+            n2 = np.sum((subgroup == 2) | (subgroup == 4))
+            n1 = np.sum((subgroup == 1) | (subgroup == 4))
+            n_total = n
+        w3 = n_total / n3
+        w2 = n_total / n2
+        w1 = n_total / n1
 
         inf_func = (
             w3 * did_results[3]["inf"] + w2 * did_results[2]["inf"] - w1 * did_results[1]["inf"]