solvax 0.1.0__tar.gz

solvax-0.1.0/.github/workflows/docs.yml

@@ -0,0 +1,19 @@
+name: docs
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+      - run: pip install -e '.[docs]'
+      - run: python -m sphinx -W -b html docs docs/_build

solvax-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,32 @@
+name: publish
+
+on:
+  push:
+    tags: ["v*"]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - run: pipx run build
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

solvax-0.1.0/.github/workflows/tests.yml

@@ -0,0 +1,30 @@
+name: tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    strategy:
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+      - run: pip install -e '.[dev]'
+      - run: pytest -n auto --cov --cov-report=xml
+      - uses: codecov/codecov-action@v4
+        if: matrix.os == 'ubuntu-latest'
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}

solvax-0.1.0/.gitignore

@@ -0,0 +1,9 @@
+__pycache__/
+*.egg-info/
+.pytest_cache/
+.coverage
+coverage.xml
+dist/
+build/
+docs/_build/
+.DS_Store

solvax-0.1.0/.readthedocs.yaml

@@ -0,0 +1,16 @@
+version: 2
+
+build:
+  os: ubuntu-24.04
+  tools:
+    python: "3.12"
+
+sphinx:
+  configuration: docs/conf.py
+  fail_on_warning: true
+
+python:
+  install:
+    - method: pip
+      path: .
+      extra_requirements: [docs]

solvax-0.1.0/CITATION.cff

@@ -0,0 +1,10 @@
+cff-version: 1.2.0
+message: "If you use solvax in your research, please cite it as below."
+title: "solvax: differentiable structured linear solvers in JAX"
+authors:
+  - family-names: Jorge
+    given-names: Rogerio
+    email: rogerio.jorge@wisc.edu
+    affiliation: University of Wisconsin-Madison
+repository-code: "https://github.com/uwplasma/SOLVAX"
+license: MIT

solvax-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 UW Plasma group, University of Wisconsin-Madison
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

solvax-0.1.0/PKG-INFO

@@ -0,0 +1,115 @@
+Metadata-Version: 2.4
+Name: solvax
+Version: 0.1.0
+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
+Requires-Python: >=3.10
+Requires-Dist: equinox
+Requires-Dist: jax
+Requires-Dist: lineax
+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: 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
+
+[![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, p-multigrid,
+Kronecker approximations, line smoothers), mixed-precision iterative
+refinement, and implicit differentiation of every solve — all
+jit/vmap/grad-transparent, on CPU and GPU.
+
+It fills a gap in the JAX ecosystem: [lineax](https://github.com/patrick-kidger/lineax)
+offers general linear-operator abstractions and standard solvers, but not
+block-structured direct elimination, coarse-operator/multigrid
+preconditioning, or Krylov subspace recycling for parameter continuation.
+`solvax` builds on lineax's operator interface and adds exactly that layer.
+
+## Install
+
+```bash
+pip install solvax
+```
+
+## Quickstart
+
+```python
+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)
+
+# 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)
+
+# 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)
+```
+
+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, alternating-direction line smoothers, p-multigrid V-cycles, nearest-Kronecker, mixed-precision wrappers |
+| `solvax.direct` | Block-tridiagonal Schur elimination (block Thomas): full, factor/solve split, truncated-storage mode |
+| `solvax.banded` | Non-pivoted banded LU with row equilibration + static pivoting; periodic variant via the Woodbury capacitance trick |
+| `solvax.krylov` | Flexible restarted GMRES (CGS2 + Givens) and GCROT-style Krylov subspace recycling for parameter continuation |
+| `solvax.implicit` | Implicit-function-theorem `linear_solve` and `root_solve` — gradients cost one extra (transposed) solve |
+| `solvax.refine` | Mixed-precision iterative refinement (float32 factor, float64 residuals) |
+| `solvax.native` | Host-side SuperLU bridge (non-differentiable, import-guarded) |
+
+Roadmap to v0.2: harmonic-Ritz recycle selection, pytree operands, complex
+dtypes, 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)
+
+# 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).

solvax-0.1.0/README.md

@@ -0,0 +1,80 @@
+# SOLVAX
+
+[![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, p-multigrid,
+Kronecker approximations, line smoothers), mixed-precision iterative
+refinement, and implicit differentiation of every solve — all
+jit/vmap/grad-transparent, on CPU and GPU.
+
+It fills a gap in the JAX ecosystem: [lineax](https://github.com/patrick-kidger/lineax)
+offers general linear-operator abstractions and standard solvers, but not
+block-structured direct elimination, coarse-operator/multigrid
+preconditioning, or Krylov subspace recycling for parameter continuation.
+`solvax` builds on lineax's operator interface and adds exactly that layer.
+
+## Install
+
+```bash
+pip install solvax
+```
+
+## Quickstart
+
+```python
+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)
+
+# 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)
+
+# 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)
+```
+
+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, alternating-direction line smoothers, p-multigrid V-cycles, nearest-Kronecker, mixed-precision wrappers |
+| `solvax.direct` | Block-tridiagonal Schur elimination (block Thomas): full, factor/solve split, truncated-storage mode |
+| `solvax.banded` | Non-pivoted banded LU with row equilibration + static pivoting; periodic variant via the Woodbury capacitance trick |
+| `solvax.krylov` | Flexible restarted GMRES (CGS2 + Givens) and GCROT-style Krylov subspace recycling for parameter continuation |
+| `solvax.implicit` | Implicit-function-theorem `linear_solve` and `root_solve` — gradients cost one extra (transposed) solve |
+| `solvax.refine` | Mixed-precision iterative refinement (float32 factor, float64 residuals) |
+| `solvax.native` | Host-side SuperLU bridge (non-differentiable, import-guarded) |
+
+Roadmap to v0.2: harmonic-Ritz recycle selection, pytree operands, complex
+dtypes, 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)
+
+# 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).

solvax-0.1.0/codecov.yml

@@ -0,0 +1,11 @@
+coverage:
+  status:
+    project:
+      default:
+        target: 95%
+        threshold: 1%
+    patch:
+      default:
+        target: 90%
+comment:
+  layout: "condensed_header, diff"

solvax-0.1.0/docs/api.md

@@ -0,0 +1,65 @@
+# API reference
+
+## Linear operators
+
+```{eval-rst}
+.. automodule:: solvax.operators
+   :members:
+   :member-order: bysource
+```
+
+## Preconditioners
+
+```{eval-rst}
+.. automodule:: solvax.precond
+   :members:
+   :member-order: bysource
+```
+
+## Structured direct solvers
+
+```{eval-rst}
+.. automodule:: solvax.direct
+   :members:
+   :member-order: bysource
+```
+
+## Banded and periodic-banded LU
+
+```{eval-rst}
+.. automodule:: solvax.banded
+   :members:
+   :member-order: bysource
+```
+
+## Krylov methods
+
+```{eval-rst}
+.. automodule:: solvax.krylov
+   :members:
+   :member-order: bysource
+```
+
+## Implicit differentiation
+
+```{eval-rst}
+.. automodule:: solvax.implicit
+   :members:
+   :member-order: bysource
+```
+
+## Mixed-precision refinement
+
+```{eval-rst}
+.. automodule:: solvax.refine
+   :members:
+   :member-order: bysource
+```
+
+## Native host bridges
+
+```{eval-rst}
+.. automodule:: solvax.native
+   :members:
+   :member-order: bysource
+```

solvax-0.1.0/docs/conf.py

@@ -0,0 +1,35 @@
+"""Sphinx configuration for the solvax documentation."""
+
+import solvax
+
+project = "solvax"
+copyright = "2026, UW Plasma group"
+author = "UW Plasma group"
+release = solvax.__version__
+
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.mathjax",
+    "sphinx.ext.intersphinx",
+    "myst_parser",
+    "sphinx_copybutton",
+    "sphinxcontrib.bibtex",
+]
+
+bibtex_bibfiles = ["references.bib"]
+myst_enable_extensions = ["dollarmath", "amsmath"]
+
+html_theme = "furo"
+html_title = "solvax"
+
+intersphinx_mapping = {
+    "jax": ("https://docs.jax.dev/en/latest/", None),
+    "python": ("https://docs.python.org/3", None),
+}
+
+autodoc_typehints = "description"
+napoleon_google_docstring = True
+# NamedTuple fields are already documented via napoleon Attributes sections;
+# ivar rendering avoids duplicate autodoc entries for the same objects.
+napoleon_use_ivar = True

solvax-0.1.0/docs/index.md

@@ -0,0 +1,51 @@
+# solvax
+
+**Differentiable structured linear solvers, preconditioners and matrix-free
+methods in JAX.**
+
+`solvax` is the solver layer that kinetic and PDE codes keep re-implementing,
+factored out once: structured direct solves, preconditioned and recycled
+Krylov methods, physics-agnostic preconditioners, mixed-precision refinement,
+and implicit differentiation — all jit/vmap/grad-transparent on CPU and GPU.
+
+## Quickstart
+
+```python
+import solvax as sx
+
+# 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)
+
+# Reuse one elimination across 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)
+```
+
+## The block-tridiagonal kernel
+
+For blocks $L_k x_{k-1} + D_k x_k + U_k x_{k+1} = b_k$, a Schur-complement
+sweep from the last block down,
+
+$$\Delta_{N-1} = D_{N-1}, \qquad
+\Delta_k = D_k - U_k \Delta_{k+1}^{-1} L_{k+1}, \qquad
+\sigma_k = b_k - U_k \Delta_{k+1}^{-1} \sigma_{k+1},$$
+
+followed by substitution upward from block 0, solves the system exactly with
+one dense LU and one matrix product per block {cite}`golub2013,demmel1995`.
+When the right-hand side vanishes above block $K$ and only the lowest $K$
+solution blocks are needed — the typical situation for spectral kinetic
+equations, where sources and velocity moments involve only the first few
+modes — storage above $K$ can be discarded on the fly, giving memory
+$O(K m^2)$ independent of $N$ {cite}`escoto2025`.
+
+```{toctree}
+:maxdepth: 1
+
+api
+```
+
+## References
+
+```{bibliography}
+```

solvax-0.1.0/docs/references.bib

@@ -0,0 +1,85 @@
+@book{golub2013,
+  author    = {Golub, Gene H. and Van Loan, Charles F.},
+  title     = {Matrix Computations},
+  edition   = {4},
+  publisher = {Johns Hopkins University Press},
+  year      = {2013},
+}
+
+@article{demmel1995,
+  author  = {Demmel, James W. and Higham, Nicholas J. and Schreiber, Robert S.},
+  title   = {Stability of block {LU} factorization},
+  journal = {Numerical Linear Algebra with Applications},
+  volume  = {2},
+  pages   = {173--190},
+  year    = {1995},
+  doi     = {10.1002/nla.1680020208},
+}
+
+@phdthesis{escoto2025,
+  author = {Escoto, F. Javier},
+  title  = {Fast and accurate calculation of the bootstrap current and radial
+            neoclassical transport in low collisionality stellarator plasmas},
+  school = {Universidad Carlos III de Madrid},
+  year   = {2025},
+  url    = {https://arxiv.org/abs/2510.27513},
+}
+
+@article{morgan2002,
+  author  = {Morgan, Ronald B.},
+  title   = {{GMRES} with deflated restarting},
+  journal = {SIAM Journal on Scientific Computing},
+  volume  = {24},
+  pages   = {20--37},
+  year    = {2002},
+  doi     = {10.1137/S1064827599364659},
+}
+
+@article{parks2006,
+  author  = {Parks, Michael L. and de Sturler, Eric and Mackey, Greg and
+             Johnson, Duane D. and Maiti, Spandan},
+  title   = {Recycling {Krylov} subspaces for sequences of linear systems},
+  journal = {SIAM Journal on Scientific Computing},
+  volume  = {28},
+  pages   = {1651--1674},
+  year    = {2006},
+  doi     = {10.1137/040607277},
+}
+
+@article{carson2018,
+  author  = {Carson, Erin and Higham, Nicholas J.},
+  title   = {Accelerating the solution of linear systems by iterative
+             refinement in three precisions},
+  journal = {SIAM Journal on Scientific Computing},
+  volume  = {40},
+  pages   = {A817--A847},
+  year    = {2018},
+  doi     = {10.1137/17M1140819},
+}
+
+@inproceedings{blondel2022,
+  author    = {Blondel, Mathieu and Berthet, Quentin and Cuturi, Marco and
+               Frostig, Roy and Hoyer, Stephan and Llinares-L{\'o}pez, Felipe
+               and Pedregosa, Fabian and Vert, Jean-Philippe},
+  title     = {Efficient and modular implicit differentiation},
+  booktitle = {Advances in Neural Information Processing Systems},
+  year      = {2022},
+  url       = {https://arxiv.org/abs/2105.15183},
+}
+
+@article{skene2026,
+  author  = {Skene, Calum S. and Burns, Keaton J.},
+  title   = {Fast automated adjoints for spectral {PDE} solvers},
+  journal = {Proceedings of the National Academy of Sciences},
+  year    = {2026},
+  url     = {https://arxiv.org/abs/2506.14792},
+}
+
+@incollection{vanloan1993,
+  author    = {Van Loan, Charles F. and Pitsianis, Nikos},
+  title     = {Approximation with {Kronecker} products},
+  booktitle = {Linear Algebra for Large Scale and Real-Time Applications},
+  publisher = {Springer},
+  year      = {1993},
+  doi       = {10.1007/978-94-015-8196-7_17},
+}

solvax-0.1.0/examples/01_block_tridiagonal_kinetic.py

@@ -0,0 +1,49 @@
+"""Solve a kinetic-style block-tridiagonal system with truncated storage.
+
+Spectral discretizations of kinetic equations (e.g. a Legendre expansion in
+pitch angle) couple only neighbouring modes l-1, l, l+1, giving a
+block-tridiagonal system whose right-hand side lives in the lowest few modes
+and whose observables (density, flow, pressure moments) touch only those same
+modes. `block_thomas_truncated` exploits both facts: memory O(K m^2)
+independent of the number of modes.
+
+Expected runtime: a few seconds on a laptop CPU.
+"""
+
+import jax
+import jax.numpy as jnp
+
+import solvax as sx
+
+jax.config.update("jax_enable_x64", True)
+
+n_modes, m = 64, 100  # Legendre-like modes x flux-surface grid points
+key = jax.random.PRNGKey(0)
+k1, k2, k3, k4 = jax.random.split(key, 4)
+
+# Streaming-like off-diagonal coupling and a collisional diagonal ~ l(l+1).
+lower = 0.3 * jax.random.normal(k1, (n_modes, m, m))
+upper = 0.3 * jax.random.normal(k2, (n_modes, m, m))
+nu = 0.5 * jnp.arange(n_modes) * (jnp.arange(n_modes) + 1) + 5.0
+diag = jax.random.normal(k3, (n_modes, m, m)) + nu[:, None, None] * jnp.eye(m)
+
+# Two drives (radial + parallel), nonzero only in modes 0..2 — solved together.
+rhs_low = jax.random.normal(k4, (3, m, 2))
+
+x_low = sx.block_thomas_truncated(lower, diag, upper, rhs_low, keep_lowest=3)
+print("lowest-mode solution block shape:", x_low.shape)
+
+# Cross-check against the full solve.
+rhs_full = jnp.zeros((n_modes, m, 2)).at[:3].set(rhs_low)
+x_full = sx.block_thomas(lower, diag, upper, rhs_full)
+err = jnp.max(jnp.abs(x_low - x_full[:3]))
+print(f"max |truncated - full| = {err:.2e}")
+
+# The whole solve is differentiable: gradient of a "flux" moment w.r.t. nu.
+def flux(nu_vec):
+    d = jax.random.normal(k3, (n_modes, m, m)) + nu_vec[:, None, None] * jnp.eye(m)
+    x = sx.block_thomas_truncated(lower, d, upper, rhs_low, keep_lowest=3)
+    return jnp.sum(x[1] ** 2)  # mode-1 moment ~ parallel flow
+
+g = jax.grad(flux)(nu)
+print("d(flux)/d(nu_0) =", float(g[0]))