Standardize solver refactoring and shared factorizations

pymatsolver/direct/mumps.py

@@ -1,12 +1,12 @@
-from pymatsolver.solvers import Base
+from pymatsolver.solvers import Base, _SharedFactorMixin
 try:
     from mumps import Context
     _available = True
 except ImportError:
     Context = None
     _available = False
 
-class Mumps(Base):
+class Mumps(_SharedFactorMixin, Base):
     """The MUMPS direct solver.
 
     This solver uses the python-mumps wrappers to factorize a sparse matrix, and use that factorization for solving.
@@ -31,7 +31,6 @@ class Mumps(Base):
     **kwargs
         Extra keyword arguments. If there are any left here a warning will be raised.
     """
-    _transposed = False
 
     def __init__(self, A, ordering=None, is_symmetric=None, is_positive_definite=False, check_accuracy=False, check_rtol=1e-6, check_atol=0, **kwargs):
         if not _available:
@@ -75,39 +74,25 @@ def get_attributes(self):
         attrs['ordering'] = self.ordering
         return attrs
 
-    def transpose(self):
-        trans_obj = Mumps.__new__(Mumps)
-        trans_obj._A = self.A
-        for attr, value in self.get_attributes().items():
-            setattr(trans_obj, attr, value)
-        trans_obj.solver = self.solver
-        trans_obj._transposed = not self._transposed
-        return trans_obj
-
-    def factor(self, A=None):
-        """(Re)factor the A matrix.
+    def _do_factor(self, reuse_analysis):
+        pivot_tol = 0.0 if self.is_positive_definite else 0.01
+        self.solver.factor(
+            ordering=self.ordering, reuse_analysis=reuse_analysis, pivot_tol=pivot_tol
+        )
 
-        Parameters
-        ----------
-        A : scipy.sparse.spmatrix
-            The matrix to be factorized. If a previous factorization has been performed, this will
-            reuse the previous factorization's analysis.
-        """
+    def _refactor(self, A):
+        # if it was previously factored then re-use the analysis.
         reuse_analysis = self._factored
-        do_factor = not self._factored
-        if A is not None and A is not self.A:
-            # if it was previously factored then re-use the analysis.
-            self._set_A(A)
-            self._A = A
-            do_factor = True
-        if do_factor:
-            pivot_tol = 0.0 if self.is_positive_definite else 0.01
-            self.solver.factor(
-                ordering=self.ordering, reuse_analysis=reuse_analysis, pivot_tol=pivot_tol
-            )
+        self._A = A
+        self._set_A(self._shared.A)
+        self._do_factor(reuse_analysis)
+
+    def _factor(self):
+        if not self._factored:
+            self._do_factor(reuse_analysis=False)
 
     def _solve_multiple(self, rhs):
-        self.factor()
+        self._factor()
         if self._transposed:
             self.solver.mumps_instance.icntl[9] = 0
         else:

pymatsolver/direct/pardiso.py

@@ -1,12 +1,12 @@
-from pymatsolver.solvers import Base
+from pymatsolver.solvers import Base, _SharedFactorMixin
 try:
     from pydiso.mkl_solver import MKLPardisoSolver
     from pydiso.mkl_solver import set_mkl_pardiso_threads, get_mkl_pardiso_max_threads
     _available = True
 except ImportError:
     _available = False
 
-class Pardiso(Base):
+class Pardiso(_SharedFactorMixin, Base):
     """The Pardiso direct solver.
 
     This solver uses the `pydiso` Intel MKL wrapper to factorize a sparse matrix, and use that
@@ -36,8 +36,6 @@ class Pardiso(Base):
         Extra keyword arguments. If there are any left here a warning will be raised.
     """
 
-    _transposed = False
-
     def __init__(self, A, n_threads=None, is_symmetric=None, is_positive_definite=False, is_hermitian=None, check_accuracy=False, check_rtol=1e-6, check_atol=0, **kwargs):
         if not _available:
             raise ImportError("Pardiso solver requires the pydiso package to be installed.")
@@ -89,32 +87,15 @@ def _matrixType(self):
             else:
                 return 13
 
-    def factor(self, A=None):
-        """(Re)factor the A matrix.
-
-        Parameters
-        ----------
-        A : scipy.sparse.spmatrix
-            The matrix to be factorized. If a previous factorization has been performed, this will
-            reuse the previous factorization's analysis.
-        """
-        if A is not None and self.A is not A:
-            self._A = A
-            self.solver.refactor(self.A)
+    def _refactor(self, A):
+        self._A = A
+        # reuses the previous factorization's analysis.
+        self.solver.refactor(self._shared.A)
 
     def _solve_multiple(self, rhs):
         sol = self.solver.solve(rhs, transpose=self._transposed)
         return sol
 
-    def transpose(self):
-        trans_obj = Pardiso.__new__(Pardiso)
-        trans_obj._A = self.A
-        for attr, value in self.get_attributes().items():
-            setattr(trans_obj, attr, value)
-        trans_obj.solver = self.solver
-        trans_obj._transposed = not self._transposed
-        return trans_obj
-
     @property
     def n_threads(self):
         """Number of threads to use for the Pardiso solver routine.

pymatsolver/iterative.py

@@ -83,7 +83,11 @@ def get_attributes(self):
         attrs["atol"] = self.atol
         return attrs
 
-    def factor(self):
+    def _refactor(self, A):
+        self._A = A
+        self._factored = False
+
+    def _factor(self):
         if self._factored:
             return
         nSize = self.A.shape[0]
@@ -97,7 +101,7 @@ def _tols(self):
 
 
     def _solve_single(self, rhs):
-        self.factor()
+        self._factor()
         sol, info = bicgstab(
             self.A, rhs,
             maxiter=self.maxiter,
@@ -107,7 +111,7 @@ def _solve_single(self, rhs):
         return sol
 
     def _solve_multiple(self, rhs):
-        self.factor()
+        self._factor()
         sol = np.empty_like(rhs)
         for icol in range(rhs.shape[1]):
             sol[:, icol] = self._solve_single(rhs[:, icol])

pymatsolver/solvers.py

@@ -5,6 +5,7 @@
 from scipy.linalg import issymmetric, ishermitian
 from abc import ABC, abstractmethod
 import copy
+import types
 
 
 class SolverAccuracyError(Exception):
@@ -368,6 +369,39 @@ def _solve_single(self, rhs):
     def _solve_multiple(self, rhs):
         ...
 
+    def factor(self, A=None):
+        """(Re)factor the matrix.
+
+        Parameters
+        ----------
+        A : optional
+            A new matrix to replace the current one. It must have the same shape and dtype as
+            the current matrix, and should share its symmetry and definiteness properties, which
+            are not re-checked. For solvers that reuse a symbolic analysis (e.g. `Pardiso` and
+            `Mumps`), it must also have the same sparsity pattern. If ``None`` (or the current
+            matrix), this performs the factorization if it has not been done yet.
+
+        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
+        objects continue to use the old matrix.
+        """
+        if A is not None and A is not self.A:
+            if A.shape != self.shape:
+                raise ValueError(f"A must have shape {self.shape}, got {A.shape}.")
+            if np.dtype(A.dtype) != self.dtype:
+                raise ValueError(f"A must have dtype {self.dtype}, got {A.dtype}.")
+            self._refactor(A)
+        else:
+            self._factor()
+
+    def _refactor(self, A):
+        self._A = A
+
+    def _factor(self):
+        pass
+
     def clean(self):
         pass
 
@@ -411,6 +445,45 @@ def get_attributes(self):
         return attrs
 
 
+class _SharedFactorMixin:
+    """Mixin for solvers whose transposed and conjugated views share one factorization.
+
+    The matrix and the underlying solver object are stored on a shared namespace, so that
+    re-factoring any view updates all of them. Classes using this mixin should inherit from
+    it before `Base` (or another `Base` subclass) so that its overrides take precedence.
+    """
+
+    _transposed = False
+
+    def __init__(self, A, **kwargs):
+        self._shared = types.SimpleNamespace(A=None, solver=None)
+        super().__init__(A, **kwargs)
+
+    @property
+    def _A(self):
+        A = self._shared.A
+        return A.T if self._transposed else A
+
+    @_A.setter
+    def _A(self, value):
+        # always store the matrix in its un-transposed orientation.
+        self._shared.A = value.T if self._transposed else value
+
+    @property
+    def solver(self):
+        """The underlying solver object shared by this solver's transposed and conjugated views."""
+        return self._shared.solver
+
+    @solver.setter
+    def solver(self, value):
+        self._shared.solver = value
+
+    def transpose(self):
+        trans_obj = copy.copy(self)
+        trans_obj._transposed = not self._transposed
+        return trans_obj
+
+
 class Diagonal(Base):
     """A solver for a diagonal matrix.
 
@@ -429,13 +502,7 @@ class Diagonal(Base):
     """
 
     def __init__(self, A, check_accuracy=False, check_rtol=1e-6, check_atol=0, **kwargs):
-        try:
-            self._diagonal = np.asarray(A.diagonal())
-            if not np.all(self._diagonal):
-                # this works because 0.0 evaluates as False!
-                raise ValueError("Diagonal matrix has a zero along the diagonal.")
-        except AttributeError:
-            raise TypeError("A must have a diagonal() method.")
+        self._set_diagonal(A)
         kwargs.pop("is_symmetric", None)
         is_hermitian = kwargs.pop("is_hermitian", None)
         is_positive_definite = kwargs.pop("is_positive_definite", None)
@@ -458,6 +525,20 @@ def __init__(self, A, check_accuracy=False, check_rtol=1e-6, check_atol=0, **kwa
                 is_hermitian = not np.any(self._diagonal.imag)
         self.is_hermitian = is_hermitian
 
+    def _set_diagonal(self, A):
+        try:
+            diagonal = np.asarray(A.diagonal())
+        except AttributeError:
+            raise TypeError("A must have a diagonal() method.")
+        if not np.all(diagonal):
+            # this works because 0.0 evaluates as False!
+            raise ValueError("Diagonal matrix has a zero along the diagonal.")
+        self._diagonal = diagonal
+
+    def _refactor(self, A):
+        self._set_diagonal(A)
+        self._A = A
+
     def _solve_single(self, rhs):
         return rhs / self._diagonal
 
@@ -490,12 +571,20 @@ def __init__(self, A, lower=True, check_accuracy=False, check_rtol=1e-6, check_a
         is_hermitian = kwargs.pop("is_hermitian", False)
         is_symmetric = kwargs.pop("is_symmetric", False)
         is_positive_definite = kwargs.pop("is_positive_definite", False)
+        A = self._prepare_matrix(A)
+        super().__init__(A, is_hermitian=is_hermitian, is_symmetric=is_symmetric, is_positive_definite=is_positive_definite, check_accuracy=check_accuracy, check_rtol=check_rtol, check_atol=check_atol, **kwargs)
+
+        self.lower = lower
+
+    @staticmethod
+    def _prepare_matrix(A):
         if not (sp.issparse(A) and A.format in ['csr', 'csc']):
             A = sp.csc_matrix(A)
         A.sum_duplicates()
-        super().__init__(A, is_hermitian=is_hermitian, is_symmetric=is_symmetric, is_positive_definite=is_positive_definite, check_accuracy=check_accuracy, check_rtol=check_rtol, check_atol=check_atol, **kwargs)
+        return A
 
-        self.lower = lower
+    def _refactor(self, A):
+        self._A = self._prepare_matrix(A)
 
     @property
     def lower(self):

pymatsolver/wrappers.py

@@ -113,6 +113,14 @@ def clean(self):
         if factorize and hasattr(self.solver, 'clean'):
             self.solver.clean()
 
+    def _refactor(self, A):
+        self._A = A
+        if factorize:
+            self.clean()
+            self.solver = fun(self.A, **self.kwargs)
+            if not hasattr(self.solver, "solve"):
+                raise TypeError(f"instance returned by {fun.__name__} must have a solve() method.")
+
     class_name = str(name if name is not None else fun.__name__)
     WrappedClass = type(
         class_name,
@@ -123,6 +131,7 @@ def clean(self):
             "_solve_multiple": _solve_multiple,
             "kwargs": kwargs,
             "clean": clean,
+            "_refactor": _refactor,
         }
     )
     WrappedClass.__doc__ = f"""Wrapped {class_name} solver.

tests/test_Wrappers.py

@@ -43,6 +43,32 @@ def __init__(self, A):
         WrappedClass(sp.eye(2))
 
 
+def test_bad_direct_function_on_refactor():
+    class Empty():
+        def __init__(self, A):
+            self.A = A
+
+    class Good(Empty):
+        def solve(self, x):
+            return x
+
+    # returns a valid solver the first call (at construction), and one missing
+    # a solve() method on the second call (triggered by a refactor).
+    calls = []
+
+    def flaky_direct_func(A):
+        calls.append(A)
+        if len(calls) == 1:
+            return Good(A)
+        return Empty(A)
+
+    WrappedClass = wrap_direct(flaky_direct_func, factorize=True)
+    Ainv = WrappedClass(sp.eye(2))
+
+    with pytest.raises(TypeError, match="instance returned by.*"):
+        Ainv.factor(sp.eye(2) * 2)
+
+
 def test_direct_clean_function():
     def direct_func(A):
         class Empty():