Architecture¶
The production implementation lives in vmex.core: one concern per
module, each with a header docstring naming its VMEC2000 counterpart file(s)
and the equations it implements. Everything is pure-JAX and shared between
the CLI, the differentiable API, plotting, and the wout writer — there is a
single set of physics kernels.
Module map¶
Module |
Role |
VMEC2000 counterpart |
|---|---|---|
|
|
|
pressure / iota / current parameterizations (power series, splines, two-power, pedestal, …) |
|
|
|
|
|
spectral <-> real-space transforms as batched matmuls |
|
|
real-space \(R, Z, \lambda\), half-mesh Jacobian |
|
|
\(B^u, B^v, |B|\), covariant B, pressure, energies |
|
|
MHD force kernels + spectral-condensation constraint force |
|
|
|
|
|
1D radial preconditioner, vectorized tridiagonal (Thomas) solve |
|
|
2D block preconditioner: matrix-free Newton step ( |
|
|
damped 2nd-order Richardson step, |
|
|
radial grids, 1D profile arrays, boundary processing, initial guess |
|
|
single-grid solve loop: |
|
|
|
|
|
NESTOR: Green’s function, |
|
|
free-boundary iteration, |
|
|
differentiable free-boundary residual via virtual casing |
(no VMEC2000 equivalent) |
|
mgrid netCDF read/write, differentiable interpolated field; external
coils live in ESSOS ( |
MAKEGRID file format, |
|
implicit differentiation of the equilibrium ( |
(no VMEC2000 equivalent) |
|
objectives (quasisymmetry ratio residual, QI residual, aspect, iota, mirror, well, DMerc, …) + least-squares driver |
(no VMEC2000 equivalent) |
|
traceable Boozer |
(no VMEC2000 equivalent) |
|
differentiable Redl bootstrap |
(no VMEC2000 equivalent; BOOTSJ-adjacent scope) |
|
infinite-n ideal-ballooning eigenvalue objective (COBRA-style) |
(no VMEC2000 equivalent; COBRA companion code) |
|
GK flux-tube geometry adapter + SPECTRAX-GK turbulence proxies |
(no VMEC2000 equivalent) |
|
Nyquist-resolution Fourier tables, |
|
|
derived wout quantities (beta, currents, |
|
|
complete |
|
|
VMEC2000-format iteration lines, stage banners, termination summary |
|
|
|
(no VMEC2000 equivalent) |
|
Boozer transform driver (thin wrapper over |
booz_xform |
|
measured CPU/GPU placement policy for the solve lanes |
(no VMEC2000 equivalent) |
|
typed zero-crash exceptions + the VMEC2000 |
|
|
the |
|
State and purity¶
The solver state is a frozen pytree
(SpectralState): spectral coefficients of
\(R, Z, \lambda\) (plus the asymmetric partners when lasym), the
Richardson velocity, time step, damping history, iteration counters, and the
restart flag. All solver functions are pure state -> state maps, which is
what makes the same kernels usable from jit, grad, and vmap.
Static configuration (resolutions, flags) is hashable and kept out of traced
signatures; mode and radial arrays are padded to the maximum multigrid
resolution so all NS_ARRAY stages share one compiled executable.
Two lanes, one physics¶
solver.solve(...)— alax.while_loopover the jitted iteration, fully traceable: the forward solver of the differentiable API.the CLI lane — a Python
whilearound the same jitted N-iteration block kernel, with host-side residual checks between blocks: exact-ftolearly exit, live VMEC2000-format printing, buffer donation, no AD bookkeeping.
Both lanes call identical physics kernels; a regression test asserts per-block state agreement to machine precision.
Device policy (CPU/GPU)¶
vmex.core.device implements a measured CPU/GPU placement policy
for the solve lanes (calibrated against benchmarks/gpu_baseline.json).
The cost driver of one iteration is the totzsps/tomnsps batched-matmul
work, proxied by
(iteration_work()). Per-iteration throughput
favours the GPU at every tested size, but the GPU pays fixed per-solve
overheads (dispatch/transfer floor plus compile or cache load), so small
decks finish faster on the CPU. The measured crossover is
GPU_MIN_ITERATION_WORK (100_000):
recommended_device() returns "cpu" below it
and "gpu" at or above it, per multigrid stage.
resolve_device() turns this into a concrete
placement with strict precedence rules: an explicit device= argument to
solve/solve_multigrid always wins; a user pin via JAX_PLATFORMS
or JAX_PLATFORM_NAME makes the automatic policy stand down entirely; and
the recommendation is applied only when the recommended platform is actually
available. device_context() wraps a stage in the
corresponding jax.default_device.
The optimization path is different:
resolve_implicit_device() always pins the
implicit-gradient work to the CPU by default. The jac="implicit"
Jacobian builds a per-dof vmapped forward-implicit-differentiation graph —
dozens of preconditioned GMRES solves with inner control flow — whose XLA
compile time grows with the dof count and whose execution is
kernel-launch-bound; measured on GPU it is slower than the CPU at every
optimization size tested, while the forward equilibrium solve inside it is a
host callback that never touches the accelerator anyway. Explicit
device= arguments and user platform pins are still honored.
Naming conventions¶
Community-expected VMEC names are kept (ns, mpol, ntor, nfp, lasym, iotaf,
presf, rmnc, zmns, lmns, bmnc, ...). Internal Fortran temporaries get
descriptive names (sqrt_g rather than gsqrt); every module docstring
cross-references the VMEC2000 source it ports.