VMEX 0.9.0 — the optimization release

The optimization release. The two workflows users reported as slow and inconclusive — quasi-isodynamic boundary design and single-stage plasma-and-coil design — now converge and meet the targets they print.

pip install -U vmex     # Python 3.11+, 3.12+ recommended

What users asked for

The refinement every optimizer trial pays is a Newton finish. It runs through one raw block factorization instead of a stalling Krylov solve. On the QI example, measured on the 36-core workstation at load ≤ 24: refinement 333.0 s and 63,316 Krylov iterations → 66.7 s, 153 GMRES iterations and 20 factorizations, with the same 17,251 descent iterations, 0 failed trials and the final cost within 7.2e-7. The example runs 565.1 → 359.1 s and its least-squares phase 366.9 → 127.5 s; the single-stage example 30.8 → 22.3 s per trial. Decks where the Newton phase stalls replay the previous refinement and keep its anchor.

The implicit adjoint is solved exactly, through that same factorization, instead of 1,100–17,000 Krylov iterations: QI eager jax.grad 203.9 → 55.5 s, QI from_loss first value and gradient 190.2 → 89.0 s. The reverse Jacobian lane factors once per point rather than once per batch of rows: seed QA 238.9 s, where the per-chunk path had not finished after 4,111 s.

The compiled lanes stopped recompiling. A concrete problem.jax_value_and_grad returns the host lane's pair and shares its solve memo, warm-start stash and counters: JAX value and gradient 39 → 1.5 s on QA and 55 → 1.6 s on QI, benchmark wall 95.3 → 51.0 s and 117.5 → 61.8 s.

The QI example converges. It seeds from input.QI_nfp2_initial instead of a circular torus, runs one max_mode = 2 stage and penalizes the mirror ratio from 1 % below its limit: 612 s, 0 failed trials (5 before), constructed QI 3.0e-3 (0.789 before), with iota, mirror and elongation limits met and the final ns = 101 solve converged. The unchanged example ran 1,456 s and then raised in that final solve.

The single-stage examples meet their stated targets or fail loudly. The fixed-boundary example meets every target on a full run: min |ι| 0.4277 ≥ 0.42, aspect 3.979 ≤ 4, B·n RMS 0.80 % ≤ 1 %, coil clearances and curvature within limits, with an independent ns = 101 check converged. The free-boundary example floors the quantity it checks rather than its smooth surrogate and reaches min |ι| 0.4271 in 614 s where it previously exited 1 at 0.4136 after 1,339 s — at the cost of quasi-axisymmetry (QA total 7.3e-4 → 4.5e-2), which that example does not constrain.

New

• A DESC bridge. vmex reads DESC .h5 equilibria and DESC input decks through the CLI and VmecInput, with round trips tested against DESC 0.17.
• The exterior field reports the accuracy it achieved. B_plasma_xyz returns what the adaptive schedule reached, a schedule that fails its own self-test warns instead of silently returning its last level, and known-answer oracles cover the vacuum identity outside, the interior identity inside and Malhotra on-surface parity. The d ≥ 2h distance rule and the default grid are documented.
• Per-evaluation counters. Every OptimizationRecord and benchmark row carries descent iterations, refinement calls, steps and matvecs, Jacobian calls and columns, certifier iterations, adjoint calls and matvecs, and host wall time in each.
• The summary plot's force panel no longer reads 1 on converged vacuum equilibria; --plot resolves the effective ripple instead of aliasing it (the previous diagnostic NEO settings were 7–20× from converged NEO, with dips to 1e-3 and 4e-6 of the reference).
• The CLI keeps the equilibrium when the final grid exhausts NITER, as vmec.f/fileout.f do, instead of discarding the run. The exit code stays the distinct ier_flag = 2 and wout.ier_flag records 2.

Breaking

The floor is Python 3.11, SciPy 1.16 and JAX 0.9.2. SciPy added least_squares(callback=) in 1.16 and 1.16 needs Python 3.11, so a 3.10 install resolved SciPy 1.15.3 and every optimization example died at its least_squares( call after minutes of compiling. Python 3.10 also resolved JAX 0.6.2, which differentiates the QA residual wrongly — RBC(1,0) and ZBS(1,0) are off by factors of 6 and 9 against central finite differences, and the example ends at QS total 5.29e-3 instead of 3.72e-4. import vmex now checks Python, SciPy, JAX and jaxlib before importing JAX and names the fix.

Note jax/jaxlib 0.11 require Python 3.12, so a 3.11 environment resolves JAX 0.10.2 at the newest.

The persistent compilation cache is off under jaxlib < 0.10 on every platform, not only macOS: jaxlib 0.9.2 segfaults loading a large cached executable on Linux too. JAX 0.10+ restores it automatically.

Where VMEX stands against other codes

itpplasma/benchmark_vmec is an independent suite that runs 344 cases through twelve equilibrium codes, VMEX among them, and published a 4,128-row result on 2026-09-14. Joining its comparison_table.csv with plots/quality.csv on residual_over_tolerance:

code	reported success	strictly converged
VMEC2000	198	179
PARVMEC	193	177
VMEX	168	168
VMEC++	184	157
jVMEC	140	140
educational_VMEC	187	94
GVEC	124	40

VMEX and jVMEC are the only codes whose reported successes are all converged: VMEX's worst residual_over_tolerance in that corpus is 0.9999, against 1e+30 for VMEC++ and 4e+22 for educational_VMEC. Read by the status column alone VMEX ranks last of the VMEC family, which is why the join matters.

The same corpus names what this release does not fix: VMEX's median wall per case is 23.1 s against VMEC++'s 3.35 s and VMEC2000's 2.04 s, one case per process — cold compile, not the iteration loop. Block-Jacobian assembly is the other. Both follow in the next release. Twenty-three cases in that corpus converge for a VMEC-family peer and not for VMEX; they are now a tracked work package in plan.md.

Three defects this release found by running things

Re-running the checks rather than reading them turned up three real faults, all
fixed here:

• The CLI threw away converged runs. On VMEX's own
  input.LandremanPaul2021_QA_lowres, which asks for FTOL 1e-13, the solver
  reached FSQR 2.63e-13, printed MORE ITERATIONS REQUIRED, exited 2 and wrote
  no WOUT — 29 s of work discarded. vmec.f and fileout.f reach wrout
  on more_iter_flag whatever LFULL3D1OUT says. The file is now written, the
  exit code stays 2, and wout.ier_flag records 2 (#356).
• A GPU regression from inside this release. An AUTO multigrid ladder
  crossing the CPU/GPU work threshold raised Received incompatible devices:
  each rung placed its state on its own device but carried the previous rung's
  residual scalars through unmoved. Bisected to #321. It went unseen because
  gpu.yml has had no registered runner since 2026-07-31; the guard added with
  the fix reproduces it on two CPU devices (#360).
• The finite-beta single-stage example crashed after 19 minutes. Its final
  ns = 31/51/101 check spent all 20,000 iterations per rung and then raised.
  The cause was trial acceptance, not the ladder: max_fsq_ratio's 1e6 default
  differentiates a trial whose residual is 1e-6 against the deck's 1e-12 target,
  and the implicit adjoint assumes a root, so the line search walked the
  boundary somewhere unsolvable. Requiring 1e2 keeps every trial at a root — the
  same ladder now converges to 9.9e-15 at ns = 101 in 5:03 instead of
  failing at 19:37 (#361).

Known limitations

• booz_xform_jax 0.3.0 and essos 0.17 are not tagged, so the magnetic-only Boozer projection and the twelve ESSOS-based examples need a git install; the examples print the exact command.
• The GPU workflow still has no registered self-hosted runner. This release was verified by running its selection by hand on a 2x RTX A4000 workstation (66 passed, 6 skipped, 35:16); until a runner is registered, tests/test_gpu_ci.py is unexecuted in CI.
• Every optimization example inherits the max_fsq_ratio = 1e6 default that #361 overrides for the finite-beta single stage. Whether that default is right for the library is an open question, not settled here.
• Force-balance polishing remains a research feature, disabled by default; recorded certified cases are axisymmetric.

The full changelog is in CHANGELOG.md, and every number above traces to a committed record listed in benchmarks/INDEX.md.