minimal fix to resolve #707

causalml/metrics/visualize.py

@@ -84,6 +84,10 @@ def get_cumlift(
         or treatment_effect_col in df.columns
     )
 
+    assert not (
+        (df[[outcome_col, treatment_col, treatment_effect_col]].isnull().values.any())
+    )
+
     df = df.copy()
     np.random.seed(random_seed)
     random_cols = []
@@ -219,6 +223,10 @@ def get_qini(
         or treatment_effect_col in df.columns
     )
 
+    assert not (
+        (df[[outcome_col, treatment_col, treatment_effect_col]].isnull().values.any())
+    )
+
     df = df.copy()
     np.random.seed(random_seed)
     random_cols = []
@@ -315,6 +323,8 @@ def get_tmlegain(
         or p_col in df.columns
     )
 
+    assert not ((df[[outcome_col, treatment_col, p_col]].isnull().values.any()))
+
     inference_col = [x for x in inference_col if x in df.columns]
 
     # Initialize TMLE
@@ -421,6 +431,8 @@ def get_tmleqini(
         or p_col in df.columns
     )
 
+    assert not ((df[[outcome_col, treatment_col, p_col]].isnull().values.any()))
+
     inference_col = [x for x in inference_col if x in df.columns]
 
     # Initialize TMLE

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)