Preserve CSR storage format when input is CSR in sparse_center_data

Author: MechCoder

sklearn/linear_model/base.py

@@ -27,6 +27,8 @@
 from ..utils import as_float_array, atleast2d_or_csr, safe_asarray
 from ..utils.extmath import safe_sparse_dot
 from ..utils.sparsefuncs import (csc_mean_variance_axis0,
+                                 csr_mean_variance_axis0,
+                                 inplace_csr_column_scale,
                                  inplace_csc_column_scale)
 from .cd_fast import sparse_std
 
@@ -50,15 +52,25 @@ def sparse_center_data(X, y, fit_intercept, normalize=False):
 
     if fit_intercept:
         X_data = X.data
-        # copy if 'normalize' is True or X is not a csc matrix
-        X = sp.csc_matrix(X, copy=normalize)
-        X_mean, X_std = csc_mean_variance_axis0(X)
+
+        # we might require not to change the csr matrix sometimes
+        # store a copy if normalize is True.
+        if sp.isspmatrix(X) and X.getformat() == 'csr':
+            X = sp.csr_matrix(X, copy=normalize)
+            sparse_mean_var_axis0 = csr_mean_variance_axis0
+            sparse_column_scale = inplace_csr_column_scale 
+        else:
+            X = sp.csc_matrix(X, copy=normalize)
+            sparse_mean_var_axis0 = csc_mean_variance_axis0
+            sparse_column_scale = inplace_csc_column_scale
+
+        X_mean, X_std = sparse_mean_var_axis0(X)
         if normalize:
             X_std = sparse_std(
                 X.shape[0], X.shape[1],
                 X_data, X.indices, X.indptr, X_mean)
             X_std[X_std == 0] = 1
-            inplace_csc_column_scale(X, 1. / X_std)
+            sparse_column_scale(X, 1. / X_std)
         else:
             X_std = np.ones(X.shape[1])
         y_mean = y.mean(axis=0)

sklearn/linear_model/tests/test_base.py

@@ -9,7 +9,7 @@
 from sklearn.utils.testing import assert_array_almost_equal
 from sklearn.utils.testing import assert_equal
 
-from sklearn.linear_model.base import LinearRegression
+from sklearn.linear_model.base import LinearRegression, sparse_center_data
 from sklearn.utils import check_random_state
 from sklearn.datasets.samples_generator import make_sparse_uncorrelated
 from sklearn.datasets.samples_generator import make_regression
@@ -110,3 +110,12 @@ def test_linear_regression_sparse_multiple_outcome(random_state=0):
     ols.fit(X, y.ravel())
     y_pred = ols.predict(X)
     assert_array_almost_equal(np.vstack((y_pred, y_pred)).T, Y_pred, decimal=3)
+
+
+def test_csr_sparse_center_data():
+    """Test output format of sparse_center_data, when input is csr"""
+    X, y = make_regression()
+    X[X < 2.5] = 0.0
+    csr = sparse.csr_matrix(X)
+    csr_, y, _, _, _ = sparse_center_data(csr, y, True)
+    assert_equal(csr_.getformat(), 'csr')