mv _prox_vec to utils

Author: PABannier

skglm/solvers/fista.py

@@ -1,16 +1,8 @@
 import numpy as np
 from scipy.sparse import issparse
-from numba import njit
 from skglm.solvers.base import BaseSolver
 from skglm.solvers.common import construct_grad, construct_grad_sparse
-
-
-@njit
-def _prox_vec(w, z, penalty, lipschitz):
-    n_features = w.shape[0]
-    for j in range(n_features):
-        w[j] = penalty.prox_1d(z[j], 1 / lipschitz, j)
-    return w
+from skglm.utils import prox_vec
 
 
 class FISTA(BaseSolver):
@@ -71,7 +63,7 @@ def solve(self, X, y, datafit, penalty, w_init=None, Xw_init=None):
             else:
                 grad = construct_grad(X, y, z, X @ z, datafit, all_features)
             z -= grad / lipschitz
-            w = _prox_vec(w, z, penalty, lipschitz)
+            w = prox_vec(w, z, penalty, lipschitz)
             Xw = X @ w
             z = w + (t_old - 1.) / t_new * (w - w_old)
 

skglm/utils.py

@@ -457,3 +457,32 @@ def extrapolate(self, w, Xw):
         C = inv_UTU_ones / np.sum(inv_UTU_ones)
         # floating point errors may cause w and Xw to disagree
         return self.arr_w_[:, 1:] @ C, self.arr_Xw_[:, 1:] @ C, True
+
+
+@njit
+def prox_vec(w, z, penalty, lipschitz):
+    """Evaluate the vectorized proximal operator for the FISTA solver.
+
+    Parameters
+    ----------
+    w : array, shape (n_features,)
+        Coefficient vector.
+
+    z : array, shape (n_features,)
+        FISTA auxiliary variable.
+
+    penalty : instance of Penalty.
+        Penalty object.
+
+    lipschitz : float
+        Global Lipschitz constant.
+
+    Returns
+    -------
+    w : array; shape (n_features,)
+        Updated coefficient vector.
+    """
+    n_features = w.shape[0]
+    for j in range(n_features):
+        w[j] = penalty.prox_1d(z[j], 1 / lipschitz, j)
+    return w