gkx 1.7.1

JAX gyrokinetic solver with Hermite-Laguerre velocity space

GKX

GKX is a JAX-native gyrokinetic solver for linear stability, nonlinear turbulence, differentiable analysis, and stellarator design. It uses Fourier perpendicular coordinates, a Hermite-Laguerre velocity basis, and field-aligned analytic, Miller, or VMEC geometry. The package runs on CPUs and GPUs, exposes a Python API for autodiff and optimization, and provides a simple executable for routine simulations.

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.

Highlights

• Electrostatic and electromagnetic gyrokinetics with kinetic or Boltzmann species.
• Linear initial-value, dominant-eigenmode, and nonlinear turbulence solvers.
• 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.
• A limited conserving Lenard-Bernstein/Dougherty-like collision model, with advanced multispecies and linearized Landau operators remaining research lanes.

Main Validation Results

The release atlas compares growth rates, frequencies, eigenfunctions, and nonlinear transport windows with established gyrokinetic reference results. Promoted cases include Cyclone ITG, Cyclone Miller, KBM, W7-X, and HSX, with ETG and kinetic-electron stress cases kept at their documented claim level.

The exact equations, normalization, grids, boundary conditions, diagnostic windows, tolerances, and artifact provenance are in the benchmark documentation and verification matrix. A visual overlay alone is not treated as parity evidence.

The advanced-collision research lane now generates the complete retained finite-Larmor Coulomb moment algebra and checks it against independent velocity-space projection, spectral convergence, drift-kinetic conservation, finite-wavelength gyro-diffusion, and the H-theorem. The panel below passes all operator-level gates, but does not claim production Landau transport yet; the finite-Larmor collisional-ITG gate and drift-kinetic collisional-zonal gate are closed, and the paper-resolution finite-wavelength zonal traces and velocity sections now pass their independent literature gate. The differentiable driven-current solve is equation-tested, and the direct Coulomb hierarchy is now converged through (P,J)=(20,5) with a maximum nested current change of 0.017%. Its collision-frequency and electric-field normalizations are now closed against the stationary Spitzer problem: all three operators saturate by t nu_ee=50, remain linear over a 100x field scan, and Coulomb reaches the high-charge analytic limit within 7.46%. The arbitrary-order original model reproduces its published 11% low-charge current deficit, while the converged improved model is within 0.31% of Coulomb over the complete ion-charge scan. The converged finite-Larmor ITG artifact is accepted; the generated equal-species operator remains a Python research path because a production input-file selector would imply broader multispecies coefficients that are not yet implemented.

Equations, thresholds, machine-readable results, literature links, and the one-command reproduction recipe are in the collision-operator documentation.

At (P,J)=(24,10), the drift-kinetic traces approach the Xiao residual and the finite-wavelength tails reproduce the published original < improved < Coulomb ordering at both kx rho_i=0.1 and 0.2. The improved model is also closer to Coulomb over t nu_ii <= 10 at both wavenumbers. Equations, velocity-space convergence, compact replay data, and the Figure 12--14 gate are documented in Operators and Terms.

Runtime and Memory

The panel reports measured cold wall time and peak memory for the tracked CPU, GPU, and comparison-code runs. Cold JAX rows include startup and compilation. Prepared Python simulations avoid recompiling a fixed geometry and numerical policy, but their CPU/GPU throughput depends on the software stack and GPU operating state. See performance for profiler artifacts, memory accounting, current reproducibility notes, and the distinction between executable, prepared, and distributed runs.

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.

Quasilinear Modeling

The current quasilinear implementation is a scoped model-development and optimization-screening result. It supports ranking and correlation studies but is not a runtime/TOML absolute-flux predictor. Absolute-flux promotion remains rejected when the declared Solovev and shaped-pressure stress outliers are retained. Model definitions, derivations, calibration splits, uncertainty, residual anatomy, and holdout gates are in the quasilinear documentation.

QA ITG Optimization

The VMEX-style examples append a GKX growth-rate, quasilinear, or nonlinear-window residual to the aspect-ratio, mean-iota, and quasisymmetry objective tuples. The baseline follows the max-mode-5 QA workflow; all transport comparisons use solved VMEC equilibria.

These rows are not promoted turbulent-flux designs. Their matched long post-transient nonlinear audits use converged post-transient heat-flux windows and do not show a statistically significant reduction relative to the strict QA baseline. They are useful negative transfer evidence for improving objective conditioning and optimizer choice.

The RBC(1,1) scan is a landscape and noise/convergence diagnostic, not a source of admitted optimized candidates. It compares linear growth, all shipped quasilinear rules, and replicated long-window nonlinear transport.

Reproducible scripts are in examples/optimization, and full objective equations, optimizer policies, comparison fingerprints, and long-window audits are in the optimization documentation.

Parallelization

Production parallelization currently covers independent k_y scans, quasilinear/UQ ensembles, and file-backed independent tasks with deterministic ordering and serial identity gates. 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 remain diagnostic. Species-first and Hermite-second decomposition, explicit Hermite halo exchange, field-moment collectives, and physical transport-window identity must pass before a nonlinear parallelization speedup is claimed. See 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.
• Production nonlinear domain decomposition, equilibrium ExB flow shear, species-coupled collisions, and linearized Landau/Sugama operators remain open.

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
• 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.