Add solver option and add scipy solver backend

pyproject.toml

@@ -14,6 +14,7 @@ license = { file = "LICENSE" }
 requires-python = ">=3.10"
 dependencies = [
     "numpy>=1.22.0",
+    "scipy>=1.8",
     "threadpoolctl>=2.1.0",
     "tqdm>=4.64.1",
 ]

src/fastl2lir/fastl2lir.py

@@ -7,15 +7,55 @@
 import numpy as np
 from threadpoolctl import threadpool_limits
 from tqdm import tqdm
+from scipy import linalg as sp_linalg
+
+
+def _solve_scipy(a, b):
+    try:
+        return sp_linalg.solve(a, b, assume_a="pos", check_finite=False)
+    except sp_linalg.LinAlgError:
+        return sp_linalg.solve(a, b, assume_a="sym", check_finite=False)
+
+
+def _solve_numpy(a, b):
+    return np.linalg.solve(a, b)
+
+
+def _get_solver_func(solver):
+    if solver == "scipy":
+        return _solve_scipy
+    if solver == "numpy":
+        return _solve_numpy
+    raise ValueError(f"Unknown solver: {solver!r}. Expected one of: 'scipy', 'numpy'.")
 
 
 class FastL2LiR(object):
     """Fast L2-regularized linear regression class."""
 
-    def __init__(self, W=np.array([]), b=np.array([]), verbose=False):
+    def __init__(self, W=np.array([]), b=np.array([]), verbose=False, solver="numpy"):
         self.__W = W
         self.__b = b
         self.__verbose = verbose
+        self.__solver = solver
+        self.__solve = _get_solver_func(solver)
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state.pop("_FastL2LiR__solve", None)
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        solver_key = "_FastL2LiR__solver"
+        if solver_key not in self.__dict__:
+            warnings.warn(
+                "Unpickled a FastL2LiR object without solver information. "
+                "Restoring it with solver='numpy' for backward compatibility.",
+                UserWarning,
+                stacklevel=2,
+            )
+            self.__dict__[solver_key] = "numpy"
+        self.__solve = _get_solver_func(self.__dict__[solver_key])
 
     @property
     def W(self):
@@ -271,19 +311,18 @@ def __sub_fit(
             # Choose the more efficient method based on matrix dimensions
             if X.shape[0] > X.shape[1]:
                 # Use primal form for tall matrices (more samples than features)
-                Wb = np.linalg.solve(
+                Wb = self.__solve(
                     np.matmul(X.T, X) + alpha * np.eye(X.shape[1], dtype=dtype),
                     np.matmul(X.T, Y),
                 )
             else:
                 # Use dual form for wide matrices (more features than samples)
                 Wb = np.matmul(
                     X.T,
-                    np.linalg.solve(
+                    self.__solve(
                         np.matmul(X, X.T) + alpha * np.eye(X.shape[0], dtype=dtype), Y
                     ),
                 )
-
             W = Wb[0:-1, :]
             b = Wb[-1, :][np.newaxis, :]  # Returning b as a 2D array
         else:
@@ -312,7 +351,7 @@ def __sub_fit(
                             ).ravel()
                         ]
                     ).reshape(feat_idx.size, feat_idx.size)
-                    Wb = np.linalg.solve(W0_sub, W1[index_outputDim, feat_idx])
+                    Wb = self.__solve(W0_sub, W1[index_outputDim, feat_idx])
                     W[index_outputDim, feat_idx[:-1]] = Wb[:-1]
                     b[0, index_outputDim] = Wb[-1]
                 W = W.T
@@ -379,7 +418,7 @@ def __sub_fit_save_select_feat(
                     newX.shape[1], dtype=dtype
                 )
                 rhs = np.matmul(selY.ravel(), newX)
-                Wb = np.linalg.solve(W0, rhs)
+                Wb = self.__solve(W0, rhs)
                 W[index_outputDim, feat_idx] = Wb[:-1]
                 b[0, index_outputDim] = Wb[-1]
             W = W.T

tests/test_fastl2lir.py

@@ -136,6 +136,83 @@ def test_chunk(self):
 
         np.testing.assert_array_almost_equal(yp_2d, data["yp_2d"])
 
+    def test_solver_numpy_matches_scipy(self):
+        """numpy solver produces same result as scipy solver (default)."""
+        data = np.load("./tests/testdata_basic.npz")
+
+        model_scipy = fastl2lir.FastL2LiR(solver="scipy")
+        model_numpy = fastl2lir.FastL2LiR(solver="numpy")
+
+        model_scipy.fit(data["x_tr"], data["y_2d"])
+        model_numpy.fit(data["x_tr"], data["y_2d"])
+
+        np.testing.assert_array_almost_equal(model_numpy.W, model_scipy.W)
+        np.testing.assert_array_almost_equal(model_numpy.b, model_scipy.b)
+
+    def test_solver_helpers_agree(self):
+        """_solve_numpy and _solve_scipy return the same result."""
+        from fastl2lir.fastl2lir import _solve_numpy, _solve_scipy
+
+        rng = np.random.default_rng(0)
+        n = 50
+        A = rng.standard_normal((n, n))
+        A = A.T @ A + np.eye(n)
+        b = rng.standard_normal((n, 10))
+
+        np.testing.assert_array_almost_equal(_solve_numpy(A, b), _solve_scipy(A, b))
+
+    def test_solver_pickle(self):
+        """FastL2LiR instance with scipy solver can be pickled and unpickled."""
+        import pickle
+
+        data = np.load("./tests/testdata_basic.npz")
+        model = fastl2lir.FastL2LiR(solver="scipy")
+        model.fit(data["x_tr"], data["y_2d"])
+        restored = pickle.loads(pickle.dumps(model))
+        np.testing.assert_array_almost_equal(model.W, restored.W)
+        np.testing.assert_array_almost_equal(model.b, restored.b)
+        np.testing.assert_array_almost_equal(
+            model.predict(data["x_te"]), restored.predict(data["x_te"])
+        )
+
+    def test_solver_getstate(self):
+        """__getstate__ does not store the solver function."""
+        model = fastl2lir.FastL2LiR(solver="scipy")
+        state = model.__getstate__()
+        self.assertNotIn("_FastL2LiR__solve", state)
+        self.assertEqual(state["_FastL2LiR__solver"], "scipy")
+
+    def test_solver_legacy_setstate(self):
+        """__setstate__ falls back to numpy for legacy pickles without solver info."""
+        model = fastl2lir.FastL2LiR.__new__(fastl2lir.FastL2LiR)
+        legacy_state = {
+            "_FastL2LiR__W": np.array([]),
+            "_FastL2LiR__b": np.array([]),
+            "_FastL2LiR__verbose": False,
+        }
+        with self.assertWarns(UserWarning):
+            model.__setstate__(legacy_state)
+        self.assertEqual(model._FastL2LiR__solver, "numpy")
+        data = np.load("./tests/testdata_basic.npz")
+        model.fit(data["x_tr"], data["y_2d"])
+
+    def test_solver_scipy_fallback(self):
+        """scipy solver falls back from assume_a='pos' to 'sym' for indefinite matrix."""
+        from fastl2lir.fastl2lir import _solve_scipy
+
+        n = 20
+        rng = np.random.default_rng(0)
+        A = np.diag([1.0] * (n - 1) + [-1.0])
+        b = rng.standard_normal((n, 5))
+        result = _solve_scipy(A, b)
+        residual = np.linalg.norm(A @ result - b) / np.linalg.norm(b)
+        self.assertLess(residual, 1e-10)
+
+    def test_solver_invalid(self):
+        """Invalid solver name raises ValueError."""
+        with self.assertRaises(ValueError):
+            fastl2lir.FastL2LiR(solver="cholesky")
+
     def test_reshape(self):
         """Test for reshaping."""
         Y_shape = (200, 10, 10, 5)
@@ -154,12 +231,11 @@ def test_reshape(self):
         np.testing.assert_array_equal(model_test.b.shape, (1,) + Y_shape[1:])
         np.testing.assert_array_equal(pred_test.shape, Y_shape)
 
-
     def test_save_select_feat_default_select_sample(self):
-        '''save_select_feat=True with default select_sample=None should not raise.'''
-        data = np.load('./tests/testdata_nfeat.npz')
+        """save_select_feat=True with default select_sample=None should not raise."""
+        data = np.load("./tests/testdata_nfeat.npz")
         model = fastl2lir.FastL2LiR()
-        model.fit(data['x_tr'], data['y_1d'], n_feat=20, save_select_feat=True)
+        model.fit(data["x_tr"], data["y_1d"], n_feat=20, save_select_feat=True)
 
 
 if __name__ == "__main__":