PCG API¶

Matrix-free preconditioned conjugate gradients on JAX pytrees.

class solvax.pcg.PCGSolution(x, residual_norm, relative_residual_norm, iterations, converged, status, residual_history)¶

Result of pcg() with fixed-shape, JIT-safe diagnostics.

Parameters:

x (Any)
residual_norm (Array)
relative_residual_norm (Array)
iterations (Array)
converged (Array)
status (Array)
residual_history (Array)

x: Any¶: Alias for field number 0

residual_norm: Array¶: Alias for field number 1

relative_residual_norm: Array¶: Alias for field number 2

iterations: Array¶: Alias for field number 3

converged: Array¶: Alias for field number 4

status: Array¶: Alias for field number 5

residual_history: Array¶: Alias for field number 6

class solvax.pcg.PCGDiagnostics(residual_norm, relative_residual_norm, iterations, converged, status, residual_history)¶

Fixed-shape PCG diagnostics carried through an implicit solve.

Parameters:

residual_norm (Array)
relative_residual_norm (Array)
iterations (Array)
converged (Array)
status (Array)
residual_history (Array)

residual_norm: Array¶: Alias for field number 0

relative_residual_norm: Array¶: Alias for field number 1

iterations: Array¶: Alias for field number 2

converged: Array¶: Alias for field number 3

status: Array¶: Alias for field number 4

residual_history: Array¶: Alias for field number 5

solvax.pcg.pcg(matvec, b, *, x0=None, precond=None, rtol=1e-08, atol=0.0, max_steps=500, single_reduction=False)¶

Solve a symmetric positive-definite pytree system with matrix-free PCG.

matvec and precond operate on a pytree with the same structure as b. Set single_reduction=True to batch the three independent inner products in an algebraically equivalent recurrence. The residual history always has shape (max_steps + 1,); entries after termination repeat the final norm, which keeps the result compatible with jit and vmap without dynamic allocation.

Parameters:

matvec (Callable[[Any], Any])
b (Any)
x0 (Any | None)
precond (Callable[[Any], Any] | None)
rtol (float)
atol (float)
max_steps (int)
single_reduction (bool)

Return type:

PCGSolution

solvax.pcg.pcg_linear_solve(matvec, b, *, x0=None, precond=None, rtol=1e-08, atol=0.0, max_steps=500, single_reduction=False, transpose_precond=None, transpose_rtol=None, transpose_atol=None, transpose_max_steps=None)¶

Implicitly differentiable PCG with retained forward diagnostics.

The primal and transpose solves use jax.lax.custom_linear_solve(), so derivatives do not trace through iteration-count branches. The operator is assumed Hermitian positive definite; the transpose preconditioner defaults to the primal preconditioner, while transpose tolerances may be controlled independently.

Parameters:

matvec (Callable[[Any], Any])
b (Any)
x0 (Any | None)
precond (Callable[[Any], Any] | None)
rtol (float)
atol (float)
max_steps (int)
single_reduction (bool)
transpose_precond (Callable[[Any], Any] | None)
transpose_rtol (float | None)
transpose_atol (float | None)
transpose_max_steps (int | None)

Return type:

PCGSolution

solvax.pcg.status_name(status)¶

Return the host-readable name for a materialized PCG status code.

Parameters:: status (int | Array)
Return type:: str