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.

Saturated ITG turbulence on a Cyclone flux tube: the perpendicular cut, and the same data as the field-aligned tube in real space. Amplitude is steady across the loop, not growing. Full-rate movie.

Install

pip install gkx
gkx --help
gkx

Python 3.11+. The wheel installs CPU JAX; for GPU or TPU, install an accelerator-enabled JAX wheel from the JAX installation guide.

gkx with no arguments runs a self-contained linear Cyclone demo — no input file, no data download. It takes well under a minute on a laptop CPU, prints the fitted gamma and omega, and writes gkx_default_linear.{toml,summary.json,timeseries.csv,eigenfunction.csv,png}.

The demo runs a deliberately coarse velocity grid (Nl = 7, Nm = 14), so it is not a converged physics result. It is resolved for the case it builds: its step stays under the estimated CFL bound and its horizon is long enough to fit, so the growth rate it prints agrees with that case's certified eigenvalue to better than 1%, and it runs without warnings. A gate holds it there.

Development checkout:

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

Run an equilibrium

Point the executable at a VMEC or VMEX wout file to get a nonlinear ITG run, its figures, a restartable NetCDF bundle, and the resolved deck that reproduces it — without writing a TOML first.

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

--estimate derives a minimum grid from the geometry and explains every entry:

equilibrium: /path/to/wout_circular_tokamak.nc
geometry: shat=+1.7190 q=2.066 nfp=1 |B| wells=1 anisotropy=0.214 -> ky_max*rho >= 2.2
    nx = 96       square perpendicular box (Lx = Ly), so Nx tracks Ny
    ny = 96       tokamak class (nfp = 1, anisotropy 0.214) asks ky_max*rho >= 2.2 (standard); reach ((Ny-1)//3)*dky = 2.21 at dky = 0.071
    nz = 48       scan found Nz weakly coupled (flux 8.47/8.78/8.39 at 24/32/48); floors: 16/2pi-period x 1, 6 x 1 |B| wells
    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)
    dt = 0.0071   explicit CFL bound cfl_fac*cfl/sum(omega_max) at this grid; the adaptive stepper raises it toward the measured ExB limit
 t_max = 400      8 x t_sat ~ 50 hard cap; run_to = "saturation" stops earlier

--estimate=cautious and --estimate=standard select the target-error tier. This is a calibrated starting point, not a convergence proof: tokamak rungs saturated with 64^2 about 8% above the converged 96/128 flux, while the stellarator ladder was still falling at 128^2. Matched Nx/Ny convergence is still yours.

A completed run groups everything under ./<wout-stem>/:

An under-resolved run warns instead of reporting the number:

warning: heat-flux ky cutoff is unresolved: the highest 10% of retained positive-ky modes reach 100% of the spectral peak (warning threshold 10%). Increase Ny at fixed Ly, then repeat matched Nx/Ny convergence; this warning is necessary, not sufficient, for resolution.
saturation: not saturated by the time horizon window=[3.06144,6] heat_flux=3.15837e-06+/-3.70499e-07 rel_sem=0.117307 tau_ac=0.521488

Run a checked-in case

Every shipped example is a runtime TOML the executable accepts directly:

gkx examples/linear/axisymmetric/cyclone.toml
gkx run-runtime-nonlinear \
  --config examples/nonlinear/axisymmetric/runtime_cyclone_nonlinear.toml \
  --steps 200 --out cyclone.out.nc
gkx plot cyclone.out.nc

The first prints the converged eigenvalue and its residual from the certified Krylov path:

runtime: adaptive solve finished with eig=0.0930891-0.282015j residual=6.08e-05 converged=True stable=True
ky=0.3000 gamma=0.093089 omega=0.282015

gkx run auto-detects linear versus nonlinear from [physics]. A deck without an [output] path and without --out writes no files. Examples live under examples/linear, examples/nonlinear, examples/optimization, examples/theory_and_demos, and benchmarks.

Configure a run

One TOML file. Every key has a default, so a working input is short. examples/common_input.toml shows every section with its defaults; it is a template rather than a runnable case, because its vmec_file is supplied by the CLI from the wout you pass. The key-by-key reference is inputs.

[terms] is the debugging lever: setting one coefficient to 0.0 removes exactly that term, which is how most physics gates isolate what they test.

Miller equilibria and VMEC/Boozer flux tubes are also built in-process through the Python API, where the metric coefficients stay differentiable — the path stellarator shape optimization uses. See geometry.

GKX also consumes a VMEX stellarator-mirror hybrid directly from memory, with no file round trip: VMEX owns the field-line closure, the Clebsch metric, and the equal-arc grid, and GKX evaluates its linear and quasilinear objectives on them. The parallel direction is a periodic FFT, so the field line must close; open-ended mirrors are a different model and are not admitted. The shipped case is a closed racetrack, solved to a converged fixed-boundary equilibrium before anything is measured on it, and the field-strength ratio its figure reports is a flux-tube modulation depth rather than a mirror ratio. Geometry, that figure, and the admission review: geometry.

Run control

Saturation also requires the two halves of the window to agree within twice their combined SEM, and holds the field energy Wphi and free energy Wg to the same stationarity, so a flat-looking flux over a still-evolving state does not count as converged.

The stop policy replayed on three tracked runs, each stopped by a different gate: an ITER-model case that clears the flux gates early but waits on Wg stationarity and stops at 66% of t_max; a circular-tokamak case where the relative SEM is binding, stopping at 62%; and an under-resolved D-shape case whose relative SEM never falls below the threshold, so it runs the full horizon and the summary reports not saturated rather than a number. Grey is the horizon that would never have been integrated. Gates: src/gkx/diagnostics/saturation.py.

Performance

Cold wall time and peak memory across the tracked cases, including JAX startup and compilation. The executable enables JAX's persistent compilation cache, so a rerun of an unchanged case is warm. Warm timings, per-case GPU ratios, and profiler artifacts: performance.

On an Apple M3 Max (JAX CPU, float32) the shipped default deck at 96x96x48 runs t_max = 200 in roughly 1.0–1.6 hours, about 97% of it time stepping. Cost is linear in degrees of freedom at about 196 ns per Nx*Ny*Nz*Nl*Nm element per step, flat from 64x64x24 to 96x96x48. Within a step, about 60% is data movement and 39% the FFTs; physics arithmetic is not separately measurable because XLA fuses it into those kernels.

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.

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 with fixed geometry and numerical policy, prepare the compiled simulation once and reuse it. A prepared object compiles one scan of a fixed length, so give it an explicit steps; the shipped decks stop at saturation instead, which decides the length mid-run and cannot be compiled ahead of time. Through the case API that is gkx.prepare(case, steps=N), and warmup() moves the compile out of the first timed solve.

from gkx.solvers_nonlinear_diagnostic_integration import prepare_nonlinear_explicit_diagnostics

simulation = prepare_nonlinear_explicit_diagnostics(
    initial_state, grid, geometry, parameters,
    dt=0.02, steps=400, resolved_diagnostics=False,
)
time, diagnostics, final_state, fields = simulation.run()

The prepared object accepts another same-shape initial state without rebuilding the scan, and a matched cache/parameter PyTree can stay dynamic for autodiff. Full API: gkx.readthedocs.io.

What GKX solves

The gyrokinetic equation for the perturbed distribution of each species, expanded in a Hermite-Laguerre velocity basis:

delta f_s = F_Maxwellian * sum_{m,l}  G_s^{m,l}  psi_m(v_par / v_th)  L_l(mu B / T)

with psi_m = H_m / sqrt(2^m m! sqrt(pi)) the normalized Hermite functions and L_l the Laguerre polynomials. Velocity space becomes two spectral indices: m resolves parallel dynamics (Landau damping, parallel heat flux), l resolves perpendicular dynamics (FLR effects, trapping). The evolved state is one array, G[species, laguerre l, hermite m, ky, kx, z], and each physical effect is a coupling on it:

Perpendicular directions are Fourier (kx, ky); the parallel direction z follows a field line. Electrons are kinetic or Boltzmann. Because m and l are the same kind of index as kx and ky, the whole problem is dense linear algebra on one array. Derivation: theory.

Velocity resolution. Truncating the Hermite ladder at m = M makes its end a reflecting wall, returning free energy as recurrence at t_rec ~ 2 sqrt(M) / (k_par v_th). Since t_rec grows only as sqrt(M), adding moments is a weak fix and the ladder has to absorb instead. Hypercollisions are the default and cut the revival to 0.0009 at M = 16; an opt-in reflectionless closure (Kanekar et al., JPP 81, 305810104 (2015)) needs no tuning but does not beat a well-tuned hypercollision. Tables and scans: numerics.

Collision operators

The models agree in the collisionless limit and separate as collisionality rises. Every shipped matrix is checked against the published closed forms: conservation and Onsager self-adjointness at 5.6e-17 and 8.3e-17 against a 5e-12 gate, published Appendix-C coefficients at 1.1e-16, and an H-theorem maximum eigenvalue of 9.0e-18. Coulomb tables are generated for like-species collisions; a multispecies request is refused rather than silently extrapolated. Equations and convergence panels: operators, metrics in collision_operator_verification.json.

Validation

Every figure is anchored to an exact root, a published coefficient, or a tracked reference run.

Landau damping against the roots of 1 + T_i/T_e + zeta Z(zeta) = 0, from GKX's own linear operator extrapolated to zero collisionality:

A collisionless truncated Hermite system has a purely real spectrum (measured 2.8e-14, gated below 1e-11), so it cannot Landau damp at all — the damping is a transient ending at recurrence, and the root is a pole reached by nu -> 0 extrapolation, not an eigenvalue.

Linear benchmark parity, as 100 * max|GKX - ref| / max|ref| over each scan 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. Detail: benchmarks and the verification matrix.

Differentiate the solver

GKX applies the full gyrokinetic RHS inside a restarted eigensolver, so storage is O(n m) rather than O(n²). The dense path is bounded by memory, not speed: at n = 494,592 a complex128 operator alone would be 3.6 TiB, while the matrix-free solve took 1,504 s.

settings = gkx.AdaptiveLinearEigensolverConfig(tolerance=1e-9, candidate_count=2)

def objective(boundary):
    values = gkx.solver_objective_vector_from_geometry(
        build_solver_geometry(boundary),
        n_laguerre=16, n_hermite=24,
        eigensolver="adaptive-propagator", adaptive_config=settings,
    )
    return values[-1]          # quasilinear transport objective

value, gradient = jax.value_and_grad(objective)(initial_shape)

Reverse mode uses dλ/dp = wᴴ(dA/dp)v / (wᴴv) plus a bordered solve for eigenvector observables — no differentiation through the iteration. The default stays dense so established results are unchanged. See eigensolver.

GKX also differentiates one production nonlinear objective: the physical heat flux averaged over a post-saturation RK window, via a block-checkpointed discrete adjoint storing 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 measured temporary state from 7.82 GB to 187 MB on CPU and 7.80 GB to 148 MB on an RTX A4000, for 1.92x and 1.77x more runtime. The exact discrete differentiation and centered finite differences agree to 1e-11 through 512 steps and 2.7e-9 at 1024, inside the 1e-6 gate, and part at 2048 where chaotic trajectory separation sets the useful window length. nonlinear autodiff.

QA shape optimization through turbulence

QA_optimization.py 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.

How GKX compares

GKX shares its Hermite-Laguerre gyro-moment velocity representation with GX, which makes GX the closest algorithmic and parity reference.

This records scope, not quality: both codes are mature and each is stronger than GKX in areas GKX does not attempt. See related codes.

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

Every figure regenerates from a checked-in script; where a figure has a machine-readable companion, that companion is the artifact of record.

Documentation and development

Full documentation is at gkx.readthedocs.io. Start with the quickstart and input reference, then physics, operators, numerics, geometry, outputs, testing, code structure, and release scope.

pytest
python tools/release/run_test_gates.py fast
ruff check .
python -m sphinx -W -b html docs docs/_build/html

The package-wide CI coverage gate is at least 95%. Physics, convergence, comparison, differentiability, and performance gates are required in addition to line coverage.

Contributing

Bug reports, decks that misbehave and physics questions are all welcome. Start with CONTRIBUTING.md: it covers the environment this code needs (jax >= 0.10.1, and the two precision variables the nightly job sets), what a change is expected to carry, and how published numbers are gated.

License

GKX is distributed under the MIT License.