CLN

Author: PABannier

skglm/solvers/fista.py

@@ -3,16 +3,25 @@
 from skglm.solvers.base import BaseSolver
 
 
+@njit
+def _prox_vec(w, z, penalty, lipschitz):
+    # XXX: TO DISCUSS: should add a vectorized prox update
+    n_features = w.shape[0]
+    for j in range(n_features):
+        w[j] = penalty.prox_1d(z[j], 1 / lipschitz, j)
+    return w
+
+
 class FISTA(BaseSolver):
     r"""ISTA solver with Nesterov acceleration (FISTA)."""
 
     def __init__(self, max_iter=100, tol=1e-4, fit_intercept=False, warm_start=False,
-                 opt_freq=50, verbose=0):
+                 opt_freq=100, 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.opt_freq = opt_freq
         self.verbose = verbose
 
     def solve(self, X, y, penalty, w_init=None, weights=None):
@@ -36,10 +45,7 @@ def solve(self, X, y, penalty, w_init=None, weights=None):
             w_old = w.copy()
             grad = (G @ z - Xty) / n_samples
             z -= grad / lipschitz
-            # TODO: TO DISCUSS!
-            # XXX: should add a full prox update
-            for j in range(n_features):
-                w[j] = penalty.prox_1d(z[j], 1 / lipschitz, j)
+            w = _prox_vec(w, z, penalty, lipschitz)
             z = w + (t_old - 1.) / t_new * (w - w_old)
 
             if n_iter % self.opt_freq == 0:
@@ -48,7 +54,7 @@ def solve(self, X, y, penalty, w_init=None, weights=None):
 
                 if self.verbose:
                     p_obj = (np.sum((y - X @ w) ** 2) / (2 * n_samples)
-                             + penalty.alpha * penalty.value(w))
+                             + penalty.value(w))
                     print(
                         f"Iteration {n_iter+1}: {p_obj:.10f}, "
                         f"stopping crit: {stop_crit:.2e}"