GKX

GKX is a JAX-native gyrokinetic solver for tokamaks and stellarators: linear stability, nonlinear turbulence, and differentiable analysis that plugs straight into stellarator optimization. It runs on CPUs and GPUs, from a one-line command or from Python.

pip install gkx && gkx

Saturated ITG turbulence on a Cyclone flux tube (32×32×24, t ≈ 300–376), shown as the perpendicular cut a gyrokineticist reads and as the field-aligned tube in real space — the same data twice, because a flux-tube movie that only shows the perpendicular plane hides the parallel elongation that defines the turbulence. Amplitude is steady across the whole loop, not growing.

The loop is an animated image rather than a <video>, which GitHub strips from Markdown: 24 frames at 4 fps, 225 kB of animated WebP, chosen over GIF because no GIF fit the repository's size target at a resolution worth showing. The full-rate movie (1.8 MB, 120 frames at 20 fps) is a release asset. Regenerate the physics with build_turbulence_movie.py. For a physics run, continue its saved saturated state with the same method and CFL policy; only the two displayed cuts are retained:

python tools/artifacts/build_turbulence_movie.py CASE.toml \
  --initial-state saturated_state.npz --snapshots movie_cuts.npz
python tools/artifacts/build_turbulence_movie.py --render-from movie_cuts.npz \
  --output turbulence.mp4

Re-encode the compact README loop from that MP4 with:

mkdir -p frames
ffmpeg -i gkx-cyclone-itg-turbulence.mp4 \
  -vf "fps=4,scale=720:-2:flags=lanczos" frames/f_%04d.png
img2webp -loop 0 -lossy -q 58 -m 6 -d 250 frames/f_*.png \
  -o docs/_static/turbulence_loop.webp

Why GKX

Installation

pip install gkx

For development:

git clone https://github.com/uwplasma/GKX
cd GKX
pip install -e .

Quickstart

Run the built-in linear initial-value example:

gkx

The equivalent gkx entry point is also installed. The default run prints setup, progress, elapsed time, and ETA, then writes its input, summary, time series, eigenfunction, and a two-panel plot in the current directory.

Run a checked-in case or plot an existing result:

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

Generate the small VMEC equilibria used by the self-contained examples:

pip install vmex
cd examples/vmec
./generate_wouts.sh

Full documentation is hosted at gkx.readthedocs.io. Start with the quickstart and input reference for linear, nonlinear, Miller, VMEC, restart, quasilinear, and plotting workflows.

Differentiable matrix-free eigenmodes

Design studies need a few physical modes and their derivatives, not a dense spectrum. 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 the largest tested truncation (n = 494,592) its operator alone would be 3.6 TiB, so it cannot represent the problem at any speed.

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. Residual, overlap, spectral-gap and conditioning gates reject ambiguous modes.

Opt-in: the default stays dense so established results are unchanged. Measured boundaries, the sparse fallback, the physics-aware shift inverse and what is not claimed are in the eigensolver documentation; the inner-solver choice behind it is in numerical defaults.

Validation

Every figure below is anchored to something external — an exact root, a published coefficient, or a tracked reference run — and regenerates from a script in tools/artifacts.

Landau damping against the exact kinetic roots

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

The measurement is harder than it looks, and the figure shows why: a collisionless truncated Hermite system has a purely real spectrum (verified to 3e-14), so it cannot Landau damp at all — what looks like damping is a transient ending at recurrence. The root is not an eigenvalue either; it is a pole reached by nu -> 0 extrapolation.

Linear benchmark parity

Max relative difference against tracked reference results: ETG and KAW below 0.1%, W7-X and HSX below 0.6%, Cyclone ITG 6.8%, KBM 20% — the KBM case is the known outlier and is documented at its claim level rather than smoothed over. Per-case detail: benchmarks and the verification matrix.

Collision Operators

GKX ships five collision models spanning the full hierarchy, selected with one TOML key or one argument to the Python factory:

[time]
collision_operator = "coulomb_finite_kperp"

The models agree in the collisionless limit and separate as collisionality rises, with Lenard-Bernstein over-damping and finite-Larmor gyroaveraging weakening the collisional damping relative to the drift-kinetic operator:

Reproduce with examples/theory_and_demos/collision_operator_comparison.py (--nu-scan draws the figure), or run examples/linear/axisymmetric/cyclone_coulomb_collisions.toml through the executable.

How GKX compares

GKX shares its Hermite-Laguerre gyro-moment velocity representation with GX, which is what makes GX the closest parity reference. The differences that matter for verification are the collision hierarchy and the differentiable geometry path:

This records scope, not quality — both codes are mature and each is stronger than GKX in areas GKX does not attempt. A well-converged Dougherty operator can also be a better physical answer than an unconverged Coulomb one; see related codes for the qualifications.

How the operators are verified

Every shipped matrix is checked against the published closed forms, not only against itself:

The finite-Larmor operator acts on gyrocenter moments, whose conservation is modified by gyroaveraging, so the ordering is the test: the defect must vanish at b = 0 and enter at first order. Coulomb tables are generated for like-species collisions; a multispecies request is refused rather than silently extrapolated. Equations, thresholds, convergence panels and the reproduction recipe: operators.

What GKX Solves

The gyrokinetic equation for the perturbed distribution of each species, expanded in a Hermite-Laguerre velocity basis. Writing the gyrocenter distribution as

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. This turns velocity space into two spectral indices: m resolves parallel dynamics (Landau damping, parallel heat flux), l resolves perpendicular dynamics (FLR effects, trapping). The evolved state is a single array

G[species, laguerre l, hermite m, ky, kx, z]

and each physical effect becomes a specific coupling on it:

Perpendicular directions are Fourier (kx, ky); the parallel direction z follows a field line (flux tube). Electrons are kinetic or Boltzmann.

Why a moment basis: m and l are the same kind of index as kx and ky, so the whole problem is dense linear algebra on one array — which is what a GPU is good at. The cost is that the parallel ladder must be terminated somewhere, which is what the closure section below is about.

Geometry

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

Configuration

One TOML file. Every key has a default, so a working input is short — the sections you actually touch most days are the first five.

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

Full key-by-key reference: inputs.

Runtime and Memory

Cold wall time and peak memory across the tracked cases. Cold times include JAX startup and compilation, which is the right number for "how long does a run take" and the wrong one for kernel speed.

GKX CPU→GPU speedup spans 0.7× to 13.1× depending on the case — the small linear cases are dominated by startup, so the GPU can be slower. Warm timings and profiler artifacts: performance.

Run to saturation

The cheapest way to make a nonlinear run faster is to not integrate past the point where the answer has stopped changing. Diagnosed nonlinear runs therefore stop at saturation by default ([time] run_to = "saturation"). GKX integrates in chunks and, after each one, measures the heat-flux trace with the spin-up phase excluded: it stops when the autocorrelation-corrected relative SEM of the windowed mean falls to saturation_rel_sem (default 5%), the window is long enough (saturation_min_window, ten autocorrelation times when unset), and the two halves of the window agree within twice their combined SEM. t_max stays the hard cap, so nothing is lost if the run never saturates, and the summary reports the window it averaged over together with mean ± SEM — the number you would have computed by hand afterwards. Set run_to = "t_max" or pass --no-until-saturated for a fixed horizon.

GX has no equivalent: it runs a fixed nstep/t_max and can only be halted early by dropping a .stop file next to the run, which is a manual intervention rather than a convergence criterion (src/run_gx.cu:128, src/diagnostics.cu:319-324).

Differentiable Python API

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)
parameters = LinearParams()
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, parameters, cfg.time
)

For repeated nonlinear calls with fixed geometry and numerical policy, prepare the compiled simulation once:

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. A matched rebuilt cache/parameter PyTree can also remain dynamic for autodiff; geometry layout is fixed, and dynamic-geometry compile reuse remains an active differentiability lane.

The planted two-mode inverse problem below recovers two gradient parameters and checks the autodiff Jacobian against finite differences. The single-mode demo in the docs intentionally demonstrates non-identifiability rather than exact parameter recovery.

See differentiable geometry, algorithms, and stellarator optimization for JVP, VJP, implicit differentiation, conditioning, covariance, and finite-difference gates.

Velocity resolution and recurrence

Truncating the Hermite ladder at m = M makes the end of it a reflecting wall: free energy streams up in m, hits the wall, and returns as recurrence at

t_rec ~ 2 sqrt(M) / (k_par v_th)

Nothing after t_rec is physics, and because t_rec grows only as sqrt(M), adding moments is a weak fix — the ladder has to absorb instead. Measured on the free-streaming hierarchy, a hard truncation returns 99.9% of the initial amplitude (it dissipates nothing), while hypercollisions cut that to 0.0002.

Hypercollisions are the default and are what you want. GKX also ships an opt-in reflectionless closure (Kanekar et al., JPP 81, 305810104 (2015)) whose coefficient tends to 1 with resolution, so it needs no tuning and touches only m = M; on measurement it does not beat a well-tuned hypercollision.

Full derivation, both metrics, resolution scans and the closure coefficient: numerics.

Quasilinear Modeling

Use it for: ranking and correlation studies, optimization screening. Not for: absolute heat flux — it is not a runtime/TOML absolute-flux predictor. Absolute-flux promotion stays rejected while the declared Solovev and shaped-pressure stress outliers are retained.

Derivations, calibration splits, uncertainty and holdout gates: quasilinear docs.

Nonlinear autodiff and QA optimization

GKX differentiates one production nonlinear objective: the physical heat flux averaged over a post-saturation RK window. A block-checkpointed discrete adjoint stores O(sqrt(N)) distribution states and works on CPU and GPU.

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)

One derivative, bounded memory. On one host, one 16x16x16 Cyclone case and one 1024-step window, block checkpointing cuts the measured temporary state from 7.82 GB to 187 MB on 36 CPU cores and from 7.80 GB to 148 MB on an RTX A4000, for 1.9x and 1.8x more runtime. The discrete adjoint and centered finite differences agree to 1e-11 through 1024 steps and part at 2048, where chaotic trajectory separation sets the useful window length; longer windows warn. The same gradient agrees between CPU and an A4000 to 1.5e-15. Regenerate every number.

The single QA_optimization.py follows VMEX's vacuum QA mode ladder and adds this heat flux as a fourth tuple. Finite a/L_T=3 and a/L_n=1 drive ITG turbulence. The analytic Jacobian composes VMEX's implicit equilibrium derivative with the exact GKX window derivative.

A small shape step with a preliminary transport effect. Eight low-order boundary coefficients move; aspect ratio and mean iota change by less than 0.05%, while the 3-D LCFS and LCFS Boozer |B| show where the equilibrium changes. The QA residual remains O(10^-3); all panels use the same |B|/<|B|> color scale.

The saturation state is detached and refreshed after accepted stages. The window is a local design derivative. Matched runs measure the accepted direction with independent trajectories.

The startup spike is excluded; the shaded window is measured. The preliminary 12.26% reduction across 24 nominal pairs has a conditional 95% CI of 10.64--13.88%. It is not statistically resolved: 4 of 48 nominal traces fail the published per-trace final-drift test, and the compact outputs contain no resolved heat-flux spectra. Signed ensemble drifts had hidden these failures.

The 24x24 and (Nl,Nm)=(6,12) means remain useful diagnostics, but 2 of 32 traces fail the drift test in each refinement. Promotion requires stationary individual traces, autocorrelation-aware batches, resolved spectral tails, and grid/timestep convergence. See the concise autodiff mathematics and the equations, scripts, matched statistics, audit, and next-run contract.

Parallelization

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 speedup is not claimed until species-first and Hermite-second decomposition, Hermite halo exchange, field-moment collectives, and transport-window identity all pass.

Details: parallelization.

Current Claim Scope

Validated release claims are bounded by the release scope:

• Standard electrostatic/electromagnetic full gyrokinetics is validated only on the promoted cases and observables in the verification matrix.
• Quasilinear outputs are diagnostics and screening models, not universal absolute nonlinear heat-flux predictions.
• Nonlinear optimization evidence requires matched, replicated, long post-transient windows; startup or reduced envelopes are not production evidence.
• W7-X zonal long-window recurrence/damping and W7-X TEM / kinetic-electron extensions are deferred.
• The Sugama, improved-Sugama, and Coulomb operators are verified against their published closed forms and structural invariants, and are validated for like-species collisions; species-coupled Coulomb coefficients remain open.
• Collision operators run on the fixed-step cached integrator. The diffrax, sharded, and Krylov eigenvalue paths reject them rather than silently substituting the built-in diagonal term.
• Production nonlinear domain decomposition and equilibrium ExB flow shear remain open.

Full feature list

• Electrostatic and electromagnetic gyrokinetics with kinetic or Boltzmann species.
• Linear initial-value, dominant-eigenmode, and nonlinear turbulence solvers.
• Matrix-free eigenmodes with certified residuals, branch continuation, and implicit derivatives — O(n m) storage instead of O(n²).
• Analytic s-alpha, Miller, imported VMEC, and differentiable VMEC/Boozer geometry.
• JAX JIT, forward/reverse autodiff, implicit eigenvalue derivatives, and UQ tools.
• Quasilinear transport diagnostics with explicit saturation-rule metadata.
• CPU/GPU execution and production parallelization for independent scans and ensembles.
• Restartable NetCDF output and gkx --plot publication-style figures.
• Five selectable collision operators, from a conserving Lenard-Bernstein model to the full linearized Coulomb (Landau) operator with finite-Larmor-radius effects.

Examples and Documentation

The repository keeps small runnable examples under:

• examples/linear: axisymmetric and stellarator linear runs.
• examples/nonlinear: nonlinear turbulence and restarts.
• examples/optimization: differentiable QA workflows.
• examples/theory_and_demos: numerical and autodiff demonstrations.
• benchmarks: comparison inputs, drivers, and compact result indexes.

Detailed user and developer documentation:

• Physics and equations
• Operators and models
• Numerics and solvers
• Matrix-free eigenmodes and the numerical defaults behind them
• Geometry
• Outputs and plotting
• Testing and validation
• Code structure
• Release and research scope

Testing

pytest
python tools/release/run_test_gates.py fast
python tools/release/run_test_gates.py wide-coverage \
  --shards 48 --timeout 300 --fail-under 95 \
  --pytest-arg=-o --pytest-arg=addopts= --pytest-arg=-m --pytest-arg="not slow"
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.

License

GKX is distributed under the MIT License.