minimal fix to resolve #707 (#720)

* minimal fix to resolve #707
* lint
* remove .bool()
* operate on series not df

---------

Co-authored-by: ras44 <rolandrmgservices@gmail.com>

causalml/metrics/visualize.py

@@ -77,11 +77,13 @@ def get_cumlift(
     Returns:
         (pandas.DataFrame): average uplifts of model estimates in cumulative population
     """
-
     assert (
-        (outcome_col in df.columns)
-        and (treatment_col in df.columns)
-        or treatment_effect_col in df.columns
+        (outcome_col in df.columns and df[outcome_col].notnull().all())
+        and (treatment_col in df.columns and df[treatment_col].notnull().all())
+        or (
+            treatment_effect_col in df.columns
+            and df[treatment_effect_col].notnull().all()
+        )
     )
 
     df = df.copy()
@@ -214,9 +216,12 @@ def get_qini(
         (pandas.DataFrame): cumulative gains of model estimates in population
     """
     assert (
-        (outcome_col in df.columns)
-        and (treatment_col in df.columns)
-        or treatment_effect_col in df.columns
+        (outcome_col in df.columns and df[outcome_col].notnull().all())
+        and (treatment_col in df.columns and df[treatment_col].notnull().all())
+        or (
+            treatment_effect_col in df.columns
+            and df[treatment_effect_col].notnull().all()
+        )
     )
 
     df = df.copy()
@@ -310,9 +315,9 @@ def get_tmlegain(
         (pandas.DataFrame): cumulative gains of model estimates based of TMLE
     """
     assert (
-        (outcome_col in df.columns)
-        and (treatment_col in df.columns)
-        or p_col in df.columns
+        (outcome_col in df.columns and df[outcome_col].notnull().all())
+        and (treatment_col in df.columns and df[treatment_col].notnull().all())
+        or (p_col in df.columns and df[p_col].notnull().all())
     )
 
     inference_col = [x for x in inference_col if x in df.columns]
@@ -416,9 +421,9 @@ def get_tmleqini(
         (pandas.DataFrame): cumulative gains of model estimates based of TMLE
     """
     assert (
-        (outcome_col in df.columns)
-        and (treatment_col in df.columns)
-        or p_col in df.columns
+        (outcome_col in df.columns and df[outcome_col].notnull().all())
+        and (treatment_col in df.columns and df[treatment_col].notnull().all())
+        or (p_col in df.columns and df[p_col].notnull().all())
     )
 
     inference_col = [x for x in inference_col if x in df.columns]

tests/test_visualize.py

@@ -0,0 +1,14 @@
+import pandas as pd
+import numpy as np
+import pytest
+from causalml.metrics.visualize import get_cumlift
+
+
+def test_visualize_get_cumlift_errors_on_nan():
+    df = pd.DataFrame(
+        [[0, np.nan, 0.5], [1, np.nan, 0.1], [1, 1, 0.4], [0, 1, 0.3], [1, 1, 0.2]],
+        columns=["w", "y", "pred"],
+    )
+
+    with pytest.raises(Exception):
+        get_cumlift(df)