MAINT Parameters validation for sklearn.utils.class_weight.compute_class_weight (scikit-learn#26512)

Co-authored-by: Guillaume Lemaitre <g.lemaitre58@gmail.com>

Author: brendanlu

sklearn/tests/test_public_functions.py

@@ -306,13 +306,13 @@ def _check_function_param_validation(
     "sklearn.tree.export_graphviz",
     "sklearn.tree.export_text",
     "sklearn.tree.plot_tree",
-    "sklearn.utils.class_weight.compute_sample_weight",
     "sklearn.utils.gen_batches",
     "sklearn.utils.gen_even_slices",
-    "sklearn.utils.graph.single_source_shortest_path_length",
     "sklearn.utils.resample",
     "sklearn.utils.safe_mask",
     "sklearn.utils.extmath.randomized_svd",
+    "sklearn.utils.class_weight.compute_class_weight",
+    "sklearn.utils.class_weight.compute_sample_weight",
     "sklearn.utils.graph.single_source_shortest_path_length",
 ]
 

sklearn/utils/class_weight.py

@@ -8,29 +8,37 @@
 from ._param_validation import StrOptions, validate_params
 
 
+@validate_params(
+    {
+        "class_weight": [dict, StrOptions({"balanced"}), None],
+        "classes": [np.ndarray],
+        "y": ["array-like"],
+    },
+    prefer_skip_nested_validation=True,
+)
 def compute_class_weight(class_weight, *, classes, y):
     """Estimate class weights for unbalanced datasets.
 
     Parameters
     ----------
-    class_weight : dict, 'balanced' or None
-        If 'balanced', class weights will be given by
-        ``n_samples / (n_classes * np.bincount(y))``.
-        If a dictionary is given, keys are classes and values
-        are corresponding class weights.
-        If None is given, the class weights will be uniform.
+    class_weight : dict, "balanced" or None
+        If "balanced", class weights will be given by
+        `n_samples / (n_classes * np.bincount(y))`.
+        If a dictionary is given, keys are classes and values are corresponding class
+        weights.
+        If `None` is given, the class weights will be uniform.
 
     classes : ndarray
         Array of the classes occurring in the data, as given by
-        ``np.unique(y_org)`` with ``y_org`` the original class labels.
+        `np.unique(y_org)` with `y_org` the original class labels.
 
     y : array-like of shape (n_samples,)
         Array of original class labels per sample.
 
     Returns
     -------
     class_weight_vect : ndarray of shape (n_classes,)
-        Array with class_weight_vect[i] the weight for i-th class.
+        Array with `class_weight_vect[i]` the weight for i-th class.
 
     References
     ----------
@@ -57,10 +65,6 @@ def compute_class_weight(class_weight, *, classes, y):
     else:
         # user-defined dictionary
         weight = np.ones(classes.shape[0], dtype=np.float64, order="C")
-        if not isinstance(class_weight, dict):
-            raise ValueError(
-                "class_weight must be dict, 'balanced', or None, got: %r" % class_weight
-            )
         unweighted_classes = []
         for i, c in enumerate(classes):
             if c in class_weight:

sklearn/utils/tests/test_class_weight.py

@@ -22,33 +22,53 @@ def test_compute_class_weight():
     assert cw[0] < cw[1] < cw[2]
 
 
-def test_compute_class_weight_not_present():
+@pytest.mark.parametrize(
+    "y_type, class_weight, classes, err_msg",
+    [
+        (
+            "numeric",
+            "balanced",
+            np.arange(4),
+            "classes should have valid labels that are in y",
+        ),
+        # Non-regression for https://github.com/scikit-learn/scikit-learn/issues/8312
+        (
+            "numeric",
+            {"label_not_present": 1.0},
+            np.arange(4),
+            r"The classes, \[0, 1, 2, 3\], are not in class_weight",
+        ),
+        (
+            "numeric",
+            "balanced",
+            np.arange(2),
+            "classes should include all valid labels",
+        ),
+        (
+            "numeric",
+            {0: 1.0, 1: 2.0},
+            np.arange(2),
+            "classes should include all valid labels",
+        ),
+        (
+            "string",
+            {"dogs": 3, "cat": 2},
+            np.array(["dog", "cat"]),
+            r"The classes, \['dog'\], are not in class_weight",
+        ),
+    ],
+)
+def test_compute_class_weight_not_present(y_type, class_weight, classes, err_msg):
     # Raise error when y does not contain all class labels
-    classes = np.arange(4)
-    y = np.asarray([0, 0, 0, 1, 1, 2])
-    with pytest.raises(ValueError):
-        compute_class_weight("balanced", classes=classes, y=y)
-    # Fix exception in error message formatting when missing label is a string
-    # https://github.com/scikit-learn/scikit-learn/issues/8312
-    with pytest.raises(
-        ValueError, match=r"The classes, \[0, 1, 2, 3\], are not in class_weight"
-    ):
-        compute_class_weight({"label_not_present": 1.0}, classes=classes, y=y)
-    # Raise error when y has items not in classes
-    classes = np.arange(2)
-    with pytest.raises(ValueError):
-        compute_class_weight("balanced", classes=classes, y=y)
-    with pytest.raises(ValueError):
-        compute_class_weight({0: 1.0, 1: 2.0}, classes=classes, y=y)
-
-    # y contains a unweighted class that is not in class_weights
-    classes = np.asarray(["cat", "dog"])
-    y = np.asarray(["dog", "cat", "dog"])
-    class_weights = {"dogs": 3, "cat": 2}
-    msg = r"The classes, \['dog'\], are not in class_weight"
-
-    with pytest.raises(ValueError, match=msg):
-        compute_class_weight(class_weights, classes=classes, y=y)
+    y = (
+        np.asarray([0, 0, 0, 1, 1, 2])
+        if y_type == "numeric"
+        else np.asarray(["dog", "cat", "dog"])
+    )
+
+    print(y)
+    with pytest.raises(ValueError, match=err_msg):
+        compute_class_weight(class_weight, classes=classes, y=y)
 
 
 def test_compute_class_weight_dict():