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).
pip install gkx
gkxPython 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/.
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)
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.
| Choice | Default | Notes |
|---|---|---|
| Geometry | s-alpha, Miller, VMEC/VMEX, imported | [geometry] model; gkx wout_*.nc for any VMEC |
| Velocity space | Hermite-Laguerre, (Nl, Nm) |
hypercollisions damp the last Hermite moments |
| Collisions | linearized Landau | conserving, H-theorem; Dougherty, Sugama alternatives (operators) |
| End damping | one rate, [time] damp_ends_rate |
same on every route; no time step enters (numerics) |
| Time stepping | rk4 with CFL controller |
imex-ars3 for nonlinear kinetic electrons at 32x32 and up |
| Eigensolver | auto (banded, else adaptive), certified |
shift-invert opt-in |
Resolution, step and solver guidance says what to set for kinetic electrons and electromagnetic runs, and how to check it.
| Case | gamma |
omega |
|---|---|---|
| KAW | 0.0004% | 0.051% |
| ETG | 0.040% | 0.074% |
| W7-X | 0.265% | 0.296% |
| HSX | 0.577% | 0.273% |
| Cyclone Miller | 5.51% | 1.25% |
| Cyclone ITG | 6.83% | 1.59% |
| KBM | 20.0% | 11.1% |
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.
(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.
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).
- Nonlinear step: about 196 ns per
Nx*Ny*Nz*Nl*Nmelement 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-ars3is 10x faster to solution than RK3 at 64x64x24, and does not pay on tutorial grids. - Parallel work: independent
k_yscans and ensembles run across devices with results identical to the serial run.
Performance · parallelization.
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 4CI 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.
MIT; see LICENSE.