pyMOTO 1.2.1__py3-none-any.whl → 1.4.0__py3-none-any.whl

pymoto/{common/solvers_dense.py → solvers/dense.py}

@@ -1,7 +1,8 @@
 import warnings
 import numpy as np
 import scipy.linalg as spla  # Dense matrix solvers
-from .solvers import matrix_is_hermitian, matrix_is_diagonal, LinearSolver
+from .matrix_checks import matrix_is_hermitian, matrix_is_diagonal
+from .solvers import LinearSolver
 
 
 class SolverDiagonal(LinearSolver):
@@ -11,24 +12,17 @@ class SolverDiagonal(LinearSolver):
         self.diag = A.diagonal()
         return self
 
-    def solve(self, rhs):
+    def solve(self, rhs, x0=None, trans='N'):
         r""" Solve using the diagonal only, by :math:`x_i = b_i / A_{ii}`
 
         The right-hand-side :math:`\mathbf{b}` can be of size ``(N)`` or ``(N, K)``, where ``N`` is the size of matrix
         :math:`\mathbf{A}` and ``K`` is the number of right-hand sides.
         """
+        d = self.diag.conj() if trans == 'H' else self.diag
         if rhs.ndim == 1:
-            return rhs / self.diag
+            return rhs / d
         else:
-            return rhs / self.diag[..., None]
-
-    def adjoint(self, rhs):
-        r""" Solve using the diagonal only, by :math:`x_i = b_i / A_{ii}^*`
-
-        The right-hand-side :math:`\mathbf{b}` can be of size ``(N)`` or ``(N, K)``, where ``N`` is the size of matrix
-        :math:`\mathbf{A}` and ``K`` is the number of right-hand sides.
-        """
-        return self.solve(rhs.conj()).conj()
+            return rhs / d[..., None]
 
 
 # Dense QR solver
@@ -42,23 +36,31 @@ class SolverDenseQR(LinearSolver):
         self.q, self.r = spla.qr(A)
         return self
 
-    def solve(self, rhs):
-        r""" Solves the linear system of equations :math:`\mathbf{A} \mathbf{x} = \mathbf{b}` by backward substitution of
-        :math:`\mathbf{x} = \mathbf{R}^{-1}\mathbf{Q}^\text{H}\mathbf{b}`.
-
-        The right-hand-side :math:`\mathbf{b}` can be of size ``(N)`` or ``(N, K)``, where ``N`` is the size of matrix
-        :math:`\mathbf{A}` and ``K`` is the number of right-hand sides.
-        """
-        return spla.solve_triangular(self.r, self.q.T.conj()@rhs)
+    def solve(self, rhs, x0=None, trans='N'):
+        r""" Solves the linear system of equations using the QR factorization.
 
-    def adjoint(self, rhs):
-        r""" Solves the linear system of equations :math:`\mathbf{A}^\text{H}\mathbf{x} = \mathbf{b}` by
-        forward substitution of :math:`\mathbf{x} = \mathbf{Q}\mathbf{R}^{-H}\mathbf{b}`.
+        ======= ================= =====================
+        `trans`     Equation      Solution of :math:`x`
+        ------- ----------------- ---------------------
+          `N`   :math:`A x = b`   :math:`R^{-1} Q^H b`
+          `T`   :math:`A^T x = b` :math:`Q^* R^{-T} b`
+          `H`   :math:`A^H x = b` :math:`Q R^{-H} b`
+        ======= ================= =====================
 
         The right-hand-side :math:`\mathbf{b}` can be of size ``(N)`` or ``(N, K)``, where ``N`` is the size of matrix
         :math:`\mathbf{A}` and ``K`` is the number of right-hand sides.
         """
-        return self.q@spla.solve_triangular(self.r, rhs, trans='C')
+        if trans == 'N':
+            # A = Q R -> inv(A) = inv(R) inv(Q) = inv(R) Q^H
+            return spla.solve_triangular(self.r, self.q.T.conj() @ rhs)
+        elif trans == 'T':
+            # A^T = R^T Q^T -> inv(A^T) = inv(Q^T) inv(R^T) = conj(Q) inv(R^T)
+            return self.q.conj() @ spla.solve_triangular(self.r, rhs, trans='T')
+        elif trans == 'H':
+            # A^H = R^H Q^H -> inv(A^H) = inv(Q^H) inv(R^H) = Q inv(R^H)
+            return self.q @ spla.solve_triangular(self.r, rhs, trans='C')
+        else:
+            raise TypeError("Only N, T, and H transposition is possible")
 
 
 # Dense LU solver
@@ -71,24 +73,34 @@ class SolverDenseLU(LinearSolver):
         self.p, self.l, self.u = spla.lu(A)
         return self
 
-    def solve(self, rhs):
-        r""" Solves the linear system of equations :math:`\mathbf{A} \mathbf{x} = \mathbf{b}` by forward and backward
-        substitution of :math:`\mathbf{x} = \mathbf{U}^{-1}\mathbf{L}^{-1}\mathbf{b}`.
+    def solve(self, rhs, x0=None, trans='N'):
+        r""" Solves the linear system of equations using the LU factorization.
 
-        The right-hand-side :math:`\mathbf{b}` can be of size ``(N)`` or ``(N, K)``, where ``N`` is the size of matrix
-        :math:`\mathbf{A}` and ``K`` is the number of right-hand sides.
-        """
-        return spla.solve_triangular(self.u, spla.solve_triangular(self.l, self.p.T@rhs, lower=True))
+        :math:`\mathbf{A} \mathbf{x} = \mathbf{b}` by forward and backward
+        substitution of :math:`\mathbf{x} = \mathbf{U}^{-1}\mathbf{L}^{-1}\mathbf{b}`.
 
-    def adjoint(self, rhs):
-        r""" Solves the linear system of equations :math:`\mathbf{A}^\text{H}\mathbf{x} = \mathbf{b}` by forward and
-        backward substitution of :math:`\mathbf{x} = \mathbf{L}^{-\text{H}}\mathbf{U}^{-\text{H}}\mathbf{b}`.
+        ======= ================= =========================
+        `trans`     Equation        Solution of :math:`x`
+        ------- ----------------- -------------------------
+          `N`   :math:`A x = b`   :math:`x = U^{-1} L^{-1}`
+          `T`   :math:`A^T x = b` :math:`x = L^{-1} U^{-1}`
+          `H`   :math:`A^H x = b` :math:`x = L^{-*} U^{-*}`
+        ======= ================= =========================
 
         The right-hand-side :math:`\mathbf{b}` can be of size ``(N)`` or ``(N, K)``, where ``N`` is the size of matrix
         :math:`\mathbf{A}` and ``K`` is the number of right-hand sides.
         """
-        return self.p@spla.solve_triangular(self.l, spla.solve_triangular(self.u, rhs, trans='C'),
-                                            lower=True, trans='C')  # TODO permutation
+        if trans == 'N':
+            # A = P L U -> x = U^-1 L^-1 P^T b
+            return spla.solve_triangular(self.u, spla.solve_triangular(self.l, self.p.T@rhs, lower=True))
+        elif trans == 'T':
+            return self.p @ spla.solve_triangular(self.l, spla.solve_triangular(self.u, rhs, trans='T'),
+                                                  lower=True, trans='T')
+        elif trans == 'H':
+            return self.p @ spla.solve_triangular(self.l, spla.solve_triangular(self.u, rhs, trans='C'),
+                                                  lower=True, trans='C')
+        else:
+            raise TypeError("Only N, T, and H transposition is possible")
 
 
 # Dense Cholesky solver
@@ -106,7 +118,7 @@ class SolverDenseCholesky(LinearSolver):
         upper triangular matrix.
         """
         try:
-            self.u = spla.cholesky(A)
+            self.U = spla.cholesky(A)
             self.success = True
         except np.linalg.LinAlgError as err:
             warnings.warn(f"{type(self).__name__}: {err} -- using {type(self.backup_solver).__name__} instead")
@@ -114,7 +126,7 @@ class SolverDenseCholesky(LinearSolver):
             self.success = False
         return self
 
-    def solve(self, rhs):
+    def solve(self, rhs, x0=None, trans='N'):
         r""" Solves the linear system of equations :math:`\mathbf{A} \mathbf{x} = \mathbf{b}` by forward and backward
         substitution of :math:`\mathbf{x} = \mathbf{U}^{-1}\mathbf{U}^{-\text{H}}\mathbf{b}`.
 
@@ -124,21 +136,16 @@ class SolverDenseCholesky(LinearSolver):
         # TODO When Cholesky factorization A = U^T U is used, symmetric complex matrices can also be solved, but this is
         #  not implemented in scipy
         if self.success:
-            return spla.solve_triangular(self.u, spla.solve_triangular(self.u, rhs, trans='C'))
-        else:
-            return self.backup_solver.solve(rhs)
-
-    def adjoint(self, rhs):
-        r""" A Hermitian matrix is self-adjoint (:math:`\mathbf{A}=\mathbf{A}^\text{H}`), so this is equal to the
-        regular solution.
-
-        The right-hand-side :math:`\mathbf{b}` can be of size ``(N)`` or ``(N, K)``, where ``N`` is the size of matrix
-        :math:`\mathbf{A}` and ``K`` is the number of right-hand sides.
-        """
-        if self.success:
-            return self.solve(rhs)
+            if trans == 'N' or trans == 'H':
+                # A = U^H U -> A^-1 = U^-1 U^-H
+                return spla.solve_triangular(self.U, spla.solve_triangular(self.U, rhs, trans='C'))
+            elif trans == 'T':
+                # A^T = U^T conj(U) -> A^-T = conj(U^-1) U^-T
+                return spla.solve_triangular(self.U, spla.solve_triangular(self.U, rhs, trans='T').conj()).conj()
+            else:
+                raise TypeError("Only N, T, and H transposition is possible")
         else:
-            return self.backup_solver.adjoint(rhs)
+            return self.backup_solver.solve(rhs, trans=trans)
 
 
 # Dense LDL solver
@@ -171,23 +178,13 @@ class SolverDenseLDL(LinearSolver):
         self.lp = self.l[self.p, :]
         return self
 
-    def solve(self, rhs):
+    def solve(self, rhs, x0=None, trans='N'):
         r""" Solves the linear system of equations :math:`\mathbf{A} \mathbf{x} = \mathbf{b}` by forward and backward
         substitution of :math:`\mathbf{x} = \mathbf{L}^{-\text{H}}\mathbf{D}^{-1}\mathbf{L}^{-1}\mathbf{b}` in the
         Hermitian case or as :math:`\mathbf{x} = \mathbf{L}^{-\text{T}}\mathbf{D}^{-1}\mathbf{L}^{-1}\mathbf{b}` in the
         symmetric case.
 
-        The right-hand-side :math:`\mathbf{b}` can be of size ``(N)`` or ``(N, K)``, where ``N`` is the size of matrix
-        :math:`\mathbf{A}` and ``K`` is the number of right-hand sides.
-        """
-        u1 = spla.solve_triangular(self.lp, rhs[self.p], lower=True, unit_diagonal=True)
-        u2 = self.dinv(u1)
-        u = np.zeros_like(rhs, dtype=u2.dtype)
-        u[self.p] = spla.solve_triangular(self.lp, u2, trans='C' if self.hermitian else 'T', lower=True, unit_diagonal=True)
-        return u
-
-    def adjoint(self, rhs):
-        r""" Solves the linear system of equations :math:`\mathbf{A}^\text{H} \mathbf{x} = \mathbf{b}` by forward and
+        The adjoint system of equations :math:`\mathbf{A}^\text{H} \mathbf{x} = \mathbf{b}` is solved by forward and
         backward substitution of
         :math:`\mathbf{x} = \mathbf{L}^{-\text{H}}\mathbf{D}^{-\text{H}}\mathbf{L}^{-1}\mathbf{b}` in the  Hermitian
         case or as :math:`\mathbf{x} = \mathbf{L}^{-\text{H}}\mathbf{D}^{-\text{H}}\mathbf{L}^{-*}\mathbf{b}`
@@ -196,11 +193,34 @@ class SolverDenseLDL(LinearSolver):
         The right-hand-side :math:`\mathbf{b}` can be of size ``(N)`` or ``(N, K)``, where ``N`` is the size of matrix
         :math:`\mathbf{A}` and ``K`` is the number of right-hand sides.
         """
-        if not self.hermitian:
-            u1 = spla.solve_triangular(self.lp, rhs[self.p].conj(), lower=True, unit_diagonal=True).conj()
-        else:
+        if trans == 'N':
+            # Hermitian matrix A: A = L D L^H -> inv(A) = inv(L^H) inv(D) inv(L)
+            # Symmetric matrix A: A = L D L^T -> inv(A) = inv(L^T) inv(D) inv(L)
             u1 = spla.solve_triangular(self.lp, rhs[self.p], lower=True, unit_diagonal=True)
-        u2 = self.dinvH(u1)
-        u = np.zeros_like(rhs, dtype=u2.dtype)
-        u[self.p] = spla.solve_triangular(self.lp, u2, trans='C', lower=True, unit_diagonal=True)
+            u2 = self.dinv(u1)
+            u = np.zeros_like(rhs, dtype=u2.dtype)
+            u[self.p] = spla.solve_triangular(self.lp, u2, trans='C' if self.hermitian else 'T', lower=True, unit_diagonal=True)
+        elif trans == 'T':
+            # Hermitian matrix A^T: A = conj(L) D^T L^T -> inv(A) = inv(L^T) inv(D^T) inv(L^*)
+            # Symmetric matrix A^T: A = L D^T L^T -> inv(A) = inv(L^T) inv(D^T) inv(L)
+            if self.hermitian:
+                u1 = spla.solve_triangular(self.lp, rhs[self.p].conj(), lower=True, unit_diagonal=True).conj()
+            else:
+                u1 = spla.solve_triangular(self.lp, rhs[self.p], lower=True, unit_diagonal=True)
+
+            u2 = self.dinvH(u1.conj()).conj()
+            u = np.zeros_like(rhs, dtype=u2.dtype)
+            u[self.p] = spla.solve_triangular(self.lp, u2, trans='T', lower=True, unit_diagonal=True)
+        elif trans == 'H':
+            # Hermitian matrix A: inv(A^H) = inv(L^H) inv(D^H) inv(L)
+            # Symmetric matrix A: inv(A^H) = inv(L^H) inv(D^H) inv(L^*)
+            if not self.hermitian:
+                u1 = spla.solve_triangular(self.lp, rhs[self.p].conj(), lower=True, unit_diagonal=True).conj()
+            else:
+                u1 = spla.solve_triangular(self.lp, rhs[self.p], lower=True, unit_diagonal=True)
+            u2 = self.dinvH(u1)
+            u = np.zeros_like(rhs, dtype=u2.dtype)
+            u[self.p] = spla.solve_triangular(self.lp, u2, trans='C', lower=True, unit_diagonal=True)
+        else:
+            raise TypeError("Only N, T, and H transposition is possible")
         return u

pymoto/solvers/iterative.py

@@ -0,0 +1,361 @@
+import warnings
+import time
+import numpy as np
+import scipy.sparse as sps
+from scipy.sparse.linalg import splu, spilu
+from .solvers import LinearSolver
+from .auto_determine import auto_determine_solver
+from pymoto import DomainDefinition
+
+
+class Preconditioner(LinearSolver):
+    """ Abstract base class for preconditioners to inexact solvers """
+    def update(self, A):
+        pass
+
+    def solve(self, rhs, x0=None, trans='N'):
+        return rhs.copy()
+
+
+class DampedJacobi(Preconditioner):
+    r""" Damped Jacobi preconditioner
+    :math:`M = \frac{1}{\omega} D`
+    Args:
+        A (optional): The matrix
+        w (optional): Weight factor :math:`0 < \omega \leq 1`
+    """
+    def __init__(self, A=None, w=1.0):
+        assert 0 < w <= 1, 'w must be between 0 and 1'
+        self.w = w
+        self.D = None
+        super().__init__(A)
+
+    def update(self, A):
+        self.D = A.diagonal()
+
+    def solve(self, rhs, x0=None, trans='N'):
+        if trans == 'N' or trans == 'T':
+            return self.w * (rhs.T/self.D).T
+        elif trans == 'H':
+            return self.w * (rhs.T/self.D.conj()).T
+        else:
+            raise TypeError("Only N, T, or H transposition is possible")
+
+
+class SOR(Preconditioner):
+    r""" Successive over-relaxation preconditioner
+    The matrix :math:`A = L + D + U` is split into a lower triangular, diagonal, and upper triangular part.
+    :math:`M = \left(\frac{D}{\omega} + L\right) \frac{\omega D^{-1}}{2-\omega} \left(\frac{D}{\omega} + U\right)`
+
+    Args:
+        A (optional): The matrix
+        w (optional): Weight factor :math:`0 < \omega < 2`
+    """
+    def __init__(self, A=None, w=1.0):
+        assert 0 < w < 2, 'w must be between 0 and 2'
+        self.w = w
+        self.L = None
+        self.U = None
+        self.Dw = None
+        super().__init__(A)
+
+    def update(self, A):
+        diag = A.diagonal()
+        diagw = sps.diags(diag)/self.w
+        self.L = splu(sps.tril(A, k=-1) + diagw)  # Lower triangular part including diagonal
+        self.U = splu(sps.triu(A, k=1) + diagw)
+
+        self.Dw = diag * (2 - self.w) / self.w
+
+    def solve(self, rhs, x0=None, trans='N'):
+        if trans == 'N':
+            # M = (D/w + L) wD^-1 / (2-w) (D/w + U)
+            # from scipy.sparse.linalg import spsolve_triangular
+            # u1 = spsolve_triangular(self.L, rhs, lower=True, overwrite_A=False)  # Solve triangular is still very slow :(
+            u1 = self.L.solve(rhs)
+            u1 *= self.Dw[:, None]
+            # u2 = spsolve_triangular(self.U, u1, lower=False, overwrite_A=False, overwrite_b=True)
+            u2 = self.U.solve(u1)
+            return u2
+        elif trans == 'T':
+            u1 = self.U.solve(rhs, trans='T')
+            u1 *= self.Dw[:, None]
+            u2 = self.L.solve(u1, trans='T')
+            return u2
+        elif trans == 'H':
+            u1 = self.U.solve(rhs, trans='H')
+            u1 *= self.Dw[:, None].conj()
+            u2 = self.L.solve(u1, trans='H')
+            return u2
+        else:
+            raise TypeError("Only N, T, or H transposition is possible")
+
+
+class ILU(Preconditioner):
+    """ Incomplete LU factorization
+
+    Args:
+        A (optional): The matrix
+        **kwargs (optional): Keyword arguments passed to `scipy.sparse.linalg.spilu`
+    """
+    def __init__(self, A=None, **kwargs):
+        self.kwargs = kwargs
+        self.ilu = None
+        super().__init__(A)
+
+    def update(self, A):
+        self.ilu = spilu(A, **self.kwargs)
+
+    def solve(self, rhs, x0=None, trans='N'):
+        return self.ilu.solve(rhs, trans=trans)
+
+
+class GeometricMultigrid(Preconditioner):
+    """ Geometric multigrid preconditioner
+
+    Args:
+        domain: The `DomainDefinition` used for the geometry
+        A (optional): The matrix
+        inner_level (optional): Inner solver for the coarse grid, for instance, a direct solver or another MG level.
+            The default is an automatically determined direct solver.
+        smoother (optional): Smoother to use to smooth the residual and solution before and after coarse level.
+            The default is `DampedJacobi(w=0.5)`.
+        smooth_steps (optional): Number of smoothing steps to execute
+    """
+    _available_cycles = ['v', 'w']
+
+    def __init__(self, domain: DomainDefinition, A=None, cycle='V', inner_level=None, smoother=None, smooth_steps=5):
+        assert domain.nelx % 2 == 0 and domain.nely % 2 == 0 and domain.nelz % 2 == 0, \
+            f"Domain sizes {domain.nelx, domain.nely, domain.nelz} must be divisible by 2"
+        self.domain = domain
+        self.A = A
+        assert cycle.lower() in self._available_cycles, f"Cycle ({cycle}) is not available. Options are {self._available_cycles}"
+        self.cycle = cycle
+        self.inner_level = None if inner_level is None else inner_level
+        self.smoother = DampedJacobi(w=0.5) if smoother is None else None
+        self.smooth_steps = smooth_steps
+        self.R = None
+        self.sub_domain = DomainDefinition(domain.nelx // 2, domain.nely // 2, domain.nelz // 2,
+                                           domain.unitx * 2, domain.unity * 2, domain.unitz * 2)
+
+        super().__init__(A)
+
+    def update(self, A):
+        if self.R is None:
+            self.setup_interpolation(A)
+        self.A = A
+        self.smoother.update(A)
+        Ac = self.R.T @ A @ self.R
+        if self.inner_level is None:
+            self.inner_level = auto_determine_solver(Ac)
+        self.inner_level.update(Ac)
+
+    def setup_interpolation(self, A):
+        assert A.shape[0] % self.domain.nnodes == 0
+        ndof = int(A.shape[0] / self.domain.nnodes)  # Number of dofs per node
+
+        w = np.ones((3, 3, 3))*0.125
+        w[1, :, :] = 0.25
+        w[:, 1, :] = 0.25
+        w[:, :, 1] = 0.25
+        w[1, 1, :] = 0.5
+        w[1, :, 1] = 0.5
+        w[:, 1, 1] = 0.5
+        w[1, 1, 1] = 1.0
+
+        rows = []
+        cols = []
+        vals = []
+        for i in [-1, 0, 1]:
+            imin, imax = max(-i, 0), min(self.sub_domain.nelx + 1 - i, self.sub_domain.nelx + 1)
+            ix = np.arange(imin, imax)
+            for j in [-1, 0, 1]:
+                jmin, jmax = max(-j, 0), min(self.sub_domain.nely + 1 - j, self.sub_domain.nely + 1)
+                iy = np.arange(jmin, jmax)
+                for k in ([-1, 0, 1] if self.domain.dim == 3 else [0]):
+                    # Coarse node cartesian indices
+                    kmin, kmax = max(-k, 0), min(self.sub_domain.nelz + 1 - k, self.sub_domain.nelz + 1)
+                    iz = np.arange(kmin, kmax)
+                    # Coarse node numbers
+                    nod_c = self.sub_domain.get_nodenumber(*np.meshgrid(ix, iy, iz, indexing='ij')).flatten()
+                    # Fine node numbers with offset
+                    ixc, iyc, izc = ix * 2 + i, iy * 2 + j, iz * 2 + k
+                    nod_f = self.domain.get_nodenumber(*np.meshgrid(ixc, iyc, izc, indexing='ij')).flatten()
+                    for d in range(ndof):
+                        rows.append(nod_f * ndof + d)
+                        cols.append(nod_c * ndof + d)
+                        vals.append(np.ones_like(rows[-1], dtype=w.dtype) * w[1+i, 1+j, 1+k])
+
+        rows = np.concatenate(rows)
+        cols = np.concatenate(cols)
+        vals = np.concatenate(vals)
+        nfine = ndof * self.domain.nnodes
+        ncoarse = ndof * self.sub_domain.nnodes
+        self.R = sps.coo_matrix((vals, (rows, cols)), shape=(nfine, ncoarse))
+        self.R = type(A)(self.R)  # Convert to correct matrix type
+
+    def solve(self, rhs, x0=None, trans='N'):
+        if trans == 'N':
+            A = self.A
+        elif trans == 'T':
+            A = self.A.T
+        elif trans == 'H':
+            A = self.A.conj().T
+        else:
+            raise TypeError("Only N, T, or H transposition is possible")
+
+        # Pre-smoothing
+        if x0 is None:
+            u_f = self.smoother.solve(rhs, trans=trans)
+        else:
+            r = rhs - self.A @ x0
+            u_f = x0 + self.smoother.solve(r, trans=trans)
+        for i in range(self.smooth_steps-1):
+            r = rhs - self.A @ u_f
+            u_f += self.smoother.solve(r, trans=trans)
+
+        r = rhs - A @ u_f
+        # Restrict residual to coarse level
+        r_c = self.R.T @ r
+
+        # Solve at coarse level
+        u_c = self.inner_level.solve(r_c)
+
+        # Interpolate and correct
+        u_f += self.R @ u_c
+
+        # Post-smoothing
+        for i in range(self.smooth_steps):
+            r = rhs - self.A @ u_f
+            u_f += self.smoother.solve(r, trans=trans)
+        return u_f
+
+
+def orth(u, normalize=True, zero_rtol=1e-15):
+    """ Create orthogonal basis from a set of vectors
+
+    Args:
+        u: Set of vectors of size (#dof, #vectors)
+        normalize: Also normalize the basis vectors
+        zero_rtol: Relative tolerance for detection of zero vectors (in case of a rank-deficient basis)
+
+    Returns:
+        v: Orthogonal basis vectors (#dof, #non-zero-vectors)
+    """
+    if u.ndim == 1:
+        return u
+    elif u.ndim > 2:
+        raise TypeError("Only valid for 1D or 2D matrix")
+
+    def dot(a, b):  # Define inner product
+        return a @ b.conj()
+
+    orth_vecs = []
+    for i in range(u.shape[-1]):
+        vi = np.copy(u[..., i])
+        beta_i = dot(vi, vi)
+        for vj in orth_vecs:
+            alpha_ij = dot(vi, vj)
+            alpha_jj = 1.0 if normalize else dot(vj, vj)
+            vi -= vj * alpha_ij / alpha_jj
+        beta_i_new = dot(vi, vi)
+        if beta_i_new / beta_i < zero_rtol:  # Detect zero vector
+            continue
+        if normalize:
+            vi /= np.sqrt(beta_i_new)
+        orth_vecs.append(vi)
+    return np.stack(orth_vecs, axis=-1)
+
+
+class CG(LinearSolver):
+    """ Preconditioned conjugate gradient method
+    Works for positive-definite self-adjoint matrices (:math:`A=A^H`)
+
+    References:
+        Ji & Li (2017), A breakdown-free BCG method. DOI 10.1007/s10543-016-0631-z
+          https://www.cs.odu.edu/~yaohang/portfolio/BIT2017.pdf
+        Shewchuck (1994), Introduction to CG method without the agonzing pain.
+          https://www.cs.cmu.edu/~quake-papers/painless-conjugate-gradient.pdf
+
+    Args:
+        A: The matrix
+        preconditioner: Preconditioner to use
+        tol: Convergence tolerance
+        maxit: Maximum number of iterations
+        restart: Restart every Nth iteration
+        verbosity: Log level
+    """
+    def __init__(self, A=None, preconditioner=Preconditioner(), tol=1e-7, maxit=10000, restart=50, verbosity=0):
+        self.preconditioner = preconditioner
+        self.A = A
+        self.tol = tol
+        self.maxit = maxit
+        self.restart = restart
+        self.verbosity = verbosity
+        super().__init__(A)
+
+    def update(self, A):
+        tstart = time.perf_counter()
+        self.A = A
+        self.preconditioner.update(A)
+        if self.verbosity >= 1:
+            print(f"Preconditioner set up in {np.round(time.perf_counter() - tstart,3)}s")
+
+    def solve(self, rhs, x0=None, trans='N'):
+        if trans == 'N':
+            A = self.A
+        elif trans == 'T':
+            A = self.A.T
+        elif trans == 'H':
+            A = self.A.conj().T
+        else:
+            raise TypeError("Only N, T, or H transposition is possible")
+
+        tstart = time.perf_counter()
+        if rhs.ndim == 1:
+            b = rhs.reshape((rhs.size, 1))
+        else:
+            b = rhs
+        x = np.zeros_like(rhs, dtype=np.result_type(rhs, A)) if x0 is None else x0.copy()
+        if x.ndim == 1:
+            x = x.reshape((x.size, 1))
+
+        r = b - A@x
+        z = self.preconditioner.solve(r, trans=trans)
+        p = orth(z, normalize=True)
+        if self.verbosity >= 2:
+            print(f"Initial residual = {np.linalg.norm(r, axis=0) / np.linalg.norm(b, axis=0)}")
+
+        for i in range(self.maxit):
+            q = A @ p
+            pq = p.conj().T @ q
+            pq_inv = np.linalg.inv(pq)
+            alpha = pq_inv @ (p.conj().T @ r)
+
+            x += p @ alpha
+            if i % 50 == 0:  # Explicit restart
+                r = b - A@x
+            else:
+                r -= q @ alpha
+
+            if self.verbosity >= 2:
+                print(f"i = {i}, residuals = {np.linalg.norm(r, axis=0) / np.linalg.norm(b, axis=0)}")
+
+            tval = np.linalg.norm(r)/np.linalg.norm(b)
+            if tval <= self.tol:
+                break
+
+            z = self.preconditioner.solve(r, trans=trans)
+
+            beta = -pq_inv @ (q.conj().T @ z)
+            p = orth(z + p@beta, normalize=False)
+
+        if tval > self.tol:
+            warnings.warn(f'Maximum iterations ({self.maxit}) reached, with final residual {tval}')
+        elif self.verbosity >= 1:
+            print(f"Converged in {i} iterations and {np.round(time.perf_counter() - tstart, 3)}s, with final residuals {np.linalg.norm(r, axis=0) / np.linalg.norm(b, axis=0)}")
+
+        if rhs.ndim == 1:
+            return x.flatten()
+        else:
+            return x

pymoto/solvers/matrix_checks.py

@@ -0,0 +1,56 @@
+import numpy as np
+import scipy.sparse as sps
+try:
+    import cvxopt
+    _has_cvxopt = True
+except ImportError:
+    _has_cvxopt = False
+
+
+def is_cvxopt_spmatrix(A):
+    """ Checks if the argument is a cvxopt sparse matrix """
+    return isinstance(A, cvxopt.spmatrix) if _has_cvxopt else False
+
+
+def matrix_is_complex(A):
+    """ Checks if the matrix is complex """
+    if is_cvxopt_spmatrix(A):
+        return A.typecode == 'z'
+    else:
+        return np.iscomplexobj(A)
+
+
+def matrix_is_diagonal(A):
+    """ Checks if the matrix is diagonal"""
+    if sps.issparse(A):
+        if isinstance(A, sps.dia_matrix):
+            return len(A.offsets) == 1 and A.offsets[0] == 0
+        else:
+            return np.allclose((A - sps.spdiags(A.diagonal(), 0, *A.shape)).data, 0.0)
+    elif is_cvxopt_spmatrix(A):
+        return max(abs(A.I - A.J)) == 0
+    else:
+        return np.allclose(A, np.diag(np.diag(A)))
+
+
+def matrix_is_symmetric(A):
+    """ Checks whether a matrix is numerically symmetric """
+    if sps.issparse(A):
+        return np.allclose((A-A.T).data, 0)
+    elif is_cvxopt_spmatrix(A):
+        return np.isclose(max(abs(A-A.T)), 0.0)
+    else:
+        return np.allclose(A, A.T)
+
+
+def matrix_is_hermitian(A):
+    """ Checks whether a matrix is numerically Hermitian """
+    if matrix_is_complex(A):
+        if sps.issparse(A):
+            return np.allclose((A-A.T.conj()).data, 0)
+        elif is_cvxopt_spmatrix(A):
+            return np.isclose(max(abs(A-A.ctrans())), 0.0)
+        else:
+            return np.allclose(A, A.T.conj())
+    else:
+        return matrix_is_symmetric(A)