solver: tweak

usadelndsoc/solver.py

@@ -10,7 +10,16 @@ import hashlib
 import tempfile
 import pickle
 from scipy.sparse import coo_matrix
-from scipy.sparse.linalg import splu, spilu, gcrotmk, LinearOperator, lgmres, gmres, cg
+from scipy.sparse.linalg import (
+    splu,
+    spilu,
+    gcrotmk,
+    LinearOperator,
+    lgmres,
+    gmres,
+    cg,
+    tfqmr,
+)
 from scipy.special import zeta
 from scipy.linalg import blas
 
@@ -85,15 +94,16 @@ class Core:
         self._U[:, 1:, 3, :, :] = np.linalg.inv(self._U[:, :-1, 1, :, :])
 
     def new_action(self):
+        self.init_U()
         return Action(self.Lx, self.Ly, self.alpha, self._mask, self._U, self._Omega)
 
     @property
     def x(self):
-        return np.linspace(-self.Lx / 2, self.Lx / 2, self._nx)
+        return np.linspace(-self.Lx / 2, self.Lx / 2, self._shape[0])
 
     @property
     def y(self):
-        return np.linspace(-self.Ly / 2, self.Ly / 2, self._ny)
+        return np.linspace(-self.Ly / 2, self.Ly / 2, self._shape[1])
 
     U = _array_property("_U")
     Omega = _array_property("_Omega")
@@ -135,6 +145,7 @@ class Solver:
     def _solve(
         self,
         eval_rhs,
+        eval_jac_mul,
         eval_jac,
         check_step,
         check_final,
@@ -145,6 +156,7 @@ class Solver:
         preconditioner=None,
         solver=None,
         restart=0,
+        finite_difference=True,
     ):
         A_size = 2 * Phi0.size
         A_shape = (A_size, A_size)
@@ -171,7 +183,13 @@ class Solver:
         skip_precond = 0
 
         if solver is None:
-            solver = "gcrotmk"
+            solver = "lgmres"
+
+        if A_size > 20000:
+            _log_solve.debug(
+                f"    Problem too large ({A_size}): cannot compute hessian"
+            )
+            preconditioner = "none"
 
         if not np.isfinite(Phi).all():
             Phi = np.zeros_like(Phi)
@@ -219,13 +237,17 @@ class Solver:
                 nonlocal count
                 count += 1
                 v = v.view(np.complex128).reshape(Phi.shape)
-                scale = rdiff * max(1, Phi_norm) / max(1, _norm(v))
-                return (eval_rhs(Phi + scale * v) - F) / scale
+                if finite_difference:
+                    scale = rdiff * max(1, Phi_norm) / max(1, _norm(v))
+                    return (eval_rhs(Phi + scale * v) - F) / scale
+                else:
+                    return eval_jac_mul(Phi, v)
 
             A = LinearOperator(matvec=op, shape=A_shape, dtype=A_dtype)
 
             l_maxiter = 2
             l_inner_m = 15
+            l_outer_k = 5
 
             if solver == "cg":
                 dx, info = cg(
@@ -236,6 +258,15 @@ class Solver:
                     atol=0,
                     maxiter=l_maxiter,
                 )
+            elif solver == "tfqmr":
+                dx, info = tfqmr(
+                    A,
+                    -F,
+                    M=M,
+                    tol=1e-3,
+                    atol=0,
+                    maxiter=l_maxiter,
+                )
             elif solver == "lgmres":
                 dx, info = lgmres(
                     A,
@@ -246,8 +277,9 @@ class Solver:
                     maxiter=l_maxiter,
                     inner_m=l_inner_m,
                     outer_v=outer_v,
+                    outer_k=l_outer_k,
                     store_outer_Av=False,
-                    prepend_outer_v=False,
+                    prepend_outer_v=True,
                 )
             elif solver == "gcrotmk":
                 dx, info = gcrotmk(
@@ -315,6 +347,7 @@ class Solver:
                 Phi[...] = 0
                 return self._solve(
                     eval_rhs,
+                    eval_jac_mul,
                     eval_jac,
                     check_step,
                     check_final,
@@ -341,7 +374,6 @@ class Solver:
         return Phi
 
     def _check_step(self, Phi, dx):
-        return True
         return (abs(Phi + dx) < 1.0).all()
 
     def _check_final(self, Phi):
@@ -351,6 +383,9 @@ class Solver:
     def _eval_rhs(self, Phi):
         return self._action.grad(Phi).ravel().view(np.float_)
 
+    def _eval_jac_mul(self, Phi, dPhi):
+        return self._action.hess_mul(Phi, dPhi).ravel().view(np.float_)
+
     def _eval_jac(self, Phi):
         # complex jacobian of S
         H = self._action.hess(Phi)
@@ -372,10 +407,13 @@ class Solver:
         Phi00 = np.zeros_like(self._Phi)
         self._action = kw.pop("action", None)
         if self._action is None:
+            _log_solve.debug("    Initialize action")
             self._action = self._core.new_action()
+            self._action.compute(Phi00)
         self._action.omega = omega
         self._Phi = self._solve(
             self._eval_rhs,
+            self._eval_jac_mul,
             self._eval_jac,
             self._check_step,
             self._check_final,
@@ -428,7 +466,9 @@ class Solver:
         self._M = None
         self._Phi[...] = 0
 
+        _log_solve.debug("    Initialize action")
         action = self._core.new_action()
+        action.compute(self._Phi)
 
         with np.nditer(omega, ["multi_index"]) as it:
             for v in it: