diff --git a/pyproject.toml b/pyproject.toml index 461cc25..27d2fea 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/src/fastl2lir/fastl2lir.py b/src/fastl2lir/fastl2lir.py index decb9f6..5a2170a 100644 --- a/src/fastl2lir/fastl2lir.py +++ b/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,7 +311,7 @@ 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), ) @@ -279,11 +319,10 @@ def __sub_fit( # 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 diff --git a/tests/test_fastl2lir.py b/tests/test_fastl2lir.py index cd144e3..e9de8b0 100644 --- a/tests/test_fastl2lir.py +++ b/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__":