465 lines
15 KiB
Python
465 lines
15 KiB
Python
import importlib
|
|
from collections import namedtuple
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from numpy.testing import assert_array_equal
|
|
|
|
from sklearn._config import config_context, get_config
|
|
from sklearn.preprocessing import StandardScaler
|
|
from sklearn.utils._set_output import (
|
|
ADAPTERS_MANAGER,
|
|
ContainerAdapterProtocol,
|
|
_get_adapter_from_container,
|
|
_get_output_config,
|
|
_safe_set_output,
|
|
_SetOutputMixin,
|
|
_wrap_data_with_container,
|
|
check_library_installed,
|
|
)
|
|
from sklearn.utils.fixes import CSR_CONTAINERS
|
|
|
|
|
|
def test_pandas_adapter():
|
|
"""Check pandas adapter has expected behavior."""
|
|
pd = pytest.importorskip("pandas")
|
|
X_np = np.asarray([[1, 0, 3], [0, 0, 1]])
|
|
columns = np.asarray(["f0", "f1", "f2"], dtype=object)
|
|
index = np.asarray([0, 1])
|
|
X_df_orig = pd.DataFrame([[1, 2], [1, 3]], index=index)
|
|
|
|
adapter = ADAPTERS_MANAGER.adapters["pandas"]
|
|
X_container = adapter.create_container(X_np, X_df_orig, columns=lambda: columns)
|
|
assert isinstance(X_container, pd.DataFrame)
|
|
assert_array_equal(X_container.columns, columns)
|
|
assert_array_equal(X_container.index, index)
|
|
|
|
# Input dataframe's index does not change
|
|
new_columns = np.asarray(["f0", "f1"], dtype=object)
|
|
X_df = pd.DataFrame([[1, 2], [1, 3]], index=[10, 12])
|
|
new_df = adapter.create_container(X_df, X_df_orig, columns=new_columns)
|
|
assert_array_equal(new_df.columns, new_columns)
|
|
assert_array_equal(new_df.index, X_df.index)
|
|
|
|
assert adapter.is_supported_container(X_df)
|
|
assert not adapter.is_supported_container(X_np)
|
|
|
|
# adapter.update_columns updates the columns
|
|
new_columns = np.array(["a", "c"], dtype=object)
|
|
new_df = adapter.rename_columns(X_df, new_columns)
|
|
assert_array_equal(new_df.columns, new_columns)
|
|
|
|
# adapter.hstack stacks the dataframes horizontally.
|
|
X_df_1 = pd.DataFrame([[1, 2, 5], [3, 4, 6]], columns=["a", "b", "e"])
|
|
X_df_2 = pd.DataFrame([[4], [5]], columns=["c"])
|
|
X_stacked = adapter.hstack([X_df_1, X_df_2])
|
|
|
|
expected_df = pd.DataFrame(
|
|
[[1, 2, 5, 4], [3, 4, 6, 5]], columns=["a", "b", "e", "c"]
|
|
)
|
|
pd.testing.assert_frame_equal(X_stacked, expected_df)
|
|
|
|
# check that we update properly the columns even with duplicate column names
|
|
# this use-case potentially happen when using ColumnTransformer
|
|
# non-regression test for gh-28260
|
|
X_df = pd.DataFrame([[1, 2], [1, 3]], columns=["a", "a"])
|
|
new_columns = np.array(["x__a", "y__a"], dtype=object)
|
|
new_df = adapter.rename_columns(X_df, new_columns)
|
|
assert_array_equal(new_df.columns, new_columns)
|
|
|
|
# check the behavior of the inplace parameter in `create_container`
|
|
# we should trigger a copy
|
|
X_df = pd.DataFrame([[1, 2], [1, 3]], index=index)
|
|
X_output = adapter.create_container(X_df, X_df, columns=["a", "b"], inplace=False)
|
|
assert X_output is not X_df
|
|
assert list(X_df.columns) == [0, 1]
|
|
assert list(X_output.columns) == ["a", "b"]
|
|
|
|
# the operation is inplace
|
|
X_df = pd.DataFrame([[1, 2], [1, 3]], index=index)
|
|
X_output = adapter.create_container(X_df, X_df, columns=["a", "b"], inplace=True)
|
|
assert X_output is X_df
|
|
assert list(X_df.columns) == ["a", "b"]
|
|
assert list(X_output.columns) == ["a", "b"]
|
|
|
|
|
|
def test_polars_adapter():
|
|
"""Check Polars adapter has expected behavior."""
|
|
pl = pytest.importorskip("polars")
|
|
X_np = np.array([[1, 0, 3], [0, 0, 1]])
|
|
columns = ["f1", "f2", "f3"]
|
|
X_df_orig = pl.DataFrame(X_np, schema=columns, orient="row")
|
|
|
|
adapter = ADAPTERS_MANAGER.adapters["polars"]
|
|
X_container = adapter.create_container(X_np, X_df_orig, columns=lambda: columns)
|
|
|
|
assert isinstance(X_container, pl.DataFrame)
|
|
assert_array_equal(X_container.columns, columns)
|
|
|
|
# Update columns with create_container
|
|
new_columns = np.asarray(["a", "b", "c"], dtype=object)
|
|
new_df = adapter.create_container(X_df_orig, X_df_orig, columns=new_columns)
|
|
assert_array_equal(new_df.columns, new_columns)
|
|
|
|
assert adapter.is_supported_container(X_df_orig)
|
|
assert not adapter.is_supported_container(X_np)
|
|
|
|
# adapter.update_columns updates the columns
|
|
new_columns = np.array(["a", "c", "g"], dtype=object)
|
|
new_df = adapter.rename_columns(X_df_orig, new_columns)
|
|
assert_array_equal(new_df.columns, new_columns)
|
|
|
|
# adapter.hstack stacks the dataframes horizontally.
|
|
X_df_1 = pl.DataFrame([[1, 2, 5], [3, 4, 6]], schema=["a", "b", "e"], orient="row")
|
|
X_df_2 = pl.DataFrame([[4], [5]], schema=["c"], orient="row")
|
|
X_stacked = adapter.hstack([X_df_1, X_df_2])
|
|
|
|
expected_df = pl.DataFrame(
|
|
[[1, 2, 5, 4], [3, 4, 6, 5]], schema=["a", "b", "e", "c"], orient="row"
|
|
)
|
|
from polars.testing import assert_frame_equal
|
|
|
|
assert_frame_equal(X_stacked, expected_df)
|
|
|
|
# check the behavior of the inplace parameter in `create_container`
|
|
# we should trigger a copy
|
|
X_df = pl.DataFrame([[1, 2], [1, 3]], schema=["a", "b"], orient="row")
|
|
X_output = adapter.create_container(X_df, X_df, columns=["c", "d"], inplace=False)
|
|
assert X_output is not X_df
|
|
assert list(X_df.columns) == ["a", "b"]
|
|
assert list(X_output.columns) == ["c", "d"]
|
|
|
|
# the operation is inplace
|
|
X_df = pl.DataFrame([[1, 2], [1, 3]], schema=["a", "b"], orient="row")
|
|
X_output = adapter.create_container(X_df, X_df, columns=["c", "d"], inplace=True)
|
|
assert X_output is X_df
|
|
assert list(X_df.columns) == ["c", "d"]
|
|
assert list(X_output.columns) == ["c", "d"]
|
|
|
|
|
|
@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
|
|
def test__container_error_validation(csr_container):
|
|
"""Check errors in _wrap_data_with_container."""
|
|
X = np.asarray([[1, 0, 3], [0, 0, 1]])
|
|
X_csr = csr_container(X)
|
|
match = "The transformer outputs a scipy sparse matrix."
|
|
with config_context(transform_output="pandas"):
|
|
with pytest.raises(ValueError, match=match):
|
|
_wrap_data_with_container("transform", X_csr, X, StandardScaler())
|
|
|
|
|
|
class EstimatorWithoutSetOutputAndWithoutTransform:
|
|
pass
|
|
|
|
|
|
class EstimatorNoSetOutputWithTransform:
|
|
def transform(self, X, y=None):
|
|
return X # pragma: no cover
|
|
|
|
|
|
class EstimatorWithSetOutput(_SetOutputMixin):
|
|
def fit(self, X, y=None):
|
|
self.n_features_in_ = X.shape[1]
|
|
return self
|
|
|
|
def transform(self, X, y=None):
|
|
return X
|
|
|
|
def get_feature_names_out(self, input_features=None):
|
|
return np.asarray([f"X{i}" for i in range(self.n_features_in_)], dtype=object)
|
|
|
|
|
|
def test__safe_set_output():
|
|
"""Check _safe_set_output works as expected."""
|
|
|
|
# Estimator without transform will not raise when setting set_output for transform.
|
|
est = EstimatorWithoutSetOutputAndWithoutTransform()
|
|
_safe_set_output(est, transform="pandas")
|
|
|
|
# Estimator with transform but without set_output will raise
|
|
est = EstimatorNoSetOutputWithTransform()
|
|
with pytest.raises(ValueError, match="Unable to configure output"):
|
|
_safe_set_output(est, transform="pandas")
|
|
|
|
est = EstimatorWithSetOutput().fit(np.asarray([[1, 2, 3]]))
|
|
_safe_set_output(est, transform="pandas")
|
|
config = _get_output_config("transform", est)
|
|
assert config["dense"] == "pandas"
|
|
|
|
_safe_set_output(est, transform="default")
|
|
config = _get_output_config("transform", est)
|
|
assert config["dense"] == "default"
|
|
|
|
# transform is None is a no-op, so the config remains "default"
|
|
_safe_set_output(est, transform=None)
|
|
config = _get_output_config("transform", est)
|
|
assert config["dense"] == "default"
|
|
|
|
|
|
class EstimatorNoSetOutputWithTransformNoFeatureNamesOut(_SetOutputMixin):
|
|
def transform(self, X, y=None):
|
|
return X # pragma: no cover
|
|
|
|
|
|
def test_set_output_mixin():
|
|
"""Estimator without get_feature_names_out does not define `set_output`."""
|
|
est = EstimatorNoSetOutputWithTransformNoFeatureNamesOut()
|
|
assert not hasattr(est, "set_output")
|
|
|
|
|
|
def test__safe_set_output_error():
|
|
"""Check transform with invalid config."""
|
|
X = np.asarray([[1, 0, 3], [0, 0, 1]])
|
|
|
|
est = EstimatorWithSetOutput()
|
|
_safe_set_output(est, transform="bad")
|
|
|
|
msg = "output config must be in"
|
|
with pytest.raises(ValueError, match=msg):
|
|
est.transform(X)
|
|
|
|
|
|
@pytest.mark.parametrize("dataframe_lib", ["pandas", "polars"])
|
|
def test_set_output_method(dataframe_lib):
|
|
"""Check that the output is a dataframe."""
|
|
lib = pytest.importorskip(dataframe_lib)
|
|
|
|
X = np.asarray([[1, 0, 3], [0, 0, 1]])
|
|
est = EstimatorWithSetOutput().fit(X)
|
|
|
|
# transform=None is a no-op
|
|
est2 = est.set_output(transform=None)
|
|
assert est2 is est
|
|
X_trans_np = est2.transform(X)
|
|
assert isinstance(X_trans_np, np.ndarray)
|
|
|
|
est.set_output(transform=dataframe_lib)
|
|
|
|
X_trans_pd = est.transform(X)
|
|
|
|
assert isinstance(X_trans_pd, lib.DataFrame)
|
|
|
|
|
|
def test_set_output_method_error():
|
|
"""Check transform fails with invalid transform."""
|
|
|
|
X = np.asarray([[1, 0, 3], [0, 0, 1]])
|
|
est = EstimatorWithSetOutput().fit(X)
|
|
est.set_output(transform="bad")
|
|
|
|
msg = "output config must be in"
|
|
with pytest.raises(ValueError, match=msg):
|
|
est.transform(X)
|
|
|
|
|
|
@pytest.mark.parametrize("transform_output", ["pandas", "polars"])
|
|
def test__get_output_config(transform_output):
|
|
"""Check _get_output_config works as expected."""
|
|
|
|
# Without a configuration set, the global config is used
|
|
global_config = get_config()["transform_output"]
|
|
config = _get_output_config("transform")
|
|
assert config["dense"] == global_config
|
|
|
|
with config_context(transform_output=transform_output):
|
|
# with estimator=None, the global config is used
|
|
config = _get_output_config("transform")
|
|
assert config["dense"] == transform_output
|
|
|
|
est = EstimatorNoSetOutputWithTransform()
|
|
config = _get_output_config("transform", est)
|
|
assert config["dense"] == transform_output
|
|
|
|
est = EstimatorWithSetOutput()
|
|
# If estimator has not config, use global config
|
|
config = _get_output_config("transform", est)
|
|
assert config["dense"] == transform_output
|
|
|
|
# If estimator has a config, use local config
|
|
est.set_output(transform="default")
|
|
config = _get_output_config("transform", est)
|
|
assert config["dense"] == "default"
|
|
|
|
est.set_output(transform=transform_output)
|
|
config = _get_output_config("transform", est)
|
|
assert config["dense"] == transform_output
|
|
|
|
|
|
class EstimatorWithSetOutputNoAutoWrap(_SetOutputMixin, auto_wrap_output_keys=None):
|
|
def transform(self, X, y=None):
|
|
return X
|
|
|
|
|
|
def test_get_output_auto_wrap_false():
|
|
"""Check that auto_wrap_output_keys=None does not wrap."""
|
|
est = EstimatorWithSetOutputNoAutoWrap()
|
|
assert not hasattr(est, "set_output")
|
|
|
|
X = np.asarray([[1, 0, 3], [0, 0, 1]])
|
|
assert X is est.transform(X)
|
|
|
|
|
|
def test_auto_wrap_output_keys_errors_with_incorrect_input():
|
|
msg = "auto_wrap_output_keys must be None or a tuple of keys."
|
|
with pytest.raises(ValueError, match=msg):
|
|
|
|
class BadEstimator(_SetOutputMixin, auto_wrap_output_keys="bad_parameter"):
|
|
pass
|
|
|
|
|
|
class AnotherMixin:
|
|
def __init_subclass__(cls, custom_parameter, **kwargs):
|
|
super().__init_subclass__(**kwargs)
|
|
cls.custom_parameter = custom_parameter
|
|
|
|
|
|
def test_set_output_mixin_custom_mixin():
|
|
"""Check that multiple init_subclasses passes parameters up."""
|
|
|
|
class BothMixinEstimator(_SetOutputMixin, AnotherMixin, custom_parameter=123):
|
|
def transform(self, X, y=None):
|
|
return X
|
|
|
|
def get_feature_names_out(self, input_features=None):
|
|
return input_features
|
|
|
|
est = BothMixinEstimator()
|
|
assert est.custom_parameter == 123
|
|
assert hasattr(est, "set_output")
|
|
|
|
|
|
def test_set_output_mro():
|
|
"""Check that multi-inheritance resolves to the correct class method.
|
|
|
|
Non-regression test gh-25293.
|
|
"""
|
|
|
|
class Base(_SetOutputMixin):
|
|
def transform(self, X):
|
|
return "Base" # noqa
|
|
|
|
class A(Base):
|
|
pass
|
|
|
|
class B(Base):
|
|
def transform(self, X):
|
|
return "B"
|
|
|
|
class C(A, B):
|
|
pass
|
|
|
|
assert C().transform(None) == "B"
|
|
|
|
|
|
class EstimatorWithSetOutputIndex(_SetOutputMixin):
|
|
def fit(self, X, y=None):
|
|
self.n_features_in_ = X.shape[1]
|
|
return self
|
|
|
|
def transform(self, X, y=None):
|
|
import pandas as pd
|
|
|
|
# transform by giving output a new index.
|
|
return pd.DataFrame(X.to_numpy(), index=[f"s{i}" for i in range(X.shape[0])])
|
|
|
|
def get_feature_names_out(self, input_features=None):
|
|
return np.asarray([f"X{i}" for i in range(self.n_features_in_)], dtype=object)
|
|
|
|
|
|
def test_set_output_pandas_keep_index():
|
|
"""Check that set_output does not override index.
|
|
|
|
Non-regression test for gh-25730.
|
|
"""
|
|
pd = pytest.importorskip("pandas")
|
|
|
|
X = pd.DataFrame([[1, 2, 3], [4, 5, 6]], index=[0, 1])
|
|
est = EstimatorWithSetOutputIndex().set_output(transform="pandas")
|
|
est.fit(X)
|
|
|
|
X_trans = est.transform(X)
|
|
assert_array_equal(X_trans.index, ["s0", "s1"])
|
|
|
|
|
|
class EstimatorReturnTuple(_SetOutputMixin):
|
|
def __init__(self, OutputTuple):
|
|
self.OutputTuple = OutputTuple
|
|
|
|
def transform(self, X, y=None):
|
|
return self.OutputTuple(X, 2 * X)
|
|
|
|
|
|
def test_set_output_named_tuple_out():
|
|
"""Check that namedtuples are kept by default."""
|
|
Output = namedtuple("Output", "X, Y")
|
|
X = np.asarray([[1, 2, 3]])
|
|
est = EstimatorReturnTuple(OutputTuple=Output)
|
|
X_trans = est.transform(X)
|
|
|
|
assert isinstance(X_trans, Output)
|
|
assert_array_equal(X_trans.X, X)
|
|
assert_array_equal(X_trans.Y, 2 * X)
|
|
|
|
|
|
class EstimatorWithListInput(_SetOutputMixin):
|
|
def fit(self, X, y=None):
|
|
assert isinstance(X, list)
|
|
self.n_features_in_ = len(X[0])
|
|
return self
|
|
|
|
def transform(self, X, y=None):
|
|
return X
|
|
|
|
def get_feature_names_out(self, input_features=None):
|
|
return np.asarray([f"X{i}" for i in range(self.n_features_in_)], dtype=object)
|
|
|
|
|
|
@pytest.mark.parametrize("dataframe_lib", ["pandas", "polars"])
|
|
def test_set_output_list_input(dataframe_lib):
|
|
"""Check set_output for list input.
|
|
|
|
Non-regression test for #27037.
|
|
"""
|
|
lib = pytest.importorskip(dataframe_lib)
|
|
|
|
X = [[0, 1, 2, 3], [4, 5, 6, 7]]
|
|
est = EstimatorWithListInput()
|
|
est.set_output(transform=dataframe_lib)
|
|
|
|
X_out = est.fit(X).transform(X)
|
|
assert isinstance(X_out, lib.DataFrame)
|
|
assert_array_equal(X_out.columns, ["X0", "X1", "X2", "X3"])
|
|
|
|
|
|
@pytest.mark.parametrize("name", sorted(ADAPTERS_MANAGER.adapters))
|
|
def test_adapter_class_has_interface(name):
|
|
"""Check adapters have the correct interface."""
|
|
assert isinstance(ADAPTERS_MANAGER.adapters[name], ContainerAdapterProtocol)
|
|
|
|
|
|
def test_check_library_installed(monkeypatch):
|
|
"""Check import error changed."""
|
|
orig_import_module = importlib.import_module
|
|
|
|
def patched_import_module(name):
|
|
if name == "pandas":
|
|
raise ImportError()
|
|
orig_import_module(name, package=None)
|
|
|
|
monkeypatch.setattr(importlib, "import_module", patched_import_module)
|
|
|
|
msg = "Setting output container to 'pandas' requires"
|
|
with pytest.raises(ImportError, match=msg):
|
|
check_library_installed("pandas")
|
|
|
|
|
|
def test_get_adapter_from_container():
|
|
"""Check the behavior fo `_get_adapter_from_container`."""
|
|
pd = pytest.importorskip("pandas")
|
|
X = pd.DataFrame({"a": [1, 2, 3], "b": [10, 20, 100]})
|
|
adapter = _get_adapter_from_container(X)
|
|
assert adapter.container_lib == "pandas"
|
|
err_msg = "The container does not have a registered adapter in scikit-learn."
|
|
with pytest.raises(ValueError, match=err_msg):
|
|
_get_adapter_from_container(X.to_numpy())
|