From 652b29fdb813dce9836b34c0fd789c99cf779f17 Mon Sep 17 00:00:00 2001
From: e-moral-sanchez <88042165+e-moral-sanchez@users.noreply.github.com>
Date: Thu, 12 Oct 2023 16:41:00 +0200
Subject: [PATCH] Fix dtype in GMRES and MINRES (#345)

In the iterative linear solvers take `dtype` from the domain or the
codomain of the linear operator to be "inverted". This avoids reading
the `dtype` property of the linear operator, which may not be properly
defined in some cases. Specifically, the classes `GMRES` and
`MinimumResidual` were fixed.
---
 psydac/linalg/solvers.py            |  6 +++---
 psydac/linalg/tests/test_solvers.py | 17 ++++++++++++++++-
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/psydac/linalg/solvers.py b/psydac/linalg/solvers.py
index d9e0d5b7e..93817ac51 100644
--- a/psydac/linalg/solvers.py
+++ b/psydac/linalg/solvers.py
@@ -994,7 +994,7 @@ def __init__(self, A, *, x0=None, tol=1e-6, maxiter=1000, verbose=False, recycle
 
         assert isinstance(A, LinearOperator)
         assert A.domain.dimension == A.codomain.dimension
-        assert A.dtype == float
+        assert A.domain.dtype == float
         domain = A.codomain
         codomain = A.domain
 
@@ -1731,7 +1731,7 @@ def __init__(self, A, *, x0=None, tol=1e-6, maxiter=100, verbose=False, recycle=
         self._tmps = {key: domain.zeros() for key in ("r", "p")}
 
         # Initialize upper Hessenberg matrix
-        self._H = np.zeros((self._options["maxiter"] + 1, self._options["maxiter"]), dtype=A.dtype)
+        self._H = np.zeros((self._options["maxiter"] + 1, self._options["maxiter"]), dtype=A.domain.dtype)
         self._Q = []
         self._info = None
 
@@ -1879,7 +1879,7 @@ def solve(self, b, out=None):
     def solve_triangular(self, T, d):
         # Backwards substitution. Assumes T is upper triangular
         k = T.shape[0]
-        y = np.zeros((k,), dtype=self._A.dtype)
+        y = np.zeros((k,), dtype=self._A.domain.dtype)
 
         for k1 in range(k):
             temp = 0.
diff --git a/psydac/linalg/tests/test_solvers.py b/psydac/linalg/tests/test_solvers.py
index a91e4f5cb..803877257 100644
--- a/psydac/linalg/tests/test_solvers.py
+++ b/psydac/linalg/tests/test_solvers.py
@@ -105,6 +105,7 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
     solv  = inverse(A, solver, tol=tol, verbose=False, recycle=True)
     solvt = solv.transpose()
     solvh = solv.H
+    solv2 = inverse(A@A, solver, tol=1e-13, verbose=True, recycle=True) # Test solver of composition of operators
 
     # Manufacture right-hand-side vector from exact solution
     be  = A @ xe
@@ -135,11 +136,21 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
     assert np.array_equal(xh.toarray(), solvh_x0.toarray())
     assert xh is not solvh_x0
 
+    if solver != 'pcg':
+        # PCG only works with operators with diagonal
+        xc = solv2 @ be2
+        solv2_x0 = solv2._options["x0"]
+        assert np.array_equal(xc.toarray(), solv2_x0.toarray())
+        assert xc is not solv2_x0
+
+
     # Verify correctness of calculation: 2-norm of error
     b = A @ x
     b2 = A @ x2
     bt = A.T @ xt
     bh = A.H @ xh
+    if solver != 'pcg':
+        bc = A @ A @ xc
 
     
     err = b - be
@@ -151,6 +162,10 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
     errh = bh - beh
     errh_norm = np.linalg.norm( errh.toarray() )
 
+    if solver != 'pcg': 
+        errc = bc - be2
+        errc_norm = np.linalg.norm( errc.toarray() )
+
     #---------------------------------------------------------------------------
     # TERMINAL OUTPUT
     #---------------------------------------------------------------------------
@@ -180,7 +195,7 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
         assert err2_norm < tol
         assert errt_norm < tol
         assert errh_norm < tol
-
+        assert solver == 'pcg' or errc_norm < tol
 
 # ===============================================================================
 # SCRIPT FUNCTIONALITY