ENH add sparse parameter to OneHotEncoder

Author: kaushik94

sklearn/preprocessing/data.py

@@ -927,6 +927,9 @@ class OneHotEncoder(BaseEstimator, TransformerMixin):
     dtype : number type, default=np.float
         Desired dtype of output.
 
+    sparse : boolean, default=True
+        Will return sparse matrix if set True else will return an array.
+
     Attributes
     ----------
     `active_features_` : array
@@ -953,7 +956,7 @@ class OneHotEncoder(BaseEstimator, TransformerMixin):
     >>> enc.fit([[0, 0, 3], [1, 1, 0], [0, 2, 1], \
 [1, 0, 2]])  # doctest: +ELLIPSIS
     OneHotEncoder(categorical_features='all', dtype=<... 'float'>,
-           n_values='auto')
+           n_values='auto', sparse=True)
     >>> enc.n_values_
     array([2, 3, 4])
     >>> enc.feature_indices_
@@ -969,10 +972,11 @@ class OneHotEncoder(BaseEstimator, TransformerMixin):
       encoding of dictionary items or strings.
     """
     def __init__(self, n_values="auto", categorical_features="all",
-                 dtype=np.float):
+                 dtype=np.float, sparse=True):
         self.n_values = n_values
         self.categorical_features = categorical_features
         self.dtype = dtype
+        self.sparse = sparse
 
     def fit(self, X, y=None):
         """Fit OneHotEncoder to X.
@@ -1033,7 +1037,7 @@ def _fit_transform(self, X):
             out = out[:, active_features]
             self.active_features_ = active_features
 
-        return out
+        return out if self.sparse else out.toarray()
 
     def fit_transform(self, X, y=None):
         """Fit OneHotEncoder to X, then transform X.
@@ -1069,7 +1073,8 @@ def _transform(self, X):
                                 dtype=self.dtype).tocsr()
         if self.n_values == 'auto':
             out = out[:, self.active_features_]
-        return out
+
+        return out if self.sparse else out.toarray()
 
     def transform(self, X):
         """Transform X using one-hot encoding.
@@ -1081,7 +1086,7 @@ def transform(self, X):
 
         Returns
         -------
-        X_out : sparse matrix, dtype=int
+        X_out : sparse matrix if sparse=True else a 2-d array, dtype=int
             Transformed input.
         """
         return _transform_selected(X, self._transform,

sklearn/preprocessing/tests/test_data.py

@@ -622,7 +622,7 @@ def test_add_dummy_feature_csr():
     assert_array_equal(X.toarray(), [[1, 1, 0], [1, 0, 1], [1, 0, 1]])
 
 
-def test_one_hot_encoder():
+def test_one_hot_encoder_sparse():
     """Test OneHotEncoder's fit and transform."""
     X = [[3, 2, 1], [0, 1, 1]]
     enc = OneHotEncoder()
@@ -674,6 +674,22 @@ def test_one_hot_encoder():
     # test negative input to transform
     enc.fit([[0], [1]])
     assert_raises(ValueError, enc.transform, [[0], [-1]])
+    
+def test_one_hot_encoder_dense():    
+    """check for sparse=False"""
+    X = [[3, 2, 1], [0, 1, 1]]
+    enc = OneHotEncoder(sparse=False)
+    # discover max values automatically
+    X_trans = enc.fit_transform(X)
+    assert_equal(X_trans.shape, (2, 5))
+    assert_array_equal(enc.active_features_,
+                       np.where([1, 0, 0, 1, 0, 1, 1, 0, 1])[0])
+    assert_array_equal(enc.feature_indices_, [0, 4, 7, 9])
+
+    # check outcome
+    assert_array_equal(X_trans,
+                       np.array([[0., 1., 0., 1., 1.],
+                                 [1., 0., 1., 0., 1.]]))
 
 
 def _check_transform_selected(X, X_expected, sel):