Skip to content

Commit 144e6b6

Browse files
jeongyoonleeclaude
andcommitted
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>
1 parent 912c266 commit 144e6b6

1 file changed

Lines changed: 24 additions & 0 deletions

File tree

tests/test_uplift_trees.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,3 +293,27 @@ def test_uplift_tree_visualization():
293293
# Plot uplift tree
294294
graph = uplift_tree_plot(uplift_model.fitted_uplift_tree, x_names)
295295
graph.create_png()
296+
297+
298+
def test_uplift_tree_pvalue_no_nan_with_sparse_groups():
299+
"""Test that p-values don't become NaN when tree nodes have zero
300+
treatment or control observations (issue #585)."""
301+
np.random.seed(RANDOM_SEED)
302+
n = 50
303+
X = np.random.randn(n, 3)
304+
# Heavily imbalanced: only 2 samples in treatment1
305+
treatment = np.array([CONTROL_NAME] * 45 + ["treatment1"] * 2 + ["treatment2"] * 3)
306+
y = np.random.randint(0, 2, n)
307+
308+
model = UpliftTreeClassifier(
309+
control_name=CONTROL_NAME,
310+
min_samples_leaf=1,
311+
min_samples_treatment=0,
312+
max_depth=5,
313+
)
314+
model.fit(X, treatment, y)
315+
preds = model.predict(X)
316+
317+
assert not np.any(
318+
np.isnan(preds)
319+
), "Predictions contain NaN (likely from NaN p-values)"

0 commit comments

Comments
 (0)