vmex 0.2.0__tar.gz

vmex-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 UW Plasma
+
+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.

vmex-0.2.0/PKG-INFO

@@ -0,0 +1,615 @@
+Metadata-Version: 2.4
+Name: vmex
+Version: 0.2.0
+Summary: JAX implementation of VMEC2000 with differentiable fixed-boundary and branch-local free-boundary research paths.
+Author: vmex contributors
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/uwplasma/VMEX
+Project-URL: Documentation, https://vmex.readthedocs.io/en/latest/
+Project-URL: Repository, https://github.com/uwplasma/VMEX
+Project-URL: Issues, https://github.com/uwplasma/VMEX/issues
+Project-URL: Changelog, https://github.com/uwplasma/VMEX/releases
+Keywords: vmec,stellarator,mhd,jax,optimization
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: scipy
+Requires-Dist: netCDF4
+Requires-Dist: matplotlib
+Requires-Dist: packaging
+Requires-Dist: booz_xform_jax
+Requires-Dist: solvax
+Requires-Dist: tomli; python_version < "3.11"
+Provides-Extra: freeb
+Requires-Dist: virtual_casing_jax; extra == "freeb"
+Provides-Extra: docs
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: sphinx-copybutton; extra == "docs"
+Requires-Dist: furo; extra == "docs"
+Provides-Extra: dev
+Requires-Dist: build; extra == "dev"
+Requires-Dist: furo; extra == "dev"
+Requires-Dist: mypy; 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: sphinx; extra == "dev"
+Requires-Dist: types-setuptools; extra == "dev"
+Dynamic: license-file
+
+# VMEX
+
+[![PyPI version](https://img.shields.io/pypi/v/vmex.svg)](https://pypi.org/project/vmex/)
+[![Python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12-blue.svg)](https://github.com/uwplasma/VMEX/blob/main/pyproject.toml)
+[![License](https://img.shields.io/github/license/uwplasma/vmex)](https://github.com/uwplasma/VMEX/blob/main/LICENSE)
+[![CI](https://img.shields.io/github/actions/workflow/status/uwplasma/vmex/ci.yml?branch=main&label=ci)](https://github.com/uwplasma/VMEX/actions/workflows/ci.yml)
+[![Docs](https://img.shields.io/readthedocs/vmex/latest?label=docs)](https://vmex.readthedocs.io/en/latest/)
+
+> **`vmec_jax` is now `vmex`.** The package was renamed: install with
+> `pip install vmex` and `import vmex`. The `vmec` CLI command still works as an
+> alias, and `import vmec_jax` keeps working (with a deprecation warning) for
+> one release. Full documentation: **[vmex.readthedocs.io](https://vmex.readthedocs.io/en/latest/)**.
+
+**VMEX** is a clean-room, JAX-native reimplementation of the
+[VMEC2000](https://princetonuniversity.github.io/STELLOPT/VMEC) ideal-MHD
+equilibrium code for stellarators and tokamaks. It reproduces VMEC2000
+iteration-for-iteration on benchmark decks — and, unlike the Fortran
+original, it is differentiable and runs on GPUs.
+
+- **VMEC2000 parity.** The solver ports VMEC2000's algorithms
+  constant-for-constant (steepest-descent moment method, radial
+  preconditioner, spectral condensation, NESTOR vacuum solve). Benchmark
+  decks converge in the *same* number of iterations and reproduce the
+  plasma energy at machine precision. An optional **2D block
+  preconditioner** cuts iterations 2.5–11x on stiff cases while leaving the
+  default path byte-identical.
+- **Differentiable.** Gradients of *fixed-boundary* equilibrium outputs with
+  respect to boundary shape and profile parameters by implicit
+  differentiation of the converged fixed point — no finite differences, no
+  unrolling — validated against central finite differences to ~1e-6 relative
+  (see the gradient table in the docs), with an O(1)-memory adjoint. **Free
+  boundary** is differentiable end-to-end through the virtual-casing vacuum
+  field (coil / `extcur` derivatives), finite-difference-validated.
+- **Drop-in.** Reads VMEC2000 `input.*` namelists and VMEC++-style JSON,
+  prints VMEC2000-format iteration output, and writes `wout_*.nc` files
+  that load unchanged in simsopt and booz_xform.
+- **Batteries included.** Plotting (`vmex --plot`), Boozer transform
+  (`vmex --booz`), spline profiles, multigrid, hot restart, free boundary
+  from mgrid files *or* directly from coils,
+  typed zero-crash errors — with the shared linear/adjoint solver layer
+  factored out into [SOLVAX](https://pypi.org/project/solvax/).
+
+![Flux surfaces, 3-D geometry, and Boozer |B| of the bundled quick-start QH case](docs/_static/figures/readme_equilibrium_showcase.png)
+
+*The bundled quick-start case (`vmex --test`): flux-surface cross sections,
+the 3-D plasma boundary coloured by `|B|`, and `|B|` in **Boozer coordinates**
+on the last closed flux surface (the near-straight diagonal contours are the
+signature of quasi-helical symmetry) for a four-field-period stellarator —
+all from the built-in `vmex.core.plotting` / `core.boozer` helpers.*
+
+## Install
+
+Install from PyPI:
+
+```bash
+pip install vmex
+```
+
+Development install from source:
+
+```bash
+git clone https://github.com/uwplasma/VMEX
+cd vmex && pip install -e .
+```
+
+## Quickstart
+
+```bash
+vmex --doctor     # check the installation and JAX backend
+vmex --test       # solve the bundled QH case, write wout + plots
+vmex input.X      # run any VMEC2000 input deck (or VMEC++-style JSON)
+```
+
+`vmex input.X` writes `wout_X.nc` next to the input (`--outdir` to
+redirect). To try it on a real deck:
+
+```bash
+curl -L -O https://raw.githubusercontent.com/uwplasma/vmex/main/examples/data/input.nfp4_QH_warm_start
+vmex input.nfp4_QH_warm_start
+```
+
+Post-process any wout file, including ones written by VMEC2000:
+
+```bash
+vmex --plot wout_nfp4_QH_warm_start.nc     # surfaces, |B|, profiles, 3D
+vmex --booz wout_nfp4_QH_warm_start.nc     # Boozer transform -> boozmn_*.nc
+vmex --plot boozmn_nfp4_QH_warm_start.nc   # Boozer |B| contours + spectrum
+```
+
+## Parity with VMEC2000
+
+VMEX is validated end-to-end against golden VMEC2000 (PARVMEC 9.0) runs:
+benchmark decks converge in **exactly** the golden iteration count — including
+DSHAPE's mid-run jacobian reset — and reproduce the plasma energy `wb` to
+1 part in 10¹⁵. Across the full benchmark suite (14 rows, all at `ns ≥ 201`),
+the iteration count matches VMEC2000 exactly on 12 rows; on the free-boundary
+CTH-like row it converges in a ~9% iteration tail, and on Nuhrenberg–Zille QHS
+it converges in *fewer* iterations (1681 vs 2829). Per-variable wout agreement
+and the full test gates live in the
+[documentation](https://vmex.readthedocs.io/en/latest/).
+
+![Force residual vs iteration for vmex, VMEC2000, and VMEC++](docs/_static/figures/readme_convergence.png)
+
+*Parity is per-iteration, not just end-to-end: the total force residual
+(`fsqr + fsqz + fsql`) of the quick-start QH case at ns=51, per iteration.
+The vmex trajectory lies exactly on top of VMEC2000's (both converge in
+502 iterations); VMEC++ follows a near-identical path (501 iterations).
+Traces: vmex `SolveResult.fsq_history`, VMEC2000 `NSTEP=1` stdout,
+VMEC++ wout `fsqt`.*
+
+### Optional 2D preconditioner: fewer iterations on stiff cases
+
+The default radial (1D) preconditioner reproduces VMEC2000 iteration-for-iteration.
+An opt-in **2D block preconditioner** (matrix-free Newton: a Jacobian-vector-product
+Hessian on SOLVAX's GMRES) cuts the iteration count **2.5–11×** at *identical*
+accuracy — the converged `wb` matches the 1D result to ~1e-10 (it changes the path,
+not the fixed point).
+
+![2D vs 1D preconditioner iteration counts on stiff cases](docs/_static/figures/readme_precond.png)
+
+**Why it is opt-in, not the default.** Fewer iterations is not the same as less
+wall-clock: each 2D Newton step (a GMRES solve of Hessian-vector products) costs far
+more than a 1D radial sweep. Measured across easy and stiff decks the wall-clock
+ranges 0.55–1.16× — a wash to *slower* (e.g. ~2× slower on a plain circular tokamak,
+a tie even on an aspect-ratio-100 stiff case) — and peak memory is ~30% higher (the
+extra GMRES/HVP compile graph). So the 1D path stays the byte-identical default, and
+the 2D preconditioner is there for cases where the 1D iteration count is the
+bottleneck or stalls.
+
+## Performance
+
+![Wall-clock comparison against VMEC2000 and VMEC++](docs/_static/figures/readme_runtime_compare.png)
+
+Full-solve wall-clock times on the bundled benchmark suite (Apple Silicon
+CPU, single thread; `benchmarks/baseline.json`; reproduce with
+`python benchmarks/run_baseline.py`):
+
+- **Warm** — kernels already compiled; the number that matters inside an
+  optimization loop or scan. Faster than VMEC2000 on **every** benchmark row
+  (1.3–2.6× on typical decks, up to ~7× on small ones) — including the
+  free-boundary rows (1.3–1.5×) since the NESTOR iteration loop was fused
+  into jitted multi-iteration lanes. Ratios measured on a shared CPU are
+  conservative lower bounds.
+- **Cold** — a fresh CLI process pays a one-time 5–25 s JAX/XLA compile, so a
+  single run is slower than Fortran. Executables cache per solver structure, so
+  scans, ladders, and optimizations recompile nothing — which is why *warm* is
+  the workflow number.
+- **GPU** — at these sizes a fixed per-solve dispatch cost dominates and the CPU
+  wins outright; per-iteration throughput favours the GPU ~3× on the largest
+  decks. The device policy picks CPU or GPU per stage.
+- **Memory** — peak (0.6–3.3 GB) is the transient XLA *compile* working set, not
+  the data: the equilibrium state is a few MB. The optimization Jacobian is
+  bounded by column chunking (`jac_chunk_size="auto"`), so it does not grow with
+  the number of design variables.
+
+## Features
+
+| | VMEX | VMEC2000 | VMEC++ |
+|---|:---:|:---:|:---:|
+| Fixed-boundary equilibria | ✅ | ✅ | ✅ |
+| Free boundary from an mgrid file | ✅ | ✅ | ✅ |
+| Free boundary directly from coils (no mgrid) | ✅ | ❌ | ❌ |
+| Free-boundary tokamaks (`ntor = 0`) | ✅ | ✅ | ❌ |
+| Non-stellarator-symmetric (`LASYM = T`) | ✅ | ✅ | ❌ |
+| Fixed-boundary fallback on missing mgrid | ✅ | ✅ | ❌ |
+| Spline profiles (cubic / Akima) | ✅ | ✅ | ❌ |
+| VMEC++-schema JSON input | ✅ | ❌ | ✅ |
+| Hot restart from a previous state | ✅ | ❌ | ✅ |
+| Typed zero-crash errors | ✅ | ❌ | ✅ |
+| Boozer transform built in (`--booz`) | ✅ | ❌ | ❌ |
+| Plotting built in (`--plot`) | ✅ | ❌ | ❌ |
+| GPU execution | ✅ | ❌ | ❌ |
+| Differentiable fixed boundary (implicit diff, O(1) memory) | ✅ | ❌ | ❌ |
+| Differentiable free boundary (virtual casing) | ✅ | ❌ | ❌ |
+| 2D block preconditioner (stiff-case speedup) | ✅ | ❌ | ❌ |
+
+### Free boundary straight from coils
+
+Free-boundary solves can run directly from a coil set: tabulate an
+[ESSOS](https://github.com/uwplasma/ESSOS) coil set onto the solver grid in
+memory (`essos.coils.Coils.to_mgrid`) and pass it as `external_field=`,
+with no MAKEGRID file involved. For gradients, the differentiable free
+boundary evaluates a JAX Biot-Savart (a plain `xyz→B` callable) at the
+boundary points of each iteration, keeping the coil degrees of freedom
+differentiable end-to-end. All coil geometry lives in ESSOS; vmex has no
+coil code of its own.
+
+![Free-boundary Landreman-Paul QA pressure scan directly from ESSOS coils](docs/_static/figures/readme_essos_beta_scan.png)
+
+*Free-boundary equilibria of the Landreman–Paul precise-QA configuration held
+by its 16 modular coils as optimized in
+[ESSOS](https://github.com/uwplasma/ESSOS) (3 KB coil JSON bundled in
+`examples/data/`). Pressure is ramped at fixed coil currents with each point
+warm-started from the previous boundary, and `PRES_SCALE` is calibrated per
+point so the **actual** volume-average beta of the converged wout
+(`betatotal`) — not a nominal input value — lands on 0, 1, 2, 3 % (all within
+0.08 %, force residual ~2e-10 at ns = 51). The plasma dilates and the magnetic
+axis Shafranov-shifts 14 cm outboard at the φ = 0 section (right panel) while
+the coils never move. Reproduce with
+`python examples/free_boundary_essos_coils.py`.*
+
+### Single-stage plasma + coil optimization
+
+VMEX can optimize the plasma boundary and the coils together, with one
+exact gradient. A single `jax.value_and_grad` differentiates through the
+fixed-boundary equilibrium (implicit adjoint), the virtual-casing surface
+field, and the Biot–Savart law of the ESSOS coil filaments, covering boundary
+Fourier modes, coil shapes, and coil currents at once. The benchmark below
+compares this against the classical two-stage approach — stage 1 shapes the
+boundary for quasi-axisymmetry, stage 2 fits coils to that frozen boundary —
+from the same seeds (a circular torus and four circular coils), with identical
+coil budgets, scored on the equilibrium each final coil set actually produces.
+The finite-β case runs the same joint optimization with a pressure profile;
+no published code demonstrates this in general form.
+
+The most effective use is to polish the two-stage result, the "stage 3" of
+[arXiv:2302.10622](https://arxiv.org/abs/2302.10622): warm-start the joint
+objective from the stage-1 boundary and stage-2 coils and let both adapt.
+In 10–30 minutes this lowers the normal-field error by 33% (vacuum) and 17%
+(finite β) below the two-stage result, with quasisymmetry and iota unchanged —
+stage 2 cannot make this correction because it holds the boundary frozen.
+A pure cold start (third column) shows the same joint descent from the crude
+seeds: after 50 iterations it reaches low B·n with compact coils, but its
+quasisymmetry is far from what a dedicated stage 1 delivers, which is why the
+polish pattern is recommended.
+
+![Cold-start single-stage vs two-stage plasma+coil optimization, vacuum and finite beta](docs/_static/figures/readme_single_stage.png)
+
+*Top: seed (grey, dashed) vs two-stage (orange) vs cold-start single-stage
+(blue) boundaries at φ = 0 and a half field period — the polish boundary is
+visually indistinguishable from two-stage (same aspect and iota), so it is not
+drawn. Middle/bottom: each approach's final LCFS coloured by |B| inside its
+own final coils.*
+
+Vacuum (measured; identical seeds and coil budgets across columns):
+
+| metric (vacuum) | two-stage | + single-stage polish | single-stage (cold) |
+|---|---|---|---|
+| QS ratio residual | 9.3e-05 | 1.6e-04 | 2.4e-02 |
+| mean iota (target 0.42) | 0.420 | 0.420 | 0.396 |
+| ⟨\|B·n\|⟩/⟨B⟩ | 2.38e-03 | **1.60e-03** | 3.05e-03 |
+| max\|B·n\|/⟨B⟩ | 1.30e-02 | **7.84e-03** | 1.18e-02 |
+| coil lengths [m] (≤ 4.40) | 4.12–4.39 | 4.11–4.40 | 3.60–3.87 |
+
+Finite β (⟨β⟩ ≈ 1.5 %, same pressure profile in all columns):
+
+| metric (finite β) | two-stage | + single-stage polish | single-stage (cold) |
+|---|---|---|---|
+| QS ratio residual | 4.4e-05 | 2.4e-04 | 2.3e-02 |
+| mean iota (target 0.42) | 0.420 | 0.422 | 0.100 |
+| ⟨\|B·n\|⟩/⟨B⟩ | 2.80e-03 | **2.34e-03** | 6.20e-03 |
+| max\|B·n\|/⟨B⟩ | 1.37e-02 | **1.27e-02** | 1.75e-02 |
+| coil lengths [m] (≤ 4.40) | 3.91–4.18 | 3.91–4.19 | 3.25–3.28 |
+
+Reproduce with `python examples/single_stage_vs_two_stage.py --case vacuum
+--phase all` (and `--case beta`). Measured on a 36-core CPU: stage 1 ≈ 7–9 min,
+stage 2 ≈ 6 min, polish ≈ 10–30 min; the optional cold-start single column is
+the long pole (≈ 1.5 h vacuum, several hours at finite β). The phases are
+resumable, so long runs can be split across sessions.
+
+## Code size
+
+VMEX delivers that superset of capabilities in little more than **half the
+code**, and is the most densely documented of the three. Solver source only (tests,
+language bindings, and vendored third-party excluded), counted with
+[`pygount`](https://pypi.org/project/pygount/) 3.2:
+
+| code base | language | files | code (SLOC) | comments / docstrings | doc-to-code |
+|---|---|---:|---:|---:|---:|
+| **VMEX** | Python | 41 | **13,326** | 6,744 | **0.51** |
+| VMEC2000 (PARVMEC) | Fortran | 115 | 24,190 | 8,425 | 0.35 |
+| VMEC++ | C++ / Python | 117 | 22,824 | 7,646 | 0.34 |
+
+VMEX is little more than half the SLOC of VMEC2000 and VMEC++, while
+*adding* differentiability, GPU execution, direct-coil free boundary, and a
+built-in Boozer transform — and it carries the highest comment/docstring
+density of the three (reproduce with
+`pygount --format=summary vmex`).
+
+## Python API
+
+```python
+from vmex.core.input import VmecInput
+from vmex.core import optimize as opt
+from vmex.core.wout import write_wout
+from vmex.core.plotting import plot_wout
+
+inp = VmecInput.from_file("input.nfp4_QH_warm_start")
+eq = opt.solve_equilibrium(inp)        # full NS_ARRAY ladder, VMEC2000 numerics
+print(eq.result.converged, eq.result.iterations, float(eq.wout.aspect))
+
+write_wout("wout_nfp4_QH_warm_start.nc", eq.wout)   # wout built lazily on eq
+plot_wout(eq.wout, "figures/")
+```
+
+Choosing an entry point: `optimize.solve_equilibrium` for Python analysis and
+objectives (state + runtime + lazy `.wout`); `multigrid.solve_multigrid` when
+you only need the converged state (the CLI's engine); `implicit.run` for
+gradients (`jax.grad`-able `ImplicitSolution`); `solver.solve` as the
+low-level single-grid building block.
+
+Optimization building blocks live in `vmex.core.optimize`
+(quasisymmetry and omnigenity residuals; aspect ratio, iota, mirror ratio,
+magnetic well, ballooning-stability targets; a least-squares driver over
+boundary Fourier coefficients) with implicit-differentiation gradients from
+`vmex.core.implicit` (`jac="implicit"`). The recommended pattern is **one
+`least_squares` call** — no `max_mode` continuation loop — with **Exponential
+Spectral Scaling** ordering the harmonics through the trust region:
+
+```python
+from vmex import optimize as opt
+
+qs = opt.QuasisymmetryRatioResidual(surfaces, helicity_m=1, helicity_n=0)
+result = opt.least_squares(
+    [(qs, 0.0, 1.0), (opt.aspect_ratio, 6.0, 1.0), (opt.mean_iota, 0.42, 1.0)],
+    inp, max_mode=5, jac="implicit",
+    use_ess=True,        # exp(-alpha*max(|m|,|n|)) trust radius per dof:
+)                        # high harmonics on short leashes — no ladder needed
+```
+
+Measured on a 36-core CPU from a near-circular torus (single call, all
+harmonics released at once; `examples/optimization/*_ess.py`; the staged
+`max_mode`-ladder variants live alongside for comparison):
+
+| class | nfp | residual | seed | achieved | max_mode | wall | status |
+|-------|-----|----------|------|----------|----------|------|--------|
+| QA | 2 | QS (1, 0)  | 2.04e-01 | **7.2e-06** | 5 | **14.5 min** | precise; aspect 6.00, iota 0.42 (ladder: 3.7e-07 in 25.5 min) |
+| QH | 4 | QS (1, −1) | 6.91e-01 | **5.83e-05** | 5 | 25.5 min (ladder) | precise; aspect 8.00, iota −1.22 |
+| QP | 2 | QS (0, 1)  | 4.46e-01 | 3.3e-02 | 5 | ~3.4 h (ladder + refinement) | hardest QS class — see caption |
+| QI | 1 | omnigenity | 4.52e-01 | **1.81e-02** | 6 | **17.3 min** | 25× via the traceable Goodman constructed-QI residual |
+
+![QA/QH/QP optimization: seed vs optimized boundary, 3-D |B| geometry, and Boozer |B| on the LCFS](docs/_static/figures/readme_optimization.png)
+
+*Each quasisymmetry class starts from a near-circular torus (grey, dashed) and
+is shaped into a quasi-symmetric stellarator (blue) by the least-squares driver
+(top row); the middle row is the optimized last-closed flux surface in 3-D
+coloured by `|B|`, and the bottom row is `|B|` in Boozer coordinates on the LCFS
+(jet line contours), whose contour geometry reads off the symmetry family —
+horizontal for QA, diagonal for QH, vertical for QP. `QS` is the quasisymmetry
+residual measured on the plotted equilibrium: QA **1.1e-6**, QH **5.8e-5**
+(note QH's near-straight diagonal contours), QP **3.3e-2**. Quasi-poloidal QP
+is the hardest class: the ladder plateaus near 5e-2, and an extended
+warm-start refinement of the shipped deck reaches 3.3e-2. Reproduce with
+`python benchmarks/make_readme_figures.py --only optimization` from the decks
+in `benchmarks/opt_decks/`.*
+
+Quasi-isodynamic (QI) shaping is intrinsically harder than quasisymmetry, so it
+gets its own row across field periods:
+
+![QI equilibria at nfp 1-4: boundary, 3-D |B| geometry, and Boozer |B| on the LCFS](docs/_static/figures/readme_qi.png)
+
+*Quasi-isodynamic (QI) equilibria at nfp 1, 2, 3, 4 (bundled decks in
+`examples/data/`): boundary cross-sections (top), 3-D `|B|` geometry (middle),
+and `|B|` in Boozer coordinates on the LCFS (jet, bottom). The label is the QI
+(omnigenity) residual — **not** QS; QI is hard, so ~1e-3–1e-2 is expected here,
+not the ~1e-5 reachable for quasisymmetry. Reproduce with
+`python benchmarks/make_readme_figures.py --only qi`.*
+
+These campaigns need implicit gradients. Finite differences stall at the
+axisymmetric seed of the QH target (a saddle point) and land in a worse basin
+for QP. Three measured optimizations keep each campaign in the minutes range:
+
+- the residual Jacobian uses a block-tridiagonal factorization of the force
+  linearization (33× faster than per-dof GMRES);
+- each trial equilibrium starts from a first-order perturbation prediction
+  (3.7× fewer solver iterations);
+- a converged-state memo avoids re-solving the point the residual just
+  converged.
+
+The implicit path runs on CPU by default, where it is fastest at production
+sizes; high-resolution forward solves can use the GPU. The device policy
+chooses per stage.
+
+### Beyond quasisymmetry: any objective, same gradients
+
+Any physics objective can drive the same machinery. Starting from the
+precise-QA deck above (QS ~1e-6, aspect 6.00, mean iota 0.42), five short
+campaigns each optimize one new objective while keeping the QA residual in the
+objective at a stiff weight:
+
+- raise the coil-simplicity proxy min L∇B (`l_grad_b_state`);
+- deepen the vacuum magnetic well;
+- raise mean iota to 0.55 at fixed aspect;
+- lower the aspect ratio to 4.8 at fixed iota;
+- push the Mercier criterion `DMerc` toward stability at ⟨β⟩ ≈ 1.25%.
+
+The first four use the implicit adjoint (`jac="implicit"`). `DMerc` has no
+traceable lane yet (it is computed from host-side Mercier tables), so that
+campaign uses finite differences at `max_mode` 2. The self-consistent Redl
+bootstrap objective has its own section below.
+
+![Objectives showcase: five one-objective campaigns off the precise-QA seed](docs/_static/figures/readme_objectives.png)
+
+| campaign | objective | seed → final | QS held? |
+|----------|-----------|--------------|----------|
+| `lgradb` | raise min L∇B to 1.3× seed (implicit adjoint) | 0.520 → 0.522 m (stiff — see note) | 9.8e-07 → 1.3e-06 |
+| `well` | deepen the vacuum magnetic well (implicit adjoint) | **−0.037 → +0.0002** (hill → well) | 9.8e-07 → 1.5e-05 |
+| `iota_up` | mean iota 0.42 → 0.55 at aspect 6 (implicit adjoint) | **0.420 → 0.535** | 9.8e-07 → 1.8e-05 |
+| `aspect_down` | aspect 6.00 → 4.8 at iota 0.42 (implicit adjoint) | **6.00 → 4.84** | 9.8e-07 → 4.2e-06 |
+| `dmerc` | interior DMerc → positive at ⟨β⟩ ≈ 1.25% (finite differences) | −16.6 → −16.5 (stiff — see note) | 6.6e-05 → 6.6e-05 |
+
+The `well`, `iota_up`, and `aspect_down` campaigns each take 2–3 minutes on a
+workstation CPU. The other two barely move, for physical reasons: with QS,
+aspect, and iota all held, the precise-QA shape is already close to its best
+attainable L∇B, and improving interior Mercier stability at fixed pressure
+requires profile or current degrees of freedom that boundary shaping alone
+does not provide.
+
+*Reproduce with `python examples/optimization/objectives_showcase.py` (an
+`--only lgradb,dmerc` flag runs subsets), then
+`python benchmarks/make_readme_figures.py --only objectives`.*
+
+### Self-consistent bootstrap current
+
+VMEX implements the **Redl** analytic bootstrap-current formula
+([Redl et al. 2021](https://doi.org/10.1063/5.0012664)) as a differentiable
+objective, and a fixed-boundary self-consistency loop that regenerates the
+toroidal current from the plasma geometry and kinetic profiles. Below,
+reproducing [Landreman, Buller & Drevlak 2022](https://arxiv.org/abs/2205.02914):
+the published precise QA and QH optima are loaded, their current profile is
+**erased**, and `self_consistent_bootstrap` recovers it from the Redl formula
+plus the paper's density/temperature profiles.
+
+![Self-consistent bootstrap current vs the published equilibria and SFINCS](docs/_static/figures/readme_bootstrap.png)
+
+*Recovered current density &#10216;J·B&#10217; (VMEC, blue) matches the analytic
+Redl profile (green), the published self-consistent equilibrium (grey), and —
+for QA — the paper's SFINCS drift-kinetic benchmark (circles). Converged in 7
+(QA) / 4 (QH) Picard iterations to bootstrap mismatch `f_boot` = 2.0e-6 / 7.5e-6;
+the recovered plasma current lands within **1.9 % (QA)** and **0.3 % (QH)** of
+the published `CURTOR`. Reproduce with
+`python examples/optimization/{QA,QH}_bootstrap_selfconsistent.py` (needs the
+paper's Zenodo dataset).*
+
+## VMEX vs DESC
+
+[DESC](https://desc-docs.readthedocs.io/) is the other JAX-native,
+differentiable, GPU-capable stellarator-equilibrium code. The key difference:
+DESC minimises the MHD force in a global Zernike–Fourier basis — its own
+equilibrium — while VMEX reproduces VMEC exactly. The two are
+complementary:
+
+| Where **VMEX** wins | Where **DESC** wins |
+|---|---|
+| **Is VMEC**: iteration-for-iteration VMEC2000 parity, standard `wout_*.nc`, VMEC-format prints | **Low-resolution accuracy**: global Zernike basis converges in fewer radial points |
+| **Drop-in**: reads VMEC2000 `input.*` and VMEC++ JSON unchanged | **Objective library**: large, mature set of built-in optimization targets |
+| **Full namelist**: non-symmetric surfaces (`LASYM = T`), NESTOR *and* virtual-casing free boundary | **Optimizers**: more built-in stochastic / constrained optimizers |
+| **O(1)-memory adjoint**: peak memory flat in the number of design variables | Adjoint gradients (both codes are differentiable) |
+
+Reach for **VMEX** to drop a differentiable code that *is* VMEC into an
+existing VMEC workflow (simsopt, `booz_xform`, near-axis tooling). Reach for
+**DESC** for its spectral accuracy at low radial resolution or its mature
+objective library.
+
+## CLI reference
+
+```text
+vmex input.X             solve (INDATA or VMEC++ JSON), write wout_X.nc
+vmex --plot wout_*.nc    diagnostic plots from a WOUT file
+vmex --booz wout_*.nc    run booz_xform_jax, write boozmn_*.nc
+vmex --plot boozmn_*.nc  Boozer contour/spectrum plots
+vmex --test              run and plot the bundled quick-start case
+vmex --doctor            installation and JAX backend diagnostics
+
+options:
+  --outdir PATH          directory for wout/boozmn/figure output
+  --mode {cli,jit}       jitted blocks with live printing (cli, default)
+                         or a single lax.while_loop (jit)
+  --ftol F               override the final-stage FTOL_ARRAY tolerance
+  --max-iter N           override the final-stage NITER_ARRAY cap
+  --coils PATH           ESSOS-style coils file: drive an LFREEB = T deck
+                         by direct Biot-Savart instead of an mgrid file
+  --mbooz/--nbooz N      Boozer spectral resolution (default 32/32)
+  --booz-surfaces S      Boozer surfaces ('all' or a list of s values)
+  --quiet                silence the VMEC-style stdout
+```
+
+`vmec` follows the selected JAX backend: with CPU-only JAX it runs on the
+CPU; with CUDA-enabled JAX it uses the GPU for the solver stages where that
+is faster (`JAX_PLATFORMS=cpu|cuda` pins it explicitly).
+
+## Documentation
+
+Full documentation — installation, quickstart, theory and numerics with
+equation-to-source cross-references, API reference, and
+performance/validation notes — at
+[vmex.readthedocs.io](https://vmex.readthedocs.io/en/latest/).
+
+## Mirror equilibria
+
+Alongside the toroidal VMEC core, `vmex.mirror` solves scalar-pressure
+equilibria for **open magnetic mirrors** and **closed stellarator–mirror
+hybrids** — the same differentiable, spline-native machinery applied to a
+straight (open) axis. Open mirrors use nonperiodic axial coordinates
+`(s, θ, ξ)` with fixed-flux end cuts, not thin-torus approximations. Coils and
+Biot–Savart fields stay in [ESSOS](https://github.com/uwplasma/ESSOS); VMEX
+consumes a supplied `xyz → B` field. The divergence-free field and
+scalar-pressure energy are
+
+```text
+√g B^θ = I'(s) − ∂_ξ λ,    √g B^ξ = Ψ'(s) + ∂_θ λ,    B^s = 0
+W = ∫ [ B²/(2μ₀) + p/(γ − 1) ] dV
+```
+
+### Fixed-boundary open mirrors
+
+A fixed-boundary solve is one call:
+
+```python
+from vmex.mirror import MirrorConfig, MirrorResolution, solve_fixed_boundary_from_radius
+
+config = MirrorConfig(resolution=MirrorResolution(ns=7, mpol=4, nxi=17))
+result = solve_fixed_boundary_from_radius(0.3, config)   # radius: scalar, (nxi,), or (ntheta, nxi)
+```
+
+![Fixed-boundary rotating-ellipse mirror: solved geometry, field lines, cross-sections, and convergence](docs/_static/figures/mirror_fixed_boundary_3d.png)
+
+The rotating-ellipse mirror converges at `ftol = 1e-12` to a normalized
+divergence of `6.6e-15`, in **6 s cold / 0.2 s warm** (peak ≈1.2 GB, CPU). Its
+implicit boundary gradient agrees with two fully reconverged finite-difference
+solves to `5.9e-10` relative — the derivative an external optimizer needs.
+
+### Free-boundary β scan
+
+`solve_beta_scan` jointly updates the spline last-closed surface, the plasma
+state, and the unbounded exterior vacuum, driven by an ESSOS two-coil field.
+The supported sequence runs from 0 % to 10 % β, and the implicit free-boundary
+derivative matches a reconverged finite difference to `1.1e-10` relative
+(adjoint residual `1.4e-9`). Pushing to a requested 50 % β grows the central
+radius 7.7 % and drops the on-axis field 23.7 % from vacuum, exercising the
+finite-β coupling end to end.
+
+![Free-boundary beta scan with ESSOS coils: field lines, LCFS, |B|, pressure, and residual histories](docs/_static/figures/mirror_free_boundary_beta50_summary.png)
+
+### Stellarator–mirror hybrid (research)
+
+A closed periodic hybrid — two exactly straight mirror legs joined by two
+curved stellarator returns on a rotation-minimizing B-spline axis — has a
+complete fixed-boundary solve and example. A finite axial current gives
+`ι = 0.085`; the case reaches a `2.4e-14` variational residual and `3.1e-14`
+normalized divergence. Its independent strong-force gate does not yet converge
+under same-geometry refinement, so it ships as a validated **research
+candidate**, not a supported benchmark — the same implicit API already
+differentiates its periodic boundary and axis controls.
+
+![Periodic B-spline stellarator–mirror hybrid: straight legs, rotating returns, field lines, and |B|](docs/_static/figures/stellarator_mirror_hybrid.png)
+
+### Run the mirror examples
+
+```bash
+python examples/mirror_fixed_boundary_nonaxisymmetric.py
+python examples/mirror_free_boundary_beta_scan.py
+python examples/stellarator_mirror_hybrid.py
+```
+
+Open-mirror `mout_*.nc` files plot with `vmex --plot mout_*.nc`. The
+[mirror-geometry documentation](https://vmex.readthedocs.io/en/latest/mirror_geometry.html)
+derives the coordinate and field models, defines the boundary conditions and
+residuals, and records the validation and derivative limits.
+
+## License
+
+MIT. If you use VMEX in published work, please cite this repository and
+the original VMEC papers (Hirshman & Whitson, *Phys. Fluids* 1983;
+Hirshman, van Rij & Merkel, *Comput. Phys. Commun.* 1986).