feature_selection: validate binary outcome in filter_LR / filter_D (#349)

causalml/feature_selection/filters.py

@@ -1,7 +1,24 @@
 """
-Filter feature selection methods for uplift modeling
-
-- Currently only for classification problem: the outcome variable of uplift model is binary.
+Filter feature selection methods for uplift modeling.
+
+These filters implement the bin-based and likelihood-based feature ranking
+described in Zhao et al. 2020 (https://arxiv.org/abs/2005.03447).
+
+.. note::
+    The current implementation only supports **binary** outcomes (values
+    coded as ``0`` and ``1``):
+
+    * ``filter_LR`` uses ``statsmodels`` ``Logit``, which assumes a binary
+      response.
+    * ``filter_D`` (KL / Chi / ED divergences) computes class probabilities
+      via ``_GetNodeSummary``, which only inspects ``y == 0`` and
+      ``y == 1`` counts and silently ignores other label values.
+
+    Passing a non-binary outcome used to fall through silently or behave
+    inconsistently per node. As of uber/causalml#349, ``filter_LR`` and
+    ``filter_D`` raise a clear ``ValueError`` if the outcome contains
+    values other than ``{0, 1}``. The ``filter_F`` (OLS) method tolerates
+    continuous outcomes and is unaffected.
 """
 
 import numpy as np
@@ -11,6 +28,41 @@
 from sklearn.impute import SimpleImputer
 
 
+def _check_binary_outcome(y, y_name="y"):
+    """Validate that an outcome vector contains only ``0`` and ``1``.
+
+    Used by ``filter_LR`` and ``filter_D`` to fail loudly when the user
+    passes a non-binary outcome — those filters silently mis-handle other
+    label sets (Logit assumes binary; ``_GetNodeSummary`` only counts
+    ``y == 0`` and ``y == 1``). See uber/causalml#349.
+
+    Args:
+        y (array-like): outcome values.
+        y_name (str): column name used in the error message.
+
+    Raises:
+        ValueError: if ``y`` contains any value other than ``0`` or ``1``.
+    """
+    arr = np.asarray(y)
+    # Drop NaNs from the validation set — they are handled separately
+    # (e.g. ``null_impute`` in ``_filter_D_one_feature``).
+    if arr.dtype.kind == "f":
+        arr = arr[~np.isnan(arr)]
+    unique = np.unique(arr)
+    # Allow either {0, 1}, {0}, or {1} — empty set is rejected too.
+    extra = set(np.atleast_1d(unique).tolist()) - {0, 1, 0.0, 1.0}
+    if extra or len(unique) == 0:
+        raise ValueError(
+            "Filter feature selection only supports binary outcomes "
+            "(values 0/1); column '{}' contains {}. See "
+            "uber/causalml#349 for the current limitation. Use "
+            "``filter_F`` (OLS) for continuous outcomes, or pre-process "
+            "your label (e.g. via ``pd.qcut``) into a binary indicator.".format(
+                y_name, sorted(unique.tolist())[:10]
+            )
+        )
+
+
 class FilterSelect:
     """A class for feature importance methods."""
 
@@ -242,6 +294,10 @@ def filter_LR(
         if order not in [1, 2, 3]:
             raise Exception("ValueError: order argument only takes value 1,2,3.")
 
+        # filter_LR uses statsmodels Logit which silently mis-handles
+        # non-binary outcomes; validate up-front per uber/causalml#349.
+        _check_binary_outcome(data[y_name], y_name=y_name)
+
         all_result = pd.DataFrame()
         for x_name_i in features:
             one_result = self._filter_LR_one_feature(
@@ -550,6 +606,12 @@ def filter_D(
                 a data frame containing the feature importance statistics
         """
 
+        # The bin-based divergence filters (KL/ED/Chi) compute per-bin
+        # class probabilities via ``_GetNodeSummary``, which only counts
+        # ``y == 0`` and ``y == 1``. A non-binary outcome would silently
+        # produce nonsense scores. Validate up-front per uber/causalml#349.
+        _check_binary_outcome(data[y_name], y_name=y_name)
+
         all_result = pd.DataFrame()
 
         for x_name_i in features:

tests/test_feature_selection.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 from causalml.feature_selection.filters import FilterSelect
 
 from .const import RANDOM_SEED, CONVERSION
@@ -44,6 +45,52 @@ def test_filter_lr(generate_classification_data):
     assert imp["score"].values[0] >= imp["score"].values[1]
 
 
+@pytest.mark.parametrize("method", ["LR", "KL", "ED", "Chi"])
+def test_filter_rejects_non_binary_outcome(generate_classification_data, method):
+    """Regression test for uber/causalml#349.
+
+    Before the fix, passing a non-binary outcome to ``filter_LR`` /
+    ``filter_D`` (KL/ED/Chi) would silently mis-handle it: ``Logit`` would
+    raise a ``PerfectSeparationError`` deep inside statsmodels, and the
+    bin-based filters' ``_GetNodeSummary`` would only count ``y == 0``
+    and ``y == 1`` rows — producing meaningless scores when the label set
+    is e.g. ``{0, 1, 2, 3}``.
+
+    After the fix, the public entry points raise a clear ``ValueError``
+    that names the offending column and points the user at the limitation.
+    """
+    np.random.seed(RANDOM_SEED)
+    df, X_names = generate_classification_data()
+    y_name = CONVERSION
+
+    # Replace the binary outcome with a multi-class label set that includes
+    # 0 and 1 (so a naive value-counts check could miss the gap).
+    df = df.copy()
+    df[y_name] = np.random.randint(0, 4, size=df.shape[0])
+
+    filter_obj = FilterSelect()
+    with pytest.raises(ValueError, match="binary"):
+        filter_obj.get_importance(
+            df, X_names, y_name, method, treatment_group="treatment1"
+        )
+
+
+def test_filter_f_accepts_continuous_outcome(generate_classification_data):
+    """``filter_F`` uses OLS and is documented to tolerate continuous y;
+    the binary-outcome guard introduced for #349 must not affect it."""
+    np.random.seed(RANDOM_SEED)
+    df, X_names = generate_classification_data()
+    y_name = CONVERSION
+    df = df.copy()
+    df[y_name] = np.random.rand(df.shape[0])  # continuous outcome
+
+    filter_obj = FilterSelect()
+    imp = filter_obj.get_importance(
+        df, X_names, y_name, "F", treatment_group="treatment1"
+    )
+    assert imp.shape[0] == len(X_names)
+
+
 def test_filter_kl(generate_classification_data):
     # generate uplift classification data
     np.random.seed(RANDOM_SEED)