GKX

GKX is a JAX-native gyrokinetic solver for tokamak and stellarator flux tubes. It takes a VMEC equilibrium or an analytic geometry, computes linear stability and nonlinear turbulence in a Hermite-Laguerre velocity basis, and differentiates the whole path end to end on CPUs and GPUs.

• Run what you already have: a VMEC or VMEX wout, a Miller or s-alpha tokamak, or one TOML deck; the executable sizes the grid and explains why.
• Differentiate it: implicit eigenvalue derivatives and a checkpointed discrete adjoint of a nonlinear heat-flux window, through the equilibrium.
• Trust the eigenvalue: every returned eigenpair is certified against the original operator, and under-resolved runs warn instead of reporting a number.
• Model collisions properly: five operators, up to gyrokinetic Coulomb, each checked against its published closed form.
• Trace every number: an evidence ledger ties each published figure to its artifact and generator, and CI recomputes them.

Saturated ITG turbulence on a Cyclone flux tube, as a perpendicular cut and as the field-aligned tube (full-rate movie). Methods, the decisions behind them and the limits of each measurement are in methods and decisions.

Install

pip install gkx
gkx

Python 3.11+. The wheel installs CPU JAX; for GPUs install an accelerator JAX wheel from the JAX installation guide.

gkx with no arguments runs a linear Cyclone demo in under a minute on a laptop CPU, prints gamma and omega, and writes gkx_default_linear.*. Its velocity grid is coarse, but its growth rate agrees with that case's certified eigenvalue to better than 1%, and a gate holds it there.

Run an equilibrium

Point the executable at a VMEC or VMEX wout for a nonlinear ITG run, its figures, a restartable NetCDF bundle and the resolved deck that reproduces it:

gkx wout_circular_tokamak.nc --estimate   # size the grid, explain it, exit
gkx wout_circular_tokamak.nc              # run to saturation
gkx plot wout_circular_tokamak/gkx.out.nc # replot a saved bundle

--estimate derives each grid entry from the geometry and says why:

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 calibrated starting point, not a convergence proof; matched Nx/Ny convergence is still yours. A run stops when the heat flux, field energy and free energy are all stationary, and an unresolved ky cutoff or a missed saturation is reported as such rather than as a number. The saturation gates are in saturation.py.

Every deck under examples/ runs with gkx <deck>.toml. The input reference is at inputs.

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 one scan once and warmup() moves the compile out of the first timed solve. Full API: gkx.readthedocs.io.

Linear physics

Growth rates and frequencies against GX for Cyclone ITG (a), W7-X ITG (b) and KBM (d), and the W7-X eigenfunction at ky rho_i = 0.3 (c), which overlaps GX's to 0.9999999994. The Cyclone and KBM references are provisional: GX's own end-damping defect (below) acts at those resolutions.

Parity across the tracked scans, as 100 * max|GKX - ref| / max|ref| against the references in tools/benchmark_atlas_manifest.toml:

KBM is the known outlier, published at its claim level rather than smoothed over. This is agreement against those tracked scans, not a claim of identical physics options or feature coverage in the reference codes. CI recomputes every percentage from its scan. Detail: benchmarks, verification matrix.

Eigenvalues come from a matrix-free restarted eigensolver that applies the full gyrokinetic RHS, O(n m) rather than O(n^2). At n = 494,592 the dense complex128 operator alone would be 3.6 TiB. Twist-and-shift chains are solved on the modes they couple, and every returned eigenpair is checked against the unprojected operator: solvers.

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 Wg and field energy Wphi, for Cyclone, Cyclone Miller, W7-X, HSX and KBM: the time average of the pointwise relative difference along each trajectory, all below the 10% gate. Evidence: verification matrix.

Collision operators

[time] collision_operator selects Lenard-Bernstein/Dougherty (the default), drift-kinetic Sugama and improved Sugama, drift-kinetic linearized Coulomb (Frei, Ernst & Ricci 2022, Eqs. C9a-C9f), or gyrokinetic Coulomb at finite k_perp (Frei et al. 2021, Eqs. 3.47-3.50). Coulomb tables are generated for like-species collisions; a multispecies request is refused rather than silently extrapolated. Their identities are in the proof tests below. Equations and convergence panels: operators.

Proof tests

Identities the discretization must satisfy exactly, and analytic limits it must reach. Bars are measured by scripts/figures.py; ticks are the tolerances the named tests assert in CI.

The Landau roots solve 1 + T_i/T_e + zeta Z(zeta) = 0. GKX reaches them by extrapolating its own linear operator to zero collisionality: a collisionless truncated Hermite system has a real spectrum, so the damping there is a transient that ends at recurrence, not an eigenvalue.

Differentiate the solver

Eigenvalue derivatives use dλ/dp = wᴴ(dA/dp)v / (wᴴv) plus a bordered solve for eigenvector observables, with no differentiation through the iteration (eigensolver).

GKX also differentiates one production nonlinear objective: the physical heat flux averaged over a post-saturation RK window, through a block-checkpointed discrete adjoint that stores O(sqrt(N)) states.

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 a 1024-step window, checkpointing cuts temporary state from 7.82 GB to 187 MB on CPU and from 7.80 GB to 148 MB on an RTX A4000, for 1.92x and 1.77x the runtime. The derivative agrees with centered finite differences to 1e-11 through 512 steps and 2.7e-9 at 1024, inside the 1e-6 gate; they part at 2048 steps, where chaotic separation sets the useful window. nonlinear autodiff.

QA shape optimization through turbulence

QA_optimization.py in examples/ adds this heat flux as a fourth objective to VMEX's vacuum QA ladder, composing VMEX's implicit equilibrium derivative with the exact GKX window derivative.

Eight low-order boundary coefficients move; aspect ratio changes by +0.0115% and mean iota by -0.044%, while the QA residual goes from 5.88e-4 to 1.54e-3.

These historical traces predate the periodic hypercollision correction and must be regenerated; they are not evidence for the current operator. The preliminary 12.26% reduction across 24 nominal pairs has a conditional 95% CI of 10.64-13.88%, and is not statistically resolved: 4 of 48 nominal traces fail the published per-trace final-drift test. Promotion requires stationary individual traces, autocorrelation-aware batches, resolved spectral tails, and grid/timestep convergence; nonlinear optimization evidence requires matched, replicated, long post-saturation windows. Every row is in qa_transport_summary.csv; the campaign is in stellarator optimization.

Performance

Cold wall time and peak memory across the tracked cases, including JAX startup and compilation; the executable caches compilations, so a rerun is warm. On an Apple M3 Max (CPU, float32), cost is about 196 ns per Nx*Ny*Nz*Nl*Nm element per step, flat from 64x64x24 to 96x96x48. Warm timings, GPU ratios and profiles: performance.

Parallelism is production for independent k_y scans, quasilinear/UQ ensembles, and file-backed tasks, all deterministically ordered and serial-identity gated. Sensitivity sweeps can use the same deterministic independent-work reconstruction, but they need a dedicated matched scaling artifact before any speedup claim is promoted; nonlinear whole-state and domain decomposition stay diagnostic only. Details: parallelization.

Claim scope

Release claims are bounded by the release scope.

Quasilinear outputs are for ranking, correlation studies, and optimization screening. They are not a runtime/TOML absolute-flux predictor: absolute-flux promotion stays rejected while the declared Solovev and shaped-pressure stress outliers are retained, the best tracked candidate misses the 0.35 transport gate, and the positive-growth mixing-length rule predicts zero for HSX and W7-X where the tracked nonlinear windows are finite. Derivations, calibration splits, and holdout gates: quasilinear.

Collision operators are validated for like-species collisions and run on the fixed-step cached integrator. W7-X zonal long-window recurrence/damping and W7-X TEM / kinetic-electron extensions are deferred. Production nonlinear domain decomposition and equilibrium ExB flow shear remain open.

Reproducing the figures

The four README figures under docs/_static/readme/ regenerate with one command, configured by scripts/figures.toml:

JAX_ENABLE_X64=true python scripts/figures.py          # or name one: linear, nonlinear, proof_tests

Each PNG has a JSON companion holding every plotted number. The proof-test figure takes about 7 CPU-minutes (the Landau collisionality scan); the others take seconds and read only tracked files. The autodiff, QA and performance figures name their generators in the evidence ledger; the turbulence movie's recipe is in scripts/artifacts/build_turbulence_movie.py.

Documentation and development

Full documentation is at gkx.readthedocs.io: start with the quickstart, then physics, numerics, geometry, outputs and testing.

git clone https://github.com/uwplasma/GKX
cd GKX
pip install -e ".[dev]"
pytest
python scripts/checks/run_test_gates.py fast

CI requires at least 95% line coverage plus the physics, convergence, comparison, differentiability and performance gates. Bug reports, misbehaving decks and physics questions are welcome: see CONTRIBUTING.md and cite GKX with CITATION.cff.

License

GKX is distributed under the MIT License.