Metadata-Version: 2.4
Name: solvax
Version: 0.11.2
Summary: Differentiable structured linear solvers, preconditioners and matrix-free methods in JAX
Project-URL: Homepage, https://github.com/uwplasma/SOLVAX
Project-URL: Documentation, https://solvax.readthedocs.io
Author-email: UW Plasma <rogerio.jorge@wisc.edu>
License: MIT
License-File: LICENSE
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Topic :: Scientific/Engineering :: Physics
Classifier: Typing :: Typed
Requires-Python: >=3.10
Requires-Dist: equinox
Requires-Dist: jax
Provides-Extra: adv
Requires-Dist: adv-jax-math>=1.1; extra == 'adv'
Provides-Extra: bench
Requires-Dist: lineax; extra == 'bench'
Requires-Dist: scipy; extra == 'bench'
Provides-Extra: dev
Requires-Dist: numpy; extra == 'dev'
Requires-Dist: pytest; extra == 'dev'
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: pytest-xdist; extra == 'dev'
Requires-Dist: ruff; extra == 'dev'
Requires-Dist: scipy; extra == 'dev'
Provides-Extra: docs
Requires-Dist: furo; extra == 'docs'
Requires-Dist: myst-parser; extra == 'docs'
Requires-Dist: sphinx; extra == 'docs'
Requires-Dist: sphinx-copybutton; extra == 'docs'
Requires-Dist: sphinxcontrib-bibtex; extra == 'docs'
Provides-Extra: native
Requires-Dist: scipy; extra == 'native'
Description-Content-Type: text/markdown

# SOLVAX

[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.21651844.svg)](https://doi.org/10.5281/zenodo.21651844)

[![tests](https://github.com/uwplasma/SOLVAX/actions/workflows/tests.yml/badge.svg)](https://github.com/uwplasma/SOLVAX/actions/workflows/tests.yml)
[![codecov](https://codecov.io/gh/uwplasma/SOLVAX/branch/main/graph/badge.svg)](https://codecov.io/gh/uwplasma/SOLVAX)
[![PyPI](https://img.shields.io/pypi/v/solvax)](https://pypi.org/project/solvax/)
[![docs](https://readthedocs.org/projects/solvax/badge/?version=latest)](https://solvax.readthedocs.io)
[![license](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)

**Differentiable structured linear solvers, preconditioners and matrix-free methods in JAX.**

`solvax` provides the solver infrastructure that kinetic and PDE codes keep
re-implementing: structured direct solves (batched dense LU, block-tridiagonal
Schur elimination with truncated storage), preconditioned and recycled Krylov
methods, physics-agnostic preconditioners (coarse-operator LU, semicoarsened
geometric and p-multigrid, Kronecker approximations, symmetric additive and
line smoothers),
mixed-precision iterative refinement, and implicit differentiation of every solve —
traceable under `jit`, `vmap` and `grad`, on CPU and GPU.

Two documented exceptions, because "transparent to every transform" would not be
true: the exact-window reverse rule is a `custom_vjp`, so `jax.jacfwd` and
`jax.jvp` raise on it rather than falling back to the taped path (pass the full
window, or differentiate the untruncated entry point, when forward mode is what
you need); and `solvax.native` runs SciPy's SuperLU on the host, so it refuses
to be traced at all and says so.

It complements general JAX solver libraries with block-structured direct
elimination, coarse-operator and multigrid preconditioning, and Krylov
subspace recycling for parameter continuation. SOLVAX operators are native
JAX pytrees; no external operator abstraction is required.

## Install

```bash
pip install solvax
```

## Quickstart

```python
import jax
import jax.numpy as jnp
import solvax as sx

# Solve a block-tridiagonal system L_k x_{k-1} + D_k x_k + U_k x_{k+1} = b_k
x = sx.block_thomas(lower, diag, upper, rhs)

# Matrix-free PCG on arrays or arbitrary JAX pytrees
solution = sx.pcg(matvec, rhs, precond=preconditioner, rtol=1e-10)
assert solution.converged

# Solve an expensive affine coupling map without assembling its Jacobian
coupled = sx.affine_fixed_point_gmres(coupling_sweep, initial_state)

# Same diagnostics, but gradients use an implicit primal/transpose solve
implicit_solution = sx.pcg_linear_solve(matvec, rhs, precond=preconditioner)

# Reuse one elimination across many right-hand sides
factors = sx.block_thomas_factor(lower, diag, upper)
x1 = sx.block_thomas_solve(factors, rhs1)
x2 = sx.block_thomas_solve(factors, rhs2)

# Generate each block once when reusable factors are needed without a stored
# diagonal band. `row(j)` returns the triple (L_j, D_j, U_j) of row j; the
# parameterized form `block_fn(params, j)` used further down takes the
# parameters as its first argument.
generated_factors = sx.block_thomas_factor_fn(row, n_blocks=N)

# One generated solve with O(sqrt(N) m^2) factor storage and exact JVP/VJP.
x = sx.block_thomas_checkpointed_fn(row, N, rhs)

# Memory-truncated mode: rhs nonzero only in the lowest K blocks and only the
# lowest K solution blocks needed -> O(K m^2) memory, independent of N.
x_low = sx.block_thomas_truncated(lower, diag, upper, rhs[:3], keep_lowest=3)
```

Differentiate a generated selected-head solve with respect to the compact
parameters that build its rows, at retained state independent of the block
count. The window is estimated up front from the chain's own localization
profile. The estimate is a diagnostic, not a certificate: it reports where the
chain's transfer norms drop below one, and you should confirm the accuracy you
need by widening the window until the gradient stops moving.

```python
advice = sx.localization_crossover_window(lambda k: block_fn(p, k), N, keep_lowest=3)
# advice.certified is False: an estimate, not a guarantee. It can be passed
# straight back to the solver, or unpacked as advice.window.

def objective(params):
    x_low = sx.block_thomas_truncated_fn(
        block_fn, N, rhs[:3], keep_lowest=3,
        params=params, adjoint_window=advice,
    )
    return loss(x_low)

grad = jax.grad(objective)(p)

# Confirm the window before trusting it: widen it and see if the gradient moves.
report = sx.check_localized_gradient(
    lambda w: jax.grad(lambda q: loss(sx.block_thomas_truncated_fn(
        block_fn, N, rhs[:3], keep_lowest=3, params=q, adjoint_window=w)))(p),
    window=advice.window,
)
```

Everything is differentiable (`jax.grad` through the solve) and batchable
(`jax.vmap` over stacked systems).

## What's in the box

| Module | Contents |
|---|---|
| `solvax.operators` | Matrix-free, sum, Kronecker, block-tridiagonal and bordered (constraint-row) operator containers with closed-form transposes |
| `solvax.precond` | Jacobi/block-Jacobi, coarse-operator LU, Galerkin-deflation coarse correction, symmetric additive and alternating-direction line composition, V-/W-/F-cycle multigrid over explicit or semicoarsened rediscretized hierarchies, nearest-Kronecker, mixed-precision wrappers |
| `solvax.transfer` | Separable per-axis restriction/prolongation (full weighting, linear, injection) with periodic, dirichlet and reflective closures, exact variational adjointness, and semicoarsening plans |
| `solvax.smoothers` | Point/block Jacobi, batched tridiagonal line and exact banded plane relaxation, upwind-ordered sweeps for streaming operators, and a measured smoothing factor |
| `solvax.direct` | Block-tridiagonal Schur elimination (block Thomas): full, factor/solve split, selected-head (truncated-storage) mode, exact-window localized adjoint, per-row localization profile and window advisor |
| `solvax.banded` | Non-pivoted banded LU with row equilibration + static pivoting; periodic variant via the Woodbury capacitance trick |
| `solvax.tridiagonal` | Batched scalar tridiagonal solve (reproducible Thomas / fused cuSPARSE backend) and periodic (cyclic) systems via a Sherman--Morrison correction |
| `solvax.elliptic` | Spectral Fourier--Helmholtz solve for separable periodic-by-bounded elliptic problems — the drift-plane / vorticity `lap phi = rhs` inversion, one FFT + one batched tridiagonal sweep |
| `solvax.krylov` | Flexible restarted GMRES (CGS2 + Givens) over arrays, scalars and arbitrary pytrees with optional custom inner products, and GCROT Krylov subspace recycling with FIFO or harmonic-Ritz (GCRO-DR) deflated restarting |
| `solvax.pcg` | Matrix-free pytree PCG with preconditioning, fixed-shape residual history, and explicit convergence/breakdown status |
| `solvax.fixed_point` | Safeguarded Aitken, bounded-memory (condition-filtered) Anderson, and matrix-free affine fixed-point FGMRES |
| `solvax.implicit` | Matrix-free `newton_krylov` (JFNK) plus implicit-function-theorem `linear_solve` and `root_solve` — gradients cost one extra (transposed) solve |
| `solvax.autodiff` | Bounded-memory chunked forward/reverse Jacobians (`chunked_jacfwd`/`jacrev`/`jacobian`) with automatic chunk sizing |
| `solvax.refine` | Mixed-precision iterative refinement (float32 factor, float64 residuals) |
| `solvax.native` | Host-side SuperLU bridge (non-differentiable, import-guarded) |

Complex-valued GMRES/GCROT, tridiagonal solves, and fixed-point acceleration
use Hermitian inner products and real-valued safeguards. Remaining roadmap:
multi-leaf pytree GCROT operands (GCROT takes arrays of any rank; GMRES is
pytree-native) and expanded GPU batched-LU benchmarks.

```python
# Preconditioned, recycled Krylov across a parameter scan:
sol = sx.gcrot(matvec, b, precond=coarse_inverse, m=50, k=10)
sol2 = sx.gcrot(matvec2, b2, precond=coarse_inverse, recycle=sol.recycle)

# Matrix-free Newton-Krylov (JFNK): Jacobian-vector products via jax.linearize,
# each correction solved by FGMRES over an array or structured pytree state.
root = sx.newton_krylov(residual_fn, x0, precond=approx_inverse, rtol=1e-8)

# Weakly contractive affine coupling map G(x) = L x + c, solved as (I - L) x = c:
fixed = sx.affine_fixed_point_gmres(coupling_map, x0, restart=20)

# Periodic (cyclic) scalar tridiagonal line, corners in sub[0] and sup[-1]:
x_line = sx.cyclic_tridiagonal_solve(sub, dia, sup, line_rhs)

# Differentiable solve wrapping any solver:
x = sx.linear_solve(matvec, b, solver=lambda mv, rhs: sx.gmres(mv, rhs).x)
```

## License

MIT. Developed by the [UW Plasma group](https://github.com/uwplasma).
