Krylov API

Flexible restarted GMRES and GCROT-style Krylov subspace recycling.

Right-preconditioned flexible GMRES (FGMRES) builds the Arnoldi relation

A Z_m = V_{m+1} Hbar_m, Z_m = [M_1^{-1} v_1, …, M_m^{-1} v_m],

where V_{m+1} is orthonormal and Hbar_m is (m+1) x m upper Hessenberg. Because the preconditioned vectors z_j are stored explicitly, the preconditioner may change from step to step (flexible mode); the correction is x += Z_m y with y minimizing || beta e_1 - Hbar_m y ||, solved incrementally with Givens rotations so the residual norm is available at every inner step. Orthogonalization uses classical Gram-Schmidt applied twice (CGS2), which reduces the sequential inner-product latency of modified Gram-Schmidt on accelerators while retaining O(eps) loss of orthogonality.

GCROT(m, k)-style recycling maintains an outer pair (C, U) with A U = C and C^H C = I. Each outer iteration first minimizes over the recycled space (x += U C^H r, r -= C C^H r), then runs one FGMRES(m) cycle on the deflated operator (I - C C^H) A, giving

(I - C C^H) A Z_m = V_{m+1} Hbar_m, B_m = C^H A Z_m,

so the cycle correction is dx = Z_m y - U B_m y with A dx = V_{m+1} Hbar_m y orthogonal to C. In this v0.1 the recycle space is updated with one direction per cycle — the cycle’s own optimal correction (dx, A dx), normalized and inserted FIFO — rather than the harmonic Ritz vectors of GCRO-DR. This is a deliberate simplification: it keeps the update O(nk) and shape-static, and it retains the directions that restarting would otherwise discard, but it deflates slowly-converging eigenmodes only indirectly. Recycle pairs may be passed between solves in a parameter continuation; on entry A U is recomputed and the pair is re-orthonormalized (thin QR) so A U = C holds for the current operator, as in Parks et al.

References

  • Y. Saad, Iterative Methods for Sparse Linear Systems, 2nd ed., SIAM (2003), sections 6.3-6.5 and 9.4 (GMRES, restarting, FGMRES).

  • R. B. Morgan, “GMRES with deflated restarting”, SIAM J. Sci. Comput. 24, 20 (2002) — GMRES-DR.

  • 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) — GCRO-DR, recycling across a sequence.

  • E. de Sturler, “Truncation strategies for optimal Krylov subspace methods”, SIAM J. Numer. Anal. 36, 864 (1999) — GCROT.

  • L. Giraud, J. Langou, M. Rozloznik & J. van den Eshof, “Rounding error analysis of the classical Gram-Schmidt orthogonalization process”, Numer. Math. 101, 87 (2005) — CGS2 stability.

class solvax.krylov.KrylovSolution(x, residual_norm, iterations, converged, recycle=None, recycle_drift=None)

Result of gmres() or gcrot().

Variables:

  • x (Any) – approximate solution with the same structure and leaf shapes as b.

  • residual_norm (jax.Array) – true residual norm ||b - A x|| (recomputed once after the iteration, not the Givens estimate).

  • iterations (jax.Array) – total inner (Arnoldi) iterations across all cycles, int32.

  • converged (jax.Array) – whether residual_norm <= max(atol, rtol * ||b||).

  • recycle (tuple[jax.Array, jax.Array] | None) – updated recycle pair (C, U) with fixed shapes (n, k) for gcrot(), None for gmres(). Unfilled columns are zero.

  • recycle_drift (jax.Array | None) – for a warm-started gcrot(), the mean principal- angle sine between the incoming recycle image space and its re-established span under the current operator — a direct measure of how far the operator has drifted since the pair was built (0 for an unchanged operator, up to 1 for an orthogonal rotation). 0.0 on a cold start, None for gmres().

Parameters:
x: Any

Alias for field number 0

residual_norm: Array

Alias for field number 1

iterations: Array

Alias for field number 2

converged: Array

Alias for field number 3

recycle: tuple[Array, Array] | None

Alias for field number 4

recycle_drift: Array | None

Alias for field number 5

solvax.krylov.gmres(matvec, b, *, x0=None, precond=None, inner_product=None, restart=30, rtol=1e-08, atol=0.0, max_restarts=50)

Restarted flexible GMRES with right preconditioning.

Solves A x = b for a matrix-free operator, stopping when ||b - A x|| <= max(atol, rtol * ||b||). Fully jit-able: all loop state has fixed shapes (the Krylov basis is zero-padded to the cycle size and early convergence exits via lax.while_loop).

Parameters:

  • matvec (Callable[[Any], Any]) – the operator v -> A v on an array or pytree; must be pure JAX (traceable) and preserve the input tree structure.

  • b (Any) – array or pytree right-hand side. Pytree leaves must have one common inexact dtype.

  • x0 (Any | None) – initial guess (defaults to zeros).

  • precond (Callable[[Any], Any] | None) – right preconditioner v -> M^{-1} v (defaults to the identity). May be flexible, i.e. nonlinear or changing between inner steps — the update uses the stored preconditioned vectors.

  • inner_product (Callable[[Any, Any], Array] | None) – optional (left, right) -> scalar product used for PyTree Arnoldi projections and norms. Defaults to the Euclidean product. Supplying it also selects the PyTree path for array inputs.

  • restart (int) – static Arnoldi cycle size m.

  • rtol (float) – relative tolerance on ||b||.

  • atol (float) – absolute tolerance floor.

  • max_restarts (int) – static maximum number of cycles.

Returns:

A KrylovSolution with recycle=None.

Return type:

KrylovSolution

solvax.krylov.gcrot(matvec, b, *, x0=None, precond=None, m=30, k=10, rtol=1e-08, atol=0.0, max_restarts=50, recycle=None)

GCROT(m, k)-style FGMRES with Krylov subspace recycling.

Like gmres(), but maintains a recycle pair (C, U) with A U = C that deflates the operator between restarts and can be carried across solves in a slowly-varying sequence (parameter continuation): pass solution.recycle of one solve as recycle= of the next. On warm start A U is recomputed for the current operator and the pair is re-orthonormalized by thin QR (rank-deficient columns — e.g. zero padding from a short previous solve — are dropped), so a stale pair is always consistent, merely less effective.

The recycle space grows by one direction per cycle (the cycle’s own correction; see the module docstring for why this simplification, and Parks et al. 2006 for the harmonic-Ritz alternative), stored FIFO in fixed-shape (n, k) arrays so the whole solve stays jit-able.

Parameters:

  • matvec (Callable[[Any], Any]) – the operator v -> A v on flat (n,) arrays.

  • b (Array) – right-hand side, shape (n,).

  • x0 (Array | None) – initial guess (defaults to zeros).

  • precond (Callable[[Any], Any] | None) – right preconditioner v -> M^{-1} v (identity default).

  • m (int) – static inner FGMRES cycle size.

  • k (int) – static number of recycle directions kept.

  • rtol (float) – relative tolerance on ||b||.

  • atol (float) – absolute tolerance floor.

  • max_restarts (int) – static maximum number of outer cycles.

  • recycle (tuple[Array, Array] | None) – optional (C, U) pair of shape (n, k) from a previous KrylovSolution to warm-start deflation.

Returns:

A KrylovSolution whose recycle field holds the updated (C, U) pair, zero-padded to shape (n, k).

Return type:

KrylovSolution