Skip to content

VMEX 0.8.0

Choose a tag to compare

View all tags
@rogeriojorge rogeriojorge released this 31 Aug 00:58
· 560 commits to main since this release
v0.8.0
d2d5a01
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.

VMEX 0.8.0

The headline: certified force-balance polishing, execution directives that VMEC2000 ignores, and large measured performance fixes to the differentiable core.

Force-balance polishing

  • vmex input.case --polish, !@VMEX POLISH = AUTO input directives, or vj.solve_file("input.case", polish="auto") polish a converged fixed-boundary state: the solution is lifted to axis-regular cubic B-splines and both physical force channels are driven to zero on an overdetermined collocation grid with matrix-free SOLVAX Gauss-Newton steps.
  • Acceptance is certified independently: volume L2 force error below 1e-2, radial-refinement stability within 1e-3, and a positive signed Jacobian, with tangent/adjoint derivatives through the polished root.
  • result.native_equilibrium exposes the continuous polished state for fields, Boozer transforms, virtual casing, and derivatives; result.polished_state is the sampled state behind the VMEC-compatible WOUT.
  • Polishing remains opt-in; the ordinary solve is unchanged and POLISH = FALLBACK|WARN control failure behavior.

Execution metadata without touching the physics

  • RunOptions/read_input_request separate how to run from what to solve: !@VMEX KEY = VALUE directives in the input file, a _vmex section in JSON inputs, Python keywords, and CLI flags, with precedence CLI > keyword > directive > default. VMEC2000 reads the same files untouched.

Performance, measured

  • The implicit lane no longer recompiles per optimization trial: the residual lives in one module-scope compiled lane and equal-content configs are canonicalized, so identity-keyed caches hit. Measured on the committed Apple M4 baselines: a steady optimization-campaign step drops 16.6 s to 10.8 s, and a warm objective value+gradient repeat drops 28.8 s to 3.0 s - bit-identical values and gradients in both cases.

  • benchmarks/profile_workflows.py is a complete performance and compilation observability harness: seventeen flagship workflows (solves, campaigns, free-boundary adjoints, mirrors, Boozer, epsilon-effective, Gamma_c) in five never-mixed timing regimes, with compile counts from JAX's own records, XProf trace capture, and 53 committed baseline records under benchmarks/baselines/m4/.

  • The polishing path itself no longer recompiles on every solve: per-solve values moved out of compilation-cache keys (timestamps and identity-compared objects were living inside them), so repeated polished solves reuse compiled programs in memory and on disk. Measured on the bundled shaped tokamak: warm polished solve 52.9 s to 22.8 s, first call 115 s to 80 s, certificates unchanged, validated on CPU and on a two-GPU CUDA rig.

CI and process

  • CI attempt wall clock roughly halved by resharding the heaviest test lanes, and tools/preflight.py runs the static gates, guard tests, and diff-affected test modules locally so pull requests arrive pre-verified.
Assets 2
Loading