Improve uplift_at_k validation and add top-k markers to prediction plots

docs/api/viz/plot_uplift_preds.rst

@@ -2,4 +2,7 @@
 `sklift.viz <./>`_.plot_uplift_preds
 ***********************************************
 
-.. autofunction:: sklift.viz.base.plot_uplift_preds
+Builds histograms for treatment/control/uplift predictions and supports optional
+top-k threshold markers via the ``k`` parameter.
+
+.. autofunction:: sklift.viz.base.plot_uplift_preds

sklift/metrics/metrics.py

@@ -409,7 +409,6 @@ def uplift_at_k(y_true, uplift, treatment, strategy, k=0.3):
         :func:`.qini_auc_score`: Compute normalized Area Under the Qini Curve from prediction scores.
     """
 
-    # TODO: checker all groups is not empty
     check_consistent_length(y_true, uplift, treatment)
     check_is_binary(treatment)
     check_is_binary(y_true)
@@ -444,9 +443,19 @@ def uplift_at_k(y_true, uplift, treatment, strategy, k=0.3):
         else:
             n_size = k
 
-        # ToDo: _checker_ there are observations among two groups among first k
-        score_ctrl = y_true[order][:n_size][treatment[order][:n_size] == 0].mean()
-        score_trmnt = y_true[order][:n_size][treatment[order][:n_size] == 1].mean()
+        y_top_k = y_true[order][:n_size]
+        treatment_top_k = treatment[order][:n_size]
+
+        ctrl_mask = treatment_top_k == 0
+        trmnt_mask = treatment_top_k == 1
+        if ctrl_mask.sum() == 0 or trmnt_mask.sum() == 0:
+            raise ValueError(
+                f'With strategy={strategy} and k={k}, there are no observations from one of the groups '
+                'in the selected top-k sample. Try another k or use strategy="by_group".'
+            )
+
+        score_ctrl = y_top_k[ctrl_mask].mean()
+        score_trmnt = y_top_k[trmnt_mask].mean()
 
     else:  # strategy == 'by_group':
         if k_type == 'f':

sklift/tests/test_metrics.py

@@ -43,7 +43,7 @@ def test_uplift_curve(binary, test_x_actual, test_y_actual):
     y_true, uplift, treatment = make_predictions(binary)
 
     if binary == False:
-        with pytest.raises(Exception):
+        with pytest.raises(ValueError):
             x_actual, y_actual = uplift_curve(y_true, uplift, treatment)
     else:
         x_actual, y_actual = uplift_curve(y_true, uplift, treatment)
@@ -54,7 +54,7 @@ def test_uplift_curve(binary, test_x_actual, test_y_actual):
 
 
 def test_uplift_curve_hard():
-    with pytest.raises(Exception):
+    with pytest.raises(ValueError):
         y_true, uplift, treatment = make_predictions(binary=True)
         y_true = np.zeros(y_true.shape)
 
@@ -81,7 +81,7 @@ def test_uplift_curve_hard():
 def test_perfect_uplift_curve(binary, test_x_actual, test_y_actual):
     y_true, uplift, treatment = make_predictions(binary)
     if binary == False:
-        with pytest.raises(Exception):
+        with pytest.raises(ValueError):
             x_actual, y_actual = perfect_uplift_curve(y_true, treatment)
     else:
         x_actual, y_actual = perfect_uplift_curve(y_true, treatment)
@@ -91,7 +91,7 @@ def test_perfect_uplift_curve(binary, test_x_actual, test_y_actual):
 
 
 def test_perfect_uplift_curve_hard():
-    with pytest.raises(Exception):
+    with pytest.raises(ValueError):
         y_true, uplift, treatment = make_predictions(binary=True)
         y_true = np.zeros(y_true.shape)
 
@@ -119,7 +119,7 @@ def test_uplift_auc_score():
     treatment = [0, 1]
     assert_array_almost_equal(uplift_auc_score(y_true, uplift, treatment), 1.)
 
-    with pytest.raises(Exception):
+    with pytest.raises(ValueError):
         y_true = [1, 1]
         uplift = [0.1, 0.3]
         treatment = [0, 1]
@@ -152,7 +152,7 @@ def test_qini_curve(binary, test_x_actual, test_y_actual):
     y_true, uplift, treatment = make_predictions(binary)
 
     if binary == False:
-        with pytest.raises(Exception):    
+        with pytest.raises(ValueError):    
             x_actual, y_actual = qini_curve(y_true, uplift, treatment)
     else:
         x_actual, y_actual = qini_curve(y_true, uplift, treatment)
@@ -162,7 +162,7 @@ def test_qini_curve(binary, test_x_actual, test_y_actual):
 
 
 def test_qini_curve_hard():
-    with pytest.raises(Exception):
+    with pytest.raises(ValueError):
         y_true, uplift, treatment = make_predictions(binary=True)
         y_true = np.zeros(y_true.shape)
 
@@ -197,7 +197,7 @@ def test_perfect_qini_curve(binary, negative_effect, test_x_actual, test_y_actua
 
 
 def test_perfect_qini_curve_hard():
-    with pytest.raises(Exception):
+    with pytest.raises(ValueError):
         y_true, uplift, treatment = make_predictions(binary=True)
         y_true = np.zeros(y_true.shape)
 
@@ -241,7 +241,7 @@ def test_qini_auc_score():
     treatment = [0, 1]
     assert_array_almost_equal(qini_auc_score(y_true, uplift, treatment), 1.)
 
-    with pytest.raises(Exception):
+    with pytest.raises(ValueError):
         y_true = [1, 1]
         uplift = [0.1, 0.3]
         treatment = [0, 1]
@@ -276,6 +276,15 @@ def test_uplift_at_k():
     assert_array_almost_equal(uplift_at_k(y_true, uplift, treatment, strategy='by_group', k=1), np.array([0.]))
     #assert_array_almost_equal(uplift_at_k(y_true, uplift, treatment, strategy='overall', k=2), np.array([0.]))
 
+
+def test_uplift_at_k_overall_empty_group_error():
+    y_true = np.array([1, 0, 1, 0])
+    uplift = np.array([0.9, 0.8, 0.2, 0.1])
+    treatment = np.array([1, 1, 0, 0])
+
+    with pytest.raises(ValueError):
+        uplift_at_k(y_true, uplift, treatment, strategy='overall', k=2)
+
 @pytest.mark.parametrize(
     "strategy, k",
     [
@@ -427,4 +436,4 @@ def test_make_scorer_error():
 
 
 
-
+

sklift/tests/test_viz.py

@@ -136,9 +136,20 @@ def test_plot_uplift_preds():
     assert isinstance(viz[0], mpl.axes.Axes)
     assert isinstance(viz[1], mpl.axes.Axes)
     assert isinstance(viz[2], mpl.axes.Axes)
+
+    viz_k = plot_uplift_preds(trmnt_preds, ctrl_preds, log=True, bins=5, k=0.5)
+    assert len(viz_k[0].lines) == 1
+    assert len(viz_k[1].lines) == 1
+    assert len(viz_k[2].lines) == 1
 
     with pytest.raises(ValueError):
     	plot_uplift_preds(trmnt_preds, ctrl_preds, log=True, bins=0)
+    with pytest.raises(ValueError):
+        plot_uplift_preds(trmnt_preds, ctrl_preds, log=True, bins=5, k=0)
+    with pytest.raises(ValueError):
+        plot_uplift_preds(trmnt_preds, ctrl_preds, log=True, bins=5, k=1.0)
+    with pytest.raises(ValueError):
+        plot_uplift_preds(trmnt_preds, ctrl_preds, log=True, bins=5, k=10)
 
 def test_plot_uplift_by_percentile():
     y_true, uplift, treatment = make_predictions()
@@ -198,4 +209,4 @@ def test_plot_treatment_balance_curve():
 def test_plot_treatment_balance_errors():
 	y_true, uplift, treatment = make_predictions()
 	with pytest.raises(ValueError):		
-		viz = plot_treatment_balance_curve(uplift, treatment, winsize=5) 
+		viz = plot_treatment_balance_curve(uplift, treatment, winsize=5) 

sklift/viz/base.py

@@ -11,7 +11,7 @@
 )
 
 
-def plot_uplift_preds(trmnt_preds, ctrl_preds, log=False, bins=100):
+def plot_uplift_preds(trmnt_preds, ctrl_preds, log=False, bins=100, k=None):
     """Plot histograms of treatment, control and uplift predictions.
 
     Args:
@@ -22,36 +22,69 @@ def plot_uplift_preds(trmnt_preds, ctrl_preds, log=False, bins=100):
             If an integer is given, bins + 1 bin edges are calculated and returned.
             If bins is a sequence, gives bin edges, including left edge of first bin and right edge of last bin.
             In this case, bins is returned unmodified. Default is 100.
+        k (float or int, optional): If provided, draw a vertical threshold line for the top-k segment.
+            If float, should be in (0, 1) and represents the sample proportion.
+            If int, should be in [1, n_samples] and represents the number of samples.
 
     Returns:
         Object that stores computed values.
     """
 
-    # TODO: Add k as parameter: vertical line on plots
     check_consistent_length(trmnt_preds, ctrl_preds)
+    trmnt_preds = np.asarray(trmnt_preds)
+    ctrl_preds = np.asarray(ctrl_preds)
 
     if not isinstance(bins, int) or bins <= 0:
         raise ValueError(
             f'Bins should be positive integer. Invalid value for bins: {bins}')
 
+    n_samples = trmnt_preds.shape[0]
+    n_size = None
+    if k is not None:
+        if isinstance(k, bool):
+            raise ValueError(f'Invalid value for k: {k}')
+        k_type = np.asarray(k).dtype.kind
+        if k_type == 'f':
+            if k <= 0 or k >= 1:
+                raise ValueError(f'k={k} should be a float in the (0, 1) range.')
+            n_size = max(int(n_samples * k), 1)
+        elif k_type == 'i':
+            if k <= 0 or k > n_samples:
+                raise ValueError(
+                    f'k={k} should be positive and smaller than or equal to the number of samples {n_samples}.'
+                )
+            n_size = int(k)
+        else:
+            raise ValueError(f'Invalid value for k: {k}')
+
     if log:
         trmnt_preds = np.log(trmnt_preds + 1)
         ctrl_preds = np.log(ctrl_preds + 1)
 
     fig, axes = plt.subplots(ncols=3, nrows=1, figsize=(20, 7))
     axes[0].hist(
         trmnt_preds, bins=bins, alpha=0.3, color='b', label='Treated', histtype='stepfilled')
+    if n_size is not None:
+        trmnt_k_threshold = np.sort(trmnt_preds)[::-1][n_size - 1]
+        axes[0].axvline(trmnt_k_threshold, color='r', linestyle='--', linewidth=1.5, label=f'Top-{k} threshold')
     axes[0].set_ylabel('Probability hist')
     axes[0].legend()
     axes[0].set_title('Treatment predictions')
 
     axes[1].hist(
         ctrl_preds, bins=bins, alpha=0.5, color='y', label='Not treated', histtype='stepfilled')
+    if n_size is not None:
+        ctrl_k_threshold = np.sort(ctrl_preds)[::-1][n_size - 1]
+        axes[1].axvline(ctrl_k_threshold, color='r', linestyle='--', linewidth=1.5, label=f'Top-{k} threshold')
     axes[1].legend()
     axes[1].set_title('Control predictions')
 
+    uplift_preds = trmnt_preds - ctrl_preds
     axes[2].hist(
-        trmnt_preds - ctrl_preds, bins=bins, alpha=0.5, color='green', label='Uplift', histtype='stepfilled')
+        uplift_preds, bins=bins, alpha=0.5, color='green', label='Uplift', histtype='stepfilled')
+    if n_size is not None:
+        uplift_k_threshold = np.sort(uplift_preds)[::-1][n_size - 1]
+        axes[2].axvline(uplift_k_threshold, color='r', linestyle='--', linewidth=1.5, label=f'Top-{k} threshold')
     axes[2].legend()
     axes[2].set_title('Uplift predictions')