Merge pull request scikit-learn#2576 from ankit-maverick/issue2560

Fixing issue 2560 : Imputer bug with median and dense input

Author: agramfort

sklearn/preprocessing/imputation.py

@@ -304,7 +304,7 @@ def _dense_fit(self, X, strategy, missing_values, axis):
             median_masked = np.ma.median(masked_X, axis=axis)
             # Avoid the warning "Warning: converting a masked element to nan."
             median = np.ma.getdata(median_masked)
-            median[np.ma.getmask(median_masked)] = np.nan
+            median[np.ma.getmaskarray(median_masked)] = np.nan
 
             return median
 

sklearn/preprocessing/tests/test_imputation.py

@@ -1,6 +1,7 @@
 import numpy as np
 from scipy import sparse
 
+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_false
@@ -76,6 +77,19 @@ def _check_statistics(X, X_true,
                            err_msg.format(1, True))
 
 
+def test_imputation_shape():
+    """Verify the shapes of the imputed matrix for different strategies."""
+    X = np.random.randn(10, 2)
+    X[::2] = np.nan
+
+    for strategy in ['mean', 'median', 'most_frequent']:
+        imputer = Imputer(strategy=strategy)
+        X_imputed = imputer.fit_transform(X)
+        assert_equal(X_imputed.shape, (10, 2))
+        X_imputed = imputer.fit_transform(sparse.csr_matrix(X))
+        assert_equal(X_imputed.shape, (10, 2))
+
+
 def test_imputation_mean_median_only_zero():
     """Test imputation using the mean and median strategies, when
        missing_values == 0."""