Fix uplift tree p-value NaN from division by zero (#585) (#882)

* Fix uplift tree p-value NaN from division by zero (#585)

When a tree node has zero treatment or control observations (n_t=0
or n_c=0), the p-value variance formula divides by zero, producing
NaN. Also guards against zero variance when all observations in a
group have the same outcome (p=0 or p=1). Returns p_value=1.0
(maximally non-significant) in these degenerate cases.

* Add edge-case test for p-value with sparse treatment groups

Test verifies predictions don't contain NaN when tree nodes have
zero treatment or control observations (min_samples_treatment=0,
heavily imbalanced groups).

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: jeongyoonlee

causalml/inference/tree/uplift.pyx

@@ -2095,7 +2095,14 @@ class UpliftTreeClassifier:
         else:
             p_t = cur_summary_p[suboptTreatment]
             n_t = cur_summary_n[suboptTreatment]
-        p_value = (1. - stats.norm.cdf(fabs(p_c - p_t) / sqrt(p_t * (1 - p_t) / n_t + p_c * (1 - p_c) / n_c))) * 2
+        if n_t > 0 and n_c > 0:
+            variance = p_t * (1 - p_t) / n_t + p_c * (1 - p_c) / n_c
+            if variance > 0:
+                p_value = (1. - stats.norm.cdf(fabs(p_c - p_t) / sqrt(variance))) * 2
+            else:
+                p_value = 1.0
+        else:
+            p_value = 1.0
         upliftScore = [maxDiff, p_value]
 
         bestGain = 0.0

tests/test_uplift_trees.py

@@ -365,3 +365,27 @@ def test_UpliftTreeClassifier_with_nan_in_categorical_features():
     preds = uplift_model.predict(X_test)
     assert preds is not None
     assert uplift_model.fitted_uplift_tree is not None
+
+
+def test_uplift_tree_pvalue_no_nan_with_sparse_groups():
+    """Test that p-values don't become NaN when tree nodes have zero
+    treatment or control observations (issue #585)."""
+    np.random.seed(RANDOM_SEED)
+    n = 50
+    X = np.random.randn(n, 3)
+    # Heavily imbalanced: only 2 samples in treatment1
+    treatment = np.array([CONTROL_NAME] * 45 + ["treatment1"] * 2 + ["treatment2"] * 3)
+    y = np.random.randint(0, 2, n)
+
+    model = UpliftTreeClassifier(
+        control_name=CONTROL_NAME,
+        min_samples_leaf=1,
+        min_samples_treatment=0,
+        max_depth=5,
+    )
+    model.fit(X, treatment, y)
+    preds = model.predict(X)
+
+    assert not np.any(
+        np.isnan(preds)
+    ), "Predictions contain NaN (likely from NaN p-values)"