fix missing value bug

uber · Nov 14, 2023 · 9ce214c · 9ce214c
1 parent c154afe
commit 9ce214c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/causalml/metrics/visualize.py b/causalml/metrics/visualize.py
@@ -114,10 +114,10 @@ def get_cumlift(
             sorted_df["cumsum_ct"] = sorted_df.index.values - sorted_df["cumsum_tr"]
             sorted_df["cumsum_y_tr"] = (
                 sorted_df[outcome_col] * sorted_df[treatment_col]
-            ).cumsum()
+            ).fillna(0).cumsum(skipna=True).astype(float)
             sorted_df["cumsum_y_ct"] = (
                 sorted_df[outcome_col] * (1 - sorted_df[treatment_col])
-            ).cumsum()
+            ).fillna(0).cumsum(skipna=True).astype(float)
 
             lift.append(
                 sorted_df["cumsum_y_tr"] / sorted_df["cumsum_tr"]