Repository navigation
VMEX 0.4.0
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'sBounce1D. - Separate quasi-isodynamic residuals (
vmex.core.qi):ConstructedQIResidual(traceable Goodman squash-and-shuffle target matching) andJInvariantQIResidual(J invariance across field lines at common physical pitch). The existingQIResidualinterface 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_boozerevaluates 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 deckinput.QI_nfp2_initial, derived from the workflow content of PR #82; an independent NumPyqi_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 LASYMbmns/gmnsreconstruction. - Certified implicit block linear responses (
implicit_state_tangent_multi_rhs,implicit_state_pullback_multi_rhswithsolver="block"): one factorization for batched tangent/transpose responses, residual-checked withLinearResponseReport, 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-compilerestores 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 outerjax.default_devicecontext.
Fixed
- Comparator (
tools/first_divergence.py) argument mistakes now produce path-freeC0 USAGE_ERRORhints, a directory--xvmec2000auto-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.