scikit-learn · NicolasHug · Apr 15, 2019 · Apr 4, 2019 · Apr 9, 2019 · Apr 9, 2019
diff --git a/doc/modules/impute.rst b/doc/modules/impute.rst
@@ -109,10 +109,12 @@ imputation round are returned.
     >>> from sklearn.impute import IterativeImputer
     >>> imp = IterativeImputer(max_iter=10, random_state=0)
     >>> imp.fit([[1, 2], [3, 6], [4, 8], [np.nan, 3], [7, np.nan]])  # doctest: +NORMALIZE_WHITESPACE
-    IterativeImputer(estimator=None, imputation_order='ascending',
-                     initial_strategy='mean', max_iter=10, max_value=None,
-                     min_value=None, missing_values=nan, n_nearest_features=None,
-                     random_state=0, sample_posterior=False, tol=0.001, verbose=0)
+    IterativeImputer(add_indicator=False, estimator=None,
+                     imputation_order='ascending', initial_strategy='mean',
+                     max_iter=10, max_value=None, min_value=None,
+                     missing_values=nan, n_nearest_features=None,
+                     random_state=0, sample_posterior=False, tol=0.001,
+                     verbose=0)
     >>> X_test = [[np.nan, 2], [6, np.nan], [np.nan, 6]]
     >>> # the model learns that the second feature is double the first
     >>> print(np.round(imp.transform(X_test)))

diff --git a/doc/whats_new/v0.21.rst b/doc/whats_new/v0.21.rst
@@ -255,11 +255,11 @@ Support for Python 3.4 and below has been officially dropped.
   used to be kept if there were no missing values at all. :issue:`13562` by
   :user:`Jérémie du Boisberranger <jeremiedbb>`.
 
-- |Feature| The :class:`impute.SimpleImputer` has a new parameter
-  ``'add_indicator'``, which simply stacks a :class:`impute.MissingIndicator`
-  transform into the output of the imputer's transform. That allows a predictive
-  estimator to account for missingness. :issue:`12583` by
-  :user:`Danylo Baibak <DanilBaibak>`.
+- |Feature| The :class:`impute.SimpleImputer` and :class:`IterativeImputer` have
+  a new parameter ``'add_indicator'``, which simply stacks a
+  :class:`impute.MissingIndicator` transform into the output of the imputer's
+  transform. That allows a predictive estimator to account for missingness.
+  :issue:`12583`, :issue:`13601` by :user:`Danylo Baibak <DanilBaibak>`.
 
 :mod:`sklearn.isotonic`
 .......................

diff --git a/sklearn/impute.py b/sklearn/impute.py
@@ -538,6 +538,14 @@ class IterativeImputer(BaseEstimator, TransformerMixin):
         ``sample_posterior`` is True. Use an integer for determinism.
         See :term:`the Glossary <random_state>`.
 
+    add_indicator : boolean, optional (default=False)
+        If True, a `MissingIndicator` transform will stack onto output
+        of the imputer's transform. This allows a predictive estimator
+        to account for missingness despite imputation. If a feature has no
+        missing values at fit/train time, the feature won't appear on
+        the missing indicator even if there are missing values at
+        transform/test time.
+
     Attributes
     ----------
     initial_imputer_ : object of type :class:`sklearn.impute.SimpleImputer`
@@ -558,6 +566,10 @@ class IterativeImputer(BaseEstimator, TransformerMixin):
     n_features_with_missing_ : int
         Number of features with missing values.
 
+    indicator_ : :class:`sklearn.impute.MissingIndicator`
+        Indicator used to add binary indicators for missing values.
+        ``None`` if add_indicator is False.
+
     See also
     --------
     SimpleImputer : Univariate imputation of missing values.
@@ -600,7 +612,8 @@ def __init__(self,
                  min_value=None,
                  max_value=None,
                  verbose=0,
-                 random_state=None):
+                 random_state=None,
+                 add_indicator=False):
 
         self.estimator = estimator
         self.missing_values = missing_values
@@ -614,6 +627,7 @@ def __init__(self,
         self.max_value = max_value
         self.verbose = verbose
         self.random_state = random_state
+        self.add_indicator = add_indicator
 
     def _impute_one_feature(self,
                             X_filled,
@@ -922,6 +936,13 @@ def fit_transform(self, X, y=None):
                 .format(self.tol)
             )
 
+        if self.add_indicator:
+            self.indicator_ = MissingIndicator(
+                missing_values=self.missing_values)
+            X_trans_indicator = self.indicator_.fit_transform(X)
+        else:
+            self.indicator_ = None
+
         if self.estimator is None:
             from .linear_model import BayesianRidge
             self._estimator = BayesianRidge()
@@ -995,6 +1016,9 @@ def fit_transform(self, X, y=None):
                 warnings.warn("[IterativeImputer] Early stopping criterion not"
                               " reached.", ConvergenceWarning)
         Xt[~mask_missing_values] = X[~mask_missing_values]
+
+        if self.add_indicator:
+            Xt = np.hstack((Xt, X_trans_indicator))
         return Xt
 
     def transform(self, X):
@@ -1015,6 +1039,9 @@ def transform(self, X):
         """
         check_is_fitted(self, 'initial_imputer_')
 
+        if self.add_indicator:
+            X_trans_indicator = self.indicator_.transform(X)
+
         X, Xt, mask_missing_values = self._initial_imputation(X)
 
         if self.n_iter_ == 0 or np.all(mask_missing_values):
@@ -1043,6 +1070,9 @@ def transform(self, X):
                 i_rnd += 1
 
         Xt[~mask_missing_values] = X[~mask_missing_values]
+
+        if self.add_indicator:
+            Xt = np.hstack((Xt, X_trans_indicator))
         return Xt
 
     def fit(self, X, y=None):

diff --git a/sklearn/tests/test_impute.py b/sklearn/tests/test_impute.py
@@ -1147,35 +1147,50 @@ def test_missing_indicator_sparse_no_explicit_zeros():
 
 
 @pytest.mark.parametrize("marker", [np.nan, -1, 0])
-def test_imputation_add_indicator(marker):
+@pytest.mark.parametrize("imputer_constructor",
+                         [SimpleImputer, IterativeImputer])
+def test_imputers_add_indicator(marker, imputer_constructor):
     X = np.array([
-        [marker, 1,      5,       marker, 1],
-        [2,      marker, 1,       marker, 2],
-        [6,      3,      marker,  marker, 3],
-        [1,      2,      9,       marker, 4]
+        [marker, 1,      5,      marker, 1],
+        [2,      marker, 1,      marker, 2],
+        [6,      3,      marker, marker, 3],
+        [1,      2,      9,      marker, 4]
     ])
-    X_true = np.array([
-        [3., 1., 5., 1., 1., 0., 0., 1.],
-        [2., 2., 1., 2., 0., 1., 0., 1.],
-        [6., 3., 5., 3., 0., 0., 1., 1.],
-        [1., 2., 9., 4., 0., 0., 0., 1.]
+    X_true_indicator = np.array([
+        [1., 0., 0., 1.],
+        [0., 1., 0., 1.],
+        [0., 0., 1., 1.],
+        [0., 0., 0., 1.]
     ])
+    imputer = imputer_constructor(missing_values=marker,
+                                  add_indicator=True)
 
-    imputer = SimpleImputer(missing_values=marker, add_indicator=True)
-    X_trans = imputer.fit_transform(X)
-
-    assert_allclose(X_trans, X_true)
+    X_trans = imputer.fit(X).transform(X)
+    # The test is for testing the indicator,
+    # that's why we're looking at the last 4 columns only.
+    assert_allclose(X_trans[:, -4:], X_true_indicator)
     assert_array_equal(imputer.indicator_.features_, np.array([0, 1, 2, 3]))
 
 
+@pytest.mark.parametrize("imputer_constructor",
+                         [SimpleImputer, IterativeImputer])
+def test_imputer_without_indicator(imputer_constructor):
+    X = np.array([[1, 1],
+                  [1, 1]])
+    imputer = imputer_constructor()
+    imputer.fit(X)
+
+    assert imputer.indicator_ is None
+
+
 @pytest.mark.parametrize(
     "arr_type",
     [
         sparse.csc_matrix, sparse.csr_matrix, sparse.coo_matrix,
         sparse.lil_matrix, sparse.bsr_matrix
     ]
 )
-def test_imputation_add_indicator_sparse_matrix(arr_type):
+def test_simple_imputation_add_indicator_sparse_matrix(arr_type):
     X_sparse = arr_type([
         [np.nan, 1, 5],
         [2, np.nan, 1],