writing tests

Author: PABannier

skglm/datafits/single_task.py

@@ -63,15 +63,13 @@ def initialize_sparse(
         for j in range(n_features):
             nrm2 = 0.
             xty = 0
-            x2 = 0.
             for idx in range(X_indptr[j], X_indptr[j + 1]):
                 nrm2 += X_data[idx] ** 2
                 xty += X_data[idx] * y[X_indices[idx]]
-                x2 += X_data[idx] ** 2 / len(y)
 
             self.lipschitz[j] = nrm2 / len(y)
             self.Xty[j] = xty
-            self.global_lipschitz += x2
+            self.global_lipschitz += nrm2 / len(y)
 
     def value(self, y, w, Xw):
         return np.sum((y - Xw) ** 2) / (2 * len(Xw))
@@ -233,10 +231,9 @@ def params_to_dict(self):
     def initialize(self, yXT, y):
         n_features = yXT.shape[1]
         self.lipschitz = np.zeros(n_features, dtype=yXT.dtype)
-        self.global_lipschitz = 0.
+        self.global_lipschitz = norm(yXT, ord=2) ** 2 / len(y)
         for j in range(n_features):
             self.lipschitz[j] = norm(yXT[:, j]) ** 2
-            self.global_lipschitz += norm(yXT[:, j]) ** 2
 
     def initialize_sparse(
             self, yXT_data, yXT_indptr, yXT_indices, y):

skglm/solvers/fista.py

@@ -13,14 +13,35 @@ def _prox_vec(w, z, penalty, lipschitz):
 
 
 class FISTA(BaseSolver):
-    r"""ISTA solver with Nesterov acceleration (FISTA)."""
+    r"""ISTA solver with Nesterov acceleration (FISTA).
 
-    def __init__(self, max_iter=100, tol=1e-4, fit_intercept=False, warm_start=False,
-                 opt_freq=10, verbose=0):
+    This solver implements accelerated proximal gradient descent for linear problems.
+
+    Attributes
+    ----------
+    max_iter : int, default 100
+        Maximum number of iterations.
+
+    tol : float, default 1e-4
+        Tolerance for convergence.
+
+    opt_freq : int, default 10
+        Frequency for optimality condition check.
+
+    verbose : bool, default False
+        Amount of verbosity. 0/False is silent.
+
+    References
+    ----------
+    .. [1] Beck, A. and Teboulle M.
+           "A Fast Iterative Shrinkage-Thresholding Algorithm for Linear Inverse
+           problems", 2009, SIAM J. Imaging Sci.
+           https://epubs.siam.org/doi/10.1137/080716542
+    """
+
+    def __init__(self, max_iter=100, tol=1e-4, opt_freq=10, verbose=0):
         self.max_iter = max_iter
         self.tol = tol
-        self.fit_intercept = fit_intercept
-        self.warm_start = warm_start
         self.opt_freq = opt_freq
         self.verbose = verbose
 
@@ -33,8 +54,11 @@ def solve(self, X, y, datafit, penalty, w_init=None, Xw_init=None):
         z = w_init.copy() if w_init is not None else np.zeros(n_features)
         Xw = Xw_init.copy() if Xw_init is not None else np.zeros(n_samples)
 
-        # TODO: OR line search
-        lipschitz = datafit.global_lipschitz
+        if hasattr(datafit, "global_lipschitz"):
+            lipschitz = datafit.global_lipschitz
+        else:
+            # TODO: OR line search
+            raise Exception("Line search is not yet implemented for FISTA solver.")
 
         for n_iter in range(self.max_iter):
             t_old = t_new

skglm/tests/test_fista.py

@@ -0,0 +1,47 @@
+import pytest
+
+import numpy as np
+from numpy.linalg import norm
+
+from skglm.datafits import Quadratic, Logistic, QuadraticSVC
+from skglm.estimators import Lasso, LinearSVC, SparseLogisticRegression
+from skglm.penalties import L1, IndicatorBox
+from skglm.solvers import FISTA
+from skglm.utils import make_correlated_data, compiled_clone
+
+
+n_samples, n_features = 10, 20
+X, y, _ = make_correlated_data(
+    n_samples=n_samples, n_features=n_features, random_state=0)
+y_classif = np.sign(y)
+
+alpha_max = norm(X.T @ y, ord=np.inf) / len(y)
+alpha = alpha_max / 100
+
+tol = 1e-8
+
+# TODO: use GeneralizedLinearEstimator (to test global lipschtiz constants of every datafit)
+# TODO: test sparse matrices (global lipschitz constants)
+@pytest.mark.parametrize("Datafit, Penalty, Estimator", [
+    (Quadratic, L1, Lasso),
+    (Logistic, L1, SparseLogisticRegression),
+    (QuadraticSVC, IndicatorBox, LinearSVC),
+])
+def test_fista_solver(Datafit, Penalty, Estimator):
+    _y = y if isinstance(Datafit, Quadratic) else y_classif
+    datafit = compiled_clone(Datafit())
+    _init = y @ X.T if isinstance(Datafit, QuadraticSVC) else X
+    datafit.initialize(_init, _y)
+    penalty = compiled_clone(Penalty(alpha))
+
+    solver = FISTA(max_iter=1000, tol=tol)
+    w = solver.solve(X, _y, datafit, penalty)
+
+    estimator = Estimator(alpha, tol=tol, fit_intercept=False)
+    estimator.fit(X, _y)
+
+    np.testing.assert_allclose(w, estimator.coef_.flatten(), rtol=1e-3)
+
+
+if __name__ == '__main__':
+    pass