Repository: scikit-learn/scikit-learn
Pull request: ENH Adds feature_names_out to impute module
<!--
Thanks for contributing a pull request! Please ensure you have taken a look at
the contribution guidelines: https://github.com/scikit-learn/scikit-learn/blob/main/CONTRIBUTING.md
-->

#### Reference Issues/PRs
<!--
Example: Fixes #1234. See also #3456.
Please use keywords (e.g., Fixes) to create link to the issues or pull requests
you resolved, so that they will automatically be closed when your pull request
is merged. See https://github.com/blog/1506-closing-issues-via-pull-requests
-->
Continues https://github.com/scikit-learn/scikit-learn/pull/18444
Fixes https://github.com/scikit-learn/scikit-learn/issues/21200

#### What does this implement/fix? Explain your changes.
Adds `feature_names_out` to the `impute` module.

#### Any other comments?
There is an edge case I am concerned about:

```py
from sklearn.impute import SimpleImputer
import pandas as pd
import numpy as np

marker = np.nan
X = np.array(
    [
        [marker, 2],
        [2, marker],
        [6, 3],
        [1, 2],
    ]
)
X_df = pd.DataFrame(X, columns=["a", "indicator_a"])

imputer = SimpleImputer(add_indicator=True)
imputer.fit_transform(X_df)
imputer.get_feature_names_out()
# array(['a', 'indicator_a', 'indicator_a', 'indicator_indicator_a'], dtype=object)
```

For this edge case, we can add `sklearn` to the prefix resulting in `sklearn_indicator_a`?

<!--
Please be aware that we are a loose team of volunteers so patience is
necessary; assistance handling other issues is very welcome. We value
all user contributions, no matter how minor they are. If we are slow to
review, either the pull request needs some benchmarking, tinkering,
convincing, etc. or more likely the reviewers are simply busy. In either
case, we ask for your understanding during the review process.
For more information, see our FAQ on this topic:
http://scikit-learn.org/dev/faq.html#why-is-my-pull-request-not-getting-any-attention.

Thanks for contributing!
-->


Problem statement:
Implement get_feature_names_out for SimpleImputer
### Describe the workflow you want to enable

I would like to get the feature names input to a classifier after imputation. 

```
column_trans = ColumnTransformer([('imputed', SimpleImputer(add_indicator=True), x_con)], remainder='drop')
# column_trans = ColumnTransformer([('scaler', preprocessing.StandardScaler(), x_con)], remainder='drop')
pipe = make_pipeline(column_trans, LogisticRegressionCV(cv=2, random_state=0))
pipe.fit(X, y)
pipe[:-1].get_feature_names_out()
```

### Describe your proposed solution

Implement get_feature_names_out for SimpleImputer as it is implemented for other transformers e.g. it works for StandardScaler.

### Describe alternatives you've considered, if relevant

_No response_

### Additional context

_No response_

Hints from issue discussion:

Commit message: BUG Remove use for dtype in concatenate

Patch to review:
diff --git a/doc/whats_new/v1.1.rst b/doc/whats_new/v1.1.rst
index 9c1084e393e8d..2b1740e6c6483 100644
--- a/doc/whats_new/v1.1.rst
+++ b/doc/whats_new/v1.1.rst
@@ -38,6 +38,13 @@ Changelog
     :pr:`123456` by :user:`Joe Bloggs <joeongithub>`.
     where 123456 is the *pull request* number, not the issue number.
 
+:mod:`sklearn.impute`
+.....................
+
+- |API| Adds :meth:`get_feature_names_out` to :class:`impute.SimpleImputer`,
+  :class:`impute.KNNImputer`, :class:`impute.IterativeImputer`, and
+  :class:`impute.MissingIndicator`. :pr:`21078` by `Thomas Fan`_.
+
 :mod:`sklearn.utils`
 ....................
 
diff --git a/sklearn/impute/_base.py b/sklearn/impute/_base.py
index 13aa7a1f19c35..5c1f9805d67a2 100644
--- a/sklearn/impute/_base.py
+++ b/sklearn/impute/_base.py
@@ -15,6 +15,7 @@
 from ..utils.sparsefuncs import _get_median
 from ..utils.validation import check_is_fitted
 from ..utils.validation import FLOAT_DTYPES
+from ..utils.validation import _check_feature_names_in
 from ..utils._mask import _get_mask
 from ..utils import is_scalar_nan
 
@@ -113,6 +114,14 @@ def _concatenate_indicator(self, X_imputed, X_indicator):
 
         return hstack((X_imputed, X_indicator))
 
+    def _concatenate_indicator_feature_names_out(self, names, input_features):
+        if not self.add_indicator:
+            return names
+
+        indicator_names = self.indicator_.get_feature_names_out(input_features)
+        indicator_names_with_prefix = [f"indicator_{name}" for name in indicator_names]
+        return np.concatenate([names, indicator_names_with_prefix])
+
     def _more_tags(self):
         return {"allow_nan": is_scalar_nan(self.missing_values)}
 
@@ -593,6 +602,30 @@ def inverse_transform(self, X):
         X_original[full_mask] = self.missing_values
         return X_original
 
+    def get_feature_names_out(self, input_features=None):
+        """Get output feature names for transformation.
+
+        Parameters
+        ----------
+        input_features : array-like of str or None, default=None
+            Input features.
+
+            - If `input_features` is `None`, then `feature_names_in_` is
+                used as feature names in. If `feature_names_in_` is not defined,
+                then names are generated: `[x0, x1, ..., x(n_features_in_)]`.
+            - If `input_features` is an array-like, then `input_features` must
+                match `feature_names_in_` if `feature_names_in_` is defined.
+
+        Returns
+        -------
+        feature_names_out : ndarray of str objects
+            Transformed feature names.
+        """
+        input_features = _check_feature_names_in(self, input_features)
+        non_missing_mask = np.logical_not(_get_mask(self.statistics_, np.nan))
+        names = input_features[non_missing_mask]
+        return self._concatenate_indicator_feature_names_out(names, input_features)
+
 
 class MissingIndicator(TransformerMixin, BaseEstimator):
     """Binary indicators for missing values.
@@ -913,6 +946,28 @@ def fit_transform(self, X, y=None):
 
         return imputer_mask
 
+    def get_feature_names_out(self, input_features=None):
+        """Get output feature names for transformation.
+
+        Parameters
+        ----------
+        input_features : array-like of str or None, default=None
+            Input features.
+
+            - If `input_features` is `None`, then `feature_names_in_` is
+              used as feature names in. If `feature_names_in_` is not defined,
+              then names are generated: `[x0, x1, ..., x(n_features_in_)]`.
+            - If `input_features` is an array-like, then `input_features` must
+              match `feature_names_in_` if `feature_names_in_` is defined.
+
+        Returns
+        -------
+        feature_names_out : ndarray of str objects
+            Transformed feature names.
+        """
+        input_features = _check_feature_names_in(self, input_features)
+        return input_features[self.features_]
+
     def _more_tags(self):
         return {
             "allow_nan": True,
diff --git a/sklearn/impute/_iterative.py b/sklearn/impute/_iterative.py
index d95e78ccd36d3..25365261197c9 100644
--- a/sklearn/impute/_iterative.py
+++ b/sklearn/impute/_iterative.py
@@ -10,6 +10,7 @@
 from ..preprocessing import normalize
 from ..utils import check_array, check_random_state, _safe_indexing, is_scalar_nan
 from ..utils.validation import FLOAT_DTYPES, check_is_fitted
+from ..utils.validation import _check_feature_names_in
 from ..utils._mask import _get_mask
 
 from ._base import _BaseImputer
@@ -773,3 +774,26 @@ def fit(self, X, y=None):
         """
         self.fit_transform(X)
         return self
+
+    def get_feature_names_out(self, input_features=None):
+        """Get output feature names for transformation.
+
+        Parameters
+        ----------
+        input_features : array-like of str or None, default=None
+            Input features.
+
+            - If `input_features` is `None`, then `feature_names_in_` is
+                used as feature names in. If `feature_names_in_` is not defined,
+                then names are generated: `[x0, x1, ..., x(n_features_in_)]`.
+            - If `input_features` is an array-like, then `input_features` must
+                match `feature_names_in_` if `feature_names_in_` is defined.
+
+        Returns
+        -------
+        feature_names_out : ndarray of str objects
+            Transformed feature names.
+        """
+        input_features = _check_feature_names_in(self, input_features)
+        names = self.initial_imputer_.get_feature_names_out(input_features)
+        return self._concatenate_indicator_feature_names_out(names, input_features)
diff --git a/sklearn/impute/_knn.py b/sklearn/impute/_knn.py
index c2bd1410e8ecd..ad7e3537d445f 100644
--- a/sklearn/impute/_knn.py
+++ b/sklearn/impute/_knn.py
@@ -13,6 +13,7 @@
 from ..utils import is_scalar_nan
 from ..utils._mask import _get_mask
 from ..utils.validation import check_is_fitted
+from ..utils.validation import _check_feature_names_in
 
 
 class KNNImputer(_BaseImputer):
@@ -206,6 +207,7 @@ def fit(self, X, y=None):
         _check_weights(self.weights)
         self._fit_X = X
         self._mask_fit_X = _get_mask(self._fit_X, self.missing_values)
+        self._valid_mask = ~np.all(self._mask_fit_X, axis=0)
 
         super()._fit_indicator(self._mask_fit_X)
 
@@ -242,7 +244,7 @@ def transform(self, X):
 
         mask = _get_mask(X, self.missing_values)
         mask_fit_X = self._mask_fit_X
-        valid_mask = ~np.all(mask_fit_X, axis=0)
+        valid_mask = self._valid_mask
 
         X_indicator = super()._transform_indicator(mask)
 
@@ -327,3 +329,26 @@ def process_chunk(dist_chunk, start):
             pass
 
         return super()._concatenate_indicator(X[:, valid_mask], X_indicator)
+
+    def get_feature_names_out(self, input_features=None):
+        """Get output feature names for transformation.
+
+        Parameters
+        ----------
+        input_features : array-like of str or None, default=None
+            Input features.
+
+            - If `input_features` is `None`, then `feature_names_in_` is
+                used as feature names in. If `feature_names_in_` is not defined,
+                then names are generated: `[x0, x1, ..., x(n_features_in_)]`.
+            - If `input_features` is an array-like, then `input_features` must
+                match `feature_names_in_` if `feature_names_in_` is defined.
+
+        Returns
+        -------
+        feature_names_out : ndarray of str objects
+            Transformed feature names.
+        """
+        input_features = _check_feature_names_in(self, input_features)
+        names = input_features[self._valid_mask]
+        return self._concatenate_indicator_feature_names_out(names, input_features)
diff --git a/sklearn/impute/tests/test_common.py b/sklearn/impute/tests/test_common.py
index c35245ac8c253..f000aedba5a8c 100644
--- a/sklearn/impute/tests/test_common.py
+++ b/sklearn/impute/tests/test_common.py
@@ -122,3 +122,43 @@ def test_imputers_pandas_na_integer_array_support(imputer, add_indicator):
     X_trans = imputer.fit_transform(X_df)
 
     assert_allclose(X_trans_expected, X_trans)
+
+
+# ConvergenceWarning will be raised by the IterativeImputer
+@pytest.mark.filterwarnings("ignore::sklearn.exceptions.ConvergenceWarning")
+@pytest.mark.parametrize("imputer", IMPUTERS, ids=lambda x: x.__class__.__name__)
+@pytest.mark.parametrize("add_indicator", [True, False])
+def test_imputers_feature_names_out_pandas(imputer, add_indicator):
+    """Check feature names out for imputers."""
+    pd = pytest.importorskip("pandas")
+    marker = np.nan
+    imputer = imputer.set_params(add_indicator=add_indicator, missing_values=marker)
+
+    X = np.array(
+        [
+            [marker, 1, 5, marker, 1],
+            [2, marker, 1, marker, 2],
+            [6, 3, marker, marker, 3],
+            [1, 2, 9, marker, 4],
+        ]
+    )
+    X_df = pd.DataFrame(X, columns=["a", "b", "c", "d", "e"])
+    imputer.fit(X_df)
+
+    names = imputer.get_feature_names_out()
+
+    if add_indicator:
+        expected_names = [
+            "a",
+            "b",
+            "c",
+            "e",
+            "indicator_a",
+            "indicator_b",
+            "indicator_c",
+            "indicator_d",
+        ]
+        assert_array_equal(expected_names, names)
+    else:
+        expected_names = ["a", "b", "c", "e"]
+        assert_array_equal(expected_names, names)
diff --git a/sklearn/tests/test_common.py b/sklearn/tests/test_common.py
index 4f6818081c67d..139a2bfb9702a 100644
--- a/sklearn/tests/test_common.py
+++ b/sklearn/tests/test_common.py
@@ -365,7 +365,6 @@ def test_pandas_column_name_consistency(estimator):
     "decomposition",
     "discriminant_analysis",
     "ensemble",
-    "impute",
     "isotonic",
     "kernel_approximation",
     "preprocessing",


Code review comments:
diff_hunk: @@ -113,6 +114,14 @@ def _concatenate_indicator(self, X_imputed, X_indicator):
 
         return hstack((X_imputed, X_indicator))
 
+    def _concatenate_indicator_feature_names_out(self, names, input_features):
+        if not self.add_indicator:
+            return names
+
+        indicator_names = self.indicator_.get_feature_names_out(input_features)
+        indicator_names_with_prefix = [f"indicator_{name}" for name in indicator_names]
path: sklearn/impute/_base.py
text: @user1:
I find the `indicator_` prefix not explicit enough. I would rather suggest `missing_` or `imputed_` to be specific that this is a missingness or imputation indicator column.

We could use `missing_indicator_` to be extra explicit but I have the feeling this would be necessary verybose.

@author:
If we follow the `SLEP007`, then we use `missingindicator_` as the prefix since it is the class name for the output of `MissingIndicator.get_feature_names_out`.

I updated the PR to reflect this, but I am open to changing it.

@user1:
One could argue that the fact that `SimpleImputer` uses a `MissingIndicator` instance internally is a private implementation detail so we would be free to name.

But since `MissingIndicator` is also a public transformer and should follow the SLEP007 convention by itself, I think it's fine (even if a bit verbose) to have the inputs also use `missing_indicator_` as the column name prefix for the sake of consistency and simplicity. So +1 for keeping it that way.

diff_hunk: @@ -913,6 +946,28 @@ def fit_transform(self, X, y=None):
 
         return imputer_mask
 
+    def get_feature_names_out(self, input_features=None):
+        """Get output feature names for transformation.
+
+        Parameters
+        ----------
+        input_features : array-like of str or None, default=None
+            Input features.
+
+            - If `input_features` is `None`, then `feature_names_in_` is
+              used as feature names in. If `feature_names_in_` is not defined,
+              then names are generated: `[x0, x1, ..., x(n_features_in_)]`.
+            - If `input_features` is an array-like, then `input_features` must
+              match `feature_names_in_` if `feature_names_in_` is defined.
+
+        Returns
+        -------
+        feature_names_out : ndarray of str objects
+            Transformed feature names.
+        """
+        input_features = _check_feature_names_in(self, input_features)
+        return input_features[self.features_]
path: sklearn/impute/_base.py
text: @user1:
What is our rule for when to prefix and when not? The indicator output column of `SimpleImputer` is prefixed with `"indicator_"`.

@author:
I do not think we have a rule. I went with `SimpleImputer` prefixing `"indicator_"` because it needed to distinguish between the imputed features and the indicator.

Looking at this again, I think we can go with `MissingIndcator` to add the prefix. This way `SimpleImputer` only needs to combine them.

@user2:
I also think this is fine to preserve the original names for the imputed features.