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 and first run

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 case in under a minute on a laptop CPU and prints gamma and omega; a CI gate holds its growth rate within 1% of the certified eigenvalue. The tutorial walks through twelve runs, from this one to a stellarator optimization; each is a directory in examples/.

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 and prints the reason. A run stops when heat flux, field energy and free energy are stationary, or reports that it did not. It writes figures, a restartable NetCDF bundle and the resolved deck that reproduces it (inputs, outputs).

From Python:

import gkx

case = gkx.load("examples/01_linear_tokamak/case.toml")
result = gkx.solve(
    case, ky_target=0.3, Nl=case.run.Nl, Nm=case.run.Nm, show_progress=True
)
print(result.gamma, result.omega)

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.

• The default eigensolver (auto) is a banded direct shift-invert where it fits, else Arnoldi on the time-advance map. It returns an eigenpair only after checking it against the full operator. Matrix-free shift-invert is opt-in for a chosen branch (policy and measurements).
• Under-resolved or unsettled fits warn instead of reporting a number.

Models and defaults

Resolution, step and solver guidance says what to set for kinetic electrons and electromagnetic runs, and how to check it.

Agreement with other codes

100 * max|GKX - ref| / max|ref| over the tracked scan, recomputed by CI from the evidence ledger. The Cyclone rows compare a time-trace fit, not the certified eigenvalue, and KBM is the known outlier. Against GX's repaired launch limit the certified Cyclone eigenpairs agree to 0.3% in growth rate at seven of eight k_y; against converged GS2 GKX is within 1.0% at the three k_y compared. These are agreements on shared test cases, not a ranking of codes. Open items (TEM velocity convergence, stella's mirror-term sensitivity, weakly growing modes) and the figures are in validation.

Nonlinear turbulence

(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 and field energy for Cyclone, Cyclone Miller, W7-X, HSX and KBM, all below 10%. This is a window comparison, not a validation of saturated transport.

Derivatives and optimization

Eigenvalue derivatives use dλ/dp = wᴴ(dA/dp)v / (wᴴv). The nonlinear heat flux over a post-saturation window is differentiated through a block-checkpointed discrete adjoint (figure): on a 16x16x16 Cyclone case over 1024 steps, checkpointing cuts temporary memory from 7.8 GB to 187 MB for about twice the runtime, and the gradient matches finite differences to 1e-11 through 512 steps.

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

VMEX composes its implicit equilibrium derivative with GKX's, so a growth rate, quasilinear flux or nonlinear heat flux sits next to quasisymmetry, aspect ratio and iota in one least-squares problem. examples/10_vmex_optimization/run.py extends VMEX's QA_optimization.py with the nonlinear heat flux as a fourth objective on the vacuum QA ladder (equilibria, traces).

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. A transport claim needs matched, replicated, long post-saturation windows. The quasilinear flux is a ranking and screening measure, not a runtime/TOML absolute-flux predictor (quasilinear, stellarator optimization).

Performance

• Nonlinear step: about 196 ns per Nx*Ny*Nz*Nl*Nm element per step on CPU; 15.6 ms for 64x64x24, Nl = 4, Nm = 8 on one RTX A4000.
• Eigenvalues: the sparse direct shift-invert route certifies the Cyclone eigenpair at n = 3,072 in about 5 s against 32-36 s matrix-free.
• Kinetic electrons: imex-ars3 is 10x faster to solution than RK3 at 64x64x24, and does not pay on tutorial grids.
• Parallel work: independent k_y scans and ensembles run across devices with results identical to the serial run.

Performance · parallelization.

Documentation and development

Documentation: gkx.readthedocs.io. Start with the tutorial and the guidance page; equations and algorithms are in theory, operators and numerics; modules in design.

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 (CONTRIBUTING.md). Figures in docs/_static/readme/ have JSON companions with every plotted number; their generators were removed from scripts/ and are listed in plan/log.md (2026-10-15). Cite GKX with CITATION.cff.

License

MIT; see LICENSE.