TST Better tests in ridge float64 upcasting (scikit-learn#9125)

massich · Jeremiah Johnson · commit 3805b8f3a89a · 2017-12-18T15:17:07.000-05:00
Also invert the solvers check to highlight the fact that new solver should support both 32 bit and 64 bit float by default from now on.
diff --git a/sklearn/linear_model/ridge.py b/sklearn/linear_model/ridge.py
@@ -472,10 +472,11 @@ def __init__(self, alpha=1.0, fit_intercept=True, normalize=False,
 
     def fit(self, X, y, sample_weight=None):
 
-        if self.solver in ['svd', 'sparse_cg', 'cholesky', 'lsqr']:
-            _dtype = [np.float64, np.float32]
-        else:
+        if self.solver in ('sag', 'saga'):
             _dtype = np.float64
+        else:
+            # all other solvers work at both float precision levels
+            _dtype = [np.float64, np.float32]
 
         X, y = check_X_y(X, y, ['csr', 'csc', 'coo'], dtype=_dtype,
                          multi_output=True, y_numeric=True)
diff --git a/sklearn/linear_model/tests/test_ridge.py b/sklearn/linear_model/tests/test_ridge.py
@@ -813,12 +813,12 @@ def test_dtype_match():
         ridge_64.fit(X_64, y_64)
         coef_64 = ridge_64.coef_
 
-        # Do all the checks at once, like this is easier to debug
-        assert_almost_equal(ridge_32.coef_, ridge_64.coef_, decimal=5)
-
         # Do the actual checks at once for easier debug
-        assert_equal(coef_32.dtype, X_32.dtype)
-        assert_equal(coef_64.dtype, X_64.dtype)
+        assert coef_32.dtype == X_32.dtype
+        assert coef_64.dtype == X_64.dtype
+        assert ridge_32.predict(X_32).dtype == X_32.dtype
+        assert ridge_64.predict(X_64).dtype == X_64.dtype
+        assert_almost_equal(ridge_32.coef_, ridge_64.coef_, decimal=5)
 
 
 def test_dtype_match_cholesky():
@@ -844,6 +844,8 @@ def test_dtype_match_cholesky():
     coef_64 = ridge_64.coef_
 
     # Do all the checks at once, like this is easier to debug
-    assert_equal(coef_32.dtype, X_32.dtype)
-    assert_equal(coef_64.dtype, X_64.dtype)
+    assert coef_32.dtype == X_32.dtype
+    assert coef_64.dtype == X_64.dtype
+    assert ridge_32.predict(X_32).dtype == X_32.dtype
+    assert ridge_64.predict(X_64).dtype == X_64.dtype
     assert_almost_equal(ridge_32.coef_, ridge_64.coef_, decimal=5)
