REFACTOR: fixed huge code duplication in _RidgeGCV

dohmatob · agramfort · commit 8c44bb6a74a1 · 2015-10-27T21:05:09.000+01:00
diff --git a/sklearn/linear_model/ridge.py b/sklearn/linear_model/ridge.py
@@ -817,24 +817,28 @@ def _diag_dot(self, D, B):
             D = D[(slice(None), ) + (np.newaxis, ) * (len(B.shape) - 1)]
         return D * B
 
-    def _errors(self, alpha, y, v, Q, QT_y):
-        # don't construct matrix G, instead compute action on y & diagonal
+    def _errors_and_values_helper(self, alpha, y, v, Q, QT_y):
+        """Helper function to avoid code duplication between self._errors and
+        self._values.
+
+        Notes
+        -----
+        We don't construct matrix G, instead compute action on y & diagonal.
+        """
         w = 1.0 / (v + alpha)
         c = np.dot(Q, self._diag_dot(w, QT_y))
         G_diag = self._decomp_diag(w, Q)
         # handle case where y is 2-d
         if len(y.shape) != 1:
             G_diag = G_diag[:, np.newaxis]
+        return G_diag, c
+
+    def _errors(self, alpha, y, v, Q, QT_y):
+        G_diag, c = self._errors_and_values_helper(alpha, y, v, Q, QT_y)
         return (c / G_diag) ** 2, c
 
     def _values(self, alpha, y, v, Q, QT_y):
-        # don't construct matrix G, instead compute action on y & diagonal
-        w = 1.0 / (v + alpha)
-        c = np.dot(Q, self._diag_dot(w, QT_y))
-        G_diag = self._decomp_diag(w, Q)
-        # handle case where y is 2-d
-        if len(y.shape) != 1:
-            G_diag = G_diag[:, np.newaxis]
+        G_diag, c = self._errors_and_values_helper(alpha, y, v, Q, QT_y)
         return y - (c / G_diag), c
 
     def _pre_compute_svd(self, X, y):
@@ -845,22 +849,24 @@ def _pre_compute_svd(self, X, y):
         UT_y = np.dot(U.T, y)
         return v, U, UT_y
 
-    def _errors_svd(self, alpha, y, v, U, UT_y):
+    def _errors_and_values_svd_helper(self, alpha, y, v, U, UT_y):
+        """Helper function to avoid code duplication between self._errors_svd
+        and self._values_svd.
+        """
         w = ((v + alpha) ** -1) - (alpha ** -1)
         c = np.dot(U, self._diag_dot(w, UT_y)) + (alpha ** -1) * y
         G_diag = self._decomp_diag(w, U) + (alpha ** -1)
         if len(y.shape) != 1:
             # handle case where y is 2-d
             G_diag = G_diag[:, np.newaxis]
+        return G_diag, c
+
+    def _errors_svd(self, alpha, y, v, U, UT_y):
+        G_diag, c = self._errors_and_values_svd_helper(alpha, y, v, U, UT_y)
         return (c / G_diag) ** 2, c
 
     def _values_svd(self, alpha, y, v, U, UT_y):
-        w = ((v + alpha) ** -1) - (alpha ** -1)
-        c = np.dot(U, self._diag_dot(w, UT_y)) + (alpha ** -1) * y
-        G_diag = self._decomp_diag(w, U) + (alpha ** -1)
-        if len(y.shape) != 1:
-            # handle case when y is 2-d
-            G_diag = G_diag[:, np.newaxis]
+        G_diag, c = self._errors_and_values_svd_helper(alpha, y, v, U, UT_y)
         return y - (c / G_diag), c
 
     def fit(self, X, y, sample_weight=None):
diff --git a/sklearn/linear_model/tests/test_ridge.py b/sklearn/linear_model/tests/test_ridge.py
@@ -32,6 +32,7 @@
 from sklearn.model_selection import GridSearchCV
 from sklearn.model_selection import KFold
 
+from sklearn.utils import check_random_state
 
 diabetes = datasets.load_diabetes()
 X_diabetes, y_diabetes = diabetes.data, diabetes.target
@@ -711,3 +712,45 @@ def test_ridge_fit_intercept_sparse():
     assert_warns(UserWarning, sparse.fit, X_csr, y)
     assert_almost_equal(dense.intercept_, sparse.intercept_)
     assert_array_almost_equal(dense.coef_, sparse.coef_)
+
+
+def test_errors_and_values_helper():
+    ridgecv = _RidgeGCV()
+    rng = check_random_state(42)
+    alpha = 1.
+    n = 5
+    y = rng.randn(n)
+    v = rng.randn(n)
+    Q = rng.randn(len(v), len(v))
+    QT_y = Q.T.dot(y)
+    G_diag, c = ridgecv._errors_and_values_helper(alpha, y, v, Q, QT_y)
+
+    # test that helper function behaves as expected
+    out, c_ = ridgecv._errors(alpha, y, v, Q, QT_y)
+    np.testing.assert_array_equal(out, (c / G_diag) ** 2)
+    np.testing.assert_array_equal(c, c)
+
+    out, c_ = ridgecv._values(alpha, y, v, Q, QT_y)
+    np.testing.assert_array_equal(out, y - (c / G_diag))
+    np.testing.assert_array_equal(c_, c)
+
+
+def test_errors_and_values_svd_helper():
+    ridgecv = _RidgeGCV()
+    rng = check_random_state(42)
+    alpha = 1.
+    for n, p in zip((5, 10), (12, 6)):
+        y = rng.randn(n)
+        v = rng.randn(p)
+        U = rng.randn(n, p)
+        UT_y = U.T.dot(y)
+        G_diag, c = ridgecv._errors_and_values_svd_helper(alpha, y, v, U, UT_y)
+
+        # test that helper function behaves as expected
+        out, c_ = ridgecv._errors_svd(alpha, y, v, U, UT_y)
+        np.testing.assert_array_equal(out, (c / G_diag) ** 2)
+        np.testing.assert_array_equal(c, c)
+
+        out, c_ = ridgecv._values_svd(alpha, y, v, U, UT_y)
+        np.testing.assert_array_equal(out, y - (c / G_diag))
+        np.testing.assert_array_equal(c_, c)
