Skip to content

VMEX 0.4.0

Choose a tag to compare

@rogeriojorge rogeriojorge released this 01 Aug 19:30
· 899 commits to main since this release
a488bc3

VMEX 0.4.0

Differentiable, JAX-native VMEC. This release adds action-based confinement objectives (bounce action, QI residuals, signed maximum-J), validated trapped-fraction output, dimensional similarity scaling, certified multi-right-hand-side implicit responses, explicit multi-device honoring, and a reproducible resource-profiling contract — plus the free-boundary visibility and diagnostics work that root-caused silent long runs.

Added

  • Differentiable multi-well bounce-action kernel (vmex.core.bounce): the second adiabatic invariant J = 2∮√(1−λB) dl at caller-supplied physical pitch over bounded multi-period traces, with fixed-shape well masks and NaN (not zero) invalidation. Validated against adaptive quadrature and DESC's Bounce1D.
  • Separate quasi-isodynamic residuals (vmex.core.qi): ConstructedQIResidual (traceable Goodman squash-and-shuffle target matching) and JInvariantQIResidual (J invariance across field lines at common physical pitch). The existing QIResidual interface and defaults are unchanged.
  • Signed matched-well maximum-J objective (vmex.core.maxj.MaximumJResidual): penalizes the physical sign convention dJ/dψ < 0 with reciprocal nearest-neighbour well matching, carrying signed toroidal flux (APHI/signgs-aware) through solver setup and traceable Boozer surfaces. qi_and_maximum_j_from_boozer evaluates QI + maximum-J from a single Boozer pass.
  • QI + maximum-J continuation example (examples/optimization/QI_maxJ_continuation.py) with polar J(α, s) contours and the nfp=2 QI seed deck input.QI_nfp2_initial, derived from the workflow content of PR #82; an independent NumPy qi_functions_mod-lineage oracle now validates the constructed-target construction.
  • Trapped-fraction WOUT profiles: trapped_fraction_from_state (symmetric + LASYM) and validated trapped-fraction fields in every VMEX-written WOUT, with LASYM bmns/gmns reconstruction.
  • Certified implicit block linear responses (implicit_state_tangent_multi_rhs, implicit_state_pullback_multi_rhs with solver="block"): one factorization for batched tangent/transpose responses, residual-checked with LinearResponseReport, NaN-poisoned on non-convergence — consumable by external codes.
  • Dimensional similarity scaling: vmex --scale PATH [B_scale R_scale] (default targets |B0| = 5.7 T, Aminor = 1.7 m) scales inputs and WOUTs consistently, including free-boundary mgrid and NESTOR fields; scale-then-reconverge agrees with scale-the-WOUT.
  • Reproducible resource profiler (benchmarks/profile_resources.py): fresh-process measurement separating compilation from warm execution, process and device memory, with VMEC2000/VMEC++ provenance recording.
  • Near-degenerate vacuum documentation deck with the measured LFORBAL remedy and a live VMEC2000 parity test; exact-JVP GMRES robustness gate pinning convergence on a stiff high-aspect equilibrium where the VMEC2000 finite-difference block-GMRES trajectory stalls.

Changed

  • CLI compiles multigrid rungs sequentially by default (lower one-shot peak memory); --prefetch-compile restores overlapped compilation.
  • Every CLI line flushes through to redirected files immediately, and each free-boundary lane announces itself before compiling — a file-redirected cluster log always shows progress. (Root cause of previously reported hours-silent runs: an unflushed sink plus long high-mode-count lane compiles.)
  • Smaller traced free-boundary lanes: the steady vacuum lane synthesizes geometry once per pass and shares the Jacobian sign with the vacuum-source gate (funct3d.f ordering); iteration rows are byte-identical, steady-lane program size −16% at 238 modes.
  • Explicit nondefault devices are honored end-to-end: device=jax.devices("gpu")[1] owns fixed- and free-boundary runtime/state arrays without an outer jax.default_device context.

Fixed

  • Comparator (tools/first_divergence.py) argument mistakes now produce path-free C0 USAGE_ERROR hints, a directory --xvmec2000 auto-resolves to the executable inside it, and new C1 sub-stages compare parsed scalars, profiles, and the full boundary table against VMEC2000's threed1 echo — localizing parser divergences to the exact field and mode index, privacy-safely.
  • APHI toroidal-flux semantics are pinned by tests against the VMEC2000 formulas (torflux(1) normalization; unnormalized clamped profile argument), including a live row-1 parity lane.
  • The vacuum.f skip-cadence contract is pinned: a skip step reuses the cached NESTOR matrix/factor and reproduces the full-step field at machine precision.