GKX

GKX is a JAX gyrokinetic solver for tokamak and stellarator flux tubes. It reads a VMEC or VMEX equilibrium, or builds an analytic Miller or s-alpha geometry, and computes linear stability and nonlinear turbulence in a Hermite-Laguerre velocity basis. The whole path, equilibrium included, is differentiable on CPUs and GPUs.

• Run what you have: a VMEC/VMEX wout, a Miller or s-alpha tokamak, or one TOML deck. The executable sizes the grid and says why.
• Trust the eigenvalue: every returned eigenpair is checked against the full operator; an under-resolved run warns instead of reporting a number.
• Differentiate it: implicit eigenvalue derivatives, and a checkpointed adjoint of the nonlinear heat flux over a saturated window.
• Optimize with it: the heat flux, growth rate or quasilinear flux is an objective in VMEX stellarator shape optimization.

Saturated ITG turbulence on a Cyclone flux tube (full-rate movie).

Install

pip install gkx
gkx

Python 3.11+. The wheel installs CPU JAX; for GPUs add an accelerator JAX wheel (JAX install guide). gkx with no arguments runs a linear Cyclone demo in under a minute on a laptop CPU and prints gamma and omega; its growth rate is within 1% of the certified eigenvalue for that case, and a CI gate holds it there.

Run a case

gkx wout_circular_tokamak.nc --estimate   # size the grid, explain it, exit
gkx wout_circular_tokamak.nc              # nonlinear ITG run to saturation
gkx plot wout_circular_tokamak/gkx.out.nc # replot a saved bundle
gkx examples/01_linear_tokamak/case.toml  # any TOML deck

--estimate derives each grid entry from the geometry:

geometry: shat=+1.7190 q=2.066 nfp=1 |B| wells=1 anisotropy=0.214 -> ky_max*rho >= 2.2
    ny = 96       tokamak class asks ky_max*rho >= 2.2; reach ((Ny-1)//3)*dky = 2.21 at dky = 0.071
    nl = 4        Laguerre FLR floor with hypercollisions; the scan converged at Nl=4
    nm = 8        hypercollisions: t_quiet ~ 5.5*sqrt(Nm) recurrence sets the published floor (4,8)
 t_max = 400      8 x t_sat ~ 50 hard cap; run_to = "saturation" stops earlier

The estimate is a starting point, not a convergence proof. A run stops when heat flux, field energy and free energy are all stationary; a missed saturation is reported as such. Each run writes figures, a restartable NetCDF bundle and the resolved deck that reproduces it. Inputs · outputs.

From Python

import jax.numpy as jnp

from gkx import CycloneBaseCase, LinearParams, integrate_linear_from_config
from gkx.core_grid import build_spectral_grid
from gkx.geometry import SAlphaGeometry

cfg = CycloneBaseCase()
grid = build_spectral_grid(cfg.grid)
geometry = SAlphaGeometry.from_config(cfg.geometry)
state = jnp.zeros((2, 2, grid.ky.size, grid.kx.size, grid.z.size), dtype=jnp.complex64)
state = state.at[0, 0, 0, 0, :].set(1.0e-3)
trajectory, potential = integrate_linear_from_config(
    state, grid, geometry, LinearParams(), cfg.time
)

For repeated nonlinear calls, gkx.prepare(case, steps=N) compiles once and warmup() moves the compile out of the first timed solve.

Examples

examples/ is a numbered gallery. Each directory has a run.py with editable parameters at the top, a case.toml that runs in seconds to a minute on a laptop, and, where it applies, a literature-resolution case_full.toml.

python examples/03_nonlinear_tokamak/run.py

Linear stability

Growth rate and frequency scans for Cyclone ITG (a), W7-X ITG (b) and KBM (d), and the W7-X eigenfunction at ky rho_i = 0.3 (c), whose overlap with GX's is 0.9999999994. Eigenvalues come from a matrix-free restarted eigensolver that applies the full gyrokinetic right-hand side, or from a sparse direct shift-invert solve; both certify the returned pair against the unprojected operator (solvers).

python examples/01_linear_tokamak/run.py      # Cyclone ky scan and eigenfunction

Benchmarks and cross-code comparison

(a) Certified GKX eigenpairs (Nl = 16, Nm = 48) on the Cyclone s-alpha and Miller decks against GX runs of the same decks, with GX's end-damping launch limit repaired. The growth rate agrees to within 0.3% at seven of eight k_y, and to 1.45% at Miller ky rho_i = 0.15; the frequency agrees to within 0.13% everywhere. (b) The growth rate from three codes: GKX is within 0.3% of converged GS2 at Miller ky rho_i 0.30 and 0.40, and 1.0% at s-alpha 0.30. Miller 0.15 is the one point where the three codes spread by more than 1%. Record: cross-code benchmark.

• GS2 and GX with kinetic electrons (Cyclone Miller) agree to 1.2% and 0.2% at ky rho_i 0.30 and 0.50; the GKX eigenpair for this case is next.
• stella gives a growth rate about 1.4 times the other codes on the same Cyclone input. In the build tested, stella's growth rate moves by 0.3% when its mirror term is switched off, while in GKX the mirror force is 14% of gamma; that term is the lead being followed.
• s-alpha ky rho_i = 0.55 is weakly growing and not converged in either GS2's energy grid or GKX's Laguerre ladder; the finest rungs of the two agree to 0.3%.

The tracked parity scans below, as 100 * max|GKX - ref| / max|ref|, are recomputed by CI from the files in the evidence ledger. The Cyclone rows compare a time-trace fit, not the certified eigenvalue above, and KBM is the known outlier.

These are agreements on shared test cases, not a ranking of codes, which differ in models and options. Detail: benchmarks, verification matrix.

Nonlinear turbulence

gkx wout_circular_tokamak.nc

(a) A circular tokamak run straight from a VMEC wout, replicated with two seeds and a second time step: window means 18.7, 19.3 and 18.9 over t = 350-700, a 3.5% spread against a 15% gate. (b) Mean relative difference from GX runs of the same decks in heat flux, free energy Wg and field energy Wphi, for Cyclone, Cyclone Miller, W7-X, HSX and KBM, all below 10%.

Collisions and proof tests

[time] collision_operator selects Lenard-Bernstein/Dougherty (default), drift-kinetic Sugama and improved Sugama, drift-kinetic linearized Coulomb (Frei, Ernst & Ricci 2022) or gyrokinetic Coulomb at finite k_perp (Frei et al. 2021), for like-species collisions.

python examples/07_collisions/run.py

The figure shows identities the discretization must satisfy exactly and analytic limits it must reach, against the tolerances CI asserts in test_collision_physics.py and test_core_numerics.py: collision conservation, self-adjointness and the H-theorem to 1e-16, the Coulomb matrix against the published coefficients, Spitzer-Härm gamma_E(Z) to 0.6%, and the Landau roots of 1 + T_i/T_e + zeta Z(zeta) = 0 to 0.25%.

Differentiate the solver

Eigenvalue derivatives use dλ/dp = wᴴ(dA/dp)v / (wᴴv), with a bordered solve for eigenvector observables. The nonlinear heat flux over a post-saturation window is differentiated through a block-checkpointed discrete adjoint.

def loss(shape):
    return gkx.nonlinear_heat_flux_window(
        saturated, grid, geometry(shape), params, dt, steps, terms=terms
    )

heat_flux, gradient = jax.value_and_grad(loss)(shape0)

On a 16x16x16 Cyclone case over 1024 steps, checkpointing cuts temporary memory from 7.8 GB to 187 MB on CPU (148 MB on an RTX A4000) for about twice the runtime. The gradient matches finite differences to 1e-11 through 512 steps; beyond about 2000 steps chaotic separation limits any window derivative. Nonlinear autodiff · eigensolver.

Stellarator optimization with VMEX

VMEX composes its implicit equilibrium derivative with the GKX derivative, so a turbulence objective sits next to quasisymmetry, aspect ratio and iota in one least-squares problem. VMEX ships three scripts in examples/optimization/ (pip install "vmex[turbulence]"):

Each is a single run; the objective trades against quasisymmetry and iota, as each script's docstring records. The quasilinear flux is a ranking and screening measure, not a runtime/TOML absolute-flux predictor (quasilinear).

GKX's own QA_optimization.py in examples/10_vmex_optimization adds the nonlinear heat flux as a fourth objective to VMEX's vacuum QA ladder.

The boundary change is small (aspect ratio +0.0115%, mean iota -0.044%). The preliminary 12.26% reduction across 24 nominal pairs (conditional 95% CI 10.64-13.88%) is not statistically resolved: 4 of 48 nominal traces fail the per-trace drift test, and these traces predate the periodic hypercollision correction. A transport claim needs matched, replicated, long post-saturation windows. Stellarator optimization.

Performance

Cold wall time and peak memory across the tracked cases, including JAX startup and compilation; the executable caches compilations, so reruns are warm.

• Nonlinear step: about 196 ns per Nx*Ny*Nz*Nl*Nm element per step on CPU, flat from 64x64x24 to 96x96x48. On one RTX A4000 an RK3 step at 64x64x24, Nl = 4, Nm = 8 takes 15.6 ms; it is memory-traffic bound.
• Eigenvalues: the sparse direct shift-invert route certifies the Cyclone eigenpair at n = 3,072 in about 5 s against 32-36 s for the matrix-free route, and its growth-rate gradient is 6-9 times faster.
• Parallel work: independent k_y scans, quasilinear and UQ ensembles run across devices with results identical to the serial run.

Performance · parallelization.

Reproducing the figures

JAX_ENABLE_X64=true PYTHONPATH=src:. python scripts/figures.py   # or: linear, nonlinear, proof_tests, crosscode

scripts/figures.toml lists each figure's inputs; each PNG in docs/_static/readme/ has a JSON companion with every plotted number. The proof-test figure takes about 7 CPU-minutes; the others seconds.

Documentation and development

Documentation: gkx.readthedocs.io, from the quickstart to design decisions.

git clone https://github.com/uwplasma/GKX
cd GKX
pip install -e ".[dev]"
pytest -n 4

CI requires 95% line coverage plus the physics, convergence, comparison, differentiability and performance gates. See CONTRIBUTING.md; cite GKX with CITATION.cff.

License

MIT; see LICENSE.