MNT - migrate sparse matrix operation to utils.sparse_ops (#176)

Co-authored-by: jasper <[email protected]>

Author: jasperlamm

skglm/datafits/single_task.py

@@ -4,8 +4,7 @@
 from numba import float64, int64, bool_
 
 from skglm.datafits.base import BaseDatafit
-from skglm.utils.sparse_ops import spectral_norm
-from skglm.solvers.prox_newton import _sparse_xj_dot
+from skglm.utils.sparse_ops import spectral_norm, _sparse_xj_dot
 
 
 class Quadratic(BaseDatafit):

skglm/solvers/prox_newton.py

@@ -6,7 +6,7 @@
 from skglm.solvers.base import BaseSolver
 
 from sklearn.exceptions import ConvergenceWarning
-
+from skglm.utils.sparse_ops import _sparse_xj_dot
 
 EPS_TOL = 0.3
 MAX_CD_ITER = 20
@@ -413,15 +413,6 @@ def _construct_grad_sparse(X_data, X_indptr, X_indices, y, w, Xw, datafit, ws):
     return grad
 
 
-@njit(fastmath=True)
-def _sparse_xj_dot(X_data, X_indptr, X_indices, j, other):
-    # Compute X[:, j] @ other in case X sparse
-    res = 0.
-    for i in range(X_indptr[j], X_indptr[j+1]):
-        res += X_data[i] * other[X_indices[i]]
-    return res
-
-
 @njit(fastmath=True)
 def _sparse_weighted_dot(X_data, X_indptr, X_indices, j, other, weights):
     # Compute X[:, j] @ (weights * other) in case X sparse

skglm/utils/sparse_ops.py

@@ -92,3 +92,12 @@ def _XT_dot_vec(X_data, X_indptr, X_indices, vec):
             result[j] += X_data[idx] * vec[X_indices[idx]]
 
     return result
+
+
+@njit(fastmath=True)
+def _sparse_xj_dot(X_data, X_indptr, X_indices, j, other):
+    # Compute X[:, j] @ other in case X sparse
+    res = 0.
+    for i in range(X_indptr[j], X_indptr[j+1]):
+        res += X_data[i] * other[X_indices[i]]
+    return res