scikit-learn · makoeppel · Dec 15, 2022 · Dec 15, 2022 · Dec 19, 2022 · Dec 19, 2022
diff --git a/doc/whats_new/v1.3.rst b/doc/whats_new/v1.3.rst
@@ -164,6 +164,10 @@ Changelog
  :pr:`24935` by :user:`Seladus <seladus>`, :user:`Guillaume Lemaitre <glemaitre>`, and
  :user:`Dea María Léon <deamarialeon>`, :pr:`25257` by :user:`Gleb Levitski <glevv>`.

+- |Enhancement| Improve error handling in :func:`feature_selection.mutual_info_classif`
+  that now checks if instances are left after masking of unique labels.
+  :pr:`25192` by :user:`makoeppel`.
+
 Code and Documentation Contributors
 -----------------------------------


diff --git a/sklearn/feature_selection/_mutual_info.py b/sklearn/feature_selection/_mutual_info.py
@@ -9,7 +9,7 @@
 from ..neighbors import NearestNeighbors, KDTree
 from ..preprocessing import scale
 from ..utils import check_random_state
-from ..utils.validation import check_array, check_X_y
+from ..utils.validation import check_array, check_X_y, _num_samples
 from ..utils.multiclass import check_classification_targets


@@ -135,6 +135,13 @@ def _compute_mi_cd(c, d, n_neighbors):
    c = c[mask]
    radius = radius[mask]

+    if _num_samples(c) == 0:
+        raise ValueError(
+            "Found array with 0 samples after masking"
+            " points with unique labels. Ensure that at least"
+            " two instances share the same label."
+        )
+
    kd = KDTree(c)
    m_all = kd.query_radius(c, radius, count_only=True, return_distance=False)
    m_all = np.array(m_all)

diff --git a/sklearn/feature_selection/tests/test_mutual_info.py b/sklearn/feature_selection/tests/test_mutual_info.py
@@ -236,3 +236,19 @@ def test_mutual_information_symmetry_classif_regression(correlated, global_rando
    )

    assert mi_classif == pytest.approx(mi_regression)
+
+
+def test_mutual_info_error_handling_for_unique_labels():
+    """Check that the correct ValueError is raised when calling `mutual_info_classif`
+    with only unique labels.
+    """
+
+    a = [[1, 0, 1], [0, 1, 1]]
+    b = [0, 1]
+    err_msg = (
+        "Found array with 0 samples after masking"
+        " points with unique labels. Ensure that at least"
+        " two instances share the same label."
+    )
+    with pytest.raises(ValueError, match=err_msg):
+        mutual_info_classif(a, b)