compile inside solver

Author: mathurinm

examples/plot_sparse_recovery.py

@@ -18,7 +18,6 @@
 from skglm.utils.data import make_correlated_data
 from skglm.solvers import AndersonCD
 from skglm.datafits import Quadratic
-from skglm.utils.jit_compilation import compiled_clone
 from skglm.penalties import L1, MCPenalty, L0_5, L2_3, SCAD
 
 cmap = plt.get_cmap('tab10')
@@ -74,7 +73,7 @@
 for idx, estimator in enumerate(penalties.keys()):
     print(f'Running {estimator}...')
     estimator_path = solver.path(
-        X, y, compiled_clone(datafit), compiled_clone(penalties[estimator]),
+        X, y, datafit, penalties[estimator],
         alphas=alphas)
 
     f1_temp = np.zeros(n_alphas)

examples/plot_survival_analysis.py

@@ -15,6 +15,15 @@
 # Let's first generate synthetic data on which to run the Cox estimator,
 # using ``skglm`` data utils.
 #
+import warnings
+import time
+from lifelines import CoxPHFitter
+import pandas as pd
+import numpy as np
+from skglm.solvers import ProxNewton
+from skglm.penalties import L1
+from skglm.datafits import Cox
+import matplotlib.pyplot as plt
 from skglm.utils.data import make_dummy_survival_data
 
 n_samples, n_features = 500, 100
@@ -34,7 +43,6 @@
 # * ``s`` indicates the observations censorship and follows a Bernoulli(0.5) distribution
 #
 # Let's inspect the data quickly:
-import matplotlib.pyplot as plt
 
 fig, axes = plt.subplots(
     1, 3,
@@ -59,18 +67,14 @@
 # Todo so, we need to combine a Cox datafit and a :math:`\ell_1` penalty
 # and solve the resulting problem using skglm Proximal Newton solver ``ProxNewton``.
 # We set the intensity of the :math:`\ell_1` regularization to ``alpha=1e-2``.
-from skglm.datafits import Cox
-from skglm.penalties import L1
-from skglm.solvers import ProxNewton
 
-from skglm.utils.jit_compilation import compiled_clone
 
 # regularization intensity
 alpha = 1e-2
 
 # skglm internals: init datafit and penalty
-datafit = compiled_clone(Cox())
-penalty = compiled_clone(L1(alpha))
+datafit = Cox()
+penalty = L1(alpha)
 
 datafit.initialize(X, y)
 
@@ -90,9 +94,6 @@
 # %%
 # Let's solve the problem with ``lifelines`` through its ``CoxPHFitter``
 # estimator and compare the objectives found by the two packages.
-import numpy as np
-import pandas as pd
-from lifelines import CoxPHFitter
 
 # format data
 stacked_y_X = np.hstack((y, X))
@@ -126,8 +127,6 @@
 # let's compare their execution time. To get the evolution of the suboptimality
 # (objective - optimal objective) we run both estimators with increasing number of
 # iterations.
-import time
-import warnings
 
 warnings.filterwarnings('ignore')
 
@@ -230,7 +229,7 @@
 # We only need to pass in ``use_efron=True`` to the ``Cox`` datafit.
 
 # ensure using Efron estimate
-datafit = compiled_clone(Cox(use_efron=True))
+datafit = Cox(use_efron=True)
 datafit.initialize(X, y)
 
 # solve the problem

skglm/estimators.py

@@ -18,7 +18,6 @@
 from sklearn.utils._param_validation import Interval, StrOptions
 from sklearn.multiclass import OneVsRestClassifier, check_classification_targets
 
-from skglm.utils.jit_compilation import compiled_clone
 from skglm.solvers import AndersonCD, MultiTaskBCD, GroupBCD
 from skglm.datafits import (Cox, Quadratic, Logistic, QuadraticSVC,
                             QuadraticMultiTask, QuadraticGroup)

skglm/solvers/base.py

@@ -1,5 +1,10 @@
+import warnings
 from abc import abstractmethod, ABC
+
+import numpy as np
+
 from skglm.utils.validation import check_attrs
+from skglm.utils.jit_compilation import compiled_clone
 
 
 class BaseSolver(ABC):
@@ -101,6 +106,13 @@ def solve(self, X, y, datafit, penalty, w_init=None, Xw_init=None,
         >>> ...
         >>> coefs, obj_out, stop_crit = solver.solve(X, y, datafit, penalty)
         """
+        if "jitclass" in str(type(datafit)):
+            warnings.warn(
+                "Do not pass a compiled datafit, compilation is done inside solver now")
+        else:
+            datafit = compiled_clone(datafit, to_float32=X.dtype == np.float32)
+            penalty = compiled_clone(penalty, to_float32=X.dtype == np.float32)
+
         if run_checks:
             self._validate(X, y, datafit, penalty)