Merge pull request scikit-learn#2938 from jnothman/clean_impute

[MRG+1] some clean-up in Imputer, particularly in calculation of sparse median

Author: GaelVaroquaux

sklearn/preprocessing/imputation.py

@@ -33,42 +33,32 @@ def _get_mask(X, value_to_mask):
         return X == value_to_mask
 
 
-def _get_median(negative_elements, n_zeros, positive_elements):
-    """Compute the median of the array formed by negative_elements,
-       n_zeros zeros and positive_elements. This function is used
-       to support sparse matrices."""
-    negative_elements = np.sort(negative_elements, kind='heapsort')
-    positive_elements = np.sort(positive_elements, kind='heapsort')
-
-    n_elems = len(negative_elements) + n_zeros + len(positive_elements)
+def _get_median(data, n_zeros):
+    """Compute the median of data with n_zeros additional zeros.
+
+    This function is used to support sparse matrices; it modifies data in-place
+    """
+    n_elems = len(data) + n_zeros
     if not n_elems:
         return np.nan
+    n_negative = np.count_nonzero(data < 0)
+    middle, is_odd = divmod(n_elems, 2)
+    data.sort()
 
-    median_position = (n_elems - 1) / 2.0
+    if is_odd:
+        return _get_elem_at_rank(middle, data, n_negative, n_zeros)
 
-    if round(median_position) == median_position:
-        median = _get_elem_at_rank(negative_elements, n_zeros,
-                                   positive_elements, median_position)
-    else:
-        a = _get_elem_at_rank(negative_elements, n_zeros,
-                              positive_elements, math.floor(median_position))
-        b = _get_elem_at_rank(negative_elements, n_zeros,
-                              positive_elements, math.ceil(median_position))
-        median = (a + b) / 2.0
-
-    return median
-
-
-def _get_elem_at_rank(negative_elements, n_zeros, positive_elements, k):
-    """Compute the kth largest element of the array formed by
-       negative_elements, n_zeros zeros and positive_elements."""
-    len_neg = len(negative_elements)
-    if k < len_neg:
-        return negative_elements[k]
-    elif k >= len_neg + n_zeros:
-        return positive_elements[k - len_neg - n_zeros]
-    else:
+    return (_get_elem_at_rank(middle - 1, data, n_negative, n_zeros) +
+            _get_elem_at_rank(middle, data, n_negative, n_zeros)) / 2.
+
+
+def _get_elem_at_rank(rank, data, n_negative, n_zeros):
+    """Find the value in data augmented with n_zeros for the given rank"""
+    if rank < n_negative:
+        return data[rank]
+    if rank - n_negative < n_zeros:
         return 0
+    return data[rank - n_zeros]
 
 
 def _most_frequent(array, extra_value, n_repeat):
@@ -137,8 +127,8 @@ class Imputer(BaseEstimator, TransformerMixin):
 
     Attributes
     ----------
-    `statistics_` : array of shape (n_features,) or (n_samples,)
-        The statistics along the imputation axis.
+    `statistics_` : array of shape (n_features,)
+        The imputation fill value for each feature if axis == 0.
 
     Notes
     -----
@@ -211,7 +201,7 @@ def _sparse_fit(self, X, strategy, missing_values, axis):
 
         # Count the zeros
         if missing_values == 0:
-            n_zeros_axis = np.zeros(X.shape[not axis])
+            n_zeros_axis = np.zeros(X.shape[not axis], dtype=int)
         else:
             n_zeros_axis = X.shape[axis] - np.diff(X.indptr)
 
@@ -257,19 +247,15 @@ def _sparse_fit(self, X, strategy, missing_values, axis):
             mask_valids = np.hsplit(np.logical_not(mask_missing_values),
                                     X.indptr[1:-1])
 
-            columns = [col[mask.astype(np.bool)]
+            # astype necessary for bug in numpy.hsplit before v1.9
+            columns = [col[mask.astype(bool, copy=False)]
                        for col, mask in zip(columns_all, mask_valids)]
 
             # Median
             if strategy == "median":
                 median = np.empty(len(columns))
                 for i, column in enumerate(columns):
-
-                    negatives = column[column < 0]
-                    positives = column[column > 0]
-                    median[i] = _get_median(negatives,
-                                            n_zeros_axis[i],
-                                            positives)
+                    median[i] = _get_median(column, n_zeros_axis[i])
 
                 return median
 

sklearn/preprocessing/tests/test_imputation.py

@@ -43,8 +43,6 @@ def _check_statistics(X, X_true,
         assert_raises(ValueError, imputer.transform, X.copy().transpose())
     else:
         X_trans = imputer.transform(X.copy().transpose())
-        assert_array_equal(imputer.statistics_, statistics,
-                           err_msg.format(1, False))
         assert_array_equal(X_trans, X_true.transpose(),
                            err_msg.format(1, False))
 
@@ -72,8 +70,6 @@ def _check_statistics(X, X_true,
         if sparse.issparse(X_trans):
             X_trans = X_trans.toarray()
 
-        assert_array_equal(imputer.statistics_, statistics,
-                           err_msg.format(1, True))
         assert_array_equal(X_trans, X_true.transpose(),
                            err_msg.format(1, True))
 
@@ -109,16 +105,20 @@ def test_imputation_mean_median_only_zero():
     ])
     statistics_mean = [np.nan, 3, np.nan, np.nan, 7]
 
+    # Behaviour of median with NaN is undefined, e.g. different results in
+    # np.median and np.ma.median
+    X_for_median = X[:, [0, 1, 2, 4]]
     X_imputed_median = np.array([
-        [2, 5,  5],
-        [1, np.nan,  3],
-        [2, 5, 5],
-        [6, 5,  13],
+        [2, 5],
+        [1, 3],
+        [2, 5],
+        [6, 13],
     ])
-    statistics_median = [np.nan, 2, np.nan, 5, 5]
+    statistics_median = [np.nan, 2, np.nan, 5]
 
     _check_statistics(X, X_imputed_mean, "mean", statistics_mean, 0)
-    _check_statistics(X, X_imputed_median, "median", statistics_median, 0)
+    _check_statistics(X_for_median, X_imputed_median, "median",
+                      statistics_median, 0)
 
 
 def test_imputation_mean_median():
@@ -191,6 +191,36 @@ def test_imputation_mean_median():
                           true_statistics, test_missing_values)
 
 
+def test_imputation_median_special_cases():
+    """Test median imputation with sparse boundary cases
+    """
+    X = np.array([
+        [0, np.nan, np.nan],  # odd: implicit zero
+        [5, np.nan, np.nan],  # odd: explicit nonzero
+        [0, 0, np.nan],    # even: average two zeros
+        [-5, 0, np.nan],   # even: avg zero and neg
+        [0, 5, np.nan],    # even: avg zero and pos
+        [4, 5, np.nan],    # even: avg nonzeros
+        [-4, -5, np.nan],  # even: avg negatives
+        [-1, 2, np.nan],   # even: crossing neg and pos
+    ]).transpose()
+
+    X_imputed_median = np.array([
+        [0, 0, 0],
+        [5, 5, 5],
+        [0, 0, 0],
+        [-5, 0, -2.5],
+        [0, 5, 2.5],
+        [4, 5, 4.5],
+        [-4, -5, -4.5],
+        [-1, 2, .5],
+    ]).transpose()
+    statistics_median = [0, 5, 0, -2.5, 2.5, 4.5, -4.5, .5]
+
+    _check_statistics(X, X_imputed_median, "median",
+                      statistics_median, 'NaN')
+
+
 def test_imputation_most_frequent():
     """Test imputation using the most-frequent strategy."""
     X = np.array([