allow SolverLU to be transposed without refactoring

pymatsolver/direct/__init__.py

@@ -1,10 +1,10 @@
-from scipy.sparse.linalg import spsolve, splu
+from scipy.sparse.linalg import spsolve
 
 from ..wrappers import WrapDirect
+from .lu import SolverLU
 from .pardiso import Pardiso
 from .mumps import Mumps
 
 Solver = WrapDirect(spsolve, factorize=False, name="Solver")
-SolverLU = WrapDirect(splu, factorize=True, name="SolverLU")
 
 __all__ = ["Solver", "SolverLU", "Pardiso", "Mumps"]

pymatsolver/direct/lu.py

@@ -0,0 +1,55 @@
+import numpy as np
+import scipy.sparse as sp
+from scipy.sparse.linalg import splu
+
+from pymatsolver.solvers import Base, _SharedFactorMixin
+from pymatsolver.wrappers import _valid_kwargs_for_func
+
+
+class SolverLU(_SharedFactorMixin, Base):
+    """Wrapped SolverLU solver, using `scipy.sparse.linalg.splu`.
+
+    The transposed solver (``.T``) re-uses this solver's LU factorization instead of
+    factoring ``A.T``. Transposed and conjugated solvers share this factorization, and are
+    updated along with it when it is re-factored.
+
+    Parameters
+    ----------
+    A
+        The matrix to use for the solver.
+    is_symmetric : bool, optional
+        Whether the matrix is symmetric. By default, it will perform some simple tests to check for symmetry, and
+        default to ``False`` if those fail.
+    is_positive_definite : bool, optional
+        Whether the matrix is positive definite.
+    is_hermitian : bool, optional
+        Whether the matrix is hermitian.
+    check_accuracy : bool, optional
+        Whether to check the accuracy of the solution.
+    check_rtol : float, optional
+        The relative tolerance to check against for accuracy.
+    check_atol : float, optional
+        The absolute tolerance to check against for accuracy.
+    **kwargs
+        Extra keyword arguments which will attempted to be passed to `scipy.sparse.linalg.splu`.
+    """
+
+    def __init__(self, A, is_symmetric=None, is_positive_definite=False, is_hermitian=None, check_accuracy=False, check_rtol=1e-6, check_atol=0, **kwargs):
+        super().__init__(A, is_symmetric=is_symmetric, is_positive_definite=is_positive_definite, is_hermitian=is_hermitian, check_accuracy=check_accuracy, check_rtol=check_rtol, check_atol=check_atol)
+        self.kwargs = _valid_kwargs_for_func(splu, **kwargs)
+        self.solver = splu(sp.csc_matrix(self.A), **self.kwargs)
+
+    def transpose(self):
+        if self.is_symmetric:
+            return self
+        return super().transpose()
+
+    def _solve_single(self, rhs):
+        rhs = rhs.astype(self.dtype)
+        return self.solver.solve(rhs, trans='T' if self._transposed else 'N')
+
+    _solve_multiple = _solve_single
+
+    def _refactor(self, A):
+        self._A = A
+        self.solver = splu(sp.csc_matrix(self._shared.A), **self.kwargs)

pymatsolver/solvers.py

@@ -383,8 +383,8 @@ def factor(self, A=None):
 
         Notes
         -----
-        Transposed and conjugated solvers created from a `Pardiso` or `Mumps` solver share its
-        factorization, and are updated along with it. For other solvers, those previously created
+        Transposed and conjugated solvers created from a `SolverLU`, `Pardiso`, or `Mumps`
+        solver share its factorization, and are updated along with it. For other solvers, those previously created
         objects continue to use the old matrix.
         """
         if A is not None and A is not self.A:

tests/test_Wrappers.py

@@ -119,3 +119,52 @@ def iterative_solver(A, x):
 
     Ainv = Wrapped(sp.eye(4))
     npt.assert_equal(Ainv @ np.arange(4), np.arange(4))
+
+
+@pytest.mark.parametrize('dtype', [np.float64, np.complex128])
+@pytest.mark.parametrize('n_rhs', [None, 3])
+@pytest.mark.parametrize('conjugate', [False, True])
+def test_splu_transpose_shares_factorization(dtype, n_rhs, conjugate):
+    rng = np.random.default_rng(1234)
+    n = 20
+    A = sp.random(n, n, density=0.2, random_state=rng) + n * sp.eye(n)
+    if dtype == np.complex128:
+        A = A + 1j * sp.random(n, n, density=0.2, random_state=rng)
+    A = A.tocsc().astype(dtype)
+    rhs = rng.random(n if n_rhs is None else (n, n_rhs)).astype(dtype)
+
+    Ainv = SolverLU(A, is_symmetric=False)
+    AinvT = Ainv.T
+    # the transposed solver reuses the same factorization rather than calling splu again.
+    assert AinvT.solver is Ainv.solver
+    assert AinvT.T.solver is Ainv.solver
+
+    if conjugate:
+        AinvT = AinvT.conj()
+        A_t = A.T.conjugate()
+    else:
+        A_t = A.T
+
+    expected = SolverLU(A_t.tocsc(), is_symmetric=False) @ rhs
+    npt.assert_allclose(AinvT @ rhs, expected, rtol=1e-8, atol=1e-10)
+    # the original is unaffected.
+    npt.assert_allclose(A @ (Ainv @ rhs), rhs, rtol=1e-8, atol=1e-10)
+
+
+def test_splu_transpose_symmetric_returns_self():
+    A = sp.diags([2.0, 4.0, 2.0], [-1, 0, 1], shape=(10, 10), format='csc')
+    Ainv = SolverLU(A)
+    assert Ainv.is_symmetric
+    assert Ainv.T is Ainv
+
+
+def test_splu_transpose_clean_and_kwargs():
+    A = sp.diags([1.0, 4.0, 2.0], [-1, 0, 1], shape=(10, 10), format='csc')
+    Ainv = SolverLU(A, is_symmetric=False, permc_spec='NATURAL')
+    AinvT = Ainv.T
+    Ainv.clean()
+    AinvT.clean()
+    rhs = np.ones(10)
+    npt.assert_allclose(A.T @ (AinvT @ rhs), rhs, atol=1e-10)
+    with pytest.warns(UnusedArgumentWarning):
+        SolverLU(A, not_a_kwarg=1)

tests/test_refactor.py

@@ -17,7 +17,7 @@
     pymatsolver.Mumps,
 ]
 
-SHARED_SOLVERS = [pymatsolver.Pardiso, pymatsolver.Mumps]
+SHARED_SOLVERS = [pymatsolver.SolverLU, pymatsolver.Pardiso, pymatsolver.Mumps]
 
 TOL = 1e-4