test with sparse data

giorgiop · giorgiop · commit 2c2897a5e06a · 2015-10-23T11:56:50.000+02:00
diff --git a/sklearn/linear_model/tests/test_base.py b/sklearn/linear_model/tests/test_base.py
@@ -10,9 +10,11 @@
 from sklearn.utils.testing import assert_equal
 
 from sklearn.linear_model.base import LinearRegression
-from sklearn.linear_model.base import center_data, sparse_center_data, _rescale_data
+from sklearn.linear_model.base import center_data
+from sklearn.linear_model.base import sparse_center_data
+from sklearn.linear_model.base import _rescale_data
+from sklearn.utils.extmath import safe_sparse_dot
 from sklearn.utils import check_random_state
-from sklearn.utils.testing import assert_raise_message
 from sklearn.utils.testing import assert_greater
 from sklearn.datasets.samples_generator import make_sparse_uncorrelated
 from sklearn.datasets.samples_generator import make_regression
@@ -47,31 +49,33 @@ def test_linear_regression_sample_weights():
     rng = np.random.RandomState(0)
     n_samples, n_features = 6, 50  # over-determined system
 
-    for fit_intercept in [True, False]:
-        y = rng.randn(n_samples)
-        X = rng.randn(n_samples, n_features)
-        sample_weight = 1.0 + rng.rand(n_samples)
-
-        reg = LinearRegression(fit_intercept=fit_intercept)
-        reg.fit(X, y, sample_weight=sample_weight)
-        coefs1 = reg.coef_
-        intercept1 = reg.intercept_
-
-        assert_equal(reg.coef_.shape, (X.shape[1], ))
-        assert_greater(reg.score(X, y, sample_weight=sample_weight), 0.9)
-        assert_greater(r2_score(y, reg.predict(X),
-                       sample_weight=sample_weight), 0.9)  # same as above
-
-        # Sample weight can be implemented via a simple rescaling
-        # for the square loss.
-        scaled_y = y * np.sqrt(sample_weight)
-        scaled_X = X * np.sqrt(sample_weight)[:, np.newaxis]
-        reg.fit(scaled_X, scaled_y)
-        coefs2 = reg.coef_
-        intercept2 = reg.intercept_
-
-        assert_array_almost_equal(coefs1, coefs2)
-        assert_array_almost_equal(intercept1, intercept2)
+    y = rng.randn(n_samples)
+    X_ = rng.randn(n_samples, n_features)
+    sample_weight = 1.0 + rng.rand(n_samples)
+
+    for X in [X_, sparse.csr_matrix(X_)]:
+        for fit_intercept in [True, False]:
+            reg = LinearRegression(fit_intercept=fit_intercept)
+            reg.fit(X, y, sample_weight=sample_weight)
+            coefs1 = reg.coef_
+            intercept1 = reg.intercept_
+
+            assert_equal(reg.coef_.shape, (X.shape[1], ))
+            assert_greater(reg.score(X, y, sample_weight=sample_weight), 0.9)
+            assert_greater(r2_score(y, reg.predict(X),
+                           sample_weight=sample_weight), 0.9)  # same as above
+
+            # Sample weight can be implemented via a simple rescaling
+            # for the square loss.
+            scaled_y = y * np.sqrt(sample_weight)
+            scaled_X = \
+                safe_sparse_dot(np.diag(np.sqrt(sample_weight)), X)
+            reg.fit(scaled_X, scaled_y)
+            coefs2 = reg.coef_
+            intercept2 = reg.intercept_
+
+            assert_array_almost_equal(coefs1, coefs2)
+            assert_array_almost_equal(intercept1, intercept2)
 
 
 def test_raises_value_error_if_sample_weights_greater_than_1d():
