Operator API

Lightweight, physics-agnostic linear-operator containers.

Thin equinox modules wrapping the action of a linear map, interoperable with the plain-callable convention used throughout solvax: everything in solvax.krylov takes a matvec callable, and every operator here is itself callable (op(v) == op.matvec(v)), so operators can be passed directly as matvec= or precond= arguments. The single invariant all operators maintain is adjoint consistency,

<A x, y> = <x, A^T y> for all x, y,

with A^T produced by the .T property.

Structure exploited per container:

  • MatrixFreeOperator wraps an arbitrary linear callable; its transpose falls back on jax.linear_transpose() when no hand-written transpose is supplied.

  • SumOperator applies (A_1 + ... + A_p) v = sum_i A_i v; the transpose distributes over the sum.

  • KroneckerOperator applies (A (x) B) v through the reshape identity (A (x) B) vec(X) = vec(B X A^T) (equivalently, with the row-major flattening used by jnp.reshape, A X B^T), so the (pr) x (qs) product is never formed: cost is two small matrix products instead of one huge one.

  • BlockTridiagonalOperator stores dense bands (L_k, D_k, U_k) in the layout of solvax.direct and applies all blocks with batched einsums; to_blocks() feeds solvax.direct.block_thomas_factor directly, pairing the operator with its natural preconditioner.

  • BorderedOperator is the saddle-point (KKT-like) structure

    K = [[A, B],

    [C, 0]], K [x, y] = [A x + B y, C x],

    arising when constraint or source rows border a physics block A. schur_projected_precond() turns an approximate inverse of A alone into a preconditioner for the full bordered system via the (small, dense) Schur complement S = C A^{-1} B:

    y = S^{-1} (C A^{-1} r_x - r_y), x = A^{-1} (r_x - B y),

    which is exactly K^{-1} when A^{-1} is exact — so a preconditioner built for the physics block preconditions the constrained system.

References

  • M. Benzi, G. H. Golub & J. Liesen, “Numerical solution of saddle point problems”, Acta Numerica 14, 1 (2005) — block preconditioners and Schur complement methods for bordered/KKT systems.

  • C. F. Van Loan, “The ubiquitous Kronecker product”, J. Comput. Appl. Math. 123, 85 (2000) — the vec/reshape identity.

  • Y. Saad, Iterative Methods for Sparse Linear Systems, 2nd ed., SIAM (2003), chapter 9 — preconditioned Krylov methods with matrix-free operators.

class solvax.operators.MatrixFreeOperator(apply, transpose_apply=None, *, shape)

A linear operator defined only by its action v -> A v.

Variables:
Parameters:
matvec(v)

Apply the operator: A @ v.

Parameters:

v (Array)

Return type:

Array

property T: MatrixFreeOperator

The transposed operator A^T.

When transpose_apply was supplied it is used directly (and the forward apply becomes the transpose of the transpose, so op.T.T is free). Otherwise the transpose is derived with jax.linear_transpose(): each application then costs roughly one forward evaluation of apply (its jaxpr is transposed rule by rule), plus a re-trace per call outside jit — supply an explicit transpose_apply when the adjoint is hot.

Returns:

A MatrixFreeOperator with shape (n_in, n_out).

class solvax.operators.SumOperator(terms)

A sum of linear operators, applied term by term: (sum_i A_i) v.

Typical use: a structured principal part (e.g. a BlockTridiagonalOperator that a direct solve can precondition) plus matrix-free perturbations or coupling terms.

Variables:

terms (tuple) – tuple of operators, plain linear callables, or dense matrices, all with the same (n_out, n_in) extent.

Parameters:

terms (tuple)

property shape: tuple[int, int]

The common (n_out, n_in) shape, read off the first shaped term.

matvec(v)

Apply every term and sum: sum_i A_i v.

Parameters:

v (Array)

Return type:

Array

property T: SumOperator

(sum A_i)^T = sum A_i^T.

Plain callables (no .T of their own) are wrapped in MatrixFreeOperator with this operator’s shape and transposed via jax.linear_transpose() — see MatrixFreeOperator.T for the cost.

Type:

The transpose, distributed over the terms

class solvax.operators.KroneckerOperator(a, b)

The Kronecker product A (x) B applied without forming it.

Uses the reshape identity (A (x) B) vec(X) = vec(B X A^T); with the row-major flattening of jnp.reshape this reads, for A of shape (p, q) and B of shape (r, s),

(A (x) B) v = (A X B^T).reshape(-1), X = v.reshape(q, s),

so one matvec costs O(pqs + prs) instead of the O(pqrs) of the assembled (pr) x (qs) matrix.

Variables:

  • a (Any) – left factor — an operator (anything with matvec-style __call__, shape, .T) or a dense matrix. Wrap plain callables in MatrixFreeOperator.

  • b (Any) – right factor, same protocol as a.

Parameters:
property shape: tuple[int, int]

(p * r, q * s) for A of shape (p, q), B of (r, s).

matvec(v)

Apply (A (x) B) v via two small products on the reshaped vector.

Parameters:

v (Array)

Return type:

Array

property T: KroneckerOperator

(A (x) B)^T = A^T (x) B^T.

materialize()

Assemble the dense product with jnp.kron() — small sizes only.

Returns:

Dense matrix of shape (p * r, q * s).

Return type:

Array

class solvax.operators.BlockTridiagonalOperator(lower, diag, upper)

Block-tridiagonal operator over dense per-block bands.

Row k of the action is L_k x_{k-1} + D_k x_k + U_k x_{k+1}, computed for all blocks at once with batched einsums over shifted slices (no Python loop over blocks). The band layout matches solvax.direct: lower[0] and upper[-1] are carried but ignored, so to_blocks() feeds solvax.direct.block_thomas_factor() unchanged — the natural direct preconditioner for this operator.

Variables:

  • lower (jax.Array) – sub-diagonal blocks L_k, shape (n_blocks, m, m); lower[0] is ignored.

  • diag (jax.Array) – diagonal blocks D_k, shape (n_blocks, m, m).

  • upper (jax.Array) – super-diagonal blocks U_k, shape (n_blocks, m, m); upper[-1] is ignored.

Parameters:
property shape: tuple[int, int]

(n_blocks * m, n_blocks * m).

matvec(v)

Apply the operator to a flat (n_blocks * m,) vector.

Parameters:

v (Array)

Return type:

Array

property T: BlockTridiagonalOperator

bands swap and every block transposes.

(A^T)_{k,k-1} = U_{k-1}^T and (A^T)_{k,k+1} = L_{k+1}^T, implemented as a roll of the transposed bands (the wrapped-around lower[0] / upper[-1] entries land in the ignored slots).

Type:

The transpose

to_blocks()

Return (lower, diag, upper) for solvax.direct.block_thomas_factor().

Return type:

tuple[Array, Array, Array]

materialize()

Assemble the dense (n_blocks m) x (n_blocks m) matrix — small sizes only.

Returns:

Dense matrix with the three block bands scattered in place.

Return type:

Array

class solvax.operators.BorderedOperator(a, b_cols, c_rows)

The bordered (KKT-like) operator [[A, B], [C, 0]].

Acts on the concatenated vector [x, y] as [A x + B y, C x] — the structure of a physics block A augmented with constraint rows C and source/coupling columns B (Benzi, Golub & Liesen 2005). Pair with schur_projected_precond() to recycle a preconditioner for A on the full constrained system.

Variables:

  • a (Any) – the (n, n) principal block — an operator or dense matrix (wrap plain callables in MatrixFreeOperator).

  • b_cols (jax.Array) – border columns B, shape (n, p).

  • c_rows (jax.Array) – border rows C, shape (q, n).

Parameters:
property shape: tuple[int, int]

(n_out + q, n_in + p).

matvec(v)

Apply to the concatenated vector: [A x + B y, C x].

Parameters:

v (Array)

Return type:

Array

property T: BorderedOperator

[[A, B], [C, 0]]^T = [[A^T, C^T], [B^T, 0]].

The border rows become the transposed columns and vice versa.

materialize()

Assemble the dense bordered matrix — small sizes only.

Returns:

Dense matrix [[A, B], [C, 0]].

Return type:

Array

solvax.operators.schur_projected_precond(a_inv, b_cols, c_rows, d_block=None)

Preconditioner for a bordered system from an approximate inverse of A.

Given a_inv ~ A^{-1} for the principal block alone, forms the small dense Schur complement of the border once (p applications of a_inv to the columns of B, then an LU factorization of the q x p result, which must be square) and returns the projection

y = S^{-1} (C a_inv(r_x) - r_y), x = a_inv(r_x - B y),

i.e. the exact inverse of the bordered (KKT-like) matrix [[A, B], [C, D]] with A^{-1} replaced by a_inv throughout (Benzi, Golub & Liesen, Acta Numerica 2005, section 5). The border-border block D defaults to zero (the constraint/saddle-point case [[A, B], [C, 0]]); pass a nonzero d_block for a general bordered system (e.g. a quasineutrality / potential border that couples the border unknowns to themselves), in which case the Schur complement is S = C A^{-1} B - D. With a_inv exact, the preconditioned operator is the identity and GMRES converges in one iteration; with an approximate a_inv, the border is still eliminated exactly through the projected Schur system, so a preconditioner built for the physics block A preconditions the full bordered system. Each application costs two calls to a_inv plus one small triangular solve.

Parameters:

  • a_inv (Callable) – callable r -> A^{-1} r (approximate is fine) on flat (n,) arrays; must be linear and pure JAX.

  • b_cols (Array) – border columns B, shape (n, p).

  • c_rows (Array) – border rows C, shape (p, n) — the Schur complement must be square.

  • d_block (Array | None) – optional border-border block D, shape (p, p); when None (default) the border is the pure [[A, B], [C, 0]] saddle point.

Returns:

A callable [r_x, r_y] -> [x, y] on concatenated (n + p,) vectors, suitable as precond= for solvax.krylov.gmres() on the matching bordered operator.

Return type:

Callable