MNT move _check_categorical_features to utils.validation (scikit-learn#33946)

Author: lorentzenchr

sklearn/ensemble/_hist_gradient_boosting/gradient_boosting.py

@@ -10,7 +10,6 @@
 from numbers import Integral, Real
 from time import time
 
-import narwhals.stable.v2 as nw
 import numpy as np
 
 from sklearn._loss.loss import (
@@ -46,6 +45,7 @@
 from sklearn.utils._param_validation import Interval, RealNotInt, StrOptions
 from sklearn.utils.multiclass import check_classification_targets
 from sklearn.utils.validation import (
+    _check_categorical_features,
     _check_monotonic_cst,
     _check_sample_weight,
     _check_y,
@@ -267,7 +267,7 @@ def _preprocess_X(self, X, *, reset):
             return self._preprocessor.transform(X)
 
         # At this point, reset is False, which runs during `fit`.
-        self.is_categorical_ = self._check_categorical_features(X)
+        self.is_categorical_ = _check_categorical_features(X, self.categorical_features)
 
         if self.is_categorical_ is None:
             self._preprocessor = None
@@ -353,116 +353,6 @@ def _check_categories(self):
             known_categories[feature_idx] = np.arange(len(categories), dtype=X_DTYPE)
         return known_categories
 
-    def _check_categorical_features(self, X):
-        """Check and validate categorical features in X
-
-        Parameters
-        ----------
-        X : {array-like, pandas DataFrame} of shape (n_samples, n_features)
-            Input data.
-
-        Return
-        ------
-        is_categorical : ndarray of shape (n_features,) or None, dtype=bool
-            Indicates whether a feature is categorical. If no feature is
-            categorical, this is None.
-        """
-        if nw.dependencies.is_into_dataframe(X):
-            X = nw.from_native(X)
-            dtypes = X.schema.dtypes()
-            X_is_dataframe = True
-            categorical_columns_mask = np.asarray(
-                [d in (nw.Categorical, nw.Enum) for d in dtypes]
-            )
-        else:
-            X_is_dataframe = False
-            categorical_columns_mask = None
-
-        categorical_features = self.categorical_features
-
-        categorical_by_dtype = (
-            isinstance(categorical_features, str)
-            and categorical_features == "from_dtype"
-        )
-        no_categorical_dtype = categorical_features is None or (
-            categorical_by_dtype and not X_is_dataframe
-        )
-
-        if no_categorical_dtype:
-            return None
-
-        if categorical_by_dtype and X_is_dataframe:
-            categorical_features = categorical_columns_mask
-        else:
-            categorical_features = np.asarray(categorical_features)
-
-        if categorical_features.size == 0:
-            return None
-
-        if categorical_features.dtype.kind not in ("i", "b", "U", "O"):
-            raise ValueError(
-                "categorical_features must be an array-like of bool, int or "
-                f"str, got: {categorical_features.dtype.name}."
-            )
-
-        if categorical_features.dtype.kind == "O":
-            types = set(type(f) for f in categorical_features)
-            if types != {str}:
-                raise ValueError(
-                    "categorical_features must be an array-like of bool, int or "
-                    f"str, got: {', '.join(sorted(t.__name__ for t in types))}."
-                )
-
-        n_features = X.shape[1]
-        # At this point `validate_data` was not called yet because we use the original
-        # dtypes to discover the categorical features. Thus `feature_names_in_`
-        # is not defined yet.
-        feature_names_in_ = getattr(X, "columns", None)
-
-        if categorical_features.dtype.kind in ("U", "O"):
-            # check for feature names
-            if feature_names_in_ is None:
-                raise ValueError(
-                    "categorical_features should be passed as an array of "
-                    "integers or as a boolean mask when the model is fitted "
-                    "on data without feature names."
-                )
-            is_categorical = np.zeros(n_features, dtype=bool)
-            feature_names = list(feature_names_in_)
-            for feature_name in categorical_features:
-                try:
-                    is_categorical[feature_names.index(feature_name)] = True
-                except ValueError as e:
-                    raise ValueError(
-                        f"categorical_features has an item value '{feature_name}' "
-                        "which is not a valid feature name of the training "
-                        f"data. Observed feature names: {feature_names}"
-                    ) from e
-        elif categorical_features.dtype.kind == "i":
-            # check for categorical features as indices
-            if (
-                np.max(categorical_features) >= n_features
-                or np.min(categorical_features) < 0
-            ):
-                raise ValueError(
-                    "categorical_features set as integer "
-                    "indices must be in [0, n_features - 1]"
-                )
-            is_categorical = np.zeros(n_features, dtype=bool)
-            is_categorical[categorical_features] = True
-        else:
-            if categorical_features.shape[0] != n_features:
-                raise ValueError(
-                    "categorical_features set as a boolean mask "
-                    "must have shape (n_features,), got: "
-                    f"{categorical_features.shape}"
-                )
-            is_categorical = categorical_features
-
-        if not np.any(is_categorical):
-            return None
-        return is_categorical
-
     def _check_interaction_cst(self, n_features):
         """Check and validation for interaction constraints."""
         if self.interaction_cst is None:

sklearn/utils/tests/test_validation.py

@@ -65,6 +65,7 @@
 from sklearn.utils.validation import (
     FLOAT_DTYPES,
     _allclose_dense_sparse,
+    _check_categorical_features,
     _check_feature_names_in,
     _check_method_params,
     _check_pos_label_consistency,
@@ -2439,3 +2440,76 @@ def test_check_array_on_sparse_inputs_with_array_api_enabled():
 def test_check_array_allow_nd_errors(X, estimator, expected_error_message):
     with pytest.raises(ValueError, match=expected_error_message):
         check_array(X, estimator=estimator)
+
+
+@pytest.mark.parametrize(
+    ["categorical_features", "expected_msg"],
+    [
+        (
+            [b"hello", b"world"],
+            re.escape(
+                "categorical_features must be an array-like of bool, int or str, "
+                "got: bytes40."
+            ),
+        ),
+        (
+            np.array([b"hello", 1.3], dtype=object),
+            re.escape(
+                "categorical_features must be an array-like of bool, int or str, "
+                "got: bytes, float."
+            ),
+        ),
+        (
+            [0, -1],
+            re.escape(
+                "categorical_features set as integer indices must be in "
+                "[0, n_features - 1]"
+            ),
+        ),
+        (
+            [True, True, False, False, True],
+            re.escape(
+                "categorical_features set as a boolean mask must have shape "
+                "(n_features,)"
+            ),
+        ),
+    ],
+)
+def test_check_categorical_features_raises(categorical_features, expected_msg):
+    """Test that check_categorical_features raises expected errors."""
+    rng = np.random.RandomState(0)
+    n_samples, n_features = 10, 10
+    X = rng.randint(0, 3, size=(n_samples, n_features))
+
+    with pytest.raises(ValueError, match=expected_msg):
+        _check_categorical_features(X, categorical_features)
+
+
+@pytest.mark.parametrize(
+    ["categorical_features", "on_array"],
+    [
+        ([False, True, True, False], True),
+        ([1, 2], True),
+        (["b", "c"], False),
+        ("from_dtype", False),
+    ],
+)
+@pytest.mark.parametrize("constructor_name", ["array", "pandas", "polars"])
+def test_check_categorical_features(categorical_features, on_array, constructor_name):
+    """Test that check_categorical_features returns as expected on simple data."""
+    rng = np.random.RandomState(0)
+    n_samples, n_features = 30, 4
+    X = rng.randint(0, 3, size=(n_samples, n_features))
+    if constructor_name == "array" and not on_array:
+        return
+    elif constructor_name == "polars":
+        X = X.astype(str)
+    X = _convert_container(
+        X,
+        constructor_name,
+        columns_name=["a", "b", "c", "d"],
+        categorical_feature_names=["b", "c"],
+    )
+
+    result = _check_categorical_features(X, categorical_features)
+    assert_allclose(result, [False, True, True, False])

sklearn/utils/validation.py

@@ -2468,6 +2468,129 @@ def _generate_get_feature_names_out(estimator, n_features_out, input_features=No
     )
 
 
+def _check_categorical_features(X, categorical_features):
+    """Check and validate categorical features in X
+
+    Parameters
+    ----------
+    X : {array-like, pandas DataFrame} of shape (n_samples, n_features)
+        Input data.
+
+    categorical_features : array-like of {bool, int, str} of shape (n_features) \
+            or shape (n_categorical_features,), default='from_dtype'
+        Indicates the categorical features in `X`.
+
+        - None : no feature will be considered categorical.
+        - boolean array-like : boolean mask indicating categorical features.
+        - integer array-like : integer indices indicating categorical
+          features.
+        - str array-like: names of categorical features (assuming the training
+          data has feature names).
+        - `"from_dtype"`: dataframe columns with dtype "Categorical" and "Enum" are
+          considered to be categorical features. The input must be a dataframe that
+          is supported by narwhals (or supports it): :func:`narwhals.from_native` must
+          work. This is the case, for instance, for pandas and polars DataFrames.
+
+    Return
+    ------
+    is_categorical : ndarray of shape (n_features,) or None, dtype=bool
+        Indicates whether a feature is categorical. If no feature is
+        categorical, this is None.
+    """
+    if nw.dependencies.is_into_dataframe(X):
+        X = nw.from_native(X)
+        dtypes = X.schema.dtypes()
+        X_is_dataframe = True
+        categorical_columns_mask = np.asarray(
+            [d in (nw.Categorical, nw.Enum) for d in dtypes]
+        )
+    else:
+        X_is_dataframe = False
+        categorical_columns_mask = None
+
+    categorical_by_dtype = (
+        isinstance(categorical_features, str) and categorical_features == "from_dtype"
+    )
+    no_categorical_dtype = categorical_features is None or (
+        categorical_by_dtype and not X_is_dataframe
+    )
+
+    if no_categorical_dtype:
+        return None
+
+    if categorical_by_dtype and X_is_dataframe:
+        categorical_features = categorical_columns_mask
+    else:
+        categorical_features = np.asarray(categorical_features)
+
+    if categorical_features.size == 0:
+        return None
+
+    if categorical_features.dtype.kind not in ("i", "b", "U", "O"):
+        raise ValueError(
+            "categorical_features must be an array-like of bool, int or "
+            f"str, got: {categorical_features.dtype.name}."
+        )
+
+    if categorical_features.dtype.kind == "O":
+        types = set(type(f) for f in categorical_features)
+        if types != {str}:
+            raise ValueError(
+                "categorical_features must be an array-like of bool, int or "
+                f"str, got: {', '.join(sorted(t.__name__ for t in types))}."
+            )
+
+    n_features = X.shape[1]
+    # At this point `validate_data` was not called yet because we use the original
+    # dtypes to discover the categorical features. Thus `feature_names_in_`
+    # is not defined yet.
+    feature_names_in_ = getattr(X, "columns", None)
+
+    if categorical_features.dtype.kind in ("U", "O"):
+        # check for feature names
+        if feature_names_in_ is None:
+            raise ValueError(
+                "categorical_features should be passed as an array of "
+                "integers or as a boolean mask when the model is fitted "
+                "on data without feature names."
+            )
+        is_categorical = np.zeros(n_features, dtype=bool)
+        feature_names = list(feature_names_in_)
+        for feature_name in categorical_features:
+            try:
+                is_categorical[feature_names.index(feature_name)] = True
+            except ValueError as e:
+                raise ValueError(
+                    f"categorical_features has an item value '{feature_name}' "
+                    "which is not a valid feature name of the training "
+                    f"data. Observed feature names: {feature_names}"
+                ) from e
+    elif categorical_features.dtype.kind == "i":
+        # check for categorical features as indices
+        if (
+            np.max(categorical_features) >= n_features
+            or np.min(categorical_features) < 0
+        ):
+            raise ValueError(
+                "categorical_features set as integer "
+                "indices must be in [0, n_features - 1]"
+            )
+        is_categorical = np.zeros(n_features, dtype=bool)
+        is_categorical[categorical_features] = True
+    else:
+        if categorical_features.shape[0] != n_features:
+            raise ValueError(
+                "categorical_features set as a boolean mask "
+                "must have shape (n_features,), got: "
+                f"{categorical_features.shape}"
+            )
+        is_categorical = categorical_features
+
+    if not np.any(is_categorical):
+        return None
+    return is_categorical
+
+
 def _check_monotonic_cst(estimator, monotonic_cst=None):
     """Check the monotonic constraints and return the corresponding array.