fixing forgotten mask

owkin · Aug 16, 2024 · c5cb2c7 · c5cb2c7
1 parent b5e2f44
commit c5cb2c7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/fedeca/strategies/fed_smd.py b/fedeca/strategies/fed_smd.py
@@ -193,7 +193,7 @@ def compute_local_moments_per_group(
         raw_data = pd.DataFrame(Xprop, columns=propensity_cols)
 
         weights_df = pd.DataFrame(
-            np.repeat(weights[:, None], Xprop.shape[1]), columns=propensity_cols
+            np.repeat(weights, Xprop.shape[1], axis=1), columns=propensity_cols
         )
 
         results = {}
@@ -204,7 +204,7 @@ def compute_local_moments_per_group(
             # Here we pass weights
             results[res_name]["weighted"] = {
                 f"moment{k}": compute_uncentered_moment(
-                    raw_data[mask_treatment], k, weights
+                    raw_data[mask_treatment], k, weights[mask_treatment]
                 )
                 for k in range(1, 3)
             }