get_feature_names_out for EstimatorTransformer

sklego/meta/estimator_transformer.py

@@ -27,6 +27,7 @@ def fit(self, X, y, **kwargs):
         """Fits the estimator"""
         X, y = check_X_y(X, y, estimator=self, dtype=FLOAT_DTYPES, multi_output=True)
         self.multi_output_ = len(y.shape) > 1
+        self.output_len_ = y.shape[1] if self.multi_output_ else 1
         self.estimator_ = clone(self.estimator)
         self.estimator_.fit(X, y, **kwargs)
         return self
@@ -38,6 +39,34 @@ def transform(self, X):
         Returns array of shape `(X.shape[0], )` if estimator is not multi output.
         For multi output estimators an array of shape `(X.shape[0], y.shape[1])` is returned.
         """
-        check_is_fitted(self, "estimator_")
+        check_is_fitted(self)
         output = getattr(self.estimator_, self.predict_func)(X)
         return output if self.multi_output_ else output.reshape(-1, 1)
+
+    def get_feature_names_out(self, feature_names_out=None) -> list:
+        """
+        Defines descriptive names for each output of the (fitted) estimator.
+        :param feature_names_out: Redundant parameter for which the contents are ignored in this function.
+        feature_names_out is defined here because EstimatorTransformer can be part of a larger complex pipeline.
+        Some components may depend on defined feature_names_out and some not, but it is passed to all components
+        in the pipeline if `Pipeline.get_feature_names_out` is called. feature_names_out is therefore necessary
+        to define here to avoid `TypeError`s when used within a scikit-learn `Pipeline` object.
+        :return: List of descriptive names for each output variable from the fitted estimator.
+        """
+        check_is_fitted(self)
+        estimator_name_lower = self.estimator_.__class__.__name__.lower()
+        if self.multi_output_:
+            feature_names = [f"{estimator_name_lower}_{i}" for i in range(self.output_len_)]
+        else:
+            feature_names = [estimator_name_lower]
+        return feature_names
+
+    def __sklearn_is_fitted(self) -> bool:
+        """
+        Custom additional requirements that need to be satisfied to pass check_is_fitted.
+        :return: Boolean indicating if the additional requirements
+        for determining check_is_fitted are satisfied.
+        """
+        has_fit_attr = all(hasattr(self, attr) for attr in ["multi_output_", "output_len_"])
+        check_is_fitted(self.estimator_)
+        return has_fit_attr

tests/test_meta/test_estimatortransformer.py

@@ -8,7 +8,7 @@
 from sklearn.linear_model import LinearRegression, Ridge
 from sklearn.pipeline import Pipeline, FeatureUnion
 from sklearn.utils import check_X_y
-
+from sklearn.exceptions import NotFittedError
 from sklego.common import flatten
 from sklego.meta import EstimatorTransformer
 from tests.conftest import transformer_checks, general_checks
@@ -111,3 +111,57 @@ def test_kwargs(patched_clone, random_xy_dataset_clf):
     np.testing.assert_array_equal(
         sample_weights, estimator.fit.call_args[1]['sample_weight']
     )
+
+
+def test_get_feature_names_out(random_xy_dataset_clf):
+    X, y = random_xy_dataset_clf
+    pipeline = EstimatorTransformer(LinearRegression())
+
+    # We shouldn't be able to call get_feature_names_out before estimator is fitted.
+    with pytest.raises(NotFittedError):
+        pipeline.get_feature_names_out()
+
+    pipeline.fit(X, y)
+    feature_names = pipeline.get_feature_names_out()
+
+    estimator_name_lower = pipeline.estimator.__class__.__name__.lower()
+    expected_feature_names = [estimator_name_lower]
+    np.testing.assert_array_equal(feature_names, expected_feature_names)
+
+
+def test_get_feature_names_out_multitarget(random_xy_dataset_multitarget):
+    X, y = random_xy_dataset_multitarget
+    pipeline = EstimatorTransformer(LinearRegression())
+
+    pipeline.fit(X, y)
+    feature_names = pipeline.get_feature_names_out()
+
+    estimator_name_lower = pipeline.estimator.__class__.__name__.lower()
+    expected_feature_names = [f"{estimator_name_lower}_{i}" for i in range(pipeline.output_len_)]
+    np.testing.assert_array_equal(feature_names, expected_feature_names)
+
+
+def test_get_feature_names_out_featureunion(random_xy_dataset_clf):
+    X, y = random_xy_dataset_clf
+    pipeline = Pipeline(
+        [
+            (
+                "ml_features",
+                FeatureUnion(
+                    [
+                        ("model_1", EstimatorTransformer(LinearRegression())),
+                        ("model_2", EstimatorTransformer(Ridge())),
+                    ]
+                ),
+            )
+        ]
+    )
+
+    # We shouldn't be able to call get_feature_names_out before estimator is fitted.
+    with pytest.raises(NotFittedError):
+        pipeline.get_feature_names_out()
+
+    pipeline.fit(X, y)
+    feature_names = pipeline.get_feature_names_out()
+    expected_feature_names = ["model_1__linearregression", "model_2__ridge"]
+    np.testing.assert_array_equal(feature_names, expected_feature_names)