Preconditioner API

Right preconditioners for the Krylov solvers in solvax.krylov.

Every builder in this module returns a callable precond(v) -> v applying an approximate inverse M^{-1}, suitable for the precond= argument of solvax.krylov.gmres() and solvax.krylov.gcrot() (right preconditioning: the solver iterates on A M^{-1} and applies x = M^{-1} y internally, so the residual being minimized is that of the original system). All factorizations happen once, at build time, and are closed over — applying the preconditioner is factor-solve only.

The catalogue, roughly in order of increasing structure exploited:

• jacobi() / block_jacobi() — (block-)diagonal scaling, M = diag(A) or the block diagonal with batched LU (Saad, ch. 10).

• coarse_operator() — the physics-based pattern: precondition a hard operator with an exact/structured solve of a simplified one (physics-coarsened, coupling-dropped), e.g. a fluid/moment approximation of a kinetic Jacobian (Chen & Chacón) or the “preconditioner matrix” handed to LU in production PETSc codes.

• line_smoother() — alternating-direction block Jacobi: damped line solves along different tensor axes of a structured grid, composed as

  x <- x + omega_i * M_i^{-1} (b - A x),

  the classic remedy for anisotropic coupling (Trottenberg et al., ch. 5).

• additive_preconditioner() — a positive weighted sum of symmetric positive component inverse actions, suitable for PCG and additive line or Schwarz preconditioning on arrays and arbitrary pytrees.

• p_multigrid() — a V-cycle over caller-supplied levels (pre-smooth, restrict residual, recurse, prolong correction, post-smooth), physics-agnostic: all matvecs, transfers, and smoothers are injected. Covers h- and p-/spectral coarsening alike.

• galerkin_deflation() — balance a symmetric smoother around an adjoint-transfer Galerkin coarse correction, preserving the symmetry required by conjugate-gradient methods.

• mixed_precision() — run any preconditioner in low precision; flexible GMRES tolerates the inexactness and the outer residual is still accumulated in working precision (Carson & Higham).

• kronecker_nkp() / nearest_kronecker() — inverse of a Kronecker product A ⊗ B at the cost of two small solve sets, with the factors extracted automatically from a dense matrix via the Van Loan-Pitsianis rearrangement (nearest Kronecker product).

A preconditioner only has to cluster the spectrum of A M^{-1} — an O(1)-accurate inverse of the dominant physics usually beats an expensive exact inverse of the wrong terms. Because solvax.krylov.gmres() is flexible (FGMRES), the callables returned here may themselves be inner iterations or change between applications.

References

• Y. Saad, Iterative Methods for Sparse Linear Systems, 2nd ed., SIAM (2003), chapters 9-10 — preconditioned Krylov methods, (block) Jacobi.

• G. Chen & L. Chacón, “An implicit energy-conserving particle-in-cell scheme” / moment-based preconditioning of kinetic Jacobians, https://arxiv.org/abs/1309.6243 — solve a fluid (physics-coarsened) operator exactly to precondition the kinetic one; the same strategy as the PETSc Pmat (preconditioner-matrix) idiom of production codes.

• U. Trottenberg, C. W. Oosterlee & A. Schüller, Multigrid, Academic Press (2001) — smoothers, line relaxation, V-cycles.

• L. Fischer et al., https://arxiv.org/abs/2110.07663 and M. Thompson et al., https://arxiv.org/abs/2108.01751 — p-multigrid / spectral coarsening with caller-supplied transfer operators.

• E. Carson & N. J. Higham, “Accelerating the solution of linear systems by iterative refinement in three precisions”, SIAM J. Sci. Comput. 40(2), A817 (2018) — low-precision inner solves, high-precision outer.

• C. F. Van Loan & N. Pitsianis, “Approximation with Kronecker products”, in Linear Algebra for Large Scale and Real-Time Applications, Kluwer (1993) — the rearrangement turning nearest-Kronecker-product approximation into a rank-1 SVD.

solvax.precond.jacobi(diagonal)

Diagonal (point-Jacobi) preconditioner M^{-1} = diag(A)^{-1}.

The cheapest useful preconditioner: rescales each equation by its diagonal entry, which equilibrates row magnitudes and collapses the spectrum of diagonally dominant operators toward 1.

Parameters:

diagonal (Array) – the diagonal of A, shape (n,).

Returns:

A callable precond(v) -> diagonal**-1 * v.

Return type:

Callable[[Array], Array]

solvax.precond.block_jacobi(blocks)

Block-Jacobi preconditioner from the dense diagonal blocks of A.

Each m x m diagonal block is LU-factored once (batched, with partial pivoting); applying the preconditioner reshapes the vector to (n_blocks, m) and runs one batched lu_solve. Exact for block-diagonal A — with a single block equal to the full matrix this is a direct solve.

Parameters:

blocks (Array) – diagonal blocks of A, shape (n_blocks, m, m); the preconditioned vectors have length n_blocks * m.

Returns:

A callable applying blockdiag(blocks)^{-1} to flat vectors.

Return type:

Callable[[Array], Array]

solvax.precond.coarse_operator(solve)

Precondition with an exact solve of a simplified operator.

This trivial adaptor documents the central physics-based pattern: when the full operator A is too hard to invert (kinetic, dense, matrix-free), build a simplified operator A_s — physics-coarsened (e.g. a fluid/moment closure of a kinetic Jacobian, Chen & Chacón, https://arxiv.org/abs/1309.6243), or coupling-dropped (keep the block-tridiagonal / banded core, discard long-range terms) — factor it exactly with the structured solvers in solvax.direct or solvax.banded, and hand v -> A_s^{-1} v to the Krylov method. The preconditioned operator is A A_s^{-1} = I + (A - A_s) A_s^{-1}, so convergence is governed by how much physics A_s captures, not by the conditioning of A. This mirrors the “preconditioner matrix” (Pmat) strategy of production PETSc codes, where the LU of a simplified operator preconditions the true Jacobian.

Parameters:

solve (Callable[[Array], Array]) – any callable v -> A_s^{-1} v applying the exact (or structured) inverse of the simplified operator — e.g. a closure over solvax.direct.block_thomas_factor() / solvax.direct.block_thomas_solve() factors, or the banded factors of solvax.banded.

Returns:

The callable itself, usable directly as precond=.

Return type:

Callable[[Array], Array]

solvax.precond.line_smoother(matvec, line_solves, *, omega=0.8, sweeps=1)

Alternating-direction block-Jacobi (line) smoother.

Given exact solves along different tensor axes of a structured grid vector — e.g. tridiagonal x-line and y-line solves built from solvax.banded.lu_factor_banded() factors, each already closing over its axis reshape — compose them as under-relaxed corrections

x <- x + omega_i * solve_i(r), r = b - A x,

starting from x = 0, cycling through the directions sweeps times. Line relaxation solves the strongly coupled direction exactly, which is the standard cure for anisotropic operators where point smoothers stall; alternating directions covers anisotropy of unknown or mixed orientation (Trottenberg et al., ch. 5).

Parameters:

• matvec (Callable[[Array], Array]) – the full operator v -> A v, used to refresh the residual between line corrections.

• line_solves (Sequence[Callable[[Array], Array]]) – callables r -> M_i^{-1} r on flat vectors, one per direction, applied in order.

• omega (float | Sequence[float]) – under-relaxation weight(s); a scalar is broadcast, or one weight per entry of line_solves.

• sweeps (int) – number of passes over all directions (static Python int).

Returns:

A callable precond(b) -> x approximating A^{-1} b.

Return type:

Callable[[Array], Array]

solvax.precond.additive_preconditioner(preconditioners, *, weights=None)

Combine symmetric positive inverse actions without breaking PCG.

Returns sum_i weights[i] * preconditioners[i](residual). The default is the arithmetic mean, which keeps the action’s scale independent of the number of components. Positive weights preserve self-adjoint positive definiteness when every component has those properties, making this the PCG-safe counterpart to multiplicative line_smoother() for additive line, block, or Schwarz preconditioners.

Inputs and component results may be arrays or matching arbitrary pytrees. The combination is pure JAX tree arithmetic, so JIT, differentiation, and the input leaves’ device placement are preserved. The caller owns the component symmetry and positivity; this function validates only weights and pytree structure.

Parameters:

• preconditioners (Sequence[Callable[[Any], Any]]) – nonempty sequence of fixed inverse actions.

• weights (Sequence[float] | None) – optional finite positive weights, one per action. Defaults to equal weights summing to one.

Returns:

A callable with the same pytree input and output structure.

Return type:

Callable[[Any], Any]

solvax.precond.additive_tridiagonal_line_preconditioner(diagonal, directions, *, periodic_last_axis=None)

Build an additive inverse from batched tridiagonal grid-line solves.

Each (axis, lower, upper) entry defines ordinary tridiagonal lines along one array axis. Optionally, periodic_last_axis adds a cyclic solve along the final axis. Component inverses are combined with additive_preconditioner(), so their arithmetic mean is fixed, symmetric, differentiable, and suitable for PCG when each line operator is symmetric positive definite.

Parameters:

• diagonal (Array) – shared cell diagonal with the same shape as a residual.

• directions (Sequence[tuple[int, Array, Array]]) – axis and lower/upper bands for each nonperiodic line set.

• periodic_last_axis (tuple[Array, Array] | None) – optional lower/upper bands for cyclic final-axis lines, including their two corner couplings.

Returns:

A JIT- and differentiation-transparent additive inverse action.

Return type:

Callable[[Array], Array]

solvax.precond.p_multigrid(matvecs, restricts, prolongs, coarse_solve, *, smoothers, cycles=1)

Multigrid V-cycle preconditioner over caller-supplied levels.

Levels are ordered finest first; level l (0 <= l < L) carries a fine matvec, a smoother, and transfers to/from level l + 1, and the coarsest level L is handled by coarse_solve. One V-cycle on level l with operator A_l, restriction R_l and prolongation P_l computes

x <- S_l(0, b) (pre-smooth from zero) e <- V-cycle_{l+1}(R_l (b - A_l x)) (coarse-grid correction) x <- S_l(x + P_l e, b) (post-smooth)

with the recursion bottoming out at x = coarse_solve(b). The recursion is plain Python over the static level list, so the whole cycle stays jit-able. This library is physics-agnostic: nothing is assumed about the transfers, so the same cycle covers geometric h-coarsening and p-/spectral coarsening (lowering polynomial or Legendre/Hermite resolution) alike — see Trottenberg et al. for the classical theory and https://arxiv.org/abs/2110.07663 (Fischer et al.) and https://arxiv.org/abs/2108.01751 (Thompson et al.) for p-multigrid with spectral level hierarchies.

Parameters:

• matvecs (Sequence[Callable[[Array], Array]]) – fine-level operators v -> A_l v, finest first, length L (the coarsest level has no matvec).

• restricts (Sequence[Callable[[Array], Array]]) – transfers r_l -> r_{l+1}, length L.

• prolongs (Sequence[Callable[[Array], Array]]) – transfers e_{l+1} -> e_l, length L.

• coarse_solve (Callable[[Array], Array]) – exact (or strong) solve on the coarsest level, b -> A_L^{-1} b.

• smoothers (Sequence[Array | Callable]) – one per fine level. Either a jax.Array holding diag(A_l) — giving one damped-Jacobi sweep x + (2/3) diag^{-1} (b - A_l x) — or a callable smoother(matvec, x, b) -> x improving the iterate x.

• cycles (int) – number of V-cycles per application (static Python int); cycles after the first act on the residual.

Returns:

A callable precond(b) -> x approximating A_0^{-1} b.

Return type:

Callable[[Array], Array]

solvax.precond.galerkin_deflation(matvec, smoother, prolong, coarse_solve, coarse_template)

Build a symmetry-preserving Galerkin deflation preconditioner.

Restriction is the exact transpose of prolong. For symmetric matvec, smoother, and coarse_solve, the balanced operator

S + (I - S A) P A_c^-1 P.T (I - A S)

is symmetric and therefore suitable for PCG. The caller constructs and factors the Galerkin coarse operator A_c = P.T A P once; applications require one smoothing solve, one coarse solve, and one balancing solve.

Parameters:

• matvec (Callable[[Array], Array]) – fine operator v -> A v.

• smoother (Callable[[Array], Array]) – symmetric fine approximate inverse v -> S v.

• prolong (Callable[[Array], Array]) – coarse-to-fine linear transfer v -> P v.

• coarse_solve (Callable[[Array], Array]) – coarse inverse v -> A_c^-1 v.

• coarse_template (Array) – zero-like coarse array used to transpose prolong.

Returns:

A callable applying the balanced fine-plus-coarse inverse.

Return type:

Callable[[Array], Array]

solvax.precond.mixed_precision(precond, dtype=<class 'jax.numpy.float32'>)

Run any preconditioner in low precision.

Wraps precond with solvax.refine.as_low_precision(): the input vector is cast down to dtype, the preconditioner applied, and the result cast back to the input’s precision. Since a right preconditioner only needs to cluster the spectrum, low-precision application typically changes the iteration count marginally while halving memory traffic — and flexible GMRES (solvax.krylov. gmres()) is specifically robust to such inexact, step-dependent preconditioning, with residuals still accumulated in working precision (Carson & Higham, SIAM J. Sci. Comput. 40, A817 (2018)).

Parameters:

• precond (Callable[[Array], Array]) – any preconditioner callable v -> M^{-1} v.

• dtype – precision to apply it in (default float32).

Returns:

A callable with the same signature, low precision inside.

Return type:

Callable[[Array], Array]

solvax.precond.kronecker_nkp(a_factors, b_factors)

Apply (A ⊗ B)^{-1} from LU factors of the small factors.

For v = vec(V) in row-major (C) order with V of shape (na, nb), the Kronecker identity reads (A ⊗ B) vec(V) = vec(A V B^T), so the inverse is two small solve sets instead of one (na*nb)-sized one:

X = A^{-1} V B^{-T}, (A ⊗ B)^{-1} v = vec(X),

at O(na^2 nb + na nb^2) cost per application. Combine with nearest_kronecker() to build an automatic structural preconditioner for operators that are only approximately Kronecker (separable up to weak coupling): M = A ⊗ B nearest to the true operator clusters the spectrum of A M^{-1} around 1.

Parameters:

• a_factors (tuple[Array, Array]) – jax.scipy.linalg.lu_factor output for A, shape (na, na).

• b_factors (tuple[Array, Array]) – jax.scipy.linalg.lu_factor output for B, shape (nb, nb).

Returns:

A callable applying (A ⊗ B)^{-1} to flat (na * nb,) vectors.

Return type:

Callable[[Array], Array]

solvax.precond.nearest_kronecker(matrix, na, nb)

Nearest-Kronecker-product factors of a dense matrix.

Finds A (na x na) and B (nb x nb) minimizing ||M - A ⊗ B||_F via the Van Loan-Pitsianis rearrangement: the permutation R(M)[i*na + j, p*nb + q] = M[i*nb + p, j*nb + q] turns every Kronecker product into a rank-1 matrix vec(A) vec(B)^T, so the nearest one is the leading singular triplet of R(M) — A = sqrt(s_1) unvec(u_1), B = sqrt(s_1) unvec(v_1) (Van Loan & Pitsianis 1993). The factors are unique up to the inert scaling (c A) ⊗ (B / c) = A ⊗ B. Feed the LU of the result to kronecker_nkp() for an automatic structural preconditioner.

Parameters:

• matrix (Array) – dense matrix of shape (na * nb, na * nb).

• na (int) – size of the left (outer) factor.

• nb (int) – size of the right (inner) factor.

Returns:

The pair (A, B) with shapes (na, na) and (nb, nb).

Return type:

tuple[Array, Array]

Randomized

Randomized Nyström preconditioning for SPD systems.

For a symmetric positive semidefinite operator A and a regularized system (A + mu I) x = b, a rank-ell randomized Nystrom approximation A_nys = U diag(lam) U^T built from ell operator applications yields the preconditioner

P^{-1} v = U diag((lam_ell + mu) / (lam + mu)) U^T v + (v - U U^T v),

which is symmetric positive definite, costs one (n, ell) matmul pair per application, and — when ell exceeds roughly twice the mu-effective dimension of A — bounds the preconditioned condition number by a small constant in expectation, independently of the spectrum’s decay rate (Frangella, Tropp & Udell, SIAM J. Matrix Anal. Appl. 44, 718 (2023)). It is the scalable coarse-correction alternative when no grid hierarchy or structured coarse operator exists.

The sketch uses an explicit PRNG key, so construction is deterministic, jit-able, and differentiable through both the sketch and the eigenfactors; the stabilized-shift construction follows Frangella et al., Algorithm 2.1.

solvax.randomized.nystrom_preconditioner(matvec, n, rank, key, *, mu=0.0, dtype=None)

Build a randomized Nystrom preconditioner for matvec + mu I.

Parameters:

• matvec (Callable[[Array], Array]) – symmetric positive semidefinite operator v -> A v on flat (n,) arrays; must be pure JAX. Symmetry is assumed, not checked.

• n (int) – static operand dimension.

• rank (int) – static sketch size ell (1 <= rank <= n). Effective when it exceeds about twice the mu-effective dimension of A.

• key (Array) – PRNG key for the Gaussian test matrix; fixing it makes the preconditioner deterministic and differentiable.

• mu (float) – regularization shift of the system being solved (A + mu I).

• dtype – sketch dtype (defaults to float32/float64 per x64 mode).

Returns:

A symmetric-positive-definite inverse action suitable as precond= for solvax.pcg.pcg() on the system (A + mu I) x = b.

Return type:

Callable[[Array], Array]