Clean up linear solvers, add StencilDiagonalMatrix and diagonal(), require pc to be LinearOperator

psydac/api/tests/test_2d_multipatch_mapping_maxwell.py

@@ -93,10 +93,19 @@ def run_maxwell_2d(uex, f, alpha, domain, *, ncells=None, degree=None, filename=
     equation_h = discretize(equation, domain_h, [Vh, Vh], backend=PSYDAC_BACKEND_GPYCCEL)
     l2norm_h   = discretize(l2norm, domain_h, Vh, backend=PSYDAC_BACKEND_GPYCCEL)
 
-    equation_h.set_solver('pcg', pc='jacobi', tol=1e-8)
+    # Explicitly assemble the linear system
+    equation_h.assemble()
 
+    # Extract the matrix' diagonal to build a Jacobi preconditioner
+    jacobi_pc = equation_h.linear_system.lhs.diagonal(inverse=True)
+
+    # Choose a linear solver and pass any flags to it
+    equation_h.set_solver('pcg', pc=jacobi_pc, tol=1e-8)
+
+    # Solve the linear system and obtain the solution as a FEM field
     uh = equation_h.solve()
 
+    # Compute the L2 norm of the error
     l2_error = l2norm_h.assemble(F=uh)
 
     return l2_error, uh

psydac/api/tests/test_api_2d_compatible_spaces.py

@@ -399,7 +399,8 @@ def run_maxwell_time_harmonic_2d_dir(uex, f, alpha, ncells, degree):
     #+++++++++++++++++++++++++++++++
 
     # Solve linear system
-    M_inv = inverse(M, 'pcg', pc='jacobi', tol=1e-8)
+    jacobi_pc = M.diagonal(inverse=True)
+    M_inv = inverse(M, 'pcg', pc=jacobi_pc, tol=1e-8)
     sol = M_inv @ b
 
     uh       = FemField( Vh, sol )

psydac/api/tests/test_api_feec_1d.py

@@ -336,7 +336,8 @@ def discrete_energies(self, e, b):
         step_ampere_1d = dt * ( M0_dir_inv @ D0_T_dir @ M1 )
 
     elif bc_mode == 'penalization':
-        M0_M0_bc_inv = inverse(M0+M0_bc, 'pcg', pc='jacobi', **kwargs)
+        M0_M0_bc = M0 + M0_bc
+        M0_M0_bc_inv = inverse(M0_M0_bc, 'pcg', pc = M0_M0_bc.diagonal(inverse=True), **kwargs)
         step_ampere_1d = dt * ( M0_M0_bc_inv @ D0_T @ M1 )
 
     half_step_faraday_1d = (dt/2) * D0

psydac/api/tests/test_api_feec_2d.py

@@ -519,7 +519,8 @@ def discrete_energies(self, e, b):
         M1_inv = inverse(M1, 'cg', **kwargs)
         step_ampere_2d = dt * (M1_inv @ D1_T @ M2)
     else:
-        M1_M1_bc_inv = inverse(M1 + M1_bc, 'pcg', pc='jacobi', **kwargs)
+        M1_M1_bc = M1 + M1_bc
+        M1_M1_bc_inv = inverse(M1_M1_bc, 'pcg', pc = M1_M1_bc.diagonal(inverse=True), **kwargs)
         step_ampere_2d = dt * (M1_M1_bc_inv @ D1_T @ M2)
 
     half_step_faraday_2d = (dt/2) * D1

psydac/linalg/basic.py

@@ -448,7 +448,6 @@ def __new__(cls, domain, codomain=None):
             return BlockLinearOperator(domain, domain, blocks)
         else:
             return super().__new__(cls)
-
 
     def __init__(self, domain, codomain=None):
 
@@ -879,11 +878,39 @@ class InverseLinearOperator(LinearOperator):
     Iterative solver for square linear system Ax=b, where x and b belong to (normed)
     vector space V.
 
+    Parameters
+    ----------
+    A : psydac.linalg.basic.LinearOperator
+        Left-hand-side matrix A of linear system.
+
+    x0 : psydac.linalg.basic.Vector
+        First guess of solution for iterative solver (optional).
+
+    tol : float
+        Absolute tolerance for L2-norm of residual r = A*x - b.
+
+    maxiter: int
+        Maximum number of iterations.
+
+    verbose : bool
+        If True, L2-norm of residual r is printed at each iteration.
     """
 
-    @property
-    def space(self):
-        return self._space
+    def __init__(self, A, **kwargs):
+
+        assert isinstance(A, LinearOperator)
+        assert A.domain.dimension == A.codomain.dimension
+        domain = A.codomain
+        codomain = A.domain
+
+        if kwargs['x0'] is None:
+            kwargs['x0'] = codomain.zeros()
+
+        self._A = A
+        self._domain = domain
+        self._codomain = codomain
+        self._options = kwargs
+        self._check_options(**self._options)
 
     @property
     def domain(self):
@@ -899,11 +926,44 @@ def dtype(self):
 
     @property
     def linop(self):
+        """
+        The linear operator L of which this object is the inverse L^{-1}.
+
+        The linear operator L can be modified in place, or replaced entirely
+        through the setter. A substitution should only be made in cases where
+        no other options are viable, as it breaks the one-to-one map between
+        the original linear operator L (passed to the constructor) and the
+        current `InverseLinearOperator` object L^{-1}. Use with extreme care!
+
+        """
         return self._A
+
+    @linop.setter
+    def linop(self, a):
+        assert isinstance(a, LinearOperator)
+        assert a.domain is self.domain
+        assert a.codomain is self.codomain
+        self._A = a
 
     @property
     def options(self):
         return self._options
+
+    def _check_options(self, **kwargs):
+        for key, value in kwargs.items():
+
+            if key == 'x0':
+                if value is not None:
+                    assert isinstance(value, Vector), "x0 must be a Vector or None"
+                    assert value.space == self.codomain, "x0 belongs to the wrong VectorSpace"
+            elif key == 'tol':
+                assert is_real(value), "tol must be a real number"
+                assert value > 0, "tol must be positive"
+            elif key == 'maxiter':
+                assert isinstance(value, int), "maxiter must be an int"
+                assert value > 0, "maxiter must be positive"
+            elif key == 'verbose':
+                assert isinstance(value, bool), "verbose must be a bool"
 
     def toarray(self):
         raise NotImplementedError('toarray() is not defined for InverseLinearOperators.')
@@ -921,69 +981,11 @@ def set_options(self, **kwargs):
         self._check_options(**kwargs)
         self._options.update(kwargs)
 
-    @abstractmethod
-    def _check_options(self, **kwargs):
-        pass
-
-    @abstractmethod
     def transpose(self, conjugate=False):
-        pass
-
-    @staticmethod
-    def jacobi(A, b, out=None):
-        """
-        Jacobi preconditioner.
-
-        A : psydac.linalg.stencil.StencilMatrix | psydac.linalg.block.BlockLinearOperator
-            Left-hand-side matrix A of linear system.
-
-        b : psydac.linalg.stencil.StencilVector | psydac.linalg.block.BlockVector
-            Right-hand-side vector of linear system.
-
-        Returns
-        -------
-        x : psydac.linalg.stencil.StencilVector | psydac.linalg.block.BlockVector
-            Preconditioner solution
-
-        """
-        from psydac.linalg.block   import BlockLinearOperator, BlockVector
-        from psydac.linalg.stencil import StencilMatrix, StencilVector
-
-        # In case A is None we return a zero vector
-        if A is None:
-            return b.space.zeros()
-
-        # Sanity checks
-        assert isinstance(A, (StencilMatrix, BlockLinearOperator))
-        assert isinstance(b, (StencilVector, BlockVector))
-        assert A.codomain.dimension == A.domain.dimension
-        assert A.codomain == b.space
-
-        #-------------------------------------------------------------
-        # Handle the case of a block linear system
-        if isinstance(A, BlockLinearOperator):
-            if out is not None:
-                for i, bi in enumerate(b.blocks):
-                    InverseLinearOperator.jacobi(A[i,i], bi, out=out[i])
-                return out
-            else:
-                x = [InverseLinearOperator.jacobi(A[i, i], bi) for i, bi in enumerate(b.blocks)]
-                y = BlockVector(b.space, blocks=x)
-                return y
-        #-------------------------------------------------------------
-
-        V = b.space
-        i = tuple(slice(s, e + 1) for s, e in zip(V.starts, V.ends))
-
-        if out is not None:
-            b.copy(out=out)
-            out[i] /= A.diagonal()
-            out.update_ghost_regions()
-        else:
-            out = b.copy()
-            out[i] /= A.diagonal()
-            out.update_ghost_regions()
-            return out
+        cls     = type(self)
+        At      = self.linop.transpose(conjugate=conjugate)
+        options = self._options
+        return cls(At, **options)
 
 #===============================================================================
 class LinearSolver(ABC):
@@ -1016,3 +1018,8 @@ def solve(self, rhs, out=None):
     @property
     def T(self):
         return self.transpose()
+
+
+def is_real(x):
+    from numbers import Number
+    return isinstance(x, Number) and np.isrealobj(x) and not isinstance(x, bool)

psydac/linalg/block.py

@@ -777,6 +777,49 @@ def __sub__(self, M):
     #--------------------------------------
     # New properties/methods
     #--------------------------------------
+    def diagonal(self, *, inverse = False, out = None):
+        """Get the coefficients on the main diagonal as another BlockLinearOperator object.
+
+        Parameters
+        ----------
+        inverse : bool
+            If True, get the inverse of the diagonal. (Default: False).
+
+        out : BlockLinearOperator
+            If provided, write the diagonal entries into this matrix. (Default: None).
+
+        Returns
+        -------
+        BlockLinearOperator
+            The matrix which contains the main diagonal of self (or its inverse).
+
+        """
+        # Determine domain and codomain of result
+        V, W = self.domain, self.codomain
+        if inverse:
+            V, W = W, V
+
+        # Check the `out` argument, if `None` create a new BlockLinearOperator
+        if out is not None:
+            assert isinstance(out, BlockLinearOperator)
+            assert out.domain is V
+            assert out.codomain is W
+
+            # Set any off-diagonal blocks to zero
+            for i, j in out.nonzero_block_indices:
+                if i != j:
+                    out[i, j] = None
+        else:
+            out = BlockLinearOperator(V, W)
+
+        # Store the diagonal (or its inverse) into `out`
+        for i, j in self.nonzero_block_indices:
+            if i == j:
+                out[i, i] = self[i, i].diagonal(inverse = inverse, out = out[i, i])
+
+        return out
+
+    # ...
     @property
     def blocks(self):
         """ Immutable 2D view (tuple of tuples) of the linear operator,
@@ -868,7 +911,8 @@ def __setitem__(self, key, value):
             assert value.codomain is self.codomain[i]
 
         self._blocks[i,j] = value
-
+
+    # ...
     def transform(self, operation):
         """
         Applies an operation on each block in this BlockLinearOperator.
@@ -881,6 +925,7 @@ def transform(self, operation):
         blocks = {ij: operation(Bij) for ij, Bij in self._blocks.items()}
         return BlockLinearOperator(self.domain, self.codomain, blocks=blocks)
 
+    # ...
     def backend(self):
         return self._backend
 
@@ -909,7 +954,6 @@ def copy(self, out=None):
 
         out.set_backend(self._backend)
 
-
     # ...
     def __imul__(self, a):
         for Bij in self._blocks.values():