text: # Copyright (C) 2015, Pauli Virtanen <pav@iki.fi>
  # Distributed under the same license as SciPy.
  
  import warnings
  import numpy as np
  from numpy.linalg import LinAlgError
  from scipy.linalg import (get_blas_funcs, qr, solve, svd, qr_insert, lstsq)
  from scipy.sparse.linalg.isolve.utils import make_system
  
  
  __all__ = ['gcrotmk']
  
  
  def _fgmres(matvec, v0, m, atol, lpsolve=None, rpsolve=None, cs=(), outer_v=(),
              prepend_outer_v=False):
      """
      FGMRES Arnoldi process, with optional projection or augmentation
  
      Parameters
      ----------
      matvec : callable
          Operation A*x
      v0 : ndarray
          Initial vector, normalized to nrm2(v0) == 1
      m : int
          Number of GMRES rounds
      atol : float
          Absolute tolerance for early exit
      lpsolve : callable
          Left preconditioner L
      rpsolve : callable
          Right preconditioner R
      CU : list of (ndarray, ndarray)
          Columns of matrices C and U in GCROT
      outer_v : list of ndarrays
          Augmentation vectors in LGMRES
      prepend_outer_v : bool, optional
          Whether augmentation vectors come before or after 
          Krylov iterates
  
      Raises
      ------
      LinAlgError
          If nans encountered
  
      Returns
      -------
      Q, R : ndarray
          QR decomposition of the upper Hessenberg H=QR
      B : ndarray
          Projections corresponding to matrix C
      vs : list of ndarray
          Columns of matrix V
      zs : list of ndarray
          Columns of matrix Z
      y : ndarray
          Solution to ||H y - e_1||_2 = min!
      res : float
          The final (preconditioned) residual norm
  
      """
  
      if lpsolve is None:
          lpsolve = lambda x: x
      if rpsolve is None:
          rpsolve = lambda x: x
  
      axpy, dot, scal, nrm2 = get_blas_funcs(['axpy', 'dot', 'scal', 'nrm2'], (v0,))
  
      vs = [v0]
      zs = []
      y = None
      res = np.nan
  
      m = m + len(outer_v)
  
      # Orthogonal projection coefficients
      B = np.zeros((len(cs), m), dtype=v0.dtype)
  
      # H is stored in QR factorized form
      Q = np.ones((1, 1), dtype=v0.dtype)
      R = np.zeros((1, 0), dtype=v0.dtype)
  
      eps = np.finfo(v0.dtype).eps
  
      breakdown = False
  
      # FGMRES Arnoldi process
      for j in range(m):
          # L A Z = C B + V H
  
          if prepend_outer_v and j < len(outer_v):
              z, w = outer_v[j]
          elif prepend_outer_v and j == len(outer_v):
              z = rpsolve(v0)
              w = None
          elif not prepend_outer_v and j >= m - len(outer_v):
              z, w = outer_v[j - (m - len(outer_v))]
          else:
              z = rpsolve(vs[-1])
              w = None
  
          if w is None:
              w = lpsolve(matvec(z))
          else:
              # w is clobbered below
              w = w.copy()
  
          w_norm = nrm2(w)
  
          # GCROT projection: L A -> (1 - C C^H) L A
          # i.e. orthogonalize against C
          for i, c in enumerate(cs):
              alpha = dot(c, w)
              B[i,j] = alpha
              w = axpy(c, w, c.shape[0], -alpha)  # w -= alpha*c
  
          # Orthogonalize against V
          hcur = np.zeros(j+2, dtype=Q.dtype)
          for i, v in enumerate(vs):
              alpha = dot(v, w)
              hcur[i] = alpha
              w = axpy(v, w, v.shape[0], -alpha)  # w -= alpha*v
          hcur[i+1] = nrm2(w)
  
          with np.errstate(over='ignore', divide='ignore'):
              # Careful with denormals
              alpha = 1/hcur[-1]
  
          if np.isfinite(alpha):
              w = scal(alpha, w)
  
          if not (hcur[-1] > eps * w_norm):
              # w essentially in the span of previous vectors,
              # or we have nans. Bail out after updating the QR
              # solution.
              breakdown = True
  
          vs.append(w)
          zs.append(z)
  
          # Arnoldi LSQ problem
  
          # Add new column to H=Q*R, padding other columns with zeros
          Q2 = np.zeros((j+2, j+2), dtype=Q.dtype, order='F')
          Q2[:j+1,:j+1] = Q
          Q2[j+1,j+1] = 1
  
          R2 = np.zeros((j+2, j), dtype=R.dtype, order='F')
          R2[:j+1,:] = R
  
          Q, R = qr_insert(Q2, R2, hcur, j, which='col',
                           overwrite_qru=True, check_finite=False)
  
          # Transformed least squares problem
          # || Q R y - inner_res_0 * e_1 ||_2 = min!
          # Since R = [R'; 0], solution is y = inner_res_0 (R')^{-1} (Q^H)[:j,0]
  
          # Residual is immediately known
          res = abs(Q[0,-1])
  
          # Check for termination
          if res < atol or breakdown:
              break
  
      if not np.isfinite(R[j,j]):
          # nans encountered, bail out
          raise LinAlgError()
  
      # -- Get the LSQ problem solution
  
      # The problem is triangular, but the condition number may be
      # bad (or in case of breakdown the last diagonal entry may be
      # zero), so use lstsq instead of trtrs.
      y, _, _, _, = lstsq(R[:j+1,:j+1], Q[0,:j+1].conj())
  
      B = B[:,:j+1]
  
      return Q, R, B, vs, zs, y, res
  
  
  def gcrotmk(A, b, x0=None, tol=1e-5, maxiter=1000, M=None, callback=None,
              m=20, k=None, CU=None, discard_C=False, truncate='oldest',
              atol=None):
      """
      Solve a matrix equation using flexible GCROT(m,k) algorithm.
  
      Parameters
      ----------
      A : {sparse matrix, dense matrix, LinearOperator}
          The real or complex N-by-N matrix of the linear system.
          Alternatively, ``A`` can be a linear operator which can
          produce ``Ax`` using, e.g.,
          ``scipy.sparse.linalg.LinearOperator``.
      b : {array, matrix}
          Right hand side of the linear system. Has shape (N,) or (N,1).
      x0  : {array, matrix}
          Starting guess for the solution.
      tol, atol : float, optional
          Tolerances for convergence, ``norm(residual) <= max(tol*norm(b), atol)``.
          The default for ``atol`` is `tol`.
  
          .. warning::
  
             The default value for `atol` will be changed in a future release.
             For future compatibility, specify `atol` explicitly.
      maxiter : int, optional
          Maximum number of iterations.  Iteration will stop after maxiter
          steps even if the specified tolerance has not been achieved.
      M : {sparse matrix, dense matrix, LinearOperator}, optional
          Preconditioner for A.  The preconditioner should approximate the
          inverse of A. gcrotmk is a 'flexible' algorithm and the preconditioner
          can vary from iteration to iteration. Effective preconditioning
          dramatically improves the rate of convergence, which implies that
          fewer iterations are needed to reach a given error tolerance.
      callback : function, optional
          User-supplied function to call after each iteration.  It is called
          as callback(xk), where xk is the current solution vector.
      m : int, optional
          Number of inner FGMRES iterations per each outer iteration.
          Default: 20
      k : int, optional
          Number of vectors to carry between inner FGMRES iterations.
          According to [2]_, good values are around m.
          Default: m
      CU : list of tuples, optional
          List of tuples ``(c, u)`` which contain the columns of the matrices
          C and U in the GCROT(m,k) algorithm. For details, see [2]_.
          The list given and vectors contained in it are modified in-place.
          If not given, start from empty matrices. The ``c`` elements in the
          tuples can be ``None``, in which case the vectors are recomputed
          via ``c = A u`` on start and orthogonalized as described in [3]_.
      discard_C : bool, optional
          Discard the C-vectors at the end. Useful if recycling Krylov subspaces
          for different linear systems.
      truncate : {'oldest', 'smallest'}, optional
          Truncation scheme to use. Drop: oldest vectors, or vectors with
          smallest singular values using the scheme discussed in [1,2].
          See [2]_ for detailed comparison.
          Default: 'oldest'
  
      Returns
      -------
      x : array or matrix
          The solution found.
      info : int
          Provides convergence information:
  
          * 0  : successful exit
          * >0 : convergence to tolerance not achieved, number of iterations
  
      References
      ----------
      .. [1] E. de Sturler, ''Truncation strategies for optimal Krylov subspace
             methods'', SIAM J. Numer. Anal. 36, 864 (1999).
      .. [2] J.E. Hicken and D.W. Zingg, ''A simplified and flexible variant
             of GCROT for solving nonsymmetric linear systems'',
             SIAM J. Sci. Comput. 32, 172 (2010).
      .. [3] M.L. Parks, E. de Sturler, G. Mackey, D.D. Johnson, S. Maiti,
             ''Recycling Krylov subspaces for sequences of linear systems'',
             SIAM J. Sci. Comput. 28, 1651 (2006).
  
      """
      A,M,x,b,postprocess = make_system(A,M,x0,b)
  
      if not np.isfinite(b).all():
          raise ValueError("RHS must contain only finite numbers")
  
      if truncate not in ('oldest', 'smallest'):
          raise ValueError("Invalid value for 'truncate': %r" % (truncate,))
  
      if atol is None:
          warnings.warn("scipy.sparse.linalg.gcrotmk called without specifying `atol`. "
                        "The default value will change in the future. To preserve "
                        "current behavior, set ``atol=tol``.",
                        category=DeprecationWarning, stacklevel=2)
          atol = tol
  
      matvec = A.matvec
      psolve = M.matvec
  
      if CU is None:
          CU = []
  
      if k is None:
          k = m
  
      axpy, dot, scal = None, None, None
  
      r = b - matvec(x)
  
      axpy, dot, scal, nrm2 = get_blas_funcs(['axpy', 'dot', 'scal', 'nrm2'], (x, r))
  
      b_norm = nrm2(b)
  
      if discard_C:
          CU[:] = [(None, u) for c, u in CU]
  
      # Reorthogonalize old vectors
      if CU:
          # Sort already existing vectors to the front
          CU.sort(key=lambda cu: cu[0] is not None)
  
          # Fill-in missing ones
          C = np.empty((A.shape[0], len(CU)), dtype=r.dtype, order='F')
          us = []
          j = 0
          while CU:
              # More memory-efficient: throw away old vectors as we go
              c, u = CU.pop(0)
              if c is None:
                  c = matvec(u)
              C[:,j] = c
              j += 1
              us.append(u)
  
          # Orthogonalize
          Q, R, P = qr(C, overwrite_a=True, mode='economic', pivoting=True)
          del C
  
          # C := Q
          cs = list(Q.T)
  
          # U := U P R^-1,  back-substitution
          new_us = []
          for j in range(len(cs)):
              u = us[P[j]]
              for i in range(j):
                  u = axpy(us[P[i]], u, u.shape[0], -R[i,j])
              if abs(R[j,j]) < 1e-12 * abs(R[0,0]):
                  # discard rest of the vectors
                  break
              u = scal(1.0/R[j,j], u)
              new_us.append(u)
  
          # Form the new CU lists
          CU[:] = list(zip(cs, new_us))[::-1]
  
      if CU:
          axpy, dot = get_blas_funcs(['axpy', 'dot'], (r,))
  
          # Solve first the projection operation with respect to the CU
          # vectors. This corresponds to modifying the initial guess to
          # be
          #
          #     x' = x + U y
          #     y = argmin_y || b - A (x + U y) ||^2
          #
          # The solution is y = C^H (b - A x)
          for c, u in CU:
              yc = dot(c, r)
              x = axpy(u, x, x.shape[0], yc)
              r = axpy(c, r, r.shape[0], -yc)
  
      # GCROT main iteration
      for j_outer in range(maxiter):
          # -- callback
          if callback is not None:
              callback(x)
  
          beta = nrm2(r)
  
          # -- check stopping condition
          beta_tol = max(atol, tol * b_norm)
  
          if beta <= beta_tol and (j_outer > 0 or CU):
              # recompute residual to avoid rounding error
              r = b - matvec(x)
              beta = nrm2(r)
  
          if beta <= beta_tol:
              j_outer = -1
              break
  
          ml = m + max(k - len(CU), 0)
  
          cs = [c for c, u in CU]
  
          try:
              Q, R, B, vs, zs, y, pres = _fgmres(matvec,
                                                 r/beta,
                                                 ml,
                                                 rpsolve=psolve,
                                                 atol=max(atol, tol*b_norm)/beta,
                                                 cs=cs)
              y *= beta
          except LinAlgError:
              # Floating point over/underflow, non-finite result from
              # matmul etc. -- report failure.
              break
  
          #
          # At this point,
          #
          #     [A U, A Z] = [C, V] G;   G =  [ I  B ]
          #                                   [ 0  H ]
          #
          # where [C, V] has orthonormal columns, and r = beta v_0. Moreover,
          #
          #     || b - A (x + Z y + U q) ||_2 = || r - C B y - V H y - C q ||_2 = min!
          #
          # from which y = argmin_y || beta e_1 - H y ||_2, and q = -B y
          #
  
          #
          # GCROT(m,k) update
          #
  
          # Define new outer vectors
  
          # ux := (Z - U B) y
          ux = zs[0]*y[0]
          for z, yc in zip(zs[1:], y[1:]):
              ux = axpy(z, ux, ux.shape[0], yc)  # ux += z*yc
          by = B.dot(y)
          for cu, byc in zip(CU, by):
              c, u = cu
              ux = axpy(u, ux, ux.shape[0], -byc)  # ux -= u*byc
  
          # cx := V H y
          hy = Q.dot(R.dot(y))
          cx = vs[0] * hy[0]
          for v, hyc in zip(vs[1:], hy[1:]):
              cx = axpy(v, cx, cx.shape[0], hyc)  # cx += v*hyc
  
          # Normalize cx, maintaining cx = A ux
          # This new cx is orthogonal to the previous C, by construction
          try:
              alpha = 1/nrm2(cx)
              if not np.isfinite(alpha):
                  raise FloatingPointError()
          except (FloatingPointError, ZeroDivisionError):
              # Cannot update, so skip it
              continue
  
          cx = scal(alpha, cx)
          ux = scal(alpha, ux)
  
          # Update residual and solution
          gamma = dot(cx, r)
          r = axpy(cx, r, r.shape[0], -gamma)  # r -= gamma*cx
          x = axpy(ux, x, x.shape[0], gamma)  # x += gamma*ux
  
          # Truncate CU
          if truncate == 'oldest':
              while len(CU) >= k and CU:
                  del CU[0]
          elif truncate == 'smallest':
              if len(CU) >= k and CU:
                  # cf. [1,2]
                  D = solve(R[:-1,:].T, B.T).T
                  W, sigma, V = svd(D)
  
                  # C := C W[:,:k-1],  U := U W[:,:k-1]
                  new_CU = []
                  for j, w in enumerate(W[:,:k-1].T):
                      c, u = CU[0]
                      c = c * w[0]
                      u = u * w[0]
                      for cup, wp in zip(CU[1:], w[1:]):
                          cp, up = cup
                          c = axpy(cp, c, c.shape[0], wp)
                          u = axpy(up, u, u.shape[0], wp)
  
                      # Reorthogonalize at the same time; not necessary
                      # in exact arithmetic, but floating point error
                      # tends to accumulate here
                      for cp, up in new_CU:
                          alpha = dot(cp, c)
                          c = axpy(cp, c, c.shape[0], -alpha)
                          u = axpy(up, u, u.shape[0], -alpha)
                      alpha = nrm2(c)
                      c = scal(1.0/alpha, c)
                      u = scal(1.0/alpha, u)
  
                      new_CU.append((c, u))
                  CU[:] = new_CU
  
          # Add new vector to CU
          CU.append((cx, ux))
  
      # Include the solution vector to the span
      CU.append((None, x.copy()))
      if discard_C:
          CU[:] = [(None, uz) for cz, uz in CU]
  
      return postprocess(x), j_outer + 1
  