ENH Add check for non binary variables in OneHotEncoder. (scikit-learn#16585)

cmarmo · thomasjpfan · glemaitre · web-flow · commit f763c614f441 · 2020-03-10T09:58:55.000+01:00
Co-authored-by: Thomas J Fan &lt;thomasjpfan@gmail.com&gt;
Co-authored-by: Guillaume Lemaitre &lt;g.lemaitre58@gmail.com&gt;
diff --git a/sklearn/preprocessing/_encoders.py b/sklearn/preprocessing/_encoders.py
@@ -225,13 +225,13 @@ class OneHotEncoder(_BaseEncoder):
         (if any).
 
     drop_idx_ : array of shape (n_features,)
-        ``drop_idx_[i]`` is the index in ``categories_[i]`` of the category to
-        be dropped for each feature.
-        ``drop_idx_[i] = -1`` if no category is to be dropped from the feature
-        with index ``i``, e.g. when `drop='if_binary'` and the feature isn't
-        binary
-
-        ``drop_idx_ = None`` if all the transformed features will be retained.
+        - ``drop_idx_[i]`` is the index in ``categories_[i]`` of the category
+          to be dropped for each feature.
+        - ``drop_idx_[i] = None`` if no category is to be dropped from the
+          feature with index ``i``, e.g. when `drop='if_binary'` and the
+          feature isn't binary.
+        - ``drop_idx_ = None`` if all the transformed features will be
+          retained.
 
     See Also
     --------
@@ -316,10 +316,10 @@ def _compute_drop_idx(self):
             return None
         elif isinstance(self.drop, str):
             if self.drop == 'first':
-                return np.zeros(len(self.categories_), dtype=np.int_)
+                return np.zeros(len(self.categories_), dtype=np.object)
             elif self.drop == 'if_binary':
-                return np.array([0 if len(cats) == 2 else -1
-                                for cats in self.categories_], dtype=np.int_)
+                return np.array([0 if len(cats) == 2 else None
+                                for cats in self.categories_], dtype=np.object)
             else:
                 msg = (
                     "Wrong input for parameter `drop`. Expected "
@@ -354,7 +354,8 @@ def _compute_drop_idx(self):
                 raise ValueError(msg)
             return np.array([np.where(cat_list == val)[0][0]
                              for (val, cat_list) in
-                             zip(self.drop, self.categories_)], dtype=np.int_)
+                             zip(self.drop, self.categories_)],
+                            dtype=np.object)
 
     def fit(self, X, y=None):
         """
@@ -421,7 +422,7 @@ def transform(self, X):
 
         n_samples, n_features = X_int.shape
 
-        if self.drop is not None:
+        if self.drop_idx_ is not None:
             to_drop = self.drop_idx_.copy()
             # We remove all the dropped categories from mask, and decrement all
             # categories that occur after them to avoid an empty column.
@@ -431,7 +432,7 @@ def transform(self, X):
                 n_cats = len(cats)
 
                 # drop='if_binary' but feature isn't binary
-                if to_drop[i] == -1:
+                if to_drop[i] is None:
                     # set to cardinality to not drop from X_int
                     to_drop[i] = n_cats
                     n_values.append(n_cats)
@@ -484,16 +485,14 @@ def inverse_transform(self, X):
 
         n_samples, _ = X.shape
         n_features = len(self.categories_)
-        if self.drop is None:
+        if self.drop_idx_ is None:
             n_transformed_features = sum(len(cats)
                                          for cats in self.categories_)
-        elif isinstance(self.drop, str) and self.drop == 'if_binary':
-            n_transformed_features = sum(1 if len(cats) == 2
-                                         else len(cats)
-                                         for cats in self.categories_)
         else:
-            n_transformed_features = sum(len(cats) - 1
-                                         for cats in self.categories_)
+            n_transformed_features = sum(
+                len(cats) - 1 if to_drop is not None else len(cats)
+                for cats, to_drop in zip(self.categories_, self.drop_idx_)
+            )
 
         # validate shape of passed X
         msg = ("Shape of the passed X data is not correct. Expected {0} "
@@ -509,7 +508,7 @@ def inverse_transform(self, X):
         found_unknown = {}
 
         for i in range(n_features):
-            if self.drop is None:
+            if self.drop_idx_ is None or self.drop_idx_[i] is None:
                 cats = self.categories_[i]
             else:
                 cats = np.delete(self.categories_[i], self.drop_idx_[i])
@@ -532,9 +531,9 @@ def inverse_transform(self, X):
                 if unknown.any():
                     found_unknown[i] = unknown
             # drop will either be None or handle_unknown will be error. If
-            # self.drop is not None, then we can safely assume that all of
+            # self.drop_idx_ is not None, then we can safely assume that all of
             # the nulls in each column are the dropped value
-            elif self.drop is not None:
+            elif self.drop_idx_ is not None:
                 dropped = np.asarray(sub.sum(axis=1) == 0).flatten()
                 if dropped.any():
                     X_tr[dropped, i] = self.categories_[i][self.drop_idx_[i]]
@@ -581,7 +580,7 @@ def get_feature_names(self, input_features=None):
         for i in range(len(cats)):
             names = [
                 input_features[i] + '_' + str(t) for t in cats[i]]
-            if self.drop is not None:
+            if self.drop_idx_ is not None and self.drop_idx_[i] is not None:
                 names.pop(self.drop_idx_[i])
             feature_names.extend(names)
 
diff --git a/sklearn/preprocessing/tests/test_encoders.py b/sklearn/preprocessing/tests/test_encoders.py
@@ -268,6 +268,22 @@ def test_one_hot_encoder_inverse_if_binary():
     assert_array_equal(ohe.inverse_transform(X_tr), X)
 
 
+# check that resetting drop option without refitting does not throw an error
+@pytest.mark.parametrize('drop', ['if_binary', 'first', None])
+@pytest.mark.parametrize('reset_drop', ['if_binary', 'first', None])
+def test_one_hot_encoder_drop_reset(drop, reset_drop):
+    X = np.array([['Male', 1],
+                  ['Female', 3],
+                  ['Female', 2]], dtype=object)
+    ohe = OneHotEncoder(drop=drop, sparse=False)
+    ohe.fit(X)
+    X_tr = ohe.transform(X)
+    feature_names = ohe.get_feature_names()
+    ohe.set_params(drop=reset_drop)
+    assert_array_equal(ohe.inverse_transform(X_tr), X)
+    assert_allclose(ohe.transform(X), X_tr)
+    assert_array_equal(ohe.get_feature_names(), feature_names)
+
 @pytest.mark.parametrize("method", ['fit', 'fit_transform'])
 @pytest.mark.parametrize("X", [
     [1, 2],
@@ -388,8 +404,9 @@ def test_one_hot_encoder_pandas():
 
 @pytest.mark.parametrize("drop, expected_names",
                          [('first', ['x0_c', 'x2_b']),
+                          ('if_binary', ['x0_c', 'x1_2', 'x2_b']),
                           (['c', 2, 'b'], ['x0_b', 'x2_a'])],
-                         ids=['first', 'manual'])
+                         ids=['first', 'binary', 'manual'])
 def test_one_hot_encoder_feature_names_drop(drop, expected_names):
     X = [['c', 2, 'a'],
          ['b', 2, 'b']]
@@ -409,7 +426,7 @@ def test_one_hot_encoder_drop_equals_if_binary():
     expected = np.array([[1., 0., 0., 1.],
                          [0., 1., 0., 0.],
                          [0., 0., 1., 1.]])
-    expected_drop_idx = np.array([-1, 0])
+    expected_drop_idx = np.array([None, 0])
 
     ohe = OneHotEncoder(drop='if_binary', sparse=False)
     result = ohe.fit_transform(X)
@@ -423,7 +440,7 @@ def test_one_hot_encoder_drop_equals_if_binary():
     expected = np.array([[1., 1.],
                          [0., 1.],
                          [0., 1.]])
-    expected_drop_idx = np.array([0, -1])
+    expected_drop_idx = np.array([0, None])
 
     ohe = OneHotEncoder(drop='if_binary', sparse=False)
     result = ohe.fit_transform(X)
@@ -662,9 +679,9 @@ def test_categories(density, drop):
         for drop_cat, drop_idx, cat_list in zip(drop,
                                                 ohe_test.drop_idx_,
                                                 ohe_test.categories_):
-            assert cat_list[drop_idx] == drop_cat
+            assert cat_list[int(drop_idx)] == drop_cat
     assert isinstance(ohe_test.drop_idx_, np.ndarray)
-    assert ohe_test.drop_idx_.dtype == np.int_
+    assert ohe_test.drop_idx_.dtype == np.object
 
 
 @pytest.mark.parametrize('Encoder', [OneHotEncoder, OrdinalEncoder])
