Validate convert to np.nan; Check for inf

raghavrv · raghavrv · commit ecc5b47be85d · 2016-12-18T14:48:01.000+01:00
diff --git a/sklearn/preprocessing/imputation.py b/sklearn/preprocessing/imputation.py
@@ -27,7 +27,7 @@
 
 def _get_mask(X, value_to_mask):
     """Compute the boolean mask X == missing_values."""
-    if value_to_mask == "NaN" or np.isnan(value_to_mask):
+    if value_to_mask is np.nan:
         return np.isnan(X)
     else:
         return X == value_to_mask
@@ -148,22 +148,29 @@ def fit(self, X, y=None):
             raise ValueError("Can only impute missing values on axis 0 and 1, "
                              " got axis={0}".format(self.axis))
 
+        # Validate missing_values and convert from "NaN" to np.nan
+        if (isinstance(self.missing_values, six.string_types) and
+                self.missing_values == "NaN"):
+            missing_values = np.nan
+        else:
+            missing_values = self.missing_values
+
         # Since two different arrays can be provided in fit(X) and
         # transform(X), the imputation data will be computed in transform()
         # when the imputation is done per sample (i.e., when axis=1).
         if self.axis == 0:
             X = check_array(X, accept_sparse='csc', dtype=np.float64,
-                            force_all_finite=False)
+                            allow_nan=True, force_all_finite=True)
 
             if sparse.issparse(X):
                 self.statistics_ = self._sparse_fit(X,
                                                     self.strategy,
-                                                    self.missing_values,
+                                                    missing_values,
                                                     self.axis)
             else:
                 self.statistics_ = self._dense_fit(X,
                                                    self.strategy,
-                                                   self.missing_values,
+                                                   missing_values,
                                                    self.axis)
 
         return self
@@ -250,7 +257,7 @@ def _sparse_fit(self, X, strategy, missing_values, axis):
 
     def _dense_fit(self, X, strategy, missing_values, axis):
         """Fit the transformer on dense data."""
-        X = check_array(X, force_all_finite=False)
+        X = check_array(X, allow_nan=True, force_all_finite=True)
         mask = _get_mask(X, missing_values)
         masked_X = ma.masked_array(X, mask=mask)
 
@@ -307,10 +314,18 @@ def transform(self, X):
         X : {array-like, sparse matrix}, shape = [n_samples, n_features]
             The input data to complete.
         """
+        # Validate missing_values and convert from "NaN" to np.nan
+        if (isinstance(self.missing_values, six.string_types) and
+                self.missing_values == "NaN"):
+            missing_values = np.nan
+        else:
+            missing_values = self.missing_values
+
         if self.axis == 0:
             check_is_fitted(self, 'statistics_')
             X = check_array(X, accept_sparse='csc', dtype=FLOAT_DTYPES,
-                            force_all_finite=False, copy=self.copy)
+                            allow_nan=True, force_all_finite=True,
+                            copy=self.copy)
             statistics = self.statistics_
             if X.shape[1] != statistics.shape[0]:
                 raise ValueError("X has %d features per sample, expected %d"
@@ -321,18 +336,19 @@ def transform(self, X):
         # when the imputation is done per sample
         else:
             X = check_array(X, accept_sparse='csr', dtype=FLOAT_DTYPES,
-                            force_all_finite=False, copy=self.copy)
+                            allow_nan=True, force_all_finite=True,
+                            copy=self.copy)
 
             if sparse.issparse(X):
                 statistics = self._sparse_fit(X,
                                               self.strategy,
-                                              self.missing_values,
+                                              missing_values,
                                               self.axis)
 
             else:
                 statistics = self._dense_fit(X,
                                              self.strategy,
-                                             self.missing_values,
+                                             missing_values,
                                              self.axis)
 
         # Delete the invalid rows/columns
@@ -352,8 +368,8 @@ def transform(self, X):
                              "missing values: %s" % missing)
 
         # Do actual imputation
-        if sparse.issparse(X) and self.missing_values != 0:
-            mask = _get_mask(X.data, self.missing_values)
+        if sparse.issparse(X) and missing_values != 0:
+            mask = _get_mask(X.data, missing_values)
             indexes = np.repeat(np.arange(len(X.indptr) - 1, dtype=np.int),
                                 np.diff(X.indptr))[mask]
 
@@ -363,7 +379,7 @@ def transform(self, X):
             if sparse.issparse(X):
                 X = X.toarray()
 
-            mask = _get_mask(X, self.missing_values)
+            mask = _get_mask(X, missing_values)
             n_missing = np.sum(mask, axis=self.axis)
             values = np.repeat(valid_statistics, n_missing)
 
diff --git a/sklearn/preprocessing/tests/test_imputation.py b/sklearn/preprocessing/tests/test_imputation.py
@@ -5,6 +5,7 @@
 from sklearn.utils.testing import assert_equal
 from sklearn.utils.testing import assert_array_equal
 from sklearn.utils.testing import assert_raises
+from sklearn.utils.testing import assert_raise_message
 from sklearn.utils.testing import assert_false
 
 from sklearn.preprocessing.imputation import Imputer
@@ -357,3 +358,12 @@ def test_imputation_copy():
 
     # Note: If X is sparse and if missing_values=0, then a (dense) copy of X is
     # made, even if copy=False.
+
+    # Raise a proper error message if input contains infinity
+    X = [[np.inf, 8, 9, np.nan], [np.nan, 10, 10, 0], [10, 11, 9, 11]]
+    assert_raise_message(ValueError, "Input contains infinity",
+                         Imputer(axis=0, missing_values="NaN").fit_transform,
+                         X)
+    assert_raise_message(ValueError, "Input contains infinity",
+                         Imputer(axis=1, missing_values=np.nan).fit_transform,
+                         X)
