FGMRES and GCROT recycling

SOLVAX provides restarted flexible GMRES for general matrix-free systems and a GCROT-style extension for sequences of related systems. GMRES supports real or complex arrays and arbitrary matching JAX pytrees; GCROT currently requires a one-dimensional array. Both use right preconditioning.

Flexible GMRES

Given \(A x=b\) and an initial guess \(x_0\), form \(r_0=b-Ax_0\) and \(v_1=r_0/\beta\), where \(\beta=\lVert r_0\rVert_2\). At Arnoldi step \(j\),

\[ z_j=M_j^{-1}v_j, \qquad w_j=A z_j. \]

Orthogonalization produces

\[ A Z_m=V_{m+1}\bar H_m, \]

where \(Z_m=[z_1,\ldots,z_m]\), \(V_{m+1}\) has orthonormal columns, and \(\bar H_m\) is upper Hessenberg. The correction is

\[ x_m=x_0+Z_my_m, \qquad y_m=\arg\min_y\lVert\beta e_1-\bar H_my\rVert_2. \]

Storing \(Z_m\) instead of assuming a fixed \(M^{-1}\) makes the method flexible: the preconditioner may change at every step [Saa03].

SOLVAX uses two passes of classical Gram-Schmidt (CGS2) and incremental Givens rotations. For complex inputs, projections are Hermitian and the rotations use a scaled unitary LAPACK-style convention. The true residual is recomputed at restart boundaries and after the final cycle.

API usage

solution = sx.gmres(
    matvec,
    b,
    x0=None,
    precond=None,
    restart=30,
    rtol=1e-8,
    atol=0.0,
    max_restarts=50,
)

Inputs

Input	Meaning
matvec	pure JAX callable v -> A v preserving the operand structure
b	array or matching-leaf pytree right-hand side, real or complex
x0	optional initial guess; zero by default
precond	right inverse action v -> M_j^{-1}v; identity by default
restart	maximum Arnoldi steps per cycle
rtol, atol	true residual stopping tolerances
max_restarts	maximum number of Arnoldi cycles

The optimized array path expects a flat one-dimensional right-hand side. For a three-dimensional field, either flatten before calling gmres and reshape the solution afterward, or supply an inner_product such as lambda left, right: jnp.vdot(left, right). Supplying inner_product selects the shape-general PyTree path and the operator, initial guess, and preconditioner must preserve that array/PyTree structure.

Output

KrylovSolution contains x, true residual_norm, total inner iterations, converged, and recycle=None.

Restart trade-off

Unrestarted GMRES minimizes over an expanding Krylov space but stores \(O(nm)\) basis data and performs \(O(nm^2)\) orthogonalization through \(m\) steps. Restarting caps memory and compiled shapes, but discards spectral information and can stagnate. Increase restart when memory permits and convergence stalls; improve the preconditioner before relying on very large restart spaces.

GCROT-style recycling

For a sequence \(A_i x_i=b_i\), retain source and image bases \((U,C)\) satisfying

\[ A U=C,\qquad C^HC=I. \]

Before each Arnoldi cycle, project the residual over the recycled space:

\[ x\leftarrow x+UC^Hr,\qquad r\leftarrow(I-CC^H)r. \]

The new Arnoldi cycle is built for the deflated operator

\[ (I-CC^H)AZ_m=V_{m+1}\bar H_m. \]

SOLVAX inserts the cycle’s normalized correction into a fixed-size FIFO recycle space. When a recycle pair is supplied for a changed operator, it recomputes \(AU\) and uses a thin QR factorization to restore \(AU=C\) for the current system.

recycle = None
for parameter in scan:
    solution = sx.gcrot(
        make_matvec(parameter),
        b,
        m=30,
        k=10,
        recycle=recycle,
        rtol=1e-10,
    )
    recycle = solution.recycle

m is the inner FGMRES cycle length and k is the fixed recycle dimension. The returned recycle arrays have shape (n, k) and may contain zero-padded unused columns.

Recycling across optimization, and the drift diagnostic

A gradient-based optimization loop solves two slowly varying families: the primal A(theta_j) x = b and the adjoint A(theta_j)^T lam = g. Each is its own sequence — carry one recycle pair per direction:

primal = sx.gcrot(matvec, b, m=30, k=10, recycle=primal_pair)
adjoint = sx.gcrot(matvec_t, g, m=30, k=10, recycle=adjoint_pair)
primal_pair, adjoint_pair = primal.recycle, adjoint.recycle

On a warm start the solution reports recycle_drift: the mean principal-angle sine between the incoming recycle image space and its re-established span under the current operator. It is zero for an unchanged operator, and grows linearly with the operator step ||A(theta_{j+1}) - A(theta_j)|| (a Davis–Kahan subspace-rotation effect), so it tells the caller directly whether the optimizer’s step size keeps the recycled space useful — small drift and falling iteration counts, or drift near one and a pair no better than a cold start. On the continuation family in benchmarks/benchmark_recycling.py, both carried pairs collapse the per-step iterations from 32 to 2–3 within four steps (a 3.3× total matvec reduction over ten steps) at measured drift ~1e-3 per step, while per-step cold GCROT shows no gain at all — the carrying is the mechanism, not the deflation itself.

Relationship to the literature

The implementation follows the recycling framework of Parks et al. and the deflated-restart motivation of Morgan [Mor02, PdSM+06]. It is not a full harmonic-Ritz GCRO-DR implementation: it retains one optimal cycle correction per restart rather than selecting harmonic Ritz vectors. This keeps the update shape-static and \(O(nk)\) but may identify difficult invariant subspaces more slowly.

Choosing between Krylov methods

Method	Matrix assumptions	Memory	Best use
PCG	Hermitian positive definite	short recurrences	SPD elliptic and energy systems
FGMRES	general square operator	restart-sized basis	isolated nonsymmetric/indefinite solve
GCROT	general sequence	basis plus recycle pair	continuation, time stepping, optimization
BiCGSTAB	general; not in SOLVAX	short recurrences	memory-limited nonsymmetric solves, but irregular residuals
direct structured	exact known structure	factors	repeated solves with exact band/block form

Preconditioning patterns

• jacobi for scaling problems;

• block_jacobi for strong within-cell coupling;

• exact banded or block-Thomas inverse of a coupling-dropped operator;

• line_smoother for anisotropy;

• p_multigrid for scale-separated elliptic error;

• mixed-precision or truncated inner solves, which flexibility permits.

See Preconditioners for formulas and examples.

Failure and performance diagnostics

• converged=False means the true final residual missed the requested tolerance; increase work or improve the preconditioner.

• A small internal least-squares estimate is not used as the final truth; inspect residual_norm.

• Near linear dependence increases orthogonalization sensitivity. CGS2 reduces but does not eliminate finite-precision loss of orthogonality [GLRvdE05].

• A stale recycle space remains algebraically refreshed for the current operator but can be ineffective. Drop it after discontinuous parameter changes.

• restart, m, k, and max_restarts are static compiled sizes.

Complex systems and gradients

Complex systems use conjugate inner products. gmres can be supplied as the primal and transposed solver to solvax.implicit.linear_solve(). For a real objective of complex state, validate the adjoint convention with finite differences; see examples/18_complex_krylov_gradient.py.

API summary

• solvax.krylov.gmres()

• solvax.krylov.gcrot()

• solvax.krylov.KrylovSolution

Runnable counterparts: examples/02_advection_preconditioning.py, examples/03_recycled_continuation.py, and examples/18_complex_krylov_gradient.py.