Solver template

usadelndsoc/solver.py

+# -*- coding:utf-8; eval: (blacken-mode) -*-
+import os
+import io
+import logging
+import numpy as np
+import concurrent.futures
+import functools
+import collections
+import hashlib
+import tempfile
+import pickle
+from scipy.sparse.linalg import splu, spilu, gcrotmk, LinearOperator, lgmres, gmres
+from scipy.special import zeta
+from scipy.linalg import blas
+
+try:
+    from scikits import umfpack
+except ImportError:
+    umfpack = None
+
+try:
+    import pyamg
+except ImportError:
+    pyamg = None
+
+from ._core import Action, MASK_NONE, MASK_TERMINAL, MASK_VACUUM
+from .matsubara import get_matsubara_sum
+
+__all__ = ["Solver", "Result", "cpr", "MASK_NONE", "MASK_TERMINAL", "MASK_VACUUM"]
+
+_log = logging.getLogger(__name__)
+_log_solve = logging.getLogger(__name__ + ".solve")
+
+
+def _core_property(name):
+    def fget(self):
+        return getattr(self._core, name)
+
+    def fset(self, value):
+        setattr(self._core, name, value)
+
+    return property(fget, fset)
+
+
+def _array_property(name):
+    def fget(self):
+        return getattr(self, name)
+
+    return property(fget)
+
+
+def _norm(z):
+    if z.dtype.char == "D":
+        if z.flags.f_contiguous:
+            return blas.dznrm2(z.ravel("F"))
+        else:
+            return blas.dznrm2(z.ravel())
+    elif z.dtype.char == "d":
+        if z.flags.f_contiguous:
+            return blas.dnrm2(z.ravel("F"))
+        else:
+            return blas.dnrm2(z.ravel())
+    else:
+        raise ValueError("Invalid data type")
+
+
+class Core:
+    def __init__(self, Lx=1, Ly=1, alpha=0, nx=1, ny=1):
+        self.Lx = 1
+        self.Ly = 1
+        self.alpha = 0
+        self.set_shape(nx, ny)
+
+    def set_shape(self, nx=1, ny=1):
+        self._shape = (nx, ny)
+        self._mask = np.zeros((nx, ny), dtype=np.uint8)
+        self._U = np.zeros((nx, ny, 4, 4, 4), dtype=complex)
+        self._Omega = np.zeros((nx, ny, 4, 4), dtype=complex)
+        self._action = None
+
+    def new_action(self):
+        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)
+
+    @property
+    def y(self):
+        return np.linspace(-self.Ly / 2, self.Ly / 2, self._ny)
+
+    U = _array_property("_U")
+    Omega = _array_property("_Omega")
+    mask = _array_property("_mask")
+    shape = property(fget=lambda self: self._shape)
+
+
+class Solver:
+    def __init__(self, nx=1, ny=1):
+        self._core = Core()
+        self.set_shape(nx, ny)
+        self._action = None
+        self._M = None
+        self.omega = 0
+
+    def set_shape(self, nx, ny=1):
+        self._core.set_shape(nx, ny)
+        self._Phi = np.zeros((nx, ny, 2, 2, 2), dtype=np.complex_, order="C")
+        self._J_tot = np.zeros((nx, ny), dtype=np.complex_, order="C")
+
+    shape = _core_property("shape")
+    Lx = _core_property("Lx")
+    Ly = _core_property("Ly")
+    x = _core_property("x")
+    y = _core_property("y")
+    mask = _core_property("mask_view")
+    Omega = _core_property("Omega")
+
+    Phi = _array_property("_Phi")
+    J_tot = _array_property("_J_tot")
+
+    def _eval_J(self, Phi):
+        return self._core.eval_J(Phi)
+
+    @property
+    def J(self):
+        return self._eval_J(self.Phi)
+
+    def _solve(
+        self,
+        eval_rhs,
+        eval_jac,
+        check_step,
+        check_final,
+        Phi0,
+        Phi00,
+        tol=1e-6,
+        maxiter=30,
+        preconditioner=None,
+        solver=None,
+        restart=0,
+    ):
+        A_size = 2 * Phi0.size
+        A_shape = (A_size, A_size)
+        A_dtype = np.dtype(float)
+
+        F0 = eval_rhs(Phi00)
+        F0_norm = _norm(F0)
+
+        if F0_norm == 0:
+            F0_norm = tol
+
+        outer_v = []
+
+        if self._M is not None and self._M.shape == A_shape:
+            M = self._M
+        else:
+            M = None
+        self._M = None
+
+        Phi = Phi0
+
+        update_skipped = False
+        tiny_count = 0
+        skip_precond = 0
+
+        if solver is None:
+            solver = "gcrotmk"
+
+        if not np.isfinite(Phi).all():
+            Phi = np.zeros_like(Phi)
+            M = self._M = None
+            restart = 1
+        elif (Phi == 0).all():
+            restart = 1
+
+        for j in range(maxiter):
+            if skip_precond > 0:
+                M = None
+                skip_precond -= 1
+            elif M is None and preconditioner != "none":
+                # Update preconditioner
+                _log_solve.debug("    Update Jacobian")
+                F, J = eval_jac(Phi)
+
+                _log_solve.debug("    Form preconditioner")
+                M = None
+                M = self._get_preconditioner(J, preconditioner=preconditioner)
+                J = None
+                _log_solve.debug("    Preconditioner formed")
+                frozen_M = False
+            else:
+                # Don't update preconditioner
+                F = eval_rhs(Phi)
+                frozen_M = True
+
+            F_norm = _norm(F)
+            _log_solve.info(f"    #{j}: residual = {F_norm / F0_norm}")
+
+            if F_norm <= tol * F0_norm:
+                if check_final(Phi):
+                    break
+
+            # Solve linearized problem
+            count = 0
+
+            Phi_norm = _norm(Phi)
+
+            def op(v, rdiff=1e-4):
+                nonlocal count
+                count += 1
+                v = v.view(np.complex128).reshape(Phi.shape, order="F")
+                scale = rdiff * max(1, Phi_norm) / max(1, _norm(v))
+                return (eval_rhs(Phi + scale * v) - F) / scale
+
+            A = LinearOperator(matvec=op, shape=A_shape, dtype=A_dtype)
+
+            l_maxiter = 2
+            l_inner_m = 15
+
+            if solver == "lgmres":
+                dx, info = lgmres(
+                    A,
+                    -F,
+                    M=M,
+                    tol=1e-3,
+                    atol=0,
+                    maxiter=l_maxiter,
+                    inner_m=l_inner_m,
+                    outer_v=outer_v,
+                    store_outer_Av=False,
+                    prepend_outer_v=False,
+                )
+            elif solver == "gcrotmk":
+                dx, info = gcrotmk(
+                    A,
+                    -F,
+                    M=M,
+                    tol=1e-3,
+                    atol=0,
+                    maxiter=l_maxiter,
+                    m=l_inner_m,
+                )
+            elif solver == "gmres":
+                dx, info = gmres(
+                    A,
+                    -F,
+                    M=M,
+                    tol=1e-3,
+                    atol=0,
+                    maxiter=l_maxiter,
+                    restart=l_inner_m,
+                )
+            else:
+                raise ValueError(f"Unknown {solver=}")
+
+            if count >= l_maxiter * l_inner_m // 2:
+                M = None
+
+            # Update
+            dx = dx.view(np.complex128).reshape(Phi.shape, order="F")
+            s = abs(dx).max()
+
+            _log_solve.debug(f"    {count} matvec, step {s}")
+
+            do_restart = False
+
+            if s >= 0.75:
+                M = None
+                dx *= 0.75 / s
+
+                if not check_step(Phi, dx):
+                    if frozen_M and not update_skipped:
+                        _log_solve.info("Bad step: force Jacobian update")
+                        update_skipped = True
+                        continue
+                    elif restart == 0:
+                        do_restart = True
+
+                tiny_count = 0
+            elif s == 0:
+                if restart == 0:
+                    do_restart = True
+                else:
+                    skip_precond = tiny_count + 2
+            elif s < 1e-10:
+                tiny_count += 1
+                if tiny_count > 5:
+                    if restart == 0:
+                        do_restart = True
+                    else:
+                        skip_precond = tiny_count + 2
+
+            if do_restart:
+                _log_solve.info("Force restart from Ansatz")
+                self._M = None
+                Phi[...] = 0
+                return self._solve(
+                    eval_rhs,
+                    eval_jac,
+                    check_step,
+                    check_final,
+                    Phi,
+                    Phi00,
+                    tol=tol,
+                    maxiter=maxiter,
+                    preconditioner=preconditioner,
+                    restart=1,
+                    solver=solver,
+                )
+
+            Phi += dx
+            update_skipped = False
+        else:
+            # Fail to solve
+            _log_solve.error("Failed to solve (maxiter)")
+            Phi[...] = np.nan
+            M = None
+
+        self._M = M
+
+        return Phi
+
+    def _check_step(self, Phi, dx):
+        return (abs(Phi + dx) < 1.0).all()
+
+    def _check_final(self, Phi):
+        # Check branch
+        if abs(Phi).max() > 1.05:
+            Phi[...] *= 0.5
+            _log_solve.info("    Force restart: wrong branch")
+            return False
+        g = (1 - Phi[0] * Phi[1]) / (1 + Phi[0] * Phi[1])
+        if g.real.min() < -0.05:
+            Phi[0] *= -0.5
+            Phi[1] *= 0.5
+            _log_solve.info("    Force restart: wrong branch")
+            return False
+        return True
+
+    def _eval_rhs(self, Phi):
+        return self._action.grad(Phi)
+
+    def _eval_jac(self, Phi):
+        return self._action.hess(Phi)
+
+    @functools.wraps(_solve)
+    def solve(self, omega, **kw):
+        _log_solve.info(f"omega = {self.omega}")
+        nx, ny = self.shape
+        Phi00 = np.zeros_like(self._Phi)
+        self._action = kw.pop("action", None)
+        if self._action is None:
+            self._action = self._core.new_action()
+        self._action.omega = omega
+        self._Phi = self._solve(
+            self._eval_rhs,
+            self._eval_jac,
+            self._check_step,
+            self._check_final,
+            self._Phi,
+            Phi00,
+            **kw,
+        )
+        return Result(self)
+
+    def _get_preconditioner(self, J, preconditioner=None):
+        if preconditioner is None:
+            sz = self.shape[0] * self.shape[1] * self.shape[2] / sum(self.shape)
+            if sz > 1000 and pyamg is not None:
+                preconditioner = "pyamg"
+            elif umfpack is not None:
+                preconditioner = "umfpack-ilu"
+            else:
+                preconditioner = "splu"
+
+        if preconditioner == "splu":
+            try:
+                J_lu = splu(J.tocsc(copy=True))
+            except RuntimeError:
+                _log_solve.error("Preconditioner failed: using none")
+                return None
+            M = LinearOperator(matvec=J_lu.solve, shape=J.shape, dtype=J.dtype)
+        elif preconditioner == "spilu":
+            try:
+                J_lu = spilu(J.tocsc(copy=True), drop_tol=1e-4, fill_factor=2.0)
+            except RuntimeError:
+                _log_solve.error("Preconditioner failed: using none")
+                return None
+            M = LinearOperator(matvec=J_lu.solve, shape=J.shape, dtype=J.dtype)
+        elif preconditioner == "pyamg" and pyamg is not None:
+            M = pyamg.ruge_stuben_solver(J.tocsr(copy=True)).aspreconditioner()
+        elif preconditioner == "umfpack" and umfpack is not None:
+            M = UmfpackILU(J.tocsc(copy=True), drop_tol=0)
+        elif preconditioner == "umfpack-ilu" and umfpack is not None:
+            M = UmfpackILU(J.tocsc(copy=True), drop_tol=1e-3)
+        else:
+            raise ValueError(f"Unknown {preconditioner=}")
+
+        return M
+
+    def solve_many(self, omega, **kw):
+        omega = np.asarray(omega)
+
+        Phi = np.zeros(self.Phi.shape + omega.shape, dtype=self.Phi.dtype)
+
+        self._M = None
+        self._Phi[...] = 0
+
+        action = self._core.new_action()
+
+        with np.nditer(omega, ["multi_index"]) as it:
+            for v in it:
+                self.solve(omega=v, action=action, **kw)
+                Phi_prev = self._Phi.copy()
+                Phi[(Ellipsis,) + it.multi_index] = self._Phi
+                if np.isnan(self._Phi).any():
+                    self._Phi[...] = Phi_prev
+
+        return Result(self, omega=omega, Phi=Phi)
+
+
+class Result:
+    def __init__(self, parent, omega=None, Phi=None, core=None):
+        self._core = core if core is not None else parent._core
+        self.shape = parent.shape
+        self.Lx = parent.Lx
+        self.Ly = parent.Ly
+        self.Lz = parent.Lz
+        self.omega = omega if omega is not None else parent.omega
+        self.Delta = parent.Delta
+        self.Phi = Phi if Phi is not None else parent.Phi
+        self.mask = parent.mask
+
+    @property
+    def G(self):
+        s = np.sign(self.omega)
+        with np.errstate(invalid="ignore"):
+            return (1 - self.Phi[0] * self.Phi[1]) / (1 + self.Phi[0] * self.Phi[1]) * s
+
+    @property
+    def F(self):
+        s = np.sign(self.omega)
+        with np.errstate(invalid="ignore"):
+            return 2 * self.Phi[0] / (1 + self.Phi[0] * self.Phi[1]) * s
+
+    @property
+    def Fc(self):
+        with np.errstate(invalid="ignore"):
+            return 2 * self.Phi[1] / (1 + self.Phi[0] * self.Phi[1])
+
+    @property
+    def g(self):
+        with np.errstate(invalid="ignore"):
+            G = self.G
+            return np.array([[G, self.F], [self.Fc, -G]])
+
+    @property
+    def J(self):
+        return self._core.eval_J(self.Phi)
+
+
+class UmfpackILU(LinearOperator):
+    def __init__(self, M, drop_tol=1e-4):
+        if M.dtype != np.float64 or M.shape[0] != M.shape[1] or M.format != "csc":
+            raise ValueError("Only square, csc, float64 supported")
+
+        super().__init__(dtype=M.dtype, shape=M.shape)
+
+        self.M = M
+        self.ctx = umfpack.UmfpackContext(
+            family=("di" if M.indptr.dtype == np.int32 else "dl")
+        )
+        if drop_tol:
+            self.ctx.control[umfpack.UMFPACK_DROPTOL] = drop_tol
+        self.ctx.numeric(self.M)
+
+    def _matvec(self, v):
+        return self.ctx.solve(umfpack.UMFPACK_A, self.M, v)