API design with penalty and datafit jitted inside fit (#44)

Co-authored-by: Badr MOUFAD <[email protected]>

Author: mathurinm

examples/plot_sparse_recovery.py

@@ -18,6 +18,7 @@
 from skglm.utils import make_correlated_data
 from skglm.solvers import cd_solver_path
 from skglm.datafits import Quadratic
+from skglm.utils import compiled_clone
 from skglm.penalties import L1, MCPenalty, L0_5, L2_3, SCAD
 
 cmap = plt.get_cmap('tab10')
@@ -71,7 +72,8 @@
 for idx, estimator in enumerate(penalties.keys()):
     print(f'Running {estimator}...')
     estimator_path = cd_solver_path(
-        X, y, datafit, penalties[estimator], alphas=alphas, ws_strategy="fixpoint")
+        X, y, compiled_clone(datafit), compiled_clone(penalties[estimator]),
+        alphas=alphas, ws_strategy="fixpoint")
 
     f1_temp = np.zeros(n_alphas)
     prediction_error_temp = np.zeros(n_alphas)

skglm/datafits/__init__.py

@@ -1,9 +1,6 @@
 from .base import BaseDatafit, BaseMultitaskDatafit  # noqa F401
 
-from .single_task import (  # noqa F401
-    Quadratic, Quadratic_32, QuadraticSVC, QuadraticSVC_32, Logistic, Logistic_32,
-    Huber, Huber_32,
-)
+from .single_task import Quadratic, QuadraticSVC, Logistic, Huber  # noqa F401
 
 from .multi_task import QuadraticMultiTask  # noqa F401
 

skglm/datafits/base.py

@@ -1,62 +1,29 @@
 from abc import abstractmethod
 
-import numba
-from numba import float32, float64
-from numba.experimental import jitclass
-
-
-def spec_to_float32(spec):
-    """Convert a numba specification to an equivalent float32 one.
-
-    Parameters
-    ----------
-    spec : list
-        A list of (name, dtype) for every attribute of a jitclass.
-
-    Returns
-    -------
-    spec32 : list
-        A list of (name, dtype) for every attribute of a jitclass, where float64
-        have been replaced by float32.
-    """
-    spec32 = []
-    for name, dtype in spec:
-        if dtype == float64:
-            dtype32 = float32
-        elif isinstance(dtype, numba.core.types.npytypes.Array):
-            dtype32 = dtype.copy(dtype=float32)
-        else:
-            raise ValueError(f"Unknown spec type {dtype}")
-        spec32.append((name, dtype32))
-    return spec32
-
-
-def jit_factory(Datafit, spec):
-    """JIT-compile a datafit class in float32 and float64 contexts.
-
-    Parameters
-    ----------
-    Datafit : datafit class, inheriting from BaseDatafit
-        A datafit class, to be compiled.
-
-    spec : list
-        A list of type specifications for every attribute of Datafit.
-
-    Returns
-    -------
-    Datafit_64 : Jitclass
-        A compiled datafit class with attribute types float64.
-
-    Datafit_32 : Jitclass
-        A compiled datafit class with attribute types float32.
-    """
-    spec32 = spec_to_float32(spec)
-    return jitclass(spec)(Datafit), jitclass(spec32)(Datafit)
-
 
 class BaseDatafit():
     """Base class for datafits."""
 
+    @abstractmethod
+    def get_spec(self):
+        """Specify the numba types of the class attributes.
+
+        Returns
+        -------
+        spec: Tuple of (attribute_name, dtype)
+            spec to be passed to Numba jitclass to compile the class.
+        """
+
+    @abstractmethod
+    def params_to_dict(self):
+        """Get the parameters to initialize an instance of the class.
+
+        Returns
+        -------
+        dict_of_params : dict
+            The parameters to instantiate an object of the class.
+        """
+
     @abstractmethod
     def initialize(self, X, y):
         """Pre-computations before fitting on X and y.
@@ -172,6 +139,26 @@ def gradient_scalar_sparse(self, X_data, X_indptr, X_indices, y, Xw, j):
 class BaseMultitaskDatafit():
     """Base class for multitask datafits."""
 
+    @abstractmethod
+    def get_spec(self):
+        """Specify the numba types of the class attributes.
+
+        Returns
+        -------
+        spec: Tuple of (attribute_name, dtype)
+            spec to be passed to Numba jitclass to compile the class.
+        """
+
+    @abstractmethod
+    def params_to_dict(self):
+        """Get the parameters to initialize an instance of the class.
+
+        Returns
+        -------
+        dict_of_params : dict
+            The parameters to instantiate an object of the class.
+        """
+
     @abstractmethod
     def initialize(self, X, Y):
         """Store useful values before fitting on X and Y.

skglm/datafits/group.py

@@ -1,19 +1,10 @@
 import numpy as np
 from numpy.linalg import norm
-from numba.experimental import jitclass
 from numba import int32, float64
 
 from skglm.datafits.base import BaseDatafit
 
 
-spec_QuadraticGroup = [
-    ('grp_ptr', int32[:]),
-    ('grp_indices', int32[:]),
-    ('lipschitz', float64[:])
-]
-
-
-@jitclass(spec_QuadraticGroup)
 class QuadraticGroup(BaseDatafit):
     r"""Quadratic datafit used with group penalties.
 
@@ -38,6 +29,18 @@ class QuadraticGroup(BaseDatafit):
     def __init__(self, grp_ptr, grp_indices):
         self.grp_ptr, self.grp_indices = grp_ptr, grp_indices
 
+    def get_spec(self):
+        spec = (
+            ('grp_ptr', int32[:]),
+            ('grp_indices', int32[:]),
+            ('lipschitz', float64[:])
+        )
+        return spec
+
+    def params_to_dict(self):
+        return dict(grp_ptr=self.grp_ptr,
+                    grp_indices=self.grp_indices)
+
     def initialize(self, X, y):
         grp_ptr, grp_indices = self.grp_ptr, self.grp_indices
         n_groups = len(grp_ptr) - 1

skglm/datafits/multi_task.py

@@ -1,18 +1,10 @@
 import numpy as np
 from numpy.linalg import norm
 from numba import float64
-from numba.experimental import jitclass
 
 from skglm.datafits.base import BaseMultitaskDatafit
 
 
-spec_quadratic = [
-    ('XtY', float64[:, :]),
-    ('lipschitz', float64[:]),
-]
-
-
-@jitclass(spec_quadratic)
 class QuadraticMultiTask(BaseMultitaskDatafit):
     """Quadratic datafit used for multi-task regression.
 
@@ -33,6 +25,16 @@ class QuadraticMultiTask(BaseMultitaskDatafit):
     def __init__(self):
         pass
 
+    def get_spec(self):
+        spec = (
+            ('XtY', float64[:, :]),
+            ('lipschitz', float64[:]),
+        )
+        return spec
+
+    def params_to_dict(self):
+        return dict()
+
     def initialize(self, X, Y):
         """Compute optimization quantities before fitting on X and Y."""
         self.XtY = X.T @ Y