lgmres.py 8.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230
  1. # Copyright (C) 2009, Pauli Virtanen <pav@iki.fi>
  2. # Distributed under the same license as SciPy.
  3. import numpy as np
  4. from numpy.linalg import LinAlgError
  5. from scipy.linalg import get_blas_funcs
  6. from .iterative import _get_atol_rtol
  7. from .utils import make_system
  8. from ._gcrotmk import _fgmres
  9. __all__ = ['lgmres']
  10. def lgmres(A, b, x0=None, *, rtol=1e-5, atol=0., maxiter=1000, M=None, callback=None,
  11. inner_m=30, outer_k=3, outer_v=None, store_outer_Av=True,
  12. prepend_outer_v=False):
  13. """
  14. Solve ``Ax = b`` with the LGMRES algorithm.
  15. The LGMRES algorithm [1]_ [2]_ is designed to avoid some problems
  16. in the convergence in restarted GMRES, and often converges in fewer
  17. iterations.
  18. Parameters
  19. ----------
  20. A : {sparse array, ndarray, LinearOperator}
  21. The real or complex N-by-N matrix of the linear system.
  22. Alternatively, ``A`` can be a linear operator which can
  23. produce ``Ax`` using, e.g.,
  24. ``scipy.sparse.linalg.LinearOperator``.
  25. b : ndarray
  26. Right hand side of the linear system. Has shape (N,) or (N,1).
  27. x0 : ndarray
  28. Starting guess for the solution.
  29. rtol, atol : float, optional
  30. Parameters for the convergence test. For convergence,
  31. ``norm(b - A @ x) <= max(rtol*norm(b), atol)`` should be satisfied.
  32. The default is ``rtol=1e-5``, the default for ``atol`` is ``0.0``.
  33. maxiter : int, optional
  34. Maximum number of iterations. Iteration will stop after maxiter
  35. steps even if the specified tolerance has not been achieved.
  36. M : {sparse array, ndarray, LinearOperator}, optional
  37. Preconditioner for A. The preconditioner should approximate the
  38. inverse of A. Effective preconditioning dramatically improves the
  39. rate of convergence, which implies that fewer iterations are needed
  40. to reach a given error tolerance.
  41. callback : function, optional
  42. User-supplied function to call after each iteration. It is called
  43. as callback(xk), where xk is the current solution vector.
  44. inner_m : int, optional
  45. Number of inner GMRES iterations per each outer iteration.
  46. outer_k : int, optional
  47. Number of vectors to carry between inner GMRES iterations.
  48. According to [1]_, good values are in the range of 1...3.
  49. However, note that if you want to use the additional vectors to
  50. accelerate solving multiple similar problems, larger values may
  51. be beneficial.
  52. outer_v : list of tuples, optional
  53. List containing tuples ``(v, Av)`` of vectors and corresponding
  54. matrix-vector products, used to augment the Krylov subspace, and
  55. carried between inner GMRES iterations. The element ``Av`` can
  56. be `None` if the matrix-vector product should be re-evaluated.
  57. This parameter is modified in-place by `lgmres`, and can be used
  58. to pass "guess" vectors in and out of the algorithm when solving
  59. similar problems.
  60. store_outer_Av : bool, optional
  61. Whether LGMRES should store also A@v in addition to vectors `v`
  62. in the `outer_v` list. Default is True.
  63. prepend_outer_v : bool, optional
  64. Whether to put outer_v augmentation vectors before Krylov iterates.
  65. In standard LGMRES, prepend_outer_v=False.
  66. Returns
  67. -------
  68. x : ndarray
  69. The converged solution.
  70. info : int
  71. Provides convergence information:
  72. - 0 : successful exit
  73. - >0 : convergence to tolerance not achieved, number of iterations
  74. - <0 : illegal input or breakdown
  75. Notes
  76. -----
  77. The LGMRES algorithm [1]_ [2]_ is designed to avoid the
  78. slowing of convergence in restarted GMRES, due to alternating
  79. residual vectors. Typically, it often outperforms GMRES(m) of
  80. comparable memory requirements by some measure, or at least is not
  81. much worse.
  82. Another advantage in this algorithm is that you can supply it with
  83. 'guess' vectors in the `outer_v` argument that augment the Krylov
  84. subspace. If the solution lies close to the span of these vectors,
  85. the algorithm converges faster. This can be useful if several very
  86. similar matrices need to be inverted one after another, such as in
  87. Newton-Krylov iteration where the Jacobian matrix often changes
  88. little in the nonlinear steps.
  89. References
  90. ----------
  91. .. [1] A.H. Baker and E.R. Jessup and T. Manteuffel, "A Technique for
  92. Accelerating the Convergence of Restarted GMRES", SIAM J. Matrix
  93. Anal. Appl. 26, 962 (2005).
  94. .. [2] A.H. Baker, "On Improving the Performance of the Linear Solver
  95. restarted GMRES", PhD thesis, University of Colorado (2003).
  96. Examples
  97. --------
  98. >>> import numpy as np
  99. >>> from scipy.sparse import csc_array
  100. >>> from scipy.sparse.linalg import lgmres
  101. >>> A = csc_array([[3, 2, 0], [1, -1, 0], [0, 5, 1]], dtype=float)
  102. >>> b = np.array([2, 4, -1], dtype=float)
  103. >>> x, exitCode = lgmres(A, b, atol=1e-5)
  104. >>> print(exitCode) # 0 indicates successful convergence
  105. 0
  106. >>> np.allclose(A.dot(x), b)
  107. True
  108. """
  109. A,M,x,b = make_system(A,M,x0,b)
  110. if not np.isfinite(b).all():
  111. raise ValueError("RHS must contain only finite numbers")
  112. matvec = A.matvec
  113. psolve = M.matvec
  114. if outer_v is None:
  115. outer_v = []
  116. axpy, dot, scal = None, None, None
  117. nrm2 = get_blas_funcs('nrm2', [b])
  118. b_norm = nrm2(b)
  119. # we call this to get the right atol/rtol and raise errors as necessary
  120. atol, rtol = _get_atol_rtol('lgmres', b_norm, atol, rtol)
  121. if b_norm == 0:
  122. x = b
  123. return (x, 0)
  124. ptol_max_factor = 1.0
  125. for k_outer in range(maxiter):
  126. r_outer = matvec(x) - b
  127. # -- callback
  128. if callback is not None:
  129. callback(x)
  130. # -- determine input type routines
  131. if axpy is None:
  132. if np.iscomplexobj(r_outer) and not np.iscomplexobj(x):
  133. x = x.astype(r_outer.dtype)
  134. axpy, dot, scal, nrm2 = get_blas_funcs(['axpy', 'dot', 'scal', 'nrm2'],
  135. (x, r_outer))
  136. # -- check stopping condition
  137. r_norm = nrm2(r_outer)
  138. if r_norm <= max(atol, rtol * b_norm):
  139. break
  140. # -- inner LGMRES iteration
  141. v0 = -psolve(r_outer)
  142. inner_res_0 = nrm2(v0)
  143. if inner_res_0 == 0:
  144. rnorm = nrm2(r_outer)
  145. raise RuntimeError("Preconditioner returned a zero vector; "
  146. f"|v| ~ {rnorm:.1g}, |M v| = 0")
  147. v0 = scal(1.0/inner_res_0, v0)
  148. ptol = min(ptol_max_factor, max(atol, rtol*b_norm)/r_norm)
  149. try:
  150. Q, R, B, vs, zs, y, pres = _fgmres(matvec,
  151. v0,
  152. inner_m,
  153. lpsolve=psolve,
  154. atol=ptol,
  155. outer_v=outer_v,
  156. prepend_outer_v=prepend_outer_v)
  157. y *= inner_res_0
  158. if not np.isfinite(y).all():
  159. # Overflow etc. in computation. There's no way to
  160. # recover from this, so we have to bail out.
  161. raise LinAlgError()
  162. except LinAlgError:
  163. # Floating point over/underflow, non-finite result from
  164. # matmul etc. -- report failure.
  165. return x, k_outer + 1
  166. # Inner loop tolerance control
  167. if pres > ptol:
  168. ptol_max_factor = min(1.0, 1.5 * ptol_max_factor)
  169. else:
  170. ptol_max_factor = max(1e-16, 0.25 * ptol_max_factor)
  171. # -- GMRES terminated: eval solution
  172. dx = zs[0]*y[0]
  173. for w, yc in zip(zs[1:], y[1:]):
  174. dx = axpy(w, dx, dx.shape[0], yc) # dx += w*yc
  175. # -- Store LGMRES augmentation vectors
  176. nx = nrm2(dx)
  177. if nx > 0:
  178. if store_outer_Av:
  179. q = Q.dot(R.dot(y))
  180. ax = vs[0]*q[0]
  181. for v, qc in zip(vs[1:], q[1:]):
  182. ax = axpy(v, ax, ax.shape[0], qc)
  183. outer_v.append((dx/nx, ax/nx))
  184. else:
  185. outer_v.append((dx/nx, None))
  186. # -- Retain only a finite number of augmentation vectors
  187. while len(outer_v) > outer_k:
  188. del outer_v[0]
  189. # -- Apply step
  190. x += dx
  191. else:
  192. # didn't converge ...
  193. return x, maxiter
  194. return x, 0