Improve behavior of LinearOperator objects

psydac/feec/multipatch/fem_linear_operators.py

@@ -58,6 +58,9 @@ def T(self):
     @property
     def dtype( self ):
         return self.domain.dtype
+
+    def copy(self):
+        return FemLinearOperator(self.fem_domain, self.fem_codomain, self.matrix, self._sparse_matrix)
 
     def toarray(self):
         return self._matrix.toarray()
@@ -135,6 +138,9 @@ def __init__( self, operators ):
 
         # matrix not defined by matrix product because it could break the Stencil Matrix structure
 
+    def copy(self):
+        return ComposedLinearOperator(self._operators)
+
     def to_sparse_matrix( self,  **kwargs):
         mat = self._operators[-1].to_sparse_matrix()
         for i in range(2, self._n+1):
@@ -160,6 +166,9 @@ class IdLinearOperator( FemLinearOperator ):
     def __init__( self, V ):
         FemLinearOperator.__init__(self, fem_domain=V)
 
+    def copy(self):
+        return IdLinearOperator(self.fem_domain)
+
     def to_sparse_matrix( self , **kwargs):
         return sparse_id( self.fem_domain.nbasis )
 
@@ -183,6 +192,9 @@ def __init__( self, B, A ):
         self._A = A
         self._B = B
 
+    def copy(self):
+        return SumLinearOperator(self._B, self._A)
+
     def to_sparse_matrix( self, **kwargs):
         return self._A.to_sparse_matrix() + self._B.to_sparse_matrix()
 
@@ -205,6 +217,9 @@ def __init__( self, c, A ):
         self._A = A
         self._c = c
 
+    def copy(self):
+        return MultLinearOperator(self._c, self._A)
+
     def to_sparse_matrix( self,  **kwargs):
         return self._c * self._A.to_sparse_matrix()
 

psydac/feec/multipatch/operators.py

@@ -472,6 +472,9 @@ def __init__(
                     storage_fn)
                 save_npz(storage_fn, self._sparse_matrix)
 
+    def copy(self):
+        raise NotImplementedError
+
     def set_homogenous_bc(self, boundary, rhs=None):
         domain = self.symbolic_domain
         Vh = self.fem_domain
@@ -721,6 +724,9 @@ def __init__(
                     storage_fn)
                 save_npz(storage_fn, self._sparse_matrix)
 
+    def copy(self):
+        raise NotImplementedError
+
     def set_homogenous_bc(self, boundary):
         domain = self.symbolic_domain
         Vh = self.fem_domain
@@ -1011,6 +1017,9 @@ def __init__(
             # matrices are not stored, we will probably compute them later
             pass
 
+    def copy(self):
+        raise NotImplementedError
+
     def to_sparse_matrix(self):
         """
         the Hodge matrix is the patch-wise multi-patch mass matrix
@@ -1098,6 +1107,9 @@ def __init__(self, V0h, V1h):
     def transpose(self, conjugate=False):
         # todo (MCP): define as the dual differential operator
         return BrokenTransposedGradient_2D(self.fem_domain, self.fem_codomain)
+
+    def copy(self):
+        return BrokenGradient_2D(self.fem_domain, self.fem_codomain)
 
 # ==============================================================================
 
@@ -1116,6 +1128,9 @@ def __init__(self, V0h, V1h):
     def transpose(self, conjugate=False):
         # todo (MCP): discard
         return BrokenGradient_2D(self.fem_codomain, self.fem_domain)
+
+    def copy(self):
+        return BrokenTransposedGradient_2D(self.fem_domain, self.fem_codomain)
 
 
 # ==============================================================================
@@ -1132,6 +1147,9 @@ def __init__(self, V1h, V2h):
     def transpose(self, conjugate=False):
         return BrokenTransposedScalarCurl_2D(
             V1h=self.fem_domain, V2h=self.fem_codomain)
+
+    def copy(self):
+        return BrokenScalarCurl_2D(self.fem_domain, self.fem_codomain)
 
 
 # ==============================================================================
@@ -1148,6 +1166,9 @@ def __init__(self, V1h, V2h):
 
     def transpose(self, conjugate=False):
         return BrokenScalarCurl_2D(V1h=self.fem_codomain, V2h=self.fem_domain)
+
+    def copy(self):
+        return BrokenTransposedScalarCurl_2D(self.fem_domain, self.fem_codomain)
 
 
 # ==============================================================================

psydac/linalg/basic.py

@@ -248,6 +248,10 @@ def codomain(self):
     def dtype(self):
         pass
 
+    @abstractmethod
+    def copy(self):
+        pass
+
     @abstractmethod
     def tosparse(self):
         pass
@@ -584,6 +588,14 @@ def operator(self):
     @property
     def dtype(self):
         return None
+
+    def copy(self):
+        return ScaledLinearOperator(self.domain, self.codomain, self.scalar, self.operator)
+
+    def set_scalar(self, c):
+        assert np.isscalar(c)
+        assert np.iscomplexobj(c) == (self.codomain._dtype == complex)
+        self._scalar = c
 
     def toarray(self):
         return self._scalar*self._operator.toarray() 
@@ -667,6 +679,9 @@ def addends(self):
     @property
     def dtype(self):
         return None
+
+    def copy(self):
+        return SumLinearOperator(self.domain, self.codomain, *self.addends)
 
     def toarray(self):
         out = np.zeros(self.shape, dtype=self.dtype)
@@ -798,6 +813,9 @@ def multiplicants(self):
     @property
     def dtype(self):
         return None
+
+    def copy(self):
+        return ComposedLinearOperator(self.domain, self.codomain, *self.multiplicants)
 
     def toarray(self):
         raise NotImplementedError('toarray() is not defined for ComposedLinearOperators.')
@@ -905,6 +923,9 @@ def operator(self):
     def factorial(self):
         """ Returns the power to which the operator is raised. """
         return self._factorial
+
+    def copy(self):
+        return PowerLinearOperator(self.domain, self.codomain, self.operator, self.factorial)
 
     def toarray(self):
         raise NotImplementedError('toarray() is not defined for PowerLinearOperators.')
@@ -1169,6 +1190,9 @@ def codomain(self):
     @property
     def dtype(self):
         return None
+
+    def copy(self):
+        return MatrixFreeLinearOperator(self.domain, self.codomain, self._dot, self._dot_transpose)
 
     def dot(self, v, out=None, **kwargs):
         assert isinstance(v, Vector)

psydac/linalg/block.py

@@ -728,6 +728,7 @@ def __mul__(self, a):
 
     # ...
     def __add__(self, M):
+        from psydac.linalg.basic import ZeroOperator
         if not isinstance(M, BlockLinearOperator):
             return LinearOperator.__add__(self, M)
 
@@ -737,9 +738,9 @@ def __add__(self, M):
         for ij in set(self._blocks.keys()) | set(M._blocks.keys()):
             Bij = self[ij]
             Mij = M[ij]
-            if   Bij is None: blocks[ij] = Mij.copy()
-            elif Mij is None: blocks[ij] = Bij.copy()
-            else            : blocks[ij] = Bij + Mij
+            if   Bij is None or isinstance(Bij, ZeroOperator): blocks[ij] = Mij.copy() if Mij is not None else None
+            elif Mij is None or isinstance(Mij, ZeroOperator): blocks[ij] = Bij.copy()
+            else                                             : blocks[ij] = Bij + Mij
         mat = BlockLinearOperator(self.domain, self.codomain, blocks=blocks)
         if len(mat._blocks) != len(self._blocks):
             mat.set_backend(self._backend)
@@ -752,6 +753,7 @@ def __add__(self, M):
 
     # ...
     def __sub__(self, M):
+        from psydac.linalg.basic import ZeroOperator
         if not isinstance(M, BlockLinearOperator):
             return LinearOperator.__sub__(self, M)
 
@@ -761,9 +763,9 @@ def __sub__(self, M):
         for ij in set(self._blocks.keys()) | set(M._blocks.keys()):
             Bij = self[ij]
             Mij = M[ij]
-            if   Bij is None: blocks[ij] = -Mij
-            elif Mij is None: blocks[ij] =  Bij.copy()
-            else            : blocks[ij] =  Bij - Mij
+            if   Bij is None                                 : blocks[ij] = -Mij
+            elif Mij is None or isinstance(Mij, ZeroOperator): blocks[ij] =  Bij.copy()
+            else                                             : blocks[ij] =  Bij - Mij
         mat = BlockLinearOperator(self.domain, self.codomain, blocks=blocks)
         if len(mat._blocks) != len(self._blocks):
             mat.set_backend(self._backend)
@@ -1254,7 +1256,8 @@ def set_backend(self, backend):
 
         if not all(isinstance(b, (StencilMatrix, StencilInterfaceMatrix)) for b in self._blocks.values()):
             for b in self._blocks.values():
-                b.set_backend(backend)
+                if getattr(b, "set_backend", None) is not None:
+                    b.set_backend(backend)
             return
 
         block_shape = (self.n_block_rows, self.n_block_cols)

psydac/linalg/fft.py

@@ -26,6 +26,9 @@ def toarray(self):
 
     def tosparse(self):
         raise NotImplementedError('tosparse() is not defined for DistributedFFTBase.')
+
+    def copy(self):
+        raise NotImplementedError('copy() is not implemented for DistributedFFTBase.')
 
     # Possible additions for the future:
     # * split off the LinearSolver class when used with the space ndarray (as used in the KroneckerLinearSolver),

psydac/linalg/kron.py

@@ -72,6 +72,9 @@ def ndim( self ):
     @property
     def mats( self ):
         return self._mats
+
+    def copy(self):
+        return KroneckerStencilMatrix(self.domain, self.codomain, *self.mats)
 
     # ...
     def dot(self, x, out=None):
@@ -257,6 +260,7 @@ def __init__(self, V, W, *args , with_pads=False):
         self._codomain = W
         self._mats     = list(args)
         self._ndim     = len(args)
+        self._with_pads = with_pads
 
     #--------------------------------------
     # Abstract interface
@@ -284,6 +288,9 @@ def ndim(self):
     @property
     def mats(self):
         return self._mats
+
+    def copy(self):
+        return KroneckerDenseMatrix(self.domain, self.codomain, *self.mats, self._with_pads)
 
     # ...
     def dot(self, x, out=None):
@@ -428,6 +435,9 @@ def __init__(self, V, W, solvers):
 
         # for now: allocate temporary arrays here (can be removed later)
         self._temp1, self._temp2 = self._allocate_temps()
+
+    def copy(self):
+        return KroneckerLinearSolver(self.domain, self.codomain, self.solvers)
 
     def _setup_solvers(self):
         """

psydac/linalg/solvers.py

@@ -131,6 +131,11 @@ def __init__(self, A, *, x0=None, tol=1e-6, maxiter=1000, verbose=False, recycle
         self._tmps = {key: self.domain.zeros() for key in ("v", "r", "p")}
         self._info = None
 
+    def copy(self):
+        return ConjugateGradient(self.linop, x0=self.get_options("x0"), tol=self.get_options("tol"), 
+                                 maxiter=self.get_options("maxiter"), verbose=self.get_options("verbose"), 
+                                 recycle=self.get_options("recycle"))
+
     def solve(self, b, out=None):
         """
         Conjugate gradient algorithm for solving linear system Ax=b.
@@ -295,6 +300,11 @@ def __init__(self, A, *, pc=None, x0=None, tol=1e-6, maxiter=1000, verbose=False
         self._tmps = {**tmps_codomain, **tmps_domain}
         self._info = None
 
+    def copy(self):
+        return PConjugateGradient(self.linop, pc=self.get_options("pc"), x0=self.get_options("x0"), 
+                                  tol=self.get_options("tol"), maxiter=self.get_options("maxiter"), 
+                                  verbose=self.get_options("verbose"), recycle=self.get_options("recycle"))
+
     def solve(self, b, out=None):
         """
         Preconditioned Conjugate Gradient (PCG) solves the symetric positive definte
@@ -454,6 +464,11 @@ def __init__(self, A, *, x0=None, tol=1e-6, maxiter=1000, verbose=False, recycle
         self._tmps = {key: self.domain.zeros() for key in ("v", "r", "p", "vs", "rs", "ps")}
         self._info = None
 
+    def copy(self):
+        return BiConjugateGradient(self.linop, x0=self.get_options("x0"), tol=self.get_options("tol"), 
+                                   maxiter=self.get_options("maxiter"), verbose=self.get_options("verbose"), 
+                                   recycle=self.get_options("recycle"))
+
     def solve(self, b, out=None):
         """
         Biconjugate gradient (BCG) algorithm for solving linear system Ax=b.
@@ -643,6 +658,11 @@ def __init__(self, A, *, x0=None, tol=1e-6, maxiter=1000, verbose=False, recycle
         self._tmps = {key: self.domain.zeros() for key in ("v", "r", "p", "vr", "r0")}
         self._info = None
 
+    def copy(self):
+        return BiConjugateGradientStabilized(self.linop, x0=self.get_options("x0"), tol=self.get_options("tol"), 
+                                             maxiter=self.get_options("maxiter"), verbose=self.get_options("verbose"), 
+                                             recycle=self.get_options("recycle"))
+
     def solve(self, b, out=None):
         """
         Biconjugate gradient stabilized method (BCGSTAB) algorithm for solving linear system Ax=b.
@@ -843,6 +863,11 @@ def __init__(self, A, *, pc=None, x0=None, tol=1e-6, maxiter=1000, verbose=False
                                                       "rp0")}
         self._info = None
 
+    def copy(self):
+        return PBiConjugateGradientStabilized(self.linop, pc=self.get_options("pc"), x0=self.get_options("x0"), 
+                                              tol=self.get_options("tol"), maxiter=self.get_options("maxiter"), 
+                                              verbose=self.get_options("verbose"), recycle=self.get_options("recycle"))
+
     def solve(self, b, out=None):
         """
         Preconditioned biconjugate gradient stabilized method (PBCGSTAB) algorithm for solving linear system Ax=b.
@@ -1085,6 +1110,11 @@ def __init__(self, A, *, x0=None, tol=1e-6, maxiter=1000, verbose=False, recycle
         self._tmps = {key: self.domain.zeros() for key in ("res_old", "res_new", "w_new", "w_work", "w_old", "v", "y")}
         self._info = None
 
+    def copy(self):
+        return MinimumResidual(self.linop, x0=self.get_options("x0"), tol=self.get_options("tol"), 
+                               maxiter=self.get_options("maxiter"), verbose=self.get_options("verbose"), 
+                               recycle=self.get_options("recycle"))
+
     def solve(self, b, out=None):
         """
         Use MINimum RESidual iteration to solve Ax=b
@@ -1399,6 +1429,11 @@ def __init__(self, A, *, x0=None, tol=None, atol=None, btol=None, maxiter=1000,
         tmps_codomain = {key: self.codomain.zeros() for key in ("v", "v_work", "h", "hbar")}
         self._tmps = {**tmps_codomain, **tmps_domain}
 
+    def copy(self):
+        return LSMR(self.linop, x0=self.get_options("x0"), tol=self.get_options("tol"), atol=self.get_options("atol"), 
+                    btol=self.get_options("btol"), maxiter=self.get_options("maxiter"), conlim=self.get_options("conlim"), 
+                    verbose=self.get_options("verbose"), recycle=self.get_options("recycle"))
+
     def get_success(self):
         return self._successful
 
@@ -1740,6 +1775,11 @@ def __init__(self, A, *, x0=None, tol=1e-6, maxiter=100, verbose=False, recycle=
         self._Q = []
         self._info = None
 
+    def copy(self):
+        return GMRES(self.linop, x0=self.get_options("x0"), tol=self.get_options("tol"), 
+                                 maxiter=self.get_options("maxiter"), verbose=self.get_options("verbose"), 
+                                 recycle=self.get_options("recycle"))
+
     def solve(self, b, out=None):
         """
         Generalized minimal residual algorithm for solving linear system Ax=b.
@@ -1911,4 +1951,3 @@ def apply_givens_rotation(self, k, sn, cn):
 
     def dot(self, b, out=None):
         return self.solve(b, out=out)
-